fix(embedding): split overlength text before inference

[jioffe502]

Summary

Embedding processes several pages together. Today, one page that is too long can cause two kinds of data loss:

• the model can reject the whole group, so normal pages beside it never become searchable; or
• the long page can be silently cut off, so words near the end never become searchable.

This PR splits only text that is too long before sending it to the embedding model. Normal inputs are unchanged. Nothing is truncated, every piece keeps its document and page information, and an incomplete result is rejected before anything is written to a vector database.

What we observed

• JP20: 0 of 3,201 text rows exceeded the 8,192-token formatted-input limit; maximum 4,555.
• BO767: 2 of 83,341 text rows exceeded the limit (0.0024%); maximum 10,647. Those two rows became four complete children, adding only two rows.
• ViDoRe V3: its 19,252 embedding rows primarily use image-bearing VL input, so they are not governed by this text-splitting policy.

This is protection for a rare production case, not a claim that these research datasets regularly overflow.

What changes for users

• Text that fits the model follows the existing path unchanged.
• Overlength text produces multiple embedded rows that exactly reconstruct the original text and token sequence.
• Every child keeps stable ordering, IDs, and source/document/page attribution.
• If any expected embedding is missing, the write fails before calling a VDB backend.
• Image inputs are unaffected.

Technical contract, scope, and validation

Accuracy contract

• Count the complete formatted input with the exact pinned tokenizer, model prefix, and special tokens.
• Use the smaller of the configured limit and the checkpoint-supported limit: 8,192 passage tokens for the pinned text model and 4,096 for the pinned VL model. Query input remains capped at 128.
• Split overlength text into deterministic contiguous token ranges before inference without normalizing or truncating the source.
• Preserve source, page, element, bounding-box, document-chunk, and embedding-child provenance.
• Pin model metadata and tokenizer assets rather than guessing for an unknown remote model.
• Send truncate="NONE" after admission so a remote endpoint cannot silently change the decision.
• Require exactly one embedding result per admitted row and reject incomplete searchable coverage before VDB dispatch.
• Fail admission if exact source-text and token-sequence reconstruction cannot be proven.

Implementation\n\n- The obsolete compatibility embedding runtime and export are removed; all embedding actors use the single canonical runtime.\n- One model specification owns the exact checkpoint revision, backend, input limits, and prompts for both warmed and cold local embedders.

• A shared formatter and pure admission planner run before local and remote text embedding; a pandas adapter expands only rows that require splitting.
• Split children carry stable parent/child IDs, order, count, exact content, and source-token ranges under metadata["embedding_split"].
• Text fallback for text_image rows without an image is admitted without changing its existing route or batch size. True image inputs do not resolve text admission.
• Default service images preload the complete pinned admission policy for offline startup.
• The backend-neutral VDB adapter rejects missing or empty embeddings before any VDB implementation is called.
• Telemetry reports input/output coverage, splits, failures, and zero truncation.

Compatibility and scope

• Overlength inputs can increase the returned row count; fitting inputs retain their existing public fields, values, dtypes, route, and batching behavior.
• An unpinned custom remote model fails setup with an actionable error. Callers must use a registered model, a local checkpoint, or an immutable model revision.
• This PR does not change Ray scheduling, batching, backpressure, model configuration, extraction chunking, retrieval settings, LanceDB index construction, or the bounded sink in feat(vdb): add optional streaming ingest with LanceDB #2567.

Validation

• Real one-H100 local-embedder smoke: short / overlength / short input produced both short embeddings and every expected split-child embedding, all 2,048-dimensional.
• Child text and token ranges reconstructed the parent exactly across two runs; IDs, order, document, and page fields were deterministic and preserved.
• The canonical VDB adapter accepted the complete result and rejected a deliberately missing child before backend dispatch.
• Full BO767 ingestion preserved every overflow child and affected page; retrieval showed no material regression (nDCG@10 delta -0.001064, Recall@5 and Recall@10 delta +0.001009).
• JP20 retrieval metrics and document/page attribution were unchanged, with zero split expansion.
• Local full suite: 3,489 passed, 171 skipped, 13 deselected, and 26 subtests passed.
• GitHub pre-commit, unit, slim-import, macOS, Windows, Docker/offline full-suite, Greptile, and final validation checks pass.
• Black, Flake8, and git diff --check pass.

Suggested review order

