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

kalm_reranker.py

@@ -283,9 +283,10 @@ class KaLMReranker:
         sorted_pairs = [validated_pairs[index] for index in length_sorted_indices]
 
         tested_batch_size = effective_batch_size
+        first_batch_scores: Optional[List[float]] = None
         while tested_batch_size > 1:
             try:
-                self._predict_batch(
+                first_batch_scores = self._predict_batch(
                     sorted_pairs[: min(len(sorted_pairs), tested_batch_size)],
                     effective_instruction,
                 )
@@ -295,9 +296,19 @@ class KaLMReranker:
                     torch.cuda.empty_cache()
                 tested_batch_size = max(1, tested_batch_size * 3 // 4)
 
-        sorted_scores: List[float] = []
+        # The while loop's condition (`> 1`) means batch size 1 is never
+        # actually probed. If every size down to 2 OOMs, it exits without a
+        # successful probe. Only skip ahead to `tested_batch_size` when the
+        # probe actually ran; otherwise fall back to starting at 0 like the
+        # loop below always did originally, or the first item(s) get dropped.
+        if first_batch_scores is None:
+            sorted_scores: List[float] = []
+            loop_start = 0
+        else:
+            sorted_scores = list(first_batch_scores)
+            loop_start = tested_batch_size
         try:
-            for start in range(0, len(sorted_pairs), tested_batch_size):
+            for start in range(loop_start, len(sorted_pairs), tested_batch_size):
                 sorted_scores.extend(
                     self._predict_batch(
                         sorted_pairs[start : start + tested_batch_size],