Text Ranking
Transformers
Safetensors
multilingual
t5gemma2
text2text-generation
reranker
encoder-decoder
FBNL
Retrieval
RAG

fix(reranker): avoid re-computing the first batch in predict()'s batch-size probe

#1
by lukann98 - opened

Existing Issue:
KaLMReranker.predict() runs a full forward pass on the first batch twice: once as a throwaway OOM-size probe, and once again for the real computation. The probe's result is computed but never used. This roughly doubles reranking latency for any call where the total number of documents fits within a single batch (the common case, since batch_size defaults to 32).

Proposed Fix:
Keep the probe's result and reuse it as the first batch's scores, starting the real loop after the first batch instead of from the beginning.

This PR introduces a bug whereby, if CUDA OOM errors occur for every batch size up to tested_batch_size=1, the first batch will never be computed. While this is likely to be a rare edge case, it is a newly introduced bug. I will therefore close my PR, thoroughly test my adapted solution, and submit a new PR.

lukann98 changed pull request status to closed

Sign up or log in to comment