1. Model-pinned input limit and tokenizer resolution.
2. Admission, exact reconstruction, and provenance tests.
3. Local/remote cardinality and truncate="NONE" behavior.
4. Fail-closed VDB coverage checks.
5. Offline image preload and documentation.

[jioffe502]

The remaining BO767 runtime gate has passed on the accuracy-first production stack containing this patch.

• Both same-stream arms wrote the identical 83,020-row canonical table from all 767 sources.
• The canonical row-identity multiset SHA-256 matched exactly.
• Both arms retained the same four deterministic split children from the same two overlength parents, with no source-allocation changes.
• The write receipt reconciled 83,343 post-split source rows, 323 deterministic policy rejections, 83,020 writes, and zero upstream errors.
• The vector index fully covered all written rows and retrieval completed successfully.
• Full CI and the fresh-cache online-preload/offline-startup gate are green.

This closes the draft gate. Timing was considered only after exact identity and receipt parity passed; the result used to release this PR is correctness, not a speed claim.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no actionable new defect or outstanding blocking finding remains.

Summary

• Resolves immutable model metadata, tokenizer assets, prompts, and supported input limits.
• Splits only overlength text into reversible, ordered children before local or remote inference.
• Aligns local, remote, and warmed embedding runtimes around the same admission policy.
• Rejects missing or empty embeddings before VectorDB backend dispatch.
• Adds focused coverage for reconstruction, cardinality, model limits, offline packaging, remote truncation behavior, and VDB validation.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Embedding rows] --> B[Resolve pinned model policy]
    B --> C{Text requires admission?}
    C -->|No: image input| F[Inference]
    C -->|Yes| D{Formatted input fits?}
    D -->|Yes| F
    D -->|No| E[Create reversible ordered split children]
    E --> F
    F --> G{Exactly one valid embedding per admitted row?}
    G -->|Yes| H[Canonical VDB records]
    G -->|No| I[Reject before backend write]
    H --> J[Vector database]

Loading

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jioffe502]

Oversized-input correctness qualification — full BO767 gate passed

Head 4c39378f closes the two cardinality holes found after the earlier qualification: the local vLLM/HF text embedders could filter a whitespace-only split child, and canonical VDB conversion could omit that child even after it was embedded. The policy now also fails closed when the selected prompt or checkpoint-supported limit is undeclared, and fitting batches preserve their original DataFrame index.

Incidence:

• JP20: 0 / 3,201 text rows over 8,192 formatted tokens; max 4,555. Frozen candidate and exact-base control had identical reported retrieval metrics.
• BO767: 2 / 83,341 text rows over limit (0.0023998%); max 10,647, one document and two pages. Splitting predicts four children and a net +2 rows.
• ViDoRe V3: 19,252 image-bearing text_image rows use the VL route; text splitting does not govern those image inputs.
• Truncating the two BO767 rows would discard 3,897 content tokens total (14.98% and 23.07% of the affected pages). Splitting discards zero tokens and adds only two rows.

Targeted validation at 4c39378f:

• 314 focused embedding, actor, and VDB tests pass.
• Full local differential against frozen 8b22fdf: identical 18 environment-dependent failures; all 12 added tests pass (3,445 vs. 3,433 passes).
• Real pinned vLLM smoke passed twice on one H100 NVL with nvidia/llama-nemotron-embed-1b-v2 revision 113abe4acafa848e77ead9c0623205e511932348.
• A deliberately whitespace-only split child survived in a short / overflow / short batch. All four rows received 2,048-dimensional embeddings; exact text/token reconstruction, provenance, deterministic order/IDs, canonical VDB acceptance, and fail-closed missing-child rejection passed.

Full BO767 normal batch path at 4c39378f:

• Canonical coverage completed: 767 PDFs, 54,730 pages, all 991 queries.
• 83,340 pre-embedding rows became exactly 83,342 post-embedding rows. The two real overflow parents (1569715.pdf, pages 68 and 98) became four children: overlength=2 split=2 truncated=0 failed=0 split_children=4.
• Both parent text hashes and concatenated child token IDs reconstruct exactly. All ranges are contiguous; every child fits the formatted limit; every child has a 2,048-dimensional vector and preserved document/page attribution.
• LanceDB contains 83,019 rows with zero null vectors. The 323-row difference exactly matches the 323 logged inputs without searchable text/image backing; no split child was lost.
• Retrieval: nDCG@10 0.7514068452, recall@5 0.8536831483, recall@10 0.9001009082. Versus the prior same-host full control: -0.001064, +0.001009, +0.001009, respectively—no material quality-regression signal.
• The affected overflow pages retained all three prior top-10 appearances: page 68 improved rank 10 → 8 for query 301; page 98 stayed rank 6 for query 632 and rank 8 for query 680.

