Skip to content

Commit 345afa9

Browse files
author
Ivan Solovyev
committed
feat: generic ONNX embedding backend + multilingual-e5-small preset
Local embedding models were fixed choices (minilm, embeddinggemma); anything else required an openai-compat endpoint. embedding_model "generic-onnx" now runs any HuggingFace-hosted ONNX encoder selected entirely from configuration (embedding_onnx_* keys / env vars), and "e5-small" ships as a bundled preset for intfloat/multilingual-e5-small: 384-dim, ~120 MB, measured MRR 0.754 vs 0.536 (minilm) on a real 200k-drawer RU-heavy palace. Same lazy deps as embeddinggemma — no new requirements. Offline tests mirror test_embeddinggemma.py (mocked hf_hub_download / Tokenizer / InferenceSession). Closes #1563, closes #1261, addresses #1663.
1 parent 06cb698 commit 345afa9

6 files changed

Lines changed: 759 additions & 5 deletions

File tree

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
88

99
## [Unreleased]
1010

11+
### Features
12+
13+
- **Config-driven local embedding models: `generic-onnx` backend + `e5-small` preset.** Local embedders were fixed choices (minilm, embeddinggemma) and anything else required standing up an `openai-compat` endpoint. `embedding_model: "generic-onnx"` now runs any HuggingFace-hosted ONNX encoder selected entirely from configuration — repo, ONNX file, pooling, instruction prefixes and the persisted EF identity come from `embedding_onnx_*` keys in `config.json` (each overridable via the matching `MEMPALACE_EMBEDDING_ONNX_*` env var). `embedding_model: "e5-small"` is a bundled preset for `intfloat/multilingual-e5-small`: 384-dim like the default collections, ~120 MB vs embeddinggemma's ~300 MB, and measured on a real 200k-drawer RU-heavy palace at MRR 0.754 vs 0.536 for minilm. Same lazy deps as embeddinggemma — no new requirements; switching models on an existing palace still requires `mempalace repair rebuild-index`. (#1563, #1261, #1663)
14+
1115
### Bug Fixes
1216

