diff --git a/src/soup_cli/utils/embed.py b/src/soup_cli/utils/embed.py index d36c132..abbf253 100644 --- a/src/soup_cli/utils/embed.py +++ b/src/soup_cli/utils/embed.py @@ -104,3 +104,116 @@ def resolve_pooling(model_id: str) -> str: f"{cleaned!r} does not declare mean-token pooling; refusing." ) return "mean" + + +# --------------------------------------------------------------------------- +# Batched encode +# --------------------------------------------------------------------------- + +_MAX_ROWS = 200_000 +_MAX_CHARS_PER_ROW = 8_192 +_MAX_BATCH_SIZE = 512 +_MAX_SEQ_TOKENS = 512 + + +def _mean_pool(hidden, attention_mask): + """Attention-masked mean over the token axis. Padding contributes 0. + + An unmasked ``hidden.mean(dim=1)`` averages padded positions too, which + silently drags every vector toward the pad embedding. + """ + mask = attention_mask.unsqueeze(-1).to(hidden.dtype) + summed = (hidden * mask).sum(dim=1) + counts = mask.sum(dim=1).clamp(min=1e-9) # never divide by zero + return summed / counts + + +def _l2_normalize(vectors): + """Row-wise L2 normalize so cosine == dot product. Zero rows stay zero.""" + import numpy as np + + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + norms = np.where(norms < 1e-12, 1.0, norms) + return (vectors / norms).astype(np.float32) + + +def _validate_texts(texts: object) -> list: + """Coerce + bound-check the input rows. + + A bare ``str`` is rejected: it is a sequence of characters, so accepting + it would silently embed one row per letter. + """ + if isinstance(texts, (str, bytes)) or not hasattr(texts, "__len__"): + raise TypeError("texts must be a sequence of str") + items = list(texts) + if not items: + raise ValueError("texts must contain at least one text") + if len(items) > _MAX_ROWS: + raise ValueError( + f"too many texts ({len(items)}); cap is {_MAX_ROWS}. Sample the " + "dataset first (`soup data sample`) — Soup refuses rather than " + "silently subsampling." + ) + for idx, item in enumerate(items): + if not isinstance(item, str): + raise TypeError( + f"texts[{idx}] must be str, got {type(item).__name__}" + ) + return [item[:_MAX_CHARS_PER_ROW] for item in items] + + +def _require_batch_size(batch_size: object) -> int: + if isinstance(batch_size, bool) or not isinstance(batch_size, int): + raise TypeError( + f"batch_size must be int, got {type(batch_size).__name__}" + ) + if batch_size < 1 or batch_size > _MAX_BATCH_SIZE: + raise ValueError( + f"batch_size must be in [1, {_MAX_BATCH_SIZE}], got {batch_size}" + ) + return batch_size + + +def embed_texts( + texts, + *, + model_id: str = DEFAULT_EMBED_MODEL, + device: str = "auto", + batch_size: int = 32, +): + """Embed ``texts`` -> an ``(n, d)`` float32 array with L2-normalized rows. + + Torch / transformers / numpy are imported lazily so this module stays + importable on the light core (a ``pip install soup-cli`` without the + ``[train]`` extra). + """ + items = _validate_texts(texts) + batch = _require_batch_size(batch_size) + # Refuse an unverified model BEFORE any download starts. + resolve_pooling(model_id) + + import numpy as np + import torch + from transformers import AutoModel, AutoTokenizer + + from soup_cli.utils.live_eval import resolve_device + + dev = resolve_device(device) + tokenizer = AutoTokenizer.from_pretrained(model_id) + model = AutoModel.from_pretrained(model_id).to(dev) + model.eval() + + chunks = [] + with torch.no_grad(): + for start in range(0, len(items), batch): + encoded = tokenizer( + items[start: start + batch], + padding=True, + truncation=True, + max_length=_MAX_SEQ_TOKENS, + return_tensors="pt", + ).to(dev) + out = model(**encoded) + pooled = _mean_pool(out.last_hidden_state, encoded["attention_mask"]) + chunks.append(pooled.float().cpu().numpy()) + return _l2_normalize(np.vstack(chunks)) diff --git a/tests/test_v07136.py b/tests/test_v07136.py index 255b9dc..3394937 100644 --- a/tests/test_v07136.py +++ b/tests/test_v07136.py @@ -117,3 +117,96 @@ class TestResolvePooling: with pytest.raises((ValueError, TypeError)): resolve_pooling(bad) + + +class TestEmbedTexts: + def test_rejects_over_cap(self): + from soup_cli.utils.embed import _MAX_ROWS, embed_texts + + with pytest.raises(ValueError, match="too many texts"): + embed_texts(["x"] * (_MAX_ROWS + 1)) + + def test_rejects_empty_list(self): + from soup_cli.utils.embed import embed_texts + + with pytest.raises(ValueError, match="at least one text"): + embed_texts([]) + + def test_rejects_non_string_row(self): + from soup_cli.utils.embed import embed_texts + + with pytest.raises(TypeError, match=r"texts\[1\] must be str"): + embed_texts(["ok", 42]) + + @pytest.mark.parametrize("bad", [0, -1, True]) + def test_rejects_bad_batch_size(self, bad): + from soup_cli.utils.embed import embed_texts + + with pytest.raises((ValueError, TypeError)): + embed_texts(["x"], batch_size=bad) + + def test_validate_texts_truncates_long_rows(self): + from soup_cli.utils.embed import _MAX_CHARS_PER_ROW, _validate_texts + + out = _validate_texts(["a" * (_MAX_CHARS_PER_ROW + 500)]) + assert len(out[0]) == _MAX_CHARS_PER_ROW + + def test_validate_texts_rejects_bare_string(self): + """A bare str is a sequence of chars — almost always a caller bug.""" + from soup_cli.utils.embed import _validate_texts + + with pytest.raises(TypeError, match="sequence of str"): + _validate_texts("not a list") + + def test_mean_pool_masks_padding(self): + """Padded positions must NOT contribute to the mean.""" + import numpy as np + + from soup_cli.utils.embed import _mean_pool + + torch = pytest.importorskip("torch") + # 3 tokens; the last is padding carrying a huge value that would + # wreck an unmasked mean (999 -> ~334 instead of 2.0). + hidden = torch.tensor([[[1.0, 1.0], [3.0, 3.0], [999.0, 999.0]]]) + mask = torch.tensor([[1, 1, 0]]) + out = _mean_pool(hidden, mask).numpy() + np.testing.assert_allclose(out, np.array([[2.0, 2.0]]), rtol=1e-6) + + def test_mean_pool_all_padding_does_not_divide_by_zero(self): + import numpy as np + + from soup_cli.utils.embed import _mean_pool + + torch = pytest.importorskip("torch") + hidden = torch.tensor([[[5.0, 5.0]]]) + mask = torch.tensor([[0]]) + out = _mean_pool(hidden, mask).numpy() + assert np.isfinite(out).all() + + def test_l2_normalize_rows(self): + import numpy as np + + from soup_cli.utils.embed import _l2_normalize + + vecs = np.array([[3.0, 4.0], [0.0, 0.0]], dtype=np.float32) + out = _l2_normalize(vecs) + np.testing.assert_allclose(np.linalg.norm(out[0]), 1.0, rtol=1e-6) + # a zero vector must not become NaN + assert np.isfinite(out[1]).all() + + def test_l2_normalize_makes_cosine_a_dot_product(self): + import numpy as np + + from soup_cli.utils.embed import _l2_normalize + + vecs = _l2_normalize(np.array([[2.0, 0.0], [0.0, 5.0]], dtype=np.float32)) + assert float(vecs[0] @ vecs[1]) == pytest.approx(0.0, abs=1e-6) + assert float(vecs[0] @ vecs[0]) == pytest.approx(1.0, abs=1e-6) + + def test_refuses_unverified_model_before_any_download(self, monkeypatch): + """resolve_pooling must gate embed_texts BEFORE a model is fetched.""" + from soup_cli.utils import embed + + monkeypatch.setattr(embed, "_fetch_pooling_config", lambda mid: None) + with pytest.raises(ValueError, match="cannot verify pooling"): + embed.embed_texts(["hello"], model_id="org/unverified-encoder")