Text Ranking
Transformers
Safetensors
multilingual
t5gemma2
text2text-generation
reranker
encoder-decoder
FBNL
Retrieval
RAG
Instructions to use KaLM-Embedding/KaLM-Reranker-V1-Nano with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use KaLM-Embedding/KaLM-Reranker-V1-Nano with Transformers:
# Load model directly from transformers import AutoProcessor, AutoModelForMultimodalLM processor = AutoProcessor.from_pretrained("KaLM-Embedding/KaLM-Reranker-V1-Nano") model = AutoModelForMultimodalLM.from_pretrained("KaLM-Embedding/KaLM-Reranker-V1-Nano") - Notebooks
- Google Colab
- Kaggle
Update kalm_reranker.py
Browse files- 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 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 |
)
|