fix(reranker): avoid re-computing the first batch in predict()'s batch-size probe to reduce additional computational effort
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.
The former PR introduced a bug when Cuda OOM errors lead to tested_batch_size=1. In the former PR the first batch was then silently dropped. This PR adds the correct logic to prevent this behavior, by checking if scores were computed for the probe batch , see comment in lines 299-303
Thank you very much. We will review it as soon as possible.