feat: support ONNX models that require token_type_ids

src/memsearch/embeddings/onnx.py

@@ -53,6 +53,11 @@ def __init__(
         self._session = ort.InferenceSession(model_path)
         self._output_names = [o.name for o in self._session.get_outputs()]
         self._has_dense_vecs = "dense_vecs" in self._output_names
+        # BERT-family exports (e.g. Xenova/all-MiniLM-L6-v2) declare a
+        # token_type_ids input; XLM-R-family exports (e.g. bge-m3) do not.
+        # Session.run() requires every declared input to be fed.
+        input_names = {i.name for i in self._session.get_inputs()}
+        self._needs_token_type_ids = "token_type_ids" in input_names
         self._model = model
 
         # Detect dimension from a probe embedding
@@ -141,6 +146,9 @@ def _encode(self, texts: list[str]) -> list[list[float]]:
             "input_ids": input_ids,
             "attention_mask": attention_mask,
         }
+        if self._needs_token_type_ids:
+            # Single-sequence embedding: segment ids are all zero.
+            feed["token_type_ids"] = np.zeros_like(input_ids)
         outputs = self._session.run(None, feed)
 
         if self._has_dense_vecs:

tests/test_embeddings_onnx_inputs.py

@@ -0,0 +1,61 @@
+"""Tests for ONNX input-feed construction (token_type_ids compatibility).
+
+Uses a stub session so no onnxruntime model download is needed.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from memsearch.embeddings.onnx import OnnxEmbedding
+
+
+class _StubEncoding:
+    def __init__(self) -> None:
+        self.ids = [1, 2, 3]
+        self.attention_mask = [1, 1, 0]
+
+
+class _StubTokenizer:
+    def encode_batch(self, texts):
+        return [_StubEncoding() for _ in texts]
+
+
+class _StubSession:
+    """Mimics ort.InferenceSession run(); records the feed it was given."""
+
+    def __init__(self) -> None:
+        self.last_feed: dict | None = None
+
+    def run(self, _output_names, feed):
+        self.last_feed = feed
+        batch = len(feed["input_ids"])
+        return [np.ones((batch, 4), dtype=np.float32)]
+
+
+def _make(needs_token_type_ids: bool) -> tuple[OnnxEmbedding, _StubSession]:
+    e = object.__new__(OnnxEmbedding)
+    session = _StubSession()
+    e._tokenizer = _StubTokenizer()
+    e._session = session
+    e._output_names = ["dense_vecs"]
+    e._has_dense_vecs = True
+    e._needs_token_type_ids = needs_token_type_ids
+    return e, session
+
+
+def test_token_type_ids_fed_as_zeros_when_model_requires() -> None:
+    e, session = _make(needs_token_type_ids=True)
+    e._encode(["hello", "world"])
+    assert session.last_feed is not None
+    assert set(session.last_feed) == {"input_ids", "attention_mask", "token_type_ids"}
+    tti = session.last_feed["token_type_ids"]
+    assert tti.shape == session.last_feed["input_ids"].shape
+    assert not tti.any()
+
+
+def test_token_type_ids_omitted_when_model_does_not_declare_it() -> None:
+    e, session = _make(needs_token_type_ids=False)
+    e._encode(["hello"])
+    assert session.last_feed is not None
+    assert set(session.last_feed) == {"input_ids", "attention_mask"}