feat(rerank): add per-model token limits for cross-encoder rerank

[vlsi]

Why

A cross-encoder scores one <s>query</s></s>document</s> sequence per candidate, so an oversized query or document inflates every pair and can exhaust memory or stall the backend when a client sends no limits. The /rerank API had no way to bound the scored sequence length, and a single global default cannot fit every reranker.

What

Add three token budgets that bound the scored sequence per pair:

• max_query_tokens: head-truncate the query.
• max_tokens_per_doc: head-truncate each document (Cohere v2 compatible).
• max_pair_tokens: cap the joined pair, trimming the longest side first.

Configured per model at startup, since a sensible ceiling depends on the model:

• EngineArgs.max_query_tokens / max_tokens_per_doc / max_pair_tokens
• env vars INFINITY_MAX_QUERY_TOKENS / INFINITY_MAX_TOKENS_PER_DOC / INFINITY_MAX_PAIR_TOKENS (per model, ;-separated)
• v2 CLI flags --max-query-tokens / --max-tokens-per-doc / --max-pair-tokens

Each defaults to None (no limit). The /rerank request still accepts the same three fields (defaulting to null), but the engine clamps each requested budget to the model's startup ceiling: a client may lower a limit to trade quality for speed, but cannot raise it above the configured maximum. The startup config owns stability; the client owns precision.

The values thread RerankInput → engine.rerank → AsyncEngineArray → batch_handler.rerank, where they are bundled into a RerankLimits tuple on each ReRankSingle and reach both cross-encoder encode_pre paths (torch and optimum) via to_input(). Truncation runs in the single preprocessing thread using the model's own tokenizer (not the token-counting copy used on another thread) to avoid a data race; per-item pair tokenisation plus tokenizer.pad keeps each item's limits correct inside a mixed batch. Two-element (query, document) tuples still work and apply no truncation.

The checked-in OpenAPI spec (docs/assets/openapi.json) is regenerated so the Swagger UI lists the three fields with their null defaults.

How to verify

• cd libs/infinity_emb && pytest tests/unit_test/test_args.py -k "clamp or reject" (ceiling/clamp logic, no model download)
• pytest tests/unit_test/transformer/crossencoder/test_torch_crossencoder.py::test_crossencoder_rerank_limits
• Start with a ceiling, e.g. infinity_emb v2 --model-id <reranker> --max-pair-tokens 512, then send a /rerank request with max_pair_tokens: 4096 and confirm it is clamped to 512.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces configurable token budgets and limits for reranking (via RerankLimits) across the engine, batch handler, and cross-encoder backends to prevent performance issues on oversized inputs. The feedback highlights three key robustness improvements: optimizing the tokenizer in truncate_texts_to_tokens with a maximum cap to prevent potential OOMs, adding a fallback for max_length in the PyTorch cross-encoder to avoid ValueError exceptions, and using getattr with a fallback for max_position_embeddings in the Optimum cross-encoder to prevent potential AttributeError exceptions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cf8a67b0fa

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[greptile-apps]

Greptile Summary

Adds per-model token-budget ceilings for the cross-encoder reranker. Three startup settings (max_query_tokens, max_tokens_per_doc, max_pair_tokens) configure a server-side maximum; the same three fields on the /rerank request body let a client lower but never exceed those ceilings. Truncation happens inside encode_pre using a dedicated helper (truncate_texts_to_tokens) that pre-truncates queries and documents independently before pair encoding.

• RerankLimits NamedTuple threads the resolved budgets from the HTTP layer all the way to the tokenizer, and its values are included in str_repr so the disk cache distinguishes capped from uncapped results for the same pair.
• Both the torch and ONNX/optimum encode_pre paths were rewritten to accept 3-tuples (query, doc, RerankLimits) while retaining 2-tuple backward compatibility; per-item pair tokenisation followed by tokenizer.pad keeps the limits correct inside a mixed batch.
• New env vars (INFINITY_MAX_QUERY_TOKENS / INFINITY_MAX_TOKENS_PER_DOC / INFINITY_MAX_PAIR_TOKENS) and v2 CLI flags expose the new options.

Confidence Score: 5/5

Safe to merge; the clamping logic, truncation helper, and cache key changes are all correct for the single-model and default-limits cases that make up the common deployment.

The core token-budget mechanism — ceiling clamping, truncate_texts_to_tokens, per-item pair encoding, RerankLimits threading — is implemented correctly. Both encoder backends guard max_pair_tokens with min(limit, model_max). The two findings are edge cases: the CLI type annotation mismatch only matters when a user needs to express a null limit for one specific model in a multi-model CLI invocation (env vars remain the correct path for that), and the cache key ambiguity requires a deliberately crafted query string to trigger.

cli.py (type annotation for the three new list options) and primitives.py (str_repr suffix format) are worth a second look before a 1.0 release, but neither blocks the feature from working correctly in normal deployments.

Important Files Changed

