fix(embeddings): truncate oversized input, isolate per-chunk failures

[vvooki-sys]

Mechanism

LocalEmbedder::embed tokenized the text and built tensors at the tokenized length. The ONNX checkpoints in the BERT / XLM-R family — including the default multilingual-e5-small — bake 512 learned position embeddings into the graph, so a longer sequence fails inside model.run():

Embedding inference failed: Evaluating #106 "/embeddings/Slice" Slice:
Invalid range 0..1083 for slicing 1,512,I64 ... on axis 1

Two amplifiers turned that into data loss:

1. embed_batch did results.push(self.embed(text)?) — the first error aborted the whole loop.
2. flush_batch treated the Err branch as fatal for the batch: warn! + batch.clear(), with batch_size = 50.

So one oversized chunk cost up to 49 healthy neighbours their embedding. Verified earlier on the live v0.5.3 engine: 5 chunks including one 4830-char member → tantivy_docs +5, embeddings +0. This is the mechanism behind the "buried entity" incident.

Measurement (real model, multilingual-e5-small)

A throwaway probe against the real ONNX model, before any code change:

oversized chunk: 5130 chars
BEFORE: 0 embeddings for 5 chunks — Embedding inference failed: ... Invalid range 0..1083 for slicing 1,512
BEFORE single oversized embed: ERR (same error)

After the fix, the same batch: 5 embeddings for 5 chunks, and end-to-end through flush_batch against a real RocksDB, 5/5 chunk ids have a 384-dim vector.

Change

local_embeddings.rs

• load sets TruncationParams { max_length: 512 } on the tokenizer, but only when the checkpoint's tokenizer.json declares no truncation of its own (ours doesn't). Truncating on the tokenizer rather than slicing tensors by hand keeps the post-processor's [CLS]/[SEP] correct — HF subtracts the added special tokens from max_length in post_process.
• Named const DEFAULT_MAX_SEQUENCE_TOKENS = 512 with a comment on where the number comes from. Reading the limit from the model was considered and rejected: the ONNX inputs are dynamic (batch, seq), the limit lives in the positional weights, and the model dir ships no config.json with max_position_embeddings — any read would be guesswork.
• New embed_batch_lenient(&self, texts: &[String]) -> Vec<Result<Vec<f32>>> in both impl branches (feature and stub). Failures are returned, not logged — the caller owns the chunk_id.
• embed_batch (all-or-nothing) is left in place; the /v1/embed-missing backfill handler still uses it.

embedding_queue.rs

• flush_batch is now per-element. A new private embed_all returns Vec<Result<Vec<f32>>> with exactly one entry per request on both paths.
• Index alignment fixed. The old batch.drain(..).zip(embeddings) silently dropped the tail whenever the embedder returned fewer vectors than requests. pad_to fills missing slots with explicit Errs, so a "lost" chunk is now reported rather than swallowed.
• The Err branch no longer calls batch.clear(). Each failure is one warn! carrying its chunk_id; the aggregate counter stays as a summary alongside, not instead.

Tests

New, #[ignore]d because they need the real model — run with:

LOOMEM_TEST_EMBED_MODEL=$HOME/.loomem/models/multilingual-e5-small cargo test -p loomem-core --lib -- --ignored

• embed_truncates_oversized_input — 10 320 chars → 384-dim, L2-normalized vector (was Err).
• lenient_batch_survives_oversized_member — 5 texts, one 5130 chars → 5/5 Ok.
• flush_stores_every_chunk_despite_oversized_member — end-to-end: real LocalEmbedder + RocksDbStore in a tempdir, asserts get_embedding for all five chunk ids.

Plus three plain unit tests for the index-alignment helpers (pad_to, all_failed) that need no model and run in CI.

Deliberately out of scope

• The OpenAI path (llm::embed_batch) is untouched. Same class of problem — one failed call means no vectors for the whole request — but the constraint there is API-side, not our tensor, and it deserves its own change. flush_batch no longer deepens it: a failed call is expanded into one reported error per chunk instead of a silent batch clear.
• handlers/admin.rs (/v1/embed-missing) is untouched. It is a god file needing its own Critical file rationale:, and after the truncation fix the backfill stops failing on long chunks anyway, because embed no longer returns Err for them.
• Retry / re-enqueue of a chunk that still fails — still the backfill handler's job, exactly as before.

