Update kalm_reranker.py

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,11 +70,7 @@ 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)
@@ -280,8 +276,7 @@ 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]
         )