1317
- **`sweep` books a failed `stat` as a failure again.** The non-regular-file gate added in 3.7.1 reads `stat.S_ISREG(f.stat().st_mode)` inside a `try`, and its `except OSError` printed `SKIP` and moved on. A dangling symlink, a symlink loop and a file unlinked between `rglob` and the gate all raise there, and before the gate existed every one of them reached `sweep()` and was booked in `failures` — so `sweep` went from reporting a transcript it could not read to reporting success. A failed probe is now an error, not a benign file type: it is logged, printed as `WARNING`, and appended to `failures`, while a probe that succeeds and says "not regular" still skips silently. (#2221)

mempalace/config.py

Lines changed: 134 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -748,10 +748,14 @@ def embedding_model(self):
748748
749749
Values: ``"minilm"`` (ChromaDB's all-MiniLM-L6-v2 — English-only),
750750
``"embeddinggemma"`` (multilingual, 100+ languages, default for
751-
new installs since onboarding writes the choice), or
752-
``"openai-compat"`` (embeddings served by an OpenAI-compatible
753-
``/v1/embeddings`` endpoint — see ``embedding_api_url`` /
754-
``embedding_api_model`` / ``embedding_api_key``). Read from env
751+
new installs since onboarding writes the choice),
752+
``"e5-small"`` (bundled ``intfloat/multilingual-e5-small`` preset —
753+
see :mod:`mempalace.generic_onnx_embedding`), ``"generic-onnx"``
754+
(any HuggingFace ONNX encoder, driven by the ``embedding_onnx_*``
755+
settings), or ``"openai-compat"`` (embeddings served by an
756+
OpenAI-compatible ``/v1/embeddings`` endpoint — see
757+
``embedding_api_url`` / ``embedding_api_model`` /
758+
``embedding_api_key``). Read from env
755759
``MEMPALACE_EMBEDDING_MODEL`` first, then ``embedding_model`` in
756760
``config.json``, then ``"minilm"`` as a back-compat fallback for
757761
palaces created before onboarding asked the question.
@@ -881,6 +885,132 @@ def embedding_api_key(self):
881885
"""
882886
return self._resolve_str_setting("MEMPALACE_EMBEDDING_API_KEY", "embedding_api_key")
883887

888+
def _resolve_prefix_setting(self, env_var: str, config_key: str) -> str:
889+
"""Resolve a prefix setting: env var > ``config.json`` > ``""``.
890+
891+
Unlike :meth:`_resolve_str_setting` the value is **not** stripped —
892+
instruction prefixes for e5-style models end with a significant
893+
space (``"query: "``). A whitespace-only env var still counts as
894+
unset so it cannot silently mask the config-file value.
895+
"""
896+
env_val = os.environ.get(env_var)
897+
if env_val is not None and env_val.strip():
898+
return env_val
899+
cfg_val = self._file_config.get(config_key)
900+
if isinstance(cfg_val, str):
901+
return cfg_val
902+
return ""
903+
904+
@property
905+
def embedding_onnx_repo(self):
906+
"""HuggingFace repo id of the ``generic-onnx`` embedding model.
907+
908+
Required when ``embedding_model == "generic-onnx"`` (e.g.
909+
``intfloat/multilingual-e5-small``). Resolved from env
910+
``MEMPALACE_EMBEDDING_ONNX_REPO`` first, then ``embedding_onnx_repo``
911+
in ``config.json``; ``None`` when unset.
912+
"""
913+
return self._resolve_str_setting("MEMPALACE_EMBEDDING_ONNX_REPO", "embedding_onnx_repo")
914+
915+
@property
916+
def embedding_onnx_file(self):
917+
"""ONNX graph filename inside the model repo (default ``model.onnx``).
918+
919+
Pick the fp32 export unless you have measured otherwise — quantized
920+
and fp16 exports are not universally faster on CPU execution
921+
providers. Env ``MEMPALACE_EMBEDDING_ONNX_FILE`` overrides.
922+
"""
923+
return (
924+
self._resolve_str_setting("MEMPALACE_EMBEDDING_ONNX_FILE", "embedding_onnx_file")
925+
or "model.onnx"
926+
)
927+
928+
@property
929+
def embedding_onnx_subfolder(self):
930+
"""Repo subfolder holding the ONNX file (default ``onnx``).
931+
932+
Set to ``"."`` for repos that keep the graph at the repo root.
933+
Env ``MEMPALACE_EMBEDDING_ONNX_SUBFOLDER`` overrides.
934+
"""
935+
return (
936+
self._resolve_str_setting(
937+
"MEMPALACE_EMBEDDING_ONNX_SUBFOLDER", "embedding_onnx_subfolder"
938+
)
939+
or "onnx"
940+
)
941+
942+
@property
943+
def embedding_onnx_pooling(self):
944+
"""Pooling strategy for the ``generic-onnx`` model: ``mean`` or ``cls``.
945+
946+
``mean`` (default) is the attention-masked token average most
947+
sentence encoders expect; ``cls`` takes the first token. Values are
948+
validated at EF construction time. Env
949+
``MEMPALACE_EMBEDDING_ONNX_POOLING`` overrides.
950+
"""
951+
val = self._resolve_str_setting(
952+
"MEMPALACE_EMBEDDING_ONNX_POOLING", "embedding_onnx_pooling"
953+
)
954+
return val.lower() if val else "mean"
955+
956+
@property
957+
def embedding_onnx_max_len(self) -> int:
958+
"""Tokenizer truncation length for the ``generic-onnx`` model.
959+
960+
Defaults to ``512`` (the ceiling for most BERT-lineage encoders).
961+
Env ``MEMPALACE_EMBEDDING_ONNX_MAX_LEN`` overrides; non-numeric or
962+
non-positive values fall back to the default.
963+
"""
964+
raw = os.environ.get("MEMPALACE_EMBEDDING_ONNX_MAX_LEN")
965+
if raw is None:
966+
raw = self._file_config.get("embedding_onnx_max_len")
967+
try:
968+
val = int(str(raw).strip())
969+
except (TypeError, ValueError):
970+
return 512
971+
return val if val > 0 else 512
972+
973+
@property
974+
def embedding_onnx_doc_prefix(self):
975+
"""Prefix prepended on the indexing path (default ``""``).
976+
977+
Instruction-tuned retrieval models expect a role marker — e5 uses
978+
``"passage: "`` for documents. The trailing space is significant and
979+
preserved; see :meth:`_resolve_prefix_setting`. Env
980+
``MEMPALACE_EMBEDDING_ONNX_DOC_PREFIX`` overrides.
981+
"""
982+
return self._resolve_prefix_setting(
983+
"MEMPALACE_EMBEDDING_ONNX_DOC_PREFIX", "embedding_onnx_doc_prefix"
984+
)
985+
986+
@property
987+
def embedding_onnx_query_prefix(self):
988+
"""Prefix prepended on the ``embed_query`` path (default ``""``).
989+
990+
Only effective for call sites that use ``embed_query``; the current
991+
embedding wrapper routes everything through ``__call__``, which
992+
applies ``embedding_onnx_doc_prefix`` symmetrically. Env
993+
``MEMPALACE_EMBEDDING_ONNX_QUERY_PREFIX`` overrides.
994+
"""
995+
return self._resolve_prefix_setting(
996+
"MEMPALACE_EMBEDDING_ONNX_QUERY_PREFIX", "embedding_onnx_query_prefix"
997+
)
998+
999+
@property
1000+
def embedding_onnx_ef_name(self):
1001+
"""Identity name persisted on the collection for ``generic-onnx``.
1002+
1003+
Defaults to a name derived from repo + ONNX file, so pointing the
1004+
config at a different model trips the embedder-identity check
1005+
instead of silently mixing vector spaces. Set it explicitly when
1006+
changing vector-affecting knobs that the derived name cannot see
1007+
(pooling, max_len, prefixes). Env
1008+
``MEMPALACE_EMBEDDING_ONNX_EF_NAME`` overrides; ``None`` when unset.
1009+
"""
1010+
return self._resolve_str_setting(
1011+
"MEMPALACE_EMBEDDING_ONNX_EF_NAME", "embedding_onnx_ef_name"
1012+
)
1013+
8841014
@property
8851015
def topic_tunnel_min_count(self):
8861016
"""Minimum number of overlapping confirmed topics required to create

mempalace/embedding.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
bound to a user-selected ONNX Runtime execution provider, or an
55
OpenAI-compatible HTTP ``/v1/embeddings`` endpoint.
66
7-
Three embedding-model options are available, selected via
7+
Five embedding-model options are available, selected via
88
``MEMPALACE_EMBEDDING_MODEL`` or ``embedding_model`` in
99
``~/.mempalace/config.json``:
1010
@@ -17,6 +17,17 @@
1717
model is lazy-downloaded from HuggingFace on first use. Switching models
1818
on an existing palace requires ``mempalace repair rebuild-index``
1919
(different vector space).
20+
* ``e5-small`` — ``intfloat/multilingual-e5-small`` (fp32), 384-dim,
21+
multilingual. A lighter alternative to embeddinggemma (~120 MB vs ~300 MB)
22+
with strong non-English retrieval; see
23+
:func:`mempalace.generic_onnx_embedding.e5_small_ef` for measured numbers.
24+
Same rebuild-index caveat when switching an existing palace.
25+
* ``generic-onnx`` — any HuggingFace-hosted ONNX encoder, fully
26+
config-driven (#1563, #1261): repo, ONNX file, pooling, prefixes and the
27+
persisted EF name come from ``embedding_onnx_*`` settings in
28+
``config.json`` (each overridable via the matching
29+
``MEMPALACE_EMBEDDING_ONNX_*`` env var). See
30+
:mod:`mempalace.generic_onnx_embedding`.
2031
* ``openai-compat`` — embeddings served by any OpenAI-compatible
2132
``/v1/embeddings`` endpoint (LM Studio, llama.cpp, vLLM, Ollama's OpenAI
2233
shim, or a self-hosted server) instead of a local ONNX model. Useful for
@@ -686,6 +697,12 @@ def get_embedding_function(device: Optional[str] = None, model: Optional[str] =
686697
threads = _resolve_intra_op_threads()
687698
if model == "embeddinggemma":
688699
ef = EmbeddinggemmaONNX(preferred_providers=providers, intra_op_num_threads=threads)
700+
elif model in ("e5-small", "generic-onnx"):
701+
from .generic_onnx_embedding import build_generic_onnx_ef
702+
703+
ef = build_generic_onnx_ef(
704+
model, preferred_providers=providers, intra_op_num_threads=threads
705+
)
689706
else:
690707
# Default: minilm (or anything we don't recognize — back-compat win).
691708
ef_cls = _build_ef_class()

0 commit comments

Comments
 (0)