Gate

• cargo fmt --check — green
• cargo clippy --workspace -- -D warnings — green
• cargo test --workspace — 0 failed
• cargo check -p loomem-core --no-default-features — green (stub branch)
• Cycle /010 multilingual gate (polish_semantic_similarity) — still passing; truncation does not touch short texts

Greptile Summary

This change truncates oversized local-embedding inputs and lets queued chunks retain successful embeddings when another chunk fails. Validation confirmed that loomem-core/src/local_embeddings.rs still permits a custom tokenizer configuration with a limit above the ONNX model's 512-token positional window, so sufficiently long text can reach inference and fail instead of being truncated.

Confidence Score: 4/5

The change is not ready to merge for custom local embedding checkpoints until configured tokenizer limits are constrained to the ONNX model capacity.

A focused executable check confirmed that an existing 1024-token tokenizer limit remains in effect despite the 512-token model window, and the resulting sequence length is forwarded into ONNX inference.

Files Needing Attention: loomem-core/src/local_embeddings.rs needs to validate or clamp existing tokenizer truncation settings before embedding input is encoded.

T-Rex Logs

What T-Rex did

• T-Rex generated a proof for the posted P1 finding and linked it to the review comment details.
• Artifacts attached to the P1 finding include a Python artifact and three proof-check logs to support the validation.
• The general-contract-validation-proof verified that the custom truncation capacity check observes the expected configuration and reaches tensor construction and model.run, as shown in the after-log.
• The validation script trex-artifacts/custom-truncation-capacity-check.py was used to read the production source and exercise the guard/dataflow inputs.
• The runtime attempt was blocked by missing llvm-config and libclang dependencies, and production remained unchanged according to the logs.

_{Ran code and verified through T-Rex}

Comments Outside Diff (1)

1. General comment

   Existing tokenizer truncation can exceed the ONNX positional capacity

   • Bug
     • At loomem-core/src/local_embeddings.rs:61-68, the 512-token truncation is installed only when tokenizer.get_truncation() is None. A custom tokenizer.json that already specifies max_length: 1024 is retained. embed uses the encoded seq_len directly for all input tensors (:94-100) and passes them to model.run (:102-109). Thus an input above a 512-position ONNX model's capacity can reach inference and fail instead of being capped.
   • Cause
     • The load path treats any existing tokenizer truncation as authoritative without validating or clamping its max_length against the fixed DEFAULT_MAX_SEQUENCE_TOKENS model window.
   • Fix
     • Clamp configured truncation to the model-supported maximum (or replace it with a 512-token truncation when it exceeds that maximum) before embeddings are encoded; retain custom settings only when their maximum is within the model capacity.

   _{Ran code and verified through T-Rex}

_{Reviews (1): Last reviewed commit: "fix(embeddings): truncate long input, is..." | Re-trigger Greptile}

Greptile also left 1 inline comment on this PR.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
loomem-site	Ready	Preview	Aug 7, 2026 11:08pm

[greptile-apps]

Configured tokenizer limit can exceed model capacity

load applies the 512-token cap only when the tokenizer has no truncation configuration. A checkpoint whose tokenizer.json already declares, for example, max_length: 1024 keeps that value, and embed forwards the resulting sequence length into the ONNX inputs. For models with the documented 512-position window, long chunks then fail in model.run rather than being truncated. Clamp existing tokenizer truncation to the model-supported maximum before encoding.

Artifacts

Evidence from the check

• The authored executable reads the production source and exercises the existing-versus-absent tokenizer truncation branches, showing the existing 1024-token value is preserved above the 512-token model capacity.

Command output from the check

• The executed Python check reports that an existing 1024-token truncation survives load, exceeds 512, and proceeds to tensor construction and model run, confirming the defect.

Command output from the check

• The attempted ignored local-embedding regression test compiled dependencies but stopped at clang-sys because llvm-config and libclang were unavailable, so a real ONNX inference failure could not be captured.

Command output from the check

• The executed git diff check returned exit code 0 for local_embeddings.rs, confirming the validation did not modify production code.

_{Ran code and verified through T-Rex}