Filename	Overview
libs/infinity_emb/infinity_emb/cli.py	New CLI options declared as list[int] but default value from the env manager is [None], creating a type mismatch; CLI users also cannot express "no limit" for a specific model in a multi-model flag sequence.
libs/infinity_emb/infinity_emb/primitives.py	Adds RerankLimits NamedTuple and threads it through ReRankSingle; str_repr correctly varies the cache key when limits are non-default, but the key format (plain concatenation) inherits the existing ambiguity of query+document concatenation.
libs/infinity_emb/infinity_emb/transformer/crossencoder/init.py	New truncate_texts_to_tokens helper; batch-level tokenisation bounded by the largest cap keeps the OOM concern in check, but the decode/re-encode round-trip means texts at exactly the cap boundary may differ slightly from their originals.
libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py	Rewrites encode_pre to handle 3-tuples; adds per-item pair tokenisation + tokenizer.pad; correctly caps pair length to min(limit, model_max) and falls back gracefully for legacy 2-tuples.
libs/infinity_emb/infinity_emb/transformer/crossencoder/optimum.py	Mirrors the torch path for the ONNX backend; per-item encoding + tokenizer.pad keeps dtype and shape consistent; model_max fallback chain is identical to torch.
libs/infinity_emb/infinity_emb/engine.py	Adds _clamp_to_ceiling helper with correct min-of-present-values semantics; ceiling clamping applied consistently in AsyncEmbeddingEngine.rerank and forwarded through AsyncEngineArray.
libs/infinity_emb/infinity_emb/args.py	Adds three new Optional[int] fields with __post_init__ validation; correctly threads them through the zip_longest multi-model expansion in EngineArgs.create.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant Server as infinity_server (FastAPI)
    participant Engine as AsyncEmbeddingEngine
    participant Array as AsyncEngineArray
    participant BH as BatchHandler
    participant Pre as encode_pre (preprocessing thread)

    Client->>Server: "POST /rerank {query, docs, max_query_tokens?, max_tokens_per_doc?, max_pair_tokens?}"
    Server->>Engine: rerank(query, docs, max_query_tokens, max_tokens_per_doc, max_pair_tokens)
    Note over Engine: _clamp_to_ceiling(requested, EngineArgs.ceiling) for each limit
    Engine->>Array: rerank(...clamped limits...)
    Array->>BH: rerank(...clamped limits...)
    Note over BH: Build RerankLimits(clamped values), Create ReRankSingle per doc
    BH->>Pre: encode_pre(list[tuple[str,str,RerankLimits]])
    Note over Pre: truncate_texts_to_tokens for queries and docs independently, then per-item pair tokenization + tokenizer.pad
    Pre-->>BH: padded batch tensors
    BH-->>Engine: scores, usage
    Engine-->>Client: RerankResponse

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Client
    participant Server as infinity_server (FastAPI)
    participant Engine as AsyncEmbeddingEngine
    participant Array as AsyncEngineArray
    participant BH as BatchHandler
    participant Pre as encode_pre (preprocessing thread)

    Client->>Server: "POST /rerank {query, docs, max_query_tokens?, max_tokens_per_doc?, max_pair_tokens?}"
    Server->>Engine: rerank(query, docs, max_query_tokens, max_tokens_per_doc, max_pair_tokens)
    Note over Engine: _clamp_to_ceiling(requested, EngineArgs.ceiling) for each limit
    Engine->>Array: rerank(...clamped limits...)
    Array->>BH: rerank(...clamped limits...)
    Note over BH: Build RerankLimits(clamped values), Create ReRankSingle per doc
    BH->>Pre: encode_pre(list[tuple[str,str,RerankLimits]])
    Note over Pre: truncate_texts_to_tokens for queries and docs independently, then per-item pair tokenization + tokenizer.pad
    Pre-->>BH: padded batch tensors
    BH-->>Engine: scores, usage
    Engine-->>Client: RerankResponse

Loading

_{Reviews (3): Last reviewed commit: "feat(rerank): add per-model token limits..." | Re-trigger Greptile}

[vlsi]

Thanks for the review. Two notes first: the limits are now configured per model at startup (EngineArgs / INFINITY_MAX_* / v2 flags) and each defaults to None rather than 256/512/768 — a request can only lower a limit, never raise it above the model's configured ceiling. The findings still applied to the threading code, so I addressed them:

• Clamp the torch pair length to the model max (P1, greptile + codex + gemini). Fixed. torch.encode_pre now uses min(max_pair_tokens, max_position_embeddings or model_max_length), matching the optimum path; an over-large or unset budget can no longer build a sequence the model cannot process in encode_core. Added a regression assertion.
• max_position_embeddings via getattr (gemini). Fixed in both paths: getattr(config, "max_position_embeddings", None) or tokenizer.model_max_length.
• truncation=False can OOM on huge inputs (gemini, greptile). Fixed. truncate_texts_to_tokens now tokenises with truncation=True, max_length=max(caps)+1, which bounds the work while still detecting overflow.
• Cache keys must include the limits (codex P2). Fixed. ReRankSingle.str_repr now includes the limits when set, so vector_disk_cache no longer serves a capped result for a later uncapped request; the no-limit path is byte-for-byte unchanged.

Force-pushed as a single squashed commit.