The tokenizer pads every row of a sub-batch to the longest sequence in it,
and attention costs batch x heads x length^2 per layer, so arrival order
decided what an embeddinggemma run cost: one long verbatim message dragged
its whole sub-batch up to its own length.
Sort by UTF-8 byte length before the split and scatter each row back to its
input index. An input that fits one sub-batch is left alone, since every row
pads to the same width either way.
Over 43,157 sweep drawers from 160 transcripts, padded token slots drop
39.7% and the quadratic attention term 45.0%. Wall clock on 52 messages with
one model instance: 1169.70 s to 490.83 s, control baseline 1286.23 s. Row
values move by at most one float32 ULP (1.2e-07, cosine 0.99999992), which
is reduction-order rounding, not a change of meaning.
Only palaces configured for embeddinggemma are affected; the default MiniLM
embedder pads to a fixed width.
ChromaDB's ONNX embedder builds its InferenceSession without a thread cap, so
ORT's intra-op pool defaults to the physical core count. OMP_NUM_THREADS is
inert against it (ORT owns its own pool), so a background `mempalace mine`
pins 4-5 cores and stacked Stop-hook fires turn the machine into a thermal
event.
Add an `embedding_threads` config knob (env MEMPALACE_EMBEDDING_THREADS or
config.json). Unset/"auto" caps the intra-op pool at half the logical CPUs so
a fresh install stays usable out of the box; a positive integer sets an exact
count; 0/negative leaves ORT uncapped for users who want max throughput.
The cap is applied via SessionOptions at session construction:
- `_MempalaceONNX` (default minilm) overrides the `model` cached_property to
rebuild the session the same way upstream does plus the cap, falling back to
upstream's uncapped build if chromadb internals shift.
- `EmbeddinggemmaONNX` builds its session through the shared
`_intra_op_session_options()` helper.
Two threads sharing a cold EmbeddinggemmaONNX via _EF_CACHE could each
build a full model session, and two factory callers could each keep a
private instance. The load is now double-check locked with the session
published last, and the factory cache has an atomic check-then-construct
behind a lock-free fast path. __call__ wraps a bare string, returns []
for None and empty input before the lazy download, and its annotation
matches the accepted types.
One session.run over a repair-scale batch (5000 docs) allocates
attention buffers far beyond available RAM and the kernel kills the
process. Mirror chromadb's ONNXMiniLM_L6_V2 and embed in sub-batches
of 32; per-chunk padding also stops one long doc inflating the whole
batch.
Co-Authored-By: mojie5 <262519016+mojie5@users.noreply.github.com>
The EmbeddinggemmaONNX lazy-load now fetches the ONNX external-weights file
(model.onnx_data) in addition to the model graph and tokenizer, so a single
warm-up issues 3 downloads, not 2. The lazy-load-once invariant is unchanged
(InferenceSession and Tokenizer.from_file are still each built exactly once).
Resolve conflicts:
- backends/chroma.py: keep both new except handlers in get_collection
(CollectionNotInitializedError from develop + EF-mismatch helper from
this branch), ordered _ChromaNotFoundError before ValueError to match
the create-branch handler order.
- uv.lock: regenerated from merged pyproject.toml.
Fix lint: ruff format mempalace/embedding.py + tests/test_embeddinggemma.py
(CI now pins ruff==0.15.9 via develop's workflow update).
Full suite: 1923 passed, 1 skipped.
Onboarding now asks the user once, on first run, whether to use the
multilingual embedding model. The default answer is yes — defaulting to
English-only made the recall promise effectively unreachable for any
non-English content (cross-lingual cos ~0.35 vs ~0.88 for the multilingual
model). The choice is written to config.json so subsequent runs pick the
right EF without re-prompting; existing installs that never set the env
var or ran onboarding stay on minilm for back-compat. MEMPALACE_EMBEDDING_MODEL
still overrides both.
Multilingual deps (huggingface_hub, tokenizers, numpy) move from the
[multilingual] extra into core. The extra is kept as a no-op alias so
existing install scripts keep working. The 300 MB ONNX model is still
lazy-downloaded on first use, not at install time.
`quick_setup` (the programmatic non-interactive path) grows an optional
`embedding_model` arg so tests and benchmark scripts can pick a model
without writing config.json by accident.
EmbeddinggemmaONNX's "missing deps" error now points at the right
recovery path (reinstall mempalace, since the deps are core) rather
than the obsolete pip install mempalace[multilingual] hint.
Tests: 9 new (3 _ask_embedding_model variants + 2 run_onboarding
persistence + 2 quick_setup + 2 set_embedding_model round-trips). The
existing 2 run_onboarding tests now patch _ask_embedding_model so they
don't print to stdout.
Three follow-ups bundled for the embeddinggemma EF added in 51702e9:
1. Offline tests for EmbeddinggemmaONNX (10 tests, 0.08s, no network).
Mocks huggingface_hub.hf_hub_download, tokenizers.Tokenizer.from_file,
and onnxruntime.InferenceSession so CI never pulls the 300 MB model.
Guarded with pytest.importorskip so the file is skipped when the
multilingual extra isn't installed. Covers: stable name(), lazy-load
runs exactly once, output shape (n, 384) after MRL truncation, L2
normalization, sim prefix applied, dispatch from
get_embedding_function(model="embeddinggemma"), cache key separates
models, helpful ImportError when deps missing, env override.
2. Friendlier ChromaDB EF-name-mismatch error. Switching
MEMPALACE_EMBEDDING_MODEL on an existing palace previously surfaced
ChromaDB's bare "Embedding function conflict: new: X vs persisted: Y"
ValueError. Now ChromaBackend.get_collection() wraps that error and
points users at the two recovery paths: revert the env var, or run
`mempalace repair rebuild-index --palace <path>`. New
_explain_ef_mismatch helper + 3 tests (unit + end-to-end).
3. Docs: CHANGELOG [Unreleased] entry covers both the new EF and the
error wrapper. README Requirements section mentions the multilingual
extra and points at the embedding.py docstring for the migration note.
mvalentsev
Sandro da Silva
Igor Lins e Silva