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

kalm_reranker.py

@@ -282,10 +282,12 @@ class KaLMReranker:
         )
         sorted_pairs = [validated_pairs[index] for index in length_sorted_indices]
 
+
         tested_batch_size = effective_batch_size
+        first_batch_scores: List[float] = []
         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 +297,9 @@ class KaLMReranker:
                     torch.cuda.empty_cache()
                 tested_batch_size = max(1, tested_batch_size * 3 // 4)
 
-        sorted_scores: List[float] = []
+        sorted_scores: List[float] = list(first_batch_scores)
         try:
-            for start in range(0, len(sorted_pairs), tested_batch_size):
+            for start in range(tested_batch_size, len(sorted_pairs), tested_batch_size):
                 sorted_scores.extend(
                     self._predict_batch(
                         sorted_pairs[start : start + tested_batch_size],