Skip to content

feat: generic ONNX embedding backend + multilingual-e5-small preset - #2280

Open
isolovyev77 wants to merge 1 commit into
MemPalace:developfrom
isolovyev77:feat/generic-onnx-embedding
Open

feat: generic ONNX embedding backend + multilingual-e5-small preset#2280
isolovyev77 wants to merge 1 commit into
MemPalace:developfrom
isolovyev77:feat/generic-onnx-embedding

Conversation

@isolovyev77

Copy link
Copy Markdown

What does this PR do?

Local embedding models are currently fixed choices — minilm and embeddinggemma are hardcoded in embedding.py, and anything else requires standing up a separate openai-compat inference server. This PR adds the missing piece requested in #1563 / #1261:

  • embedding_model: "generic-onnx" — run any HuggingFace-hosted ONNX encoder, selected entirely from configuration: repo, ONNX file, subfolder, pooling (mean/cls), tokenizer truncation length, instruction prefixes, and the persisted EF identity name all come from embedding_onnx_* keys in config.json, each overridable via the matching MEMPALACE_EMBEDDING_ONNX_* env var. No source edits, no separate server.
  • embedding_model: "e5-small" — a bundled preset for intfloat/multilingual-e5-small (384-dim, fp32, ~120 MB): a lighter multilingual alternative to embeddinggemma-300m (~1/4 the download), also relevant to Feature request: support configurable embedding model (e.g. embeddinggemma-300m for multilingual) #1663.

Why e5-small specifically

This exact class + preset has been running in production since July on a 200k-drawer, RU-heavy palace. Measured on that corpus:

model MRR R@1
minilm 0.536
e5-small, symmetric passage: 0.684
e5-small, symmetric query: 0.754 0.65

One measured detail worth flagging: canonical e5 usage is asymmetric (passage: for documents, query: for queries), but the embedding wrapper routes both sides through __call__, applying one prefix symmetrically. On a real corpus symmetric query: beats symmetric passage: (0.754 vs 0.684), so the preset pins query: on both paths. GenericONNXEmbedding keeps separate doc_prefix / query_prefix (and embed_documents / embed_query), so if call sites later split the two paths, the canonical asymmetric pair is one config change away.

Design notes

  • Same lazy-loaded deps as EmbeddinggemmaONNX (huggingface_hub, tokenizers, numpy, onnxruntime) — no new requirements.
  • Follows the EmbeddinggemmaONNX conventions: double-checked lazy load, session assigned last, _intra_op_session_options thread cap, ImportError with reinstall hint.
  • EF identity: the default collection name for generic-onnx is derived from repo + ONNX file, so pointing the config at a different model trips the embedder-identity check instead of silently mixing vector spaces; embedding_onnx_ef_name overrides it for vector-affecting knobs the derived name cannot see (pooling, max_len, prefixes). The e5 preset persists as e5_small_384.
  • Switching models on an existing palace still requires mempalace repair rebuild-index — docstrings repeat the caveat.

Closes #1563, closes #1261; addresses #1663.

How to test

uv run pytest tests/test_generic_onnx.py -v   # 21 offline tests, no network
uv run pytest tests/ --ignore=tests/benchmarks

tests/test_generic_onnx.py mirrors test_embeddinggemma.py: fully offline via mocked hf_hub_download / Tokenizer / InferenceSession (module opted into _REAL_EMBEDDING_TEST_MODULES in conftest). Covers pooling (masked mean vs cls), L2 normalization, padding invariance, per-path prefixes, token_type_ids feed detection, batching, identity-name stability, config resolution (env + file, incl. trailing-space prefix preservation), and factory dispatch/caching.

To try it live: set "embedding_model": "e5-small" in ~/.mempalace/config.json on a fresh palace (or run mempalace repair rebuild-index on an existing one).

Full suite locally (macOS, py3.12): 4320 passed; the one failure is test_mcp_server::TestWriteTools::test_checkpoint_added_by_accepted_via_dispatch, which fails identically on clean develop.

Checklist

  • Tests pass (python -m pytest tests/ -v)
  • No hardcoded paths
  • Linter passes (ruff check .)

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 MemPalace#1563, closes MemPalace#1261, addresses MemPalace#1663.
@igorls igorls added enhancement New feature or request storage labels Aug 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request storage

Projects

None yet

2 participants