fix(embedding): split overlength text before inference

[jioffe502]

Problem

The embedding stage currently submits a pandas batch directly to the model. If one text row exceeds the checkpoint-supported input limit, the local model rejects the request, the batch is returned without embeddings, and downstream VDB conversion can discard valid neighboring rows. Remote embedding can instead truncate, which changes indexed content silently.

This is an accuracy and coverage bug. Batch composition must not determine which valid documents reach the index.

Accuracy contract

This PR establishes one fail-closed contract for local and remote text embedding:

• count the complete formatted input with the exact model tokenizer, including the query/document prefix and special tokens;
• leave a batch byte-for-byte unchanged when every text input fits;
• split overlength text into deterministic contiguous ranges before inference, without truncating or changing the source text or token sequence;
• preserve source, page, element, bounding-box, document-chunk, and embedding-child provenance;
• require pinned model metadata and tokenizer assets instead of guessing for an unknown remote model;
• send truncate="NONE" after client-side admission so a remote endpoint cannot silently override the decision;
• treat an unexpected backend rejection as a batch failure rather than guessing from an HTTP status or exception that one row is invalid;
• refuse a mixed dense VDB write when searchable rows are missing embeddings.

If exact splitting cannot preserve the original source text and token sequence, admission fails before inference.

What changes

• Embedding model resolution now includes the checkpoint-declared input limit and exact query/document prompts.
• A shared admission policy runs before both local and remote embedding actors.
• Split children receive stable parent/child IDs, order, count, and source-token ranges under metadata.
• The default service images preload the complete pinned admission policy, not only tokenizer.json, so offline actor startup has the required config and prompt assets.
• The backend-neutral canonical VDB adapter rejects incomplete searchable coverage before any VDB implementation is called. A page URI inherited only as structured-element provenance is not mistaken for independently searchable image content.
• Admission telemetry reports input/output coverage, splits, failures, and zero truncation.

Compatibility and deliberate behavior changes

• Inputs that fit retain their existing public fields, values, and dtypes.
• Overlength inputs produce multiple embedded rows, so the returned DataFrame can contain more rows than the input batch.
• The effective limit is the smaller of the configured runtime limit and the checkpoint-declared limit. For the pinned defaults this is 4,096 passage tokens for the VL model and 8,192 for the text model; query input remains capped at 128.
• An unpinned custom remote model now fails actor setup with an actionable error. Callers must use a registered model, a local checkpoint, or provide an immutable model revision.
• Existing generic embedding failure behavior remains batch-level. This PR does not introduce recursive request replay or generic row-failure classification.

Scope

This PR does not change Ray scheduling, batching, backpressure, LanceDB-specific row builders, index construction, or the bounded sink in #2567. The batch library path continues through IngestVdbOperator and the abstract VDB interface. It can be reviewed and merged independently ahead of #2567.

Validation

• 218 focused admission, local/remote actor, backend-neutral VDB conversion, model-spec, packaging, and operator tests pass locally.
• The exact pinned default VL tokenizer preserved source text and token sequences for long whitespace/control-character, decomposed Unicode/non-Latin/emoji, and literal-special-token inputs; every emitted child fit the 4,096-token checkpoint limit.
• Black, Flake8, and git diff --check pass for the touched Python files.
• The two Arrow-dtype failures encountered in the broader shared local environment reproduce unchanged on current upstream/main; the pull-request CI matrix is the clean-environment full-suite gate.
• Earlier source-frozen qualification of the core admission path completed matched 12-dataset base/candidate runs across 142,967 pages per arm, with exact warm JP20 accuracy and aggregate text-ingestion throughput within observed run variance.

Suggested review order

1. Model-pinned input limit and tokenizer resolution.
2. Admission, exact split invariants, and provenance tests.
3. Local/remote actor wiring and truncate="NONE" behavior.
4. Fail-closed canonical VDB coverage checks and inherited-page-URI regression tests.
5. Offline service-image preload and user 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]

Greptile Summary

The PR adds tokenizer-exact admission and deterministic splitting for overlength embedding inputs while preserving provenance and failing closed at inference and VDB boundaries.

• Resolves pinned tokenizer, prompt, and checkpoint input-limit metadata for local and remote embedding.
• Splits oversized text into deterministic embedding children before inference and disables backend truncation.
• Rejects partial dense VDB writes when searchable rows lack embeddings.
• Updates service-image preloading, documentation, telemetry, and focused tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/models/inference/embedding_input.py	Implements pinned tokenizer admission, exact formatted-token counting, deterministic splitting, and child provenance.
nemo_retriever/src/nemo_retriever/models/inference/runtime.py	Applies admission before inference, disables remote truncation, and maps backend errors to batch-level failures with coverage telemetry.
nemo_retriever/src/nemo_retriever/common/vdb/records.py	Preserves embedding-child provenance and rejects mixed dense writes with missing searchable embeddings while excluding inherited page URIs.
nemo_retriever/src/nemo_retriever/models/embed_model_spec.py	Resolves checkpoint-declared input limits and exact prompt metadata from pinned or local model assets.
Dockerfile	Preloads the complete default embedding admission policy for offline service startup.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Input[Embedding batch] --> Resolve[Resolve pinned model policy]
    Resolve --> Admit[Count formatted tokenizer inputs]
    Admit -->|Fits| Infer[Local or remote inference]
    Admit -->|Overlength| Split[Create deterministic provenance-preserving children]
    Split --> Infer
    Infer -->|Success| Coverage[Validate searchable-row embedding coverage]
    Infer -->|Backend rejection| Fail[Fail prepared batch]
    Coverage -->|Complete| VDB[Write canonical VDB records]
    Coverage -->|Incomplete| Reject[Raise VdbUploadError]

Loading

_{Reviews (4): Last reviewed commit: "fix(embedding): tighten overflow admissi..." | Re-trigger Greptile}

[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.