fix(embedding): remap unsupported EmbeddingGemma token IDs

mempalace/embedding.py

@@ -220,6 +220,43 @@ _EMBEDDINGGEMMA_MAX_LEN = 2048
 _EMBEDDINGGEMMA_BATCH_SIZE = 32
 
 
+def _sanitize_embeddinggemma_input_ids(tokenizer, input_ids, np):
+    """Replace tokenizer-only IDs that the text ONNX model cannot embed."""
+    model_vocab_size = tokenizer.get_vocab_size(with_added_tokens=False)
+    out_of_range = (input_ids < 0) | (input_ids >= model_vocab_size)
+
+    if not np.any(out_of_range):
+        return input_ids
+
+    unknown_token_id = tokenizer.token_to_id("<unk>")
+    if unknown_token_id is None or not 0 <= unknown_token_id < model_vocab_size:
+        raise RuntimeError(
+            "EmbeddingGemma tokenizer produced token IDs outside the ONNX "
+            "text vocabulary, but no valid <unk> token is available"
+        )
+
+    invalid_ids = sorted({int(token_id) for token_id in input_ids[out_of_range]})
+    warning_key = (
+        "embeddinggemma-out-of-range-token-ids",
+        model_vocab_size,
+        tuple(invalid_ids),
+    )
+
+    if warning_key not in _WARNED:
+        logger.warning(
+            "EmbeddingGemma tokenizer produced token IDs outside the ONNX "
+            "text vocabulary (size=%d): %s; remapping to <unk> (%d)",
+            model_vocab_size,
+            invalid_ids,
+            unknown_token_id,
+        )
+        _WARNED.add(warning_key)
+
+    sanitized = input_ids.copy()
+    sanitized[out_of_range] = unknown_token_id
+    return sanitized
+
+
 class EmbeddinggemmaONNX:
     """ChromaDB-compatible EF using embeddinggemma-300m ONNX (q8, MRL→384d).
 
@@ -340,6 +377,11 @@ class EmbeddinggemmaONNX:
             texts = [_EMBEDDINGGEMMA_PREFIX + t for t in chunk]
             encs = self._tokenizer.encode_batch(texts)
             input_ids = np.asarray([e.ids for e in encs], dtype=np.int64)
+            input_ids = _sanitize_embeddinggemma_input_ids(
+                self._tokenizer,
+                input_ids,
+                np,
+            )
             attention_mask = np.asarray([e.attention_mask for e in encs], dtype=np.int64)
             outputs = self._session.run(
                 None, {"input_ids": input_ids, "attention_mask": attention_mask}

tests/test_embeddinggemma.py

@@ -70,6 +70,12 @@ class _FakeTokenizer:
         self._truncation_enabled = True
         self._truncation_max = max_length
 
+    def get_vocab_size(self, with_added_tokens=True):
+        return 262145 if with_added_tokens else 262144
+
+    def token_to_id(self, token):
+        return 3 if token == "<unk>" else None
+
     def encode_batch(self, texts):
         class _Enc:
             def __init__(self, n):
@@ -397,3 +403,59 @@ def test_config_embedding_model_default_is_minilm(monkeypatch):
 
     monkeypatch.delenv("MEMPALACE_EMBEDDING_MODEL", raising=False)
     assert MempalaceConfig().embedding_model == "minilm"
+
+
+def test_out_of_range_added_token_is_remapped_to_unknown(
+    patched_lazy_load,
+    monkeypatch,
+    caplog,
+):
+    class EncodingWithAddedToken:
+        ids = [2, 262144, 1]
+        attention_mask = [1, 1, 1]
+
+    def encode_with_added_token(_self, texts):
+        return [EncodingWithAddedToken() for _ in texts]
+
+    monkeypatch.setattr(
+        _FakeTokenizer,
+        "encode_batch",
+        encode_with_added_token,
+    )
+
+    captured = {}
+    fake_session_class = _make_fake_session()
+
+    class BoundsCheckingSession(fake_session_class):
+        def run(self, output_names, feed):
+            captured["input_ids"] = feed["input_ids"].copy()
+
+            assert np.all(feed["input_ids"] >= 0)
+            assert np.all(feed["input_ids"] < 262144)
+
+            return super().run(
+                output_names,
+                feed,
+            )
+
+    import onnxruntime
+
+    monkeypatch.setattr(
+        onnxruntime,
+        "InferenceSession",
+        lambda *_args, **_kwargs: BoundsCheckingSession(),
+    )
+
+    caplog.set_level(
+        "WARNING",
+        logger=embedding.__name__,
+    )
+
+    embedding_function = embedding.EmbeddinggemmaONNX()
+
+    result = embedding_function(["literal <image_soft_token> in source"])
+
+    assert captured["input_ids"].tolist() == [[2, 3, 1]]
+    assert np.asarray(result).shape == (1, 384)
+    assert "remapping to <unk>" in caplog.text
+    assert patched_lazy_load["hf_hub_download"] == 3