Fix Transformers 5 tokenizer compatibility

FlagEmbedding/abc/finetune/reranker/AbsDataset.py

@@ -16,7 +16,7 @@
 from typing import List
 
 from .AbsArguments import AbsRerankerDataArguments
-from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
 
 logger = logging.getLogger(__name__)
 
@@ -197,7 +197,8 @@ def __call__(self, features) -> List[BatchEncoding]:
         if isinstance(features[0], list):
             features = sum(features, [])
 
-        collated = self.tokenizer.pad(
+        collated = pad_with_compat(
+            self.tokenizer,
             features,
             padding=self.padding,
             max_length=self.query_max_len + self.passage_max_len,
@@ -391,7 +392,8 @@ def __call__(self, features, return_tensors='pt'):
                 else:
                     feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
 
-        collated = self.tokenizer.pad(
+        collated = pad_with_compat(
+            self.tokenizer,
             features,
             padding=self.padding,
             max_length=self.query_max_len + self.passage_max_len,

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -6,6 +6,7 @@
 from transformers import AutoModel, AutoTokenizer
 
 from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
 
 
 # Pooling function for LLM-based embedding models
@@ -251,7 +252,8 @@ def encode_single_device(
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -270,7 +272,8 @@ def encode_single_device(
         for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
                                 disable=len(sentences) < batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -10,6 +10,7 @@
 from transformers import AutoModel, AutoTokenizer
 
 from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
 
 
 # Pooling function for LLM-based embedding models
@@ -406,7 +407,8 @@ def encode_queries_single_device(
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -425,7 +427,8 @@ def encode_queries_single_device(
         for start_index in tqdm(range(0, len(sentences_sorted), batch_size), desc="Inference Embeddings",
                                 disable=len(sentences_sorted) < batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',
@@ -517,7 +520,8 @@ def encode_single_device(
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -536,7 +540,8 @@ def encode_single_device(
         for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
                                 disable=len(sentences) < batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from .base import BaseLLMEmbedder, last_token_pool
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
 
 
 class PseudoMoELLMEmbedder(BaseLLMEmbedder):
@@ -139,7 +140,8 @@ def encode_single_device(
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -160,7 +162,8 @@ def encode_single_device(
         all_embeddings = []
         for start_index in range(0, len(sentences), batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -6,6 +6,7 @@
 from transformers import AutoModel, AutoTokenizer
 
 from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
 
 
 class BaseEmbedder(AbsEmbedder):
@@ -232,7 +233,8 @@ def encode_single_device(
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -251,7 +253,8 @@ def encode_single_device(
         for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
                                 disable=len(sentences) < batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -10,6 +10,7 @@
 from typing import Any, List, Union, Dict, Literal, Tuple, Optional
 
 from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
 from FlagEmbedding.finetune.embedder.encoder_only.m3 import (
     EncoderOnlyEmbedderM3ModelForInference, EncoderOnlyEmbedderM3Runner
 )
@@ -396,7 +397,8 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
         flag = False
         while flag is False:
             try:
-                inputs_batch = self.tokenizer.pad(
+                inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[: batch_size],
                     padding=True,
                     return_tensors='pt',
@@ -419,7 +421,8 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
         for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
                                 disable=len(sentences) < batch_size):
             inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs_batch = self.tokenizer.pad(
+            inputs_batch = pad_with_compat(
+                self.tokenizer,
                 inputs_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -10,7 +10,7 @@
 
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
-from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
 
 
 def last_logit_pool(logits: Tensor,
@@ -163,7 +163,8 @@ def __call__(self, data):
                 else:
                     feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
 
-        return self.tokenizer.pad(
+        return pad_with_compat(
+            self.tokenizer,
             data,
             padding=True,
             pad_to_multiple_of=8,

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -10,7 +10,7 @@
 
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
-from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
 
 
 def last_logit_pool_lightweight(logits: Tensor,
@@ -78,7 +78,8 @@ def __call__(self, data):
                 else:
                     feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
 
-        collected = self.tokenizer.pad(
+        collected = pad_with_compat(
+            self.tokenizer,
             features,
             padding=True,
             pad_to_multiple_of=8,

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -5,7 +5,7 @@
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from FlagEmbedding.abc.inference import AbsReranker
-from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat
+from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
 
 
 def sigmoid(x):
@@ -162,7 +162,8 @@ def compute_score_single_gpu(
         flag = False
         while flag is False:
             try:
-                test_inputs_batch = self.tokenizer.pad(
+                test_inputs_batch = pad_with_compat(
+                    self.tokenizer,
                     all_inputs_sorted[:min(len(all_inputs_sorted), batch_size)],
                     padding=True,
                     return_tensors='pt',
@@ -179,7 +180,8 @@ def compute_score_single_gpu(
         for start_index in tqdm(range(0, len(all_inputs_sorted), batch_size), desc="Compute Scores",
                                 disable=len(all_inputs_sorted) < batch_size):
             sentences_batch = all_inputs_sorted[start_index:start_index + batch_size]
-            inputs = self.tokenizer.pad(
+            inputs = pad_with_compat(
+                self.tokenizer,
                 sentences_batch,
                 padding=True,
                 return_tensors='pt',

FlagEmbedding/utils/tokenizer_compat.py

@@ -1,8 +1,34 @@
 """Tokenizer compatibility helpers for supported Transformers versions."""
 
+from collections.abc import Mapping
 from typing import Any, List, Optional, Sequence
 
 
+def pad_with_compat(tokenizer: Any, encoded_inputs: Any, **kwargs: Any) -> Any:
+    """Pad tokenized examples across Transformers v4 and v5.
+
+    Several FlagEmbedding inference paths sort individually tokenized
+    examples before padding, producing a list of mappings.  Transformers
+    releases do not all handle that representation consistently, so normalize
+    it to the equivalent mapping-of-lists representation before calling
+    ``tokenizer.pad``.
+
+    Inputs that are already mappings, empty inputs, and other supported
+    tokenizer inputs are passed through unchanged.
+    """
+    if (
+        isinstance(encoded_inputs, (list, tuple))
+        and encoded_inputs
+        and isinstance(encoded_inputs[0], Mapping)
+    ):
+        encoded_inputs = {
+            key: [example[key] for example in encoded_inputs]
+            for key in encoded_inputs[0].keys()
+        }
+
+    return tokenizer.pad(encoded_inputs, **kwargs)
+
+
 def _decode_token_ids(tokenizer: Any, token_ids: Sequence[int]) -> str:
     """Decode token ids without dropping unknown or other special tokens."""
     return tokenizer.decode(

examples/finetune/ds_stage0.json

@@ -22,7 +22,8 @@
             "lr": "auto",
             "betas": "auto",
             "eps": "auto",
-            "weight_decay": "auto"
+            "weight_decay": "auto",
+            "torch_adam": true
         }
     },
 

examples/finetune/ds_stage1.json

@@ -13,12 +13,7 @@
         "min_loss_scale": 1
     },
     "bf16": {
-        "enabled": "auto",
-        "loss_scale": 0,
-        "initial_scale_power": 10,
-        "loss_scale_window": 1000,
-        "hysteresis": 2,
-        "min_loss_scale": 1
+        "enabled": "auto"
     },
     "optimizer": {
         "type": "AdamW",

examples/finetune/embedder/README.md

@@ -197,8 +197,9 @@ torchrun --nproc_per_node 2 \
     --temperature 0.02 \
     --sentence_pooling_method cls \
     --normalize_embeddings True \
-    --kd_loss_type kl_div
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+    --kd_loss_type kl_div \
+    --warmup_steps 0.1
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 ```
 
 ### (2) bge-m3
@@ -239,8 +240,9 @@ torchrun --nproc_per_node 2 \
     --unified_finetuning True \
     --use_self_distill True \
     --fix_encoder False \
-    --self_distill_start_step 0
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+    --self_distill_start_step 0 \
+    --warmup_steps 0.1
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 ```
 
 Here are some new arguments:
@@ -293,8 +295,9 @@ torchrun --nproc_per_node 2 \
     --temperature 0.02 \
     --sentence_pooling_method last_token \
     --normalize_embeddings True \
-    --kd_loss_type m3_kd_loss
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+    --kd_loss_type m3_kd_loss \
+    --warmup_steps 0.1
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 ```
 
 Here are some new arguments:
@@ -356,8 +359,9 @@ torchrun --nproc_per_node 2 \
     --temperature 0.02 \
     --sentence_pooling_method last_token \
     --normalize_embeddings True \
-    --kd_loss_type kl_div
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+    --kd_loss_type kl_div \
+    --warmup_steps 0.1
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 ```
 
 Here are some new arguments:

examples/finetune/embedder/decoder_only/base.sh

@@ -56,9 +56,10 @@ training_args="\
     --sentence_pooling_method last_token \
     --normalize_embeddings True \
     --kd_loss_type m3_kd_loss \
+    --warmup_steps 0.1 \
 "
 
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 
 cmd="torchrun --nproc_per_node $num_gpus \
     -m FlagEmbedding.finetune.embedder.decoder_only.base \

examples/finetune/embedder/decoder_only/base_same_dataset.sh

@@ -59,9 +59,10 @@ training_args="\
     --sentence_pooling_method last_token \
     --normalize_embeddings True \
     --kd_loss_type m3_kd_loss \
+    --warmup_steps 0.1 \
 "
 
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 
 cmd="torchrun --nproc_per_node $num_gpus \
     -m FlagEmbedding.finetune.embedder.decoder_only.base \

examples/finetune/embedder/decoder_only/icl_same_dataset.sh

@@ -63,9 +63,10 @@ training_args="\
     --sentence_pooling_method last_token \
     --normalize_embeddings True \
     --kd_loss_type kl_div \
+    --warmup_steps 0.1 \
 "
 
-# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1.
+# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1.
 
 cmd="torchrun --nproc_per_node $num_gpus \
     -m FlagEmbedding.finetune.embedder.decoder_only.icl \