diff --git a/FlagEmbedding/abc/finetune/reranker/AbsDataset.py b/FlagEmbedding/abc/finetune/reranker/AbsDataset.py index 0150c7e98..812ce6eda 100644 --- a/FlagEmbedding/abc/finetune/reranker/AbsDataset.py +++ b/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, diff --git a/FlagEmbedding/inference/embedder/decoder_only/base.py b/FlagEmbedding/inference/embedder/decoder_only/base.py index 7117c44d6..84375a912 100644 --- a/FlagEmbedding/inference/embedder/decoder_only/base.py +++ b/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', diff --git a/FlagEmbedding/inference/embedder/decoder_only/icl.py b/FlagEmbedding/inference/embedder/decoder_only/icl.py index bc962e89b..d4d6e087a 100644 --- a/FlagEmbedding/inference/embedder/decoder_only/icl.py +++ b/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', diff --git a/FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py b/FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py index 5fe0b8e90..83bf070c2 100644 --- a/FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py +++ b/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', diff --git a/FlagEmbedding/inference/embedder/encoder_only/base.py b/FlagEmbedding/inference/embedder/encoder_only/base.py index d01eb885a..e3f8ae11b 100644 --- a/FlagEmbedding/inference/embedder/encoder_only/base.py +++ b/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', diff --git a/FlagEmbedding/inference/embedder/encoder_only/m3.py b/FlagEmbedding/inference/embedder/encoder_only/m3.py index b98efcab3..dc533f554 100644 --- a/FlagEmbedding/inference/embedder/encoder_only/m3.py +++ b/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', diff --git a/FlagEmbedding/inference/reranker/decoder_only/base.py b/FlagEmbedding/inference/reranker/decoder_only/base.py index b4aa8f50d..574b48bf3 100644 --- a/FlagEmbedding/inference/reranker/decoder_only/base.py +++ b/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, diff --git a/FlagEmbedding/inference/reranker/decoder_only/lightweight.py b/FlagEmbedding/inference/reranker/decoder_only/lightweight.py index 44422c523..6bbac3269 100644 --- a/FlagEmbedding/inference/reranker/decoder_only/lightweight.py +++ b/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, diff --git a/FlagEmbedding/inference/reranker/encoder_only/base.py b/FlagEmbedding/inference/reranker/encoder_only/base.py index 15686dcca..69b9c3505 100644 --- a/FlagEmbedding/inference/reranker/encoder_only/base.py +++ b/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', diff --git a/FlagEmbedding/utils/tokenizer_compat.py b/FlagEmbedding/utils/tokenizer_compat.py index 955c43083..cdd6b61fe 100644 --- a/FlagEmbedding/utils/tokenizer_compat.py +++ b/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( diff --git a/examples/finetune/ds_stage0.json b/examples/finetune/ds_stage0.json index 495d5321c..9d03271ee 100644 --- a/examples/finetune/ds_stage0.json +++ b/examples/finetune/ds_stage0.json @@ -22,7 +22,8 @@ "lr": "auto", "betas": "auto", "eps": "auto", - "weight_decay": "auto" + "weight_decay": "auto", + "torch_adam": true } }, diff --git a/examples/finetune/ds_stage1.json b/examples/finetune/ds_stage1.json index 580d29772..4bb81a3e0 100644 --- a/examples/finetune/ds_stage1.json +++ b/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", diff --git a/examples/finetune/embedder/README.md b/examples/finetune/embedder/README.md index b1197ee5e..f3b508928 100644 --- a/examples/finetune/embedder/README.md +++ b/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: diff --git a/examples/finetune/embedder/decoder_only/base.sh b/examples/finetune/embedder/decoder_only/base.sh index 8268a2228..73b6bb44f 100644 --- a/examples/finetune/embedder/decoder_only/base.sh +++ b/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 \ diff --git a/examples/finetune/embedder/decoder_only/base_same_dataset.sh b/examples/finetune/embedder/decoder_only/base_same_dataset.sh index 197e52888..5ad9a5396 100644 --- a/examples/finetune/embedder/decoder_only/base_same_dataset.sh +++ b/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 \ diff --git a/examples/finetune/embedder/decoder_only/icl_same_dataset.sh b/examples/finetune/embedder/decoder_only/icl_same_dataset.sh index 2e080f04c..9b7843a0e 100644 --- a/examples/finetune/embedder/decoder_only/icl_same_dataset.sh +++ b/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 \ diff --git a/examples/finetune/embedder/encoder_only/base.sh b/examples/finetune/embedder/encoder_only/base.sh index c345fa8df..b8fadf6ea 100644 --- a/examples/finetune/embedder/encoder_only/base.sh +++ b/examples/finetune/embedder/encoder_only/base.sh @@ -50,9 +50,10 @@ training_args="\ --sentence_pooling_method cls \ --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.encoder_only.base \ diff --git a/examples/finetune/embedder/encoder_only/base_same_dataset.sh b/examples/finetune/embedder/encoder_only/base_same_dataset.sh index 944e463e5..3aa4e3cba 100644 --- a/examples/finetune/embedder/encoder_only/base_same_dataset.sh +++ b/examples/finetune/embedder/encoder_only/base_same_dataset.sh @@ -53,9 +53,10 @@ training_args="\ --sentence_pooling_method cls \ --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.encoder_only.base \ diff --git a/examples/finetune/embedder/encoder_only/m3.sh b/examples/finetune/embedder/encoder_only/m3.sh index 201cebb89..0222e1af8 100644 --- a/examples/finetune/embedder/encoder_only/m3.sh +++ b/examples/finetune/embedder/encoder_only/m3.sh @@ -52,9 +52,10 @@ training_args="\ --use_self_distill True \ --fix_encoder False \ --self_distill_start_step 0 \ + --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.encoder_only.m3 \ diff --git a/examples/finetune/embedder/encoder_only/m3_same_dataset.sh b/examples/finetune/embedder/encoder_only/m3_same_dataset.sh index 3b7d1296a..1a3fa8bbb 100644 --- a/examples/finetune/embedder/encoder_only/m3_same_dataset.sh +++ b/examples/finetune/embedder/encoder_only/m3_same_dataset.sh @@ -55,9 +55,10 @@ training_args="\ --use_self_distill True \ --fix_encoder False \ --self_distill_start_step 0 \ + --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.encoder_only.m3 \ diff --git a/examples/finetune/reranker/README.md b/examples/finetune/reranker/README.md index 4fde383d5..209bdba62 100644 --- a/examples/finetune/reranker/README.md +++ b/examples/finetune/reranker/README.md @@ -168,8 +168,9 @@ torchrun --nproc_per_node 2 \ --weight_decay 0.01 \ --deepspeed ../ds_stage0.json \ --logging_steps 1 \ - --save_steps 1000 -# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + --save_steps 1000 \ + --warmup_steps 0.1 +# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1. ``` ### (2) bge-reranker-v2-gemma @@ -208,8 +209,9 @@ torchrun --nproc_per_node 2 \ --weight_decay 0.01 \ --deepspeed ../ds_stage0.json \ --logging_steps 1 \ - --save_steps 1000 -# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + --save_steps 1000 \ + --warmup_steps 0.1 +# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1. ``` Here are some new arguments: @@ -266,8 +268,9 @@ torchrun --nproc_per_node 2 \ --weight_decay 0.01 \ --deepspeed ../ds_stage0.json \ --logging_steps 1 \ - --save_steps 1000 -# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + --save_steps 1000 \ + --warmup_steps 0.1 +# For transformers <= 4.57.3, replace --warmup_steps 0.1 with --warmup_ratio 0.1. ``` Here are some new arguments: diff --git a/examples/finetune/reranker/decoder_only/base.sh b/examples/finetune/reranker/decoder_only/base.sh index 2230edfa3..59b363b55 100644 --- a/examples/finetune/reranker/decoder_only/base.sh +++ b/examples/finetune/reranker/decoder_only/base.sh @@ -55,9 +55,10 @@ training_args="\ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ --save_steps 1000 \ + --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.reranker.decoder_only.base \ diff --git a/examples/finetune/reranker/decoder_only/layerwise.sh b/examples/finetune/reranker/decoder_only/layerwise.sh index 8d2b6848f..4762fb0d7 100644 --- a/examples/finetune/reranker/decoder_only/layerwise.sh +++ b/examples/finetune/reranker/decoder_only/layerwise.sh @@ -60,9 +60,10 @@ training_args="\ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ --save_steps 1000 \ + --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.reranker.decoder_only.layerwise \ diff --git a/examples/finetune/reranker/encoder_only/base.sh b/examples/finetune/reranker/encoder_only/base.sh index d0cbea6ee..4229967ce 100644 --- a/examples/finetune/reranker/encoder_only/base.sh +++ b/examples/finetune/reranker/encoder_only/base.sh @@ -44,9 +44,10 @@ training_args="\ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ --save_steps 1000 \ + --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.reranker.encoder_only.base \ diff --git a/tests/README.md b/tests/README.md index fcf55e37f..9ab0f3638 100644 --- a/tests/README.md +++ b/tests/README.md @@ -7,6 +7,7 @@ This directory contains tests for the FlagEmbedding library, including compatibi - `test_imports_v5.py`: Tests that imports work with Transformers v5, particularly the compatibility layer for `is_torch_fx_available`. - `test_finetune_trainer_compat.py`: Tests the Transformers Trainer API migration, including `processing_class`/legacy `tokenizer` construction, reranker runner arguments, and processor checkpoint saving. - `test_reranker_tokenizer_compat.py`: Tests reranker pair preparation across Transformers v4/v5, including the v5 removal of `prepare_for_model`. +- `test_tokenizer_pad_compat.py`: Tests normalization of sorted per-example tokenized records before `tokenizer.pad`, including the local BGE reranker tokenizer when available. - `test_infer_embedder_basic.py`: Tests basic functionality of BGE embedder models with a small public checkpoint. - `test_infer_reranker_basic.py`: Tests basic functionality of reranker models. @@ -32,6 +33,9 @@ pytest tests/test_finetune_trainer_compat.py # Run the reranker tokenizer compatibility tests pytest tests/test_reranker_tokenizer_compat.py +# Run the tokenizer padding compatibility tests +pytest tests/test_tokenizer_pad_compat.py + # Run with verbose output pytest -v tests/ ``` diff --git a/tests/test_reranker_tokenizer_compat.py b/tests/test_reranker_tokenizer_compat.py index 41f96819c..75decf1cf 100644 --- a/tests/test_reranker_tokenizer_compat.py +++ b/tests/test_reranker_tokenizer_compat.py @@ -1,16 +1,15 @@ import os from types import SimpleNamespace -import pytest from transformers import AutoTokenizer from FlagEmbedding.abc.finetune.reranker.AbsDataset import AbsRerankerTrainDataset from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat -RERANKER_MODEL_PATH = os.environ.get( +RERANKER_MODEL_NAME = os.environ.get( "FLAGEMBEDDING_RERANKER_MODEL", - "/share/project/Search-SWE/model/bge-reranker-base", + "BAAI/bge-reranker-base", ) @@ -26,14 +25,9 @@ def prepare_for_model(self, first_ids, second_ids, **kwargs): return {"input_ids": list(first_ids) + list(second_ids)} -@pytest.mark.skipif( - not os.path.isdir(RERANKER_MODEL_PATH), - reason="local reranker model is not available", -) def test_encoder_pair_preserves_v5_tokenizer_input_ids(): tokenizer = AutoTokenizer.from_pretrained( - RERANKER_MODEL_PATH, - local_files_only=True, + RERANKER_MODEL_NAME, ) query = "๐Ÿ˜€๐Ÿงช" @@ -62,14 +56,9 @@ def test_encoder_pair_preserves_v5_tokenizer_input_ids(): assert actual["attention_mask"] == expected["attention_mask"] -@pytest.mark.skipif( - not os.path.isdir(RERANKER_MODEL_PATH), - reason="local reranker model is not available", -) def test_encoder_training_example_works_without_prepare_for_model(): tokenizer = AutoTokenizer.from_pretrained( - RERANKER_MODEL_PATH, - local_files_only=True, + RERANKER_MODEL_NAME, ) dataset = AbsRerankerTrainDataset.__new__(AbsRerankerTrainDataset) dataset.tokenizer = tokenizer diff --git a/tests/test_tokenizer_pad_compat.py b/tests/test_tokenizer_pad_compat.py new file mode 100644 index 000000000..aaf948a56 --- /dev/null +++ b/tests/test_tokenizer_pad_compat.py @@ -0,0 +1,115 @@ +import os + +import torch +from transformers import AutoTokenizer + +from FlagEmbedding.utils.tokenizer_compat import pad_with_compat + + +RERANKER_MODEL_NAME = os.environ.get( + "FLAGEMBEDDING_RERANKER_MODEL", + "BAAI/bge-reranker-base", +) +ENCODER_MODEL_NAME = os.environ.get( + "FLAGEMBEDDING_ENCODER_MODEL", + "BAAI/bge-base-en-v1.5", +) + + +class MappingOnlyTokenizer: + """Minimal tokenizer double that rejects list-of-mapping inputs.""" + + def __init__(self): + self.received = None + self.kwargs = None + + def pad(self, encoded_inputs, **kwargs): + if not isinstance(encoded_inputs, dict): + raise AttributeError("'list' object has no attribute 'keys'") + self.received = encoded_inputs + self.kwargs = kwargs + return encoded_inputs + + +def test_pad_with_compat_normalizes_list_of_mappings(): + tokenizer = MappingOnlyTokenizer() + records = [ + {"input_ids": [101, 10, 102], "attention_mask": [1, 1, 1]}, + {"input_ids": [101, 20, 30, 102], "attention_mask": [1, 1, 1, 1]}, + ] + + result = pad_with_compat(tokenizer, records, padding=True, return_tensors="pt") + + assert result == { + "input_ids": [[101, 10, 102], [101, 20, 30, 102]], + "attention_mask": [[1, 1, 1], [1, 1, 1, 1]], + } + assert tokenizer.received == result + assert tokenizer.kwargs == {"padding": True, "return_tensors": "pt"} + + +def test_pad_with_compat_passes_mapping_inputs_through(): + tokenizer = MappingOnlyTokenizer() + encoded_inputs = { + "input_ids": [[101, 10, 102]], + "attention_mask": [[1, 1, 1]], + } + + result = pad_with_compat(tokenizer, encoded_inputs, padding=True) + + assert result is encoded_inputs + assert tokenizer.received is encoded_inputs + + +def test_local_bge_reranker_tokenizer_padding_compatibility(): + tokenizer = AutoTokenizer.from_pretrained( + RERANKER_MODEL_NAME, + ) + encoded = tokenizer( + ["What is AI?", "Explain artificial intelligence in one sentence."], + truncation=True, + padding=False, + return_tensors=None, + ) + records = [ + {key: encoded[key][index] for key in encoded.keys()} + for index in range(2) + ] + + padded = pad_with_compat( + tokenizer, + records, + padding=True, + return_tensors="pt", + ) + + assert padded["input_ids"].shape[0] == 2 + assert padded["attention_mask"].shape == padded["input_ids"].shape + assert torch.all(padded["attention_mask"][:, 0] == 1) + + +def test_local_bge_encoder_token_type_ids_are_preserved(): + tokenizer = AutoTokenizer.from_pretrained( + ENCODER_MODEL_NAME, + ) + encoded = tokenizer( + ["short text", "a longer text for padding"], + truncation=True, + padding=False, + return_tensors=None, + ) + records = [ + {key: encoded[key][index] for key in encoded.keys()} + for index in range(2) + ] + + padded = pad_with_compat( + tokenizer, + records, + padding=True, + return_tensors="pt", + ) + + assert "token_type_ids" in padded + assert padded["token_type_ids"].shape == padded["input_ids"].shape + assert padded["attention_mask"].shape == padded["input_ids"].shape