From aed7881fc8659cfaf3e3eef8fc29ddddcc2da1d3 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Wed, 15 Jul 2026 23:05:20 +0500 Subject: [PATCH] feat(embed): pooling gate that refuses unverified models (v0.71.36 Part A) --- src/soup_cli/utils/embed.py | 106 ++++++++++++++++++++++++++++++++++++ tests/test_v07136.py | 87 +++++++++++++++++++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 src/soup_cli/utils/embed.py create mode 100644 tests/test_v07136.py diff --git a/src/soup_cli/utils/embed.py b/src/soup_cli/utils/embed.py new file mode 100644 index 0000000..d36c132 --- /dev/null +++ b/src/soup_cli/utils/embed.py @@ -0,0 +1,106 @@ +"""v0.71.36 — the one embedding kernel (Data Moat II). + +``transformers.AutoModel`` + attention-masked mean-pool + L2-normalize. +This is exactly what sentence-transformers does for MiniLM-class models +(their own model card documents the equivalence), so Soup needs NO new +dependency — ``transformers`` already ships in the ``[train]`` extra. + +Torch is imported lazily inside :func:`embed_texts`; this module must stay +importable on the light core. +""" + +from __future__ import annotations + +import json +from typing import Optional + +# Models whose pooling is verified pure-mean. Short-circuits the hub fetch. +# all-mpnet-base-v2 already ships as the `ra-dit-retriever` recipe base. +POOLING_ALLOWLIST: dict[str, str] = { + "sentence-transformers/all-minilm-l6-v2": "mean", + "sentence-transformers/all-minilm-l12-v2": "mean", + "sentence-transformers/all-mpnet-base-v2": "mean", +} + +DEFAULT_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" + +_MAX_MODEL_ID_CHARS = 256 + +# Keys in 1_Pooling/config.json that are NOT plain mean pooling. If any is +# true we refuse — mean-pooling such a model emits wrong vectors silently. +_NON_MEAN_POOLING_KEYS = ( + ("pooling_mode_cls_token", "cls"), + ("pooling_mode_max_tokens", "max"), + ("pooling_mode_weightedmean_tokens", "weightedmean"), + ("pooling_mode_lasttoken", "lasttoken"), +) + + +def _require_model_id(model_id: object) -> str: + if isinstance(model_id, bool) or not isinstance(model_id, str): + raise TypeError( + f"model_id must be str, got {type(model_id).__name__}" + ) + cleaned = model_id.strip() + if not cleaned: + raise ValueError("model_id must be a non-empty string") + if "\x00" in cleaned: + raise ValueError("model_id must not contain null bytes") + if len(cleaned) > _MAX_MODEL_ID_CHARS: + raise ValueError( + f"model_id too long (max {_MAX_MODEL_ID_CHARS} chars)" + ) + return cleaned + + +def _fetch_pooling_config(model_id: str) -> Optional[dict]: + """Read ``1_Pooling/config.json`` from the hub repo. + + Returns None when the file is absent or unreadable — the caller then + REFUSES rather than assuming mean pooling. + """ + try: + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id=model_id, filename="1_Pooling/config.json" + ) + with open(path, "r", encoding="utf-8") as handle: + data = json.load(handle) + return data if isinstance(data, dict) else None + except Exception: + return None + + +def resolve_pooling(model_id: str) -> str: + """Return ``"mean"`` for a verified mean-pooled model, else raise. + + Never guesses. An unverifiable model is refused, because silently + mean-pooling a CLS model produces wrong vectors and every downstream + number is then quietly garbage. + """ + cleaned = _require_model_id(model_id) + allow = POOLING_ALLOWLIST.get(cleaned.lower()) + if allow is not None: + return allow + + config = _fetch_pooling_config(cleaned) + if config is None: + raise ValueError( + f"cannot verify pooling for {cleaned!r}: no 1_Pooling/config.json " + "and not in the verified allowlist. Soup refuses rather than " + "assume mean-pooling (wrong vectors would be silent). Use " + f"{DEFAULT_EMBED_MODEL} or another allowlisted model." + ) + for key, label in _NON_MEAN_POOLING_KEYS: + if config.get(key): + raise ValueError( + f"{cleaned!r} uses {label!r} pooling; Soup's embedding kernel " + "implements mean-pooling only. Refusing rather than emitting " + "wrong vectors." + ) + if not config.get("pooling_mode_mean_tokens"): + raise ValueError( + f"{cleaned!r} does not declare mean-token pooling; refusing." + ) + return "mean" diff --git a/tests/test_v07136.py b/tests/test_v07136.py new file mode 100644 index 0000000..3e1a3ec --- /dev/null +++ b/tests/test_v07136.py @@ -0,0 +1,87 @@ +"""v0.71.36 — Data Moat II: semantic layer + canaries + replay.""" + +from __future__ import annotations + +import re + +import pytest + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]") + + +def _clean(text: str) -> str: + """Strip ANSI + collapse whitespace. + + Rich splits flag names with color codes and wraps at terminal width; + asserting on raw output has broken CI three times (v0.71.26/32/35). + """ + return " ".join(_ANSI_RE.sub("", text).split()) + + +class TestResolvePooling: + def test_allowlisted_model_short_circuits(self): + from soup_cli.utils.embed import resolve_pooling + + assert resolve_pooling("sentence-transformers/all-MiniLM-L6-v2") == "mean" + assert resolve_pooling("sentence-transformers/all-mpnet-base-v2") == "mean" + + def test_allowlist_is_case_insensitive(self): + from soup_cli.utils.embed import resolve_pooling + + assert resolve_pooling("Sentence-Transformers/All-MiniLM-L6-v2") == "mean" + + def test_cls_pooling_config_is_refused(self, monkeypatch): + from soup_cli.utils import embed + + monkeypatch.setattr( + embed, + "_fetch_pooling_config", + lambda mid: { + "pooling_mode_mean_tokens": False, + "pooling_mode_cls_token": True, + }, + ) + with pytest.raises(ValueError, match="cls"): + embed.resolve_pooling("some/cls-model") + + def test_mean_pooling_config_is_accepted(self, monkeypatch): + from soup_cli.utils import embed + + monkeypatch.setattr( + embed, + "_fetch_pooling_config", + lambda mid: { + "pooling_mode_mean_tokens": True, + "pooling_mode_cls_token": False, + }, + ) + assert embed.resolve_pooling("some/mean-model") == "mean" + + def test_unfetchable_config_is_refused_not_assumed(self, monkeypatch): + from soup_cli.utils import embed + + monkeypatch.setattr(embed, "_fetch_pooling_config", lambda mid: None) + with pytest.raises(ValueError, match="cannot verify pooling"): + embed.resolve_pooling("some/unknown-model") + + def test_refusal_names_what_it_found(self, monkeypatch): + from soup_cli.utils import embed + + monkeypatch.setattr( + embed, + "_fetch_pooling_config", + lambda mid: { + "pooling_mode_mean_tokens": False, + "pooling_mode_max_tokens": True, + }, + ) + with pytest.raises(ValueError) as exc: + embed.resolve_pooling("some/max-model") + assert "max" in str(exc.value) + + @pytest.mark.parametrize("bad", ["", " ", None, 123, True]) + def test_bad_model_id_rejected(self, bad): + from soup_cli.utils.embed import resolve_pooling + + with pytest.raises((ValueError, TypeError)): + resolve_pooling(bad)