feat(embed): pooling gate that refuses unverified models (v0.71.36 Part A)

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"

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)