Scope remained accuracy-first: no batching, Ray scheduling, model configuration, extraction/chunking defaults, or retrieval settings changed.

Caveat: this fresh extraction produced one fewer pre-embedding row than the earlier incidence snapshot (83,340 vs. 83,341), upstream of embedding and consistent with known independent-run extraction variability. The full-control comparison is same-host and configuration-equivalent at the relevant ingestion/model settings, but it is not a same-stream replay; strict row/ranked-list identity is not claimed.

Conclusion: the oversized-input path is rare in these corpora, but the two real BO767 cases now pass end-to-end without truncation, row loss, or attribution loss. This is production protection, not evidence that normal traffic routinely overflows.

[jperez999]

Could we place this under one namespaced mapping, such as metadata["embedding_split"]? These top-level fields include the generic content key and expand the record-wide metadata contract, which makes collisions with extraction and VDB metadata more likely. A nested object would keep split provenance cohesive and easier to evolve.

[jioffe502]

Addressed in fdf6dd1. Split provenance now lives in one metadata["embedding_split"] mapping containing exact content, parent/child IDs, order/count, and token range. Because this contract is introduced by this PR, I updated the producer, consumers, docs, and tests directly without a legacy flat-key path.

[jperez999]

Could we split pure admission/planning from DataFrame expansion and provenance mutation? EmbeddingInputPolicy currently owns tokenizer policy, token planning, row copying, ID derivation, metadata mutation, and pandas construction. Separating a pure split plan from its pandas adapter would make it easier to test and reuse outside this runtime.

[jioffe502]

Addressed in fdf6dd1. EmbeddingInputPolicy.plan() now returns immutable, side-effect-free EmbeddingSplitPlan / EmbeddingSplitChild values. The separate prepare_embedding_inputs() pandas adapter owns row copying, stable ID derivation, provenance attachment, and DataFrame construction.

[jperez999]

The generic VDB record adapter now recognizes embedding-specific split children and validates embedding coverage. Is there an existing embedding-boundary helper this can reuse or move behind? Keeping VDB conversion ingestion-agnostic avoids future embedding-policy changes requiring synchronized edits in record conversion.

[jioffe502]

Addressed in fdf6dd1. records.py no longer imports split-key lists or interprets/copies individual split fields. It calls the single embedding_record_content() helper, which keeps the split layout behind the embedding schema seam. I retained the generic missing-embedding check at VDB dispatch because that is the point that prevents an incomplete partial write.

[jperez999]

Prefix formatting is duplicated between this admission policy and the local embedder implementations that prepare query and passage strings. Since admission counts tokens after prefixing, these need to stay byte-for-byte aligned with actual inference. Could we extract a shared embedding-input formatter so a future prefix-rule change cannot make preflight accounting disagree with the backend?

[jioffe502]

Addressed in fdf6dd1. format_embedding_input() is now the shared formatter used by admission, local vLLM prefix application, the local HF text embedder, and the legacy local runtime adapter. The explicit prefix_if_missing argument preserves the HF backend rule while keeping the formatted string identical to what admission counts.

[jperez999]

[P2] Consolidate the duplicate embedding runtime

nemo_retriever.common.modality.pipeline still exports a second embed_text_main_text_embed implementation. That path calls create_text_embeddings_for_df() directly, never accepts or applies embedding_input_policy, and hard-codes truncate="END", so callers of the exported compatibility entry point can still silently truncate the overlength text this PR is intended to preserve. It also duplicates the grouping, prefixing, error, and output-column logic being changed here. Could we make that module a thin compatibility wrapper around this implementation (or remove its export if it is truly dead) so there is one admission/runtime path?

[jioffe502]

Addressed in 3656624. I confirmed repository-wide that no caller imports the old common.modality.pipeline embedding export; the CPU/GPU actors and tests already use models.inference.runtime. I removed the export and deleted the duplicate 188-line implementation rather than adding a compatibility wrapper. There is now one embedding runtime, so admission policy, no-truncation behavior, grouping, errors, and output cardinality cannot diverge between these two paths. Validation: 353 focused tests passed, plus Black, Flake8, import smoke, and git diff --check.