feat: generic ONNX embedding backend + multilingual-e5-small preset - #2280
Open
isolovyev77 wants to merge 1 commit into
Open
feat: generic ONNX embedding backend + multilingual-e5-small preset#2280isolovyev77 wants to merge 1 commit into
isolovyev77 wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Local embedding models are currently fixed choices —
minilmandembeddinggemmaare hardcoded inembedding.py, and anything else requires standing up a separateopenai-compatinference 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 fromembedding_onnx_*keys inconfig.json, each overridable via the matchingMEMPALACE_EMBEDDING_ONNX_*env var. No source edits, no separate server.embedding_model: "e5-small"— a bundled preset forintfloat/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:
passage:query: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 symmetricquery:beats symmetricpassage:(0.754 vs 0.684), so the preset pinsquery:on both paths.GenericONNXEmbeddingkeeps separatedoc_prefix/query_prefix(andembed_documents/embed_query), so if call sites later split the two paths, the canonical asymmetric pair is one config change away.Design notes
EmbeddinggemmaONNX(huggingface_hub,tokenizers,numpy,onnxruntime) — no new requirements.EmbeddinggemmaONNXconventions: double-checked lazy load, session assigned last,_intra_op_session_optionsthread cap, ImportError with reinstall hint.generic-onnxis 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_nameoverrides it for vector-affecting knobs the derived name cannot see (pooling, max_len, prefixes). The e5 preset persists ase5_small_384.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/benchmarkstests/test_generic_onnx.pymirrorstest_embeddinggemma.py: fully offline via mockedhf_hub_download/Tokenizer/InferenceSession(module opted into_REAL_EMBEDDING_TEST_MODULESin conftest). Covers pooling (masked mean vs cls), L2 normalization, padding invariance, per-path prefixes,token_type_idsfeed 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.jsonon a fresh palace (or runmempalace repair rebuild-indexon 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 cleandevelop.Checklist
python -m pytest tests/ -v)ruff check .)