From 8a469da517317510dca25c17749f49879bd07e62 Mon Sep 17 00:00:00 2001 From: Jeffrey Sperling Date: Sat, 4 Jul 2026 23:26:08 -0700 Subject: [PATCH 1/3] feat: support ONNX models that require token_type_ids BERT-family ONNX exports (Xenova/all-MiniLM-L6-v2, Xenova/bge-small-en-v1.5) declare a token_type_ids input, and session.run() requires every declared input to be fed, so such models fail to embed with "Required inputs (['token_type_ids']) are missing". Detect the declared inputs once at init and feed all-zero segment ids when required (single-sequence embedding). This unlocks a small-model tier for bulk ingest: bge-small-en-v1.5 (33M, CLS-pooling-native, matching this provider's pooling) embeds a 64-text batch in ~0.02s vs ~6.6s for the 568M default on the same CPU. --- src/memsearch/embeddings/onnx.py | 8 ++++ tests/test_embeddings_onnx_inputs.py | 59 ++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 tests/test_embeddings_onnx_inputs.py diff --git a/src/memsearch/embeddings/onnx.py b/src/memsearch/embeddings/onnx.py index de99e31b..2b72a60e 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. + self._input_names = {i.name for i in self._session.get_inputs()} + self._needs_token_type_ids = "token_type_ids" in self._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..de036951 --- /dev/null +++ b/tests/test_embeddings_onnx_inputs.py @@ -0,0 +1,59 @@ +"""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 _StubTokenizer: + def encode_batch(self, texts): + class E: + ids = [1, 2, 3] + attention_mask = [1, 1, 0] + + return [E() 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"} From a7c882b0f28cc6c34e9b103efed361e76fcf01f2 Mon Sep 17 00:00:00 2001 From: Jeffrey Sperling Date: Sat, 4 Jul 2026 23:27:28 -0700 Subject: [PATCH 2/3] test: build stub encoding in __init__ to satisfy RUF012 --- tests/test_embeddings_onnx_inputs.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_embeddings_onnx_inputs.py b/tests/test_embeddings_onnx_inputs.py index de036951..7dee4d4b 100644 --- a/tests/test_embeddings_onnx_inputs.py +++ b/tests/test_embeddings_onnx_inputs.py @@ -10,13 +10,15 @@ 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): - class E: - ids = [1, 2, 3] - attention_mask = [1, 1, 0] - - return [E() for _ in texts] + return [_StubEncoding() for _ in texts] class _StubSession: From d670b458ef1b3359066557e833f74a1d93fd370a Mon Sep 17 00:00:00 2001 From: Jeffrey Sperling Date: Sun, 5 Jul 2026 00:24:05 -0700 Subject: [PATCH 3/3] Keep session input names local to __init__ --- src/memsearch/embeddings/onnx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/memsearch/embeddings/onnx.py b/src/memsearch/embeddings/onnx.py index 2b72a60e..014c17b0 100644 --- a/src/memsearch/embeddings/onnx.py +++ b/src/memsearch/embeddings/onnx.py @@ -56,8 +56,8 @@ def __init__( # 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. - self._input_names = {i.name for i in self._session.get_inputs()} - self._needs_token_type_ids = "token_type_ids" in self._input_names + 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