feat(embedding): add embeddinggemma-300m ONNX as opt-in multilingual embedder
MemPalace's default embedder (all-MiniLM-L6-v2) is English-only-trained. Cross-lingual cosine similarity on parallel-translated text averages 0.35 across DE/FR/HI/IT/KO/RU — vs 0.88 for embeddinggemma-300m ONNX (q8) with the semantic-similarity prefix. RU is the worst at 0.17, meaning a Russian memory and its identical English translation embed to nearly orthogonal vectors. Multilingual users effectively cannot retrieve their own memories. This commit adds embeddinggemma-300m as an opt-in alternative: * New EmbeddinggemmaONNX class implementing ChromaDB's EF protocol. Lazy-downloads model_quantized.onnx (~300 MB) via huggingface_hub on first use; cached under ~/.cache/huggingface/. Applies the sim prefix, runs onnxruntime inference, truncates to 384 dims via Matryoshka (MRL), L2-normalizes. * MRL truncation to 384d is intentional: matches MiniLM's vector width so collection schemas don't change, and validation showed 384d MRL actually outperforms full 768d on these similarity tasks (0.893 vs 0.881 avg) — known property of MRL training. * MEMPALACE_EMBEDDING_MODEL env (default "minilm" for back-compat). Switching models on an existing palace requires re-embedding — ChromaDB rejects reads with a mismatched EF name. Run `mempalace repair rebuild-index` after changing the value. * New optional dep group: pip install mempalace[multilingual] Adds huggingface_hub + tokenizers + numpy. Core deps unchanged. ONNX q8 validated lossless vs the Ollama gguf benchmarked previously (max delta 0.002 cos across 240 parallel pairs).
This commit is contained in:
Igor Lins e Silva
parent
8591d133da
commit
51702e9f0f
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
13
uv.lock
13
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 = [
|
||||
|
|
|
|||
Loading…
Reference in New Issue