diff --git a/mempalace/embedding.py b/mempalace/embedding.py index 8952f91..9c701bc 100644 --- a/mempalace/embedding.py +++ b/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("") + 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 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 (%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} diff --git a/tests/test_embeddinggemma.py b/tests/test_embeddinggemma.py index 54d78b3..eb512ed 100644 --- a/tests/test_embeddinggemma.py +++ b/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 == "" 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 in source"]) + + assert captured["input_ids"].tolist() == [[2, 3, 1]] + assert np.asarray(result).shape == (1, 384) + assert "remapping to " in caplog.text + assert patched_lazy_load["hf_hub_download"] == 3