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

kalm_reranker.py

@@ -62,7 +62,7 @@ class KaLMReranker:
             if self.tokenizer.eos_token_id is None:
                 raise ValueError("The tokenizer must define a pad token or an EOS token.")
             self.tokenizer.pad_token = self.tokenizer.eos_token
-        
+        # Last-token indexing below assumes right padding, matching training.
         self.tokenizer.padding_side = "right"
 
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -70,7 +70,11 @@ class KaLMReranker:
             dtype=self.dtype,
             **model_kwargs,
         )
-        
+        # Preserve model buffers in their checkpoint dtypes. In particular,
+        # T5Gemma2 keeps RoPE inverse-frequency buffers in FP32 even for BF16
+        # inference. Casting the whole module would silently change its scores.
+        # A few tied parameters can retain a nested config dtype on CPU, so only
+        # parameters that need correction are converted explicitly.
         for parameter in self.model.parameters():
             if parameter.is_floating_point() and parameter.dtype != self.dtype:
                 parameter.data = parameter.data.to(dtype=self.dtype)
@@ -276,16 +280,18 @@ class KaLMReranker:
         if not isinstance(effective_batch_size, int) or effective_batch_size <= 0:
             raise ValueError("batch_size must be a positive integer.")
 
-        
+        # Match FlagEmbedding: sort by approximate text length to reduce padding,
+        # score contiguous batches, then restore the caller's original order.
         length_sorted_indices = np.argsort(
             [-(len(query) + len(document)) for query, document in validated_pairs]
         )
         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 +301,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],