230 lines
8.3 KiB
Python
230 lines
8.3 KiB
Python
"""Offline tests for EmbeddinggemmaONNX.
|
|
|
|
The real ONNX model is ~300 MB and pulled from HuggingFace on first use, so
|
|
these tests mock huggingface_hub.hf_hub_download, tokenizers.Tokenizer, and
|
|
onnxruntime.InferenceSession to keep CI fast and network-free.
|
|
|
|
Skipped when the multilingual extra isn't installed (huggingface_hub/
|
|
tokenizers/numpy) — CI runs only core deps by default.
|
|
"""
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
np = pytest.importorskip("numpy")
|
|
pytest.importorskip("huggingface_hub")
|
|
pytest.importorskip("tokenizers")
|
|
|
|
import mempalace.embedding as embedding # noqa: E402 (after importorskip)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def isolate_embedding_state(monkeypatch):
|
|
monkeypatch.setattr(embedding, "_EF_CACHE", {})
|
|
monkeypatch.setattr(embedding, "_WARNED", set())
|
|
|
|
|
|
def _make_fake_session(out_dim=768):
|
|
"""Fake onnxruntime InferenceSession that returns a deterministic tensor.
|
|
|
|
Shape: (batch, out_dim). The values aren't important — tests check shape,
|
|
truncation, and L2-normalization, not numerical correctness.
|
|
"""
|
|
|
|
class _Output:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
class _Session:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def get_outputs(self):
|
|
return [_Output("last_hidden_state"), _Output("sentence_embedding")]
|
|
|
|
def run(self, _output_names, feed):
|
|
batch = feed["input_ids"].shape[0]
|
|
# Deterministic non-trivial values so L2-norm isn't degenerate.
|
|
sent = np.arange(batch * out_dim, dtype=np.float32).reshape(batch, out_dim) + 1.0
|
|
last_hidden = np.zeros((batch, feed["input_ids"].shape[1], out_dim), dtype=np.float32)
|
|
return [last_hidden, sent]
|
|
|
|
return _Session
|
|
|
|
|
|
class _FakeTokenizer:
|
|
"""Stand-in for tokenizers.Tokenizer with the methods _lazy_load uses."""
|
|
|
|
def __init__(self):
|
|
self._padding_enabled = False
|
|
self._truncation_enabled = False
|
|
self._truncation_max = None
|
|
|
|
def enable_padding(self):
|
|
self._padding_enabled = True
|
|
|
|
def enable_truncation(self, max_length):
|
|
self._truncation_enabled = True
|
|
self._truncation_max = max_length
|
|
|
|
def encode_batch(self, texts):
|
|
class _Enc:
|
|
def __init__(self, n):
|
|
self.ids = [0] * n
|
|
self.attention_mask = [1] * n
|
|
|
|
# Same fixed length per batch — real tokenizers pad to the longest.
|
|
max_len = max(len(t.split()) for t in texts)
|
|
return [_Enc(max_len) for _ in texts]
|
|
|
|
|
|
@pytest.fixture
|
|
def patched_lazy_load(monkeypatch):
|
|
"""Patch the third-party deps imported inside EmbeddinggemmaONNX._lazy_load.
|
|
|
|
Returns a dict of recording counters so tests can assert how many times
|
|
each was called (e.g. confirm lazy-load caches after first call).
|
|
"""
|
|
calls = {"hf_hub_download": 0, "InferenceSession": 0, "Tokenizer.from_file": 0}
|
|
|
|
def fake_download(repo, filename=None, subfolder=None, **kwargs):
|
|
calls["hf_hub_download"] += 1
|
|
return f"/tmp/fake/{subfolder or ''}/{filename}"
|
|
|
|
fake_session_cls = _make_fake_session()
|
|
|
|
def fake_session_ctor(*args, **kwargs):
|
|
calls["InferenceSession"] += 1
|
|
return fake_session_cls()
|
|
|
|
def fake_tokenizer_from_file(_path):
|
|
calls["Tokenizer.from_file"] += 1
|
|
return _FakeTokenizer()
|
|
|
|
# huggingface_hub and tokenizers are real packages (installed via the
|
|
# multilingual extra), so we patch the functions in place rather than
|
|
# injecting stub modules.
|
|
import huggingface_hub
|
|
import onnxruntime
|
|
import tokenizers
|
|
|
|
monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_download)
|
|
monkeypatch.setattr(onnxruntime, "InferenceSession", fake_session_ctor)
|
|
monkeypatch.setattr(tokenizers.Tokenizer, "from_file", staticmethod(fake_tokenizer_from_file))
|
|
|
|
return calls
|
|
|
|
|
|
def test_name_is_stable():
|
|
"""ChromaDB persists this on the collection — changing it breaks reads."""
|
|
assert embedding.EmbeddinggemmaONNX.name() == "embeddinggemma_300m"
|
|
|
|
|
|
def test_lazy_load_runs_once(patched_lazy_load):
|
|
ef = embedding.EmbeddinggemmaONNX()
|
|
ef(["one"])
|
|
ef(["two"])
|
|
ef(["three"])
|
|
assert patched_lazy_load["hf_hub_download"] == 2 # model + tokenizer, once total
|
|
assert patched_lazy_load["InferenceSession"] == 1
|
|
assert patched_lazy_load["Tokenizer.from_file"] == 1
|
|
|
|
|
|
def test_output_shape_is_truncated_to_384(patched_lazy_load):
|
|
ef = embedding.EmbeddinggemmaONNX()
|
|
out = ef(["one", "two", "three"])
|
|
arr = np.asarray(out)
|
|
assert arr.shape == (3, 384), f"expected (3, 384) after MRL truncation, got {arr.shape}"
|
|
|
|
|
|
def test_output_is_l2_normalized(patched_lazy_load):
|
|
ef = embedding.EmbeddinggemmaONNX()
|
|
out = ef(["hello world", "another sentence"])
|
|
arr = np.asarray(out)
|
|
norms = np.linalg.norm(arr, axis=1)
|
|
assert np.allclose(norms, 1.0, atol=1e-5), f"vectors not unit-norm: {norms}"
|
|
|
|
|
|
def test_prefix_is_applied(patched_lazy_load, monkeypatch):
|
|
captured = []
|
|
original_encode_batch = _FakeTokenizer.encode_batch
|
|
|
|
def fake_encode_batch(self, texts):
|
|
captured.extend(texts)
|
|
return original_encode_batch(self, texts)
|
|
|
|
monkeypatch.setattr(_FakeTokenizer, "encode_batch", fake_encode_batch)
|
|
ef = embedding.EmbeddinggemmaONNX()
|
|
ef(["raw text one", "raw text two"])
|
|
assert all(t.startswith("task: sentence similarity | query: ") for t in captured)
|
|
# And the raw text is preserved after the prefix.
|
|
assert any("raw text one" in t for t in captured)
|
|
|
|
|
|
def test_get_embedding_function_dispatches_to_embeddinggemma(monkeypatch):
|
|
"""model='embeddinggemma' must build EmbeddinggemmaONNX, not the MiniLM EF."""
|
|
monkeypatch.setattr(
|
|
embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu")
|
|
)
|
|
ef = embedding.get_embedding_function(device="cpu", model="embeddinggemma")
|
|
assert isinstance(ef, embedding.EmbeddinggemmaONNX)
|
|
assert ef.name() == "embeddinggemma_300m"
|
|
|
|
|
|
def test_cache_key_separates_models(monkeypatch):
|
|
"""Switching model must not return the cached EF for the other model.
|
|
|
|
The cache key changed from `providers` to `(model, providers)` for exactly
|
|
this reason — without it, the second call would silently reuse the wrong EF.
|
|
"""
|
|
|
|
class DummyMiniLM:
|
|
def __init__(self, preferred_providers=None):
|
|
self.kind = "minilm"
|
|
|
|
monkeypatch.setattr(embedding, "_build_ef_class", lambda: DummyMiniLM)
|
|
monkeypatch.setattr(
|
|
embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu")
|
|
)
|
|
|
|
ml = embedding.get_embedding_function(device="cpu", model="minilm")
|
|
eg = embedding.get_embedding_function(device="cpu", model="embeddinggemma")
|
|
ml_again = embedding.get_embedding_function(device="cpu", model="minilm")
|
|
|
|
assert ml is ml_again, "minilm should cache-hit on second call"
|
|
assert isinstance(eg, embedding.EmbeddinggemmaONNX), "embeddinggemma should not collide with minilm cache"
|
|
assert ml is not eg
|
|
|
|
|
|
def test_missing_deps_raise_helpful_error(monkeypatch):
|
|
"""If the user hasn't installed `mempalace[multilingual]`, the error must
|
|
name the extra rather than just spilling a bare ImportError."""
|
|
|
|
# Drop tokenizers from sys.modules and block re-import, simulating a user
|
|
# who didn't install the multilingual extra. huggingface_hub and onnxruntime
|
|
# are present (they ship with core), so the failure should land on tokenizers.
|
|
monkeypatch.setitem(sys.modules, "tokenizers", None)
|
|
|
|
ef = embedding.EmbeddinggemmaONNX()
|
|
with pytest.raises(ImportError, match=r"mempalace\[multilingual\]"):
|
|
ef(["anything"])
|
|
|
|
|
|
def test_config_embedding_model_env_override(monkeypatch):
|
|
"""MEMPALACE_EMBEDDING_MODEL env var must override the config file default."""
|
|
from mempalace.config import MempalaceConfig
|
|
|
|
monkeypatch.setenv("MEMPALACE_EMBEDDING_MODEL", "embeddinggemma")
|
|
assert MempalaceConfig().embedding_model == "embeddinggemma"
|
|
|
|
monkeypatch.setenv("MEMPALACE_EMBEDDING_MODEL", "MiniLM") # case-insensitive
|
|
assert MempalaceConfig().embedding_model == "minilm"
|
|
|
|
|
|
def test_config_embedding_model_default_is_minilm(monkeypatch):
|
|
"""Back-compat: existing installs without explicit config get minilm."""
|
|
from mempalace.config import MempalaceConfig
|
|
|
|
monkeypatch.delenv("MEMPALACE_EMBEDDING_MODEL", raising=False)
|
|
assert MempalaceConfig().embedding_model == "minilm"
|