Text Ranking
Transformers
Safetensors
multilingual
t5gemma2
text2text-generation
reranker
encoder-decoder
FBNL
Retrieval
RAG
Yuki131 commited on
Commit
4388e79
·
verified ·
1 Parent(s): 4aadb1c

Update kalm_reranker.py

Browse files
Files changed (1) hide show
  1. kalm_reranker.py +3 -8
kalm_reranker.py CHANGED
@@ -62,7 +62,7 @@ class KaLMReranker:
62
  if self.tokenizer.eos_token_id is None:
63
  raise ValueError("The tokenizer must define a pad token or an EOS token.")
64
  self.tokenizer.pad_token = self.tokenizer.eos_token
65
- # Last-token indexing below assumes right padding, matching training.
66
  self.tokenizer.padding_side = "right"
67
 
68
  self.model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -70,11 +70,7 @@ class KaLMReranker:
70
  dtype=self.dtype,
71
  **model_kwargs,
72
  )
73
- # Preserve model buffers in their checkpoint dtypes. In particular,
74
- # T5Gemma2 keeps RoPE inverse-frequency buffers in FP32 even for BF16
75
- # inference. Casting the whole module would silently change its scores.
76
- # A few tied parameters can retain a nested config dtype on CPU, so only
77
- # parameters that need correction are converted explicitly.
78
  for parameter in self.model.parameters():
79
  if parameter.is_floating_point() and parameter.dtype != self.dtype:
80
  parameter.data = parameter.data.to(dtype=self.dtype)
@@ -280,8 +276,7 @@ class KaLMReranker:
280
  if not isinstance(effective_batch_size, int) or effective_batch_size <= 0:
281
  raise ValueError("batch_size must be a positive integer.")
282
 
283
- # Match FlagEmbedding: sort by approximate text length to reduce padding,
284
- # score contiguous batches, then restore the caller's original order.
285
  length_sorted_indices = np.argsort(
286
  [-(len(query) + len(document)) for query, document in validated_pairs]
287
  )
 
62
  if self.tokenizer.eos_token_id is None:
63
  raise ValueError("The tokenizer must define a pad token or an EOS token.")
64
  self.tokenizer.pad_token = self.tokenizer.eos_token
65
+
66
  self.tokenizer.padding_side = "right"
67
 
68
  self.model = AutoModelForSeq2SeqLM.from_pretrained(
 
70
  dtype=self.dtype,
71
  **model_kwargs,
72
  )
73
+
 
 
 
 
74
  for parameter in self.model.parameters():
75
  if parameter.is_floating_point() and parameter.dtype != self.dtype:
76
  parameter.data = parameter.data.to(dtype=self.dtype)
 
276
  if not isinstance(effective_batch_size, int) or effective_batch_size <= 0:
277
  raise ValueError("batch_size must be a positive integer.")
278
 
279
+
 
280
  length_sorted_indices = np.argsort(
281
  [-(len(query) + len(document)) for query, document in validated_pairs]
282
  )