diff --git a/src/memsearch/embeddings/onnx.py b/src/memsearch/embeddings/onnx.py index de99e31b..014c17b0 100644 --- a/src/memsearch/embeddings/onnx.py +++ b/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: diff --git a/tests/test_embeddings_onnx_inputs.py b/tests/test_embeddings_onnx_inputs.py new file mode 100644 index 00000000..7dee4d4b --- /dev/null +++ b/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"}