fix(reranker): avoid re-computing the first batch in predict()'s batch-size probe to reduce additional computational effort

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],
@@ -346,4 +357,4 @@ class KaLMReranker:
         return rankings if top_k is None else rankings[:top_k]
 
 
-__all__ = ["KaLMReranker"]
+__all__ = ["KaLMReranker"]