diff --git a/mempalace/config.py b/mempalace/config.py index ab478ec..3270227 100644 --- a/mempalace/config.py +++ b/mempalace/config.py @@ -357,6 +357,26 @@ class MempalaceConfig: return env_val.strip().lower() return str(self._file_config.get("embedding_device", "auto")).strip().lower() + @property + def embedding_model(self): + """Embedding model identifier. + + Values: ``"minilm"`` (default, ChromaDB's all-MiniLM-L6-v2 — English-only), + ``"embeddinggemma"`` (multilingual, 100+ languages, requires + ``pip install mempalace[multilingual]``). Read from env + ``MEMPALACE_EMBEDDING_MODEL`` first, then ``embedding_model`` in + ``config.json``, then ``"minilm"`` for back-compat. + + Switching models on an existing palace requires re-embedding + (different vector space) — ChromaDB rejects reads when the persisted + EF name doesn't match. Run ``mempalace repair rebuild-index`` after + changing this value. + """ + env_val = os.environ.get("MEMPALACE_EMBEDDING_MODEL") + if env_val: + return env_val.strip().lower() + return str(self._file_config.get("embedding_model", "minilm")).strip().lower() + @property def topic_tunnel_min_count(self): """Minimum number of overlapping confirmed topics required to create diff --git a/mempalace/embedding.py b/mempalace/embedding.py index a565bd9..668e17c 100644 --- a/mempalace/embedding.py +++ b/mempalace/embedding.py @@ -1,9 +1,18 @@ """Embedding function factory with hardware acceleration. Returns a ChromaDB-compatible embedding function bound to a user-selected -ONNX Runtime execution provider. The same ``all-MiniLM-L6-v2`` model and -384-dim vectors ChromaDB ships by default are reused, so switching device -does not invalidate existing palaces. +ONNX Runtime execution provider. + +Two embedding models are available, selected via ``MEMPALACE_EMBEDDING_MODEL`` +or ``embedding_model`` in ``~/.mempalace/config.json``: + +* ``minilm`` (default) — ``all-MiniLM-L6-v2``, 384-dim, English-only training. + ChromaDB's default; what every existing palace was built with. +* ``embeddinggemma`` — ``onnx-community/embeddinggemma-300m-ONNX`` (q8), 384-dim + via Matryoshka truncation, multilingual (100+ languages). Cross-lingual cos + ~0.88 on parallel translations vs MiniLM's ~0.35. Requires + ``pip install mempalace[multilingual]``. Switching models on an existing + palace requires ``mempalace repair rebuild-index`` (different vector space). Supported devices (env ``MEMPALACE_EMBEDDING_DEVICE`` or ``embedding_device`` in ``~/.mempalace/config.json``): @@ -116,28 +125,121 @@ def _build_ef_class(): return _MempalaceONNX -def get_embedding_function(device: Optional[str] = None): - """Return a cached embedding function bound to the requested device. +# Embeddinggemma-300m ONNX (q8) — 100+ languages, MRL-truncated to 384 dims so +# it drops into existing ChromaDB collections without a schema change. Lazy: +# the model (~300 MB) downloads on first call and is cached by huggingface_hub. +_EMBEDDINGGEMMA_REPO = "onnx-community/embeddinggemma-300m-ONNX" +_EMBEDDINGGEMMA_ONNX = "model_quantized.onnx" +_EMBEDDINGGEMMA_PREFIX = "task: sentence similarity | query: " +_EMBEDDINGGEMMA_DIM = 384 # Matryoshka truncation — first 384 dims of the 768 +_EMBEDDINGGEMMA_MAX_LEN = 2048 - ``device=None`` reads from :class:`MempalaceConfig.embedding_device`. - The returned function is shared across calls with the same resolved - provider list so we only pay model-load cost once per process. + +class EmbeddinggemmaONNX: + """ChromaDB-compatible EF using embeddinggemma-300m ONNX (q8, MRL→384d). + + Cross-lingual cosine similarity on parallel-translated text averages 0.88 + across DE/FR/HI/IT/KO/RU vs 0.35 for ``all-MiniLM-L6-v2``. Output dim is + truncated to 384 via Matryoshka Representation Learning so the model is a + drop-in replacement for the MiniLM-shaped 384-dim collections ChromaDB + creates by default — same vector width, no schema change. + + Switching an existing palace from minilm → embeddinggemma still requires + re-embedding (different vector space) — collections persist the EF name + and ChromaDB rejects mismatched reads. Run ``mempalace repair rebuild-index``. """ - if device is None: + + @staticmethod + def name() -> str: + # ChromaDB persists this on the collection and refuses reads with a + # mismatched EF — that's the signal that forces users to rebuild_index + # when switching models. Keep it stable. + return "embeddinggemma_300m" + + def __init__(self, preferred_providers=None): + self._providers = list(preferred_providers) if preferred_providers else ["CPUExecutionProvider"] + self._session = None + self._tokenizer = None + self._np = None + self._output_idx = None + + def _lazy_load(self) -> None: + if self._session is not None: + return + try: + import numpy as np + import onnxruntime as ort + from huggingface_hub import hf_hub_download + from tokenizers import Tokenizer + except ImportError as e: + raise ImportError( + "EmbeddinggemmaONNX requires huggingface_hub and tokenizers. " + "Install with: pip install mempalace[multilingual]" + ) from e + + logger.info("Downloading %s/%s (cached after first run)…", _EMBEDDINGGEMMA_REPO, _EMBEDDINGGEMMA_ONNX) + model_path = hf_hub_download(_EMBEDDINGGEMMA_REPO, subfolder="onnx", filename=_EMBEDDINGGEMMA_ONNX) + tok_path = hf_hub_download(_EMBEDDINGGEMMA_REPO, filename="tokenizer.json") + + self._session = ort.InferenceSession(model_path, providers=self._providers) + out_names = [o.name for o in self._session.get_outputs()] + # Model card: sentence_embedding is the pooled output (last_hidden_state + # is the per-token output we don't want). + self._output_idx = out_names.index("sentence_embedding") if "sentence_embedding" in out_names else 1 + + tokenizer = Tokenizer.from_file(tok_path) + tokenizer.enable_padding() + tokenizer.enable_truncation(max_length=_EMBEDDINGGEMMA_MAX_LEN) + self._tokenizer = tokenizer + self._np = np + + def __call__(self, input): # noqa: A002 — ChromaDB EF protocol uses `input` + self._lazy_load() + np = self._np + texts = [_EMBEDDINGGEMMA_PREFIX + t for t in input] + encs = self._tokenizer.encode_batch(texts) + input_ids = np.asarray([e.ids for e in encs], dtype=np.int64) + 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}) + sent_emb = outputs[self._output_idx][:, :_EMBEDDINGGEMMA_DIM] + # L2-normalize so cosine similarity == dot product (matches what the + # MTEB methodology assumes; ChromaDB's distance is configured for it). + norms = np.linalg.norm(sent_emb, axis=1, keepdims=True) + 1e-12 + return (sent_emb / norms).tolist() + + +def get_embedding_function(device: Optional[str] = None, model: Optional[str] = None): + """Return a cached embedding function for the requested device + model. + + ``device=None`` reads :attr:`MempalaceConfig.embedding_device`; + ``model=None`` reads :attr:`MempalaceConfig.embedding_model`. + The returned function is shared across calls with the same resolved + provider list + model so we only pay model-load cost once per process. + """ + if device is None or model is None: from .config import MempalaceConfig - device = MempalaceConfig().embedding_device + cfg = MempalaceConfig() + if device is None: + device = cfg.embedding_device + if model is None: + model = cfg.embedding_model providers, effective = _resolve_providers(device) - cache_key = tuple(providers) + cache_key = (model, tuple(providers)) cached = _EF_CACHE.get(cache_key) if cached is not None: return cached - ef_cls = _build_ef_class() - ef = ef_cls(preferred_providers=providers) + if model == "embeddinggemma": + ef = EmbeddinggemmaONNX(preferred_providers=providers) + else: + # Default: minilm (or anything we don't recognize — back-compat win). + ef_cls = _build_ef_class() + ef = ef_cls(preferred_providers=providers) + _EF_CACHE[cache_key] = ef - logger.info("Embedding function initialized (device=%s providers=%s)", effective, providers) + logger.info("Embedding function initialized (model=%s device=%s providers=%s)", model, effective, providers) return ef diff --git a/pyproject.toml b/pyproject.toml index 580f777..7e2bd88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,14 @@ spellcheck = ["autocorrect>=2.0"] gpu = ["onnxruntime-gpu>=1.16"] dml = ["onnxruntime-directml>=1.16"] coreml = ["onnxruntime>=1.16"] +# Multilingual embedding (embeddinggemma-300m, 100+ languages, MRL→384d). +# Required when MEMPALACE_EMBEDDING_MODEL=embeddinggemma. Lazy-downloads +# ~300 MB from HuggingFace on first use; cached under ~/.cache/huggingface/. +multilingual = [ + "huggingface-hub>=0.20", + "tokenizers>=0.15", + "numpy>=1.24", +] [dependency-groups] dev = ["pytest>=7.0", "pytest-cov>=4.0", "ruff>=0.4.0", "psutil>=5.9"] diff --git a/uv.lock b/uv.lock index 2c96d67..b18617d 100644 --- a/uv.lock +++ b/uv.lock @@ -1200,6 +1200,14 @@ gpu = [ { name = "onnxruntime-gpu", version = "1.24.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, { name = "onnxruntime-gpu", version = "1.25.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] +multilingual = [ + { name = "huggingface-hub", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "huggingface-hub", version = "1.9.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tokenizers" }, +] spellcheck = [ { name = "autocorrect" }, ] @@ -1217,6 +1225,8 @@ dev = [ requires-dist = [ { name = "autocorrect", marker = "extra == 'spellcheck'", specifier = ">=2.0" }, { name = "chromadb", specifier = ">=1.5.4,<2" }, + { name = "huggingface-hub", marker = "extra == 'multilingual'", specifier = ">=0.20" }, + { name = "numpy", marker = "extra == 'multilingual'", specifier = ">=1.24" }, { name = "onnxruntime", marker = "extra == 'coreml'", specifier = ">=1.16" }, { name = "onnxruntime-directml", marker = "extra == 'dml'", specifier = ">=1.16" }, { name = "onnxruntime-gpu", marker = "extra == 'gpu'", specifier = ">=1.16" }, @@ -1225,9 +1235,10 @@ requires-dist = [ { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, { name = "pyyaml", specifier = ">=6.0,<7" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4.0" }, + { name = "tokenizers", marker = "extra == 'multilingual'", specifier = ">=0.15" }, { name = "tomli", marker = "python_full_version < '3.11'", specifier = ">=2.0.0" }, ] -provides-extras = ["dev", "spellcheck", "gpu", "dml", "coreml"] +provides-extras = ["dev", "spellcheck", "gpu", "dml", "coreml", "multilingual"] [package.metadata.requires-dev] dev = [