diff --git a/FlagEmbedding/abc/finetune/embedder/AbsRunner.py b/FlagEmbedding/abc/finetune/embedder/AbsRunner.py index 9466a3db..0ad49493 100644 --- a/FlagEmbedding/abc/finetune/embedder/AbsRunner.py +++ b/FlagEmbedding/abc/finetune/embedder/AbsRunner.py @@ -43,10 +43,11 @@ def __init__( os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir) and training_args.do_train - and not training_args.overwrite_output_dir + and training_args.resume_from_checkpoint is None ): raise ValueError( - f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome." + f"Output directory ({training_args.output_dir}) already exists and is not empty. " + "Please use a new output directory or set --resume_from_checkpoint." ) # Setup logging diff --git a/FlagEmbedding/abc/finetune/embedder/AbsTrainer.py b/FlagEmbedding/abc/finetune/embedder/AbsTrainer.py index 38944ed8..0f726c7c 100644 --- a/FlagEmbedding/abc/finetune/embedder/AbsTrainer.py +++ b/FlagEmbedding/abc/finetune/embedder/AbsTrainer.py @@ -1,4 +1,5 @@ import logging +import inspect from typing import Optional from abc import ABC, abstractmethod from transformers.trainer import Trainer @@ -10,6 +11,39 @@ class AbsEmbedderTrainer(ABC, Trainer): """ Abstract class for the trainer of embedder. """ + def __init__( + self, + *args, + processing_class=None, + tokenizer=None, + **kwargs, + ): + """Initialize ``Trainer`` across Transformers API versions. + + Transformers 5 renamed the ``tokenizer`` argument to + ``processing_class``. FlagEmbedding accepts both names so callers + from older examples remain compatible with either Transformers API. + """ + if processing_class is None: + processing_class = tokenizer + + trainer_parameters = inspect.signature(Trainer.__init__).parameters + if "processing_class" in trainer_parameters: + kwargs["processing_class"] = processing_class + else: + kwargs["tokenizer"] = processing_class + + super().__init__(*args, **kwargs) + + def _save_processing_class(self, output_dir: str): + """Save the tokenizer/processor using the active Transformers API.""" + processing_class = getattr(self, "processing_class", None) + if processing_class is None: + processing_class = getattr(self, "tokenizer", None) + + if processing_class is not None and self.is_world_process_zero(): + processing_class.save_pretrained(output_dir) + @abstractmethod def _save(self, output_dir: Optional[str] = None, state_dict=None): pass diff --git a/FlagEmbedding/abc/finetune/reranker/AbsRunner.py b/FlagEmbedding/abc/finetune/reranker/AbsRunner.py index b40c2d2a..875776c1 100644 --- a/FlagEmbedding/abc/finetune/reranker/AbsRunner.py +++ b/FlagEmbedding/abc/finetune/reranker/AbsRunner.py @@ -43,10 +43,11 @@ def __init__( os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir) and training_args.do_train - and not training_args.overwrite_output_dir + and training_args.resume_from_checkpoint is None ): raise ValueError( - f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome." + f"Output directory ({training_args.output_dir}) already exists and is not empty. " + "Please use a new output directory or set --resume_from_checkpoint." ) # Setup logging diff --git a/FlagEmbedding/abc/finetune/reranker/AbsTrainer.py b/FlagEmbedding/abc/finetune/reranker/AbsTrainer.py index 213d5f79..afd11b6e 100644 --- a/FlagEmbedding/abc/finetune/reranker/AbsTrainer.py +++ b/FlagEmbedding/abc/finetune/reranker/AbsTrainer.py @@ -1,4 +1,5 @@ import logging +import inspect from typing import Optional from abc import ABC, abstractmethod from transformers.trainer import Trainer @@ -10,6 +11,34 @@ class AbsRerankerTrainer(ABC, Trainer): """ Abstract class for the trainer of reranker. """ + def __init__( + self, + *args, + processing_class=None, + tokenizer=None, + **kwargs, + ): + """Initialize ``Trainer`` across Transformers API versions.""" + if processing_class is None: + processing_class = tokenizer + + trainer_parameters = inspect.signature(Trainer.__init__).parameters + if "processing_class" in trainer_parameters: + kwargs["processing_class"] = processing_class + else: + kwargs["tokenizer"] = processing_class + + super().__init__(*args, **kwargs) + + def _save_processing_class(self, output_dir: str): + """Save the tokenizer/processor using the active Transformers API.""" + processing_class = getattr(self, "processing_class", None) + if processing_class is None: + processing_class = getattr(self, "tokenizer", None) + + if processing_class is not None and self.is_world_process_zero(): + processing_class.save_pretrained(output_dir) + @abstractmethod def _save(self, output_dir: Optional[str] = None, state_dict=None): pass diff --git a/FlagEmbedding/finetune/embedder/decoder_only/base/trainer.py b/FlagEmbedding/finetune/embedder/decoder_only/base/trainer.py index bb834c66..36122a57 100644 --- a/FlagEmbedding/finetune/embedder/decoder_only/base/trainer.py +++ b/FlagEmbedding/finetune/embedder/decoder_only/base/trainer.py @@ -33,8 +33,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/embedder/decoder_only/icl/trainer.py b/FlagEmbedding/finetune/embedder/decoder_only/icl/trainer.py index c5d166f9..ee2224f3 100644 --- a/FlagEmbedding/finetune/embedder/decoder_only/icl/trainer.py +++ b/FlagEmbedding/finetune/embedder/decoder_only/icl/trainer.py @@ -33,8 +33,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/embedder/encoder_only/base/trainer.py b/FlagEmbedding/finetune/embedder/encoder_only/base/trainer.py index affa34a8..46d674ba 100644 --- a/FlagEmbedding/finetune/embedder/encoder_only/base/trainer.py +++ b/FlagEmbedding/finetune/embedder/encoder_only/base/trainer.py @@ -32,8 +32,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): f'does not support save interface') else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/embedder/encoder_only/m3/trainer.py b/FlagEmbedding/finetune/embedder/encoder_only/m3/trainer.py index a6bf5dd0..4aa179ec 100644 --- a/FlagEmbedding/finetune/embedder/encoder_only/m3/trainer.py +++ b/FlagEmbedding/finetune/embedder/encoder_only/m3/trainer.py @@ -32,8 +32,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): f'does not support save interface') else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/reranker/decoder_only/base/runner.py b/FlagEmbedding/finetune/reranker/decoder_only/base/runner.py index 14fa3dd9..e2e4c457 100644 --- a/FlagEmbedding/finetune/reranker/decoder_only/base/runner.py +++ b/FlagEmbedding/finetune/reranker/decoder_only/base/runner.py @@ -89,7 +89,7 @@ def load_trainer(self) -> DecoderOnlyRerankerTrainer: args=self.training_args, train_dataset=self.train_dataset, data_collator=self.data_collator, - tokenizer=self.tokenizer + processing_class=self.tokenizer ) return trainer diff --git a/FlagEmbedding/finetune/reranker/decoder_only/base/trainer.py b/FlagEmbedding/finetune/reranker/decoder_only/base/trainer.py index d74ad0b5..2331df85 100644 --- a/FlagEmbedding/finetune/reranker/decoder_only/base/trainer.py +++ b/FlagEmbedding/finetune/reranker/decoder_only/base/trainer.py @@ -35,8 +35,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/reranker/decoder_only/layerwise/runner.py b/FlagEmbedding/finetune/reranker/decoder_only/layerwise/runner.py index b2502004..79786c82 100644 --- a/FlagEmbedding/finetune/reranker/decoder_only/layerwise/runner.py +++ b/FlagEmbedding/finetune/reranker/decoder_only/layerwise/runner.py @@ -90,7 +90,7 @@ def load_trainer(self) -> DecoderOnlyRerankerTrainer: args=self.training_args, train_dataset=self.train_dataset, data_collator=self.data_collator, - tokenizer=self.tokenizer + processing_class=self.tokenizer ) return trainer diff --git a/FlagEmbedding/finetune/reranker/decoder_only/layerwise/trainer.py b/FlagEmbedding/finetune/reranker/decoder_only/layerwise/trainer.py index dc8c7a18..2f71682d 100644 --- a/FlagEmbedding/finetune/reranker/decoder_only/layerwise/trainer.py +++ b/FlagEmbedding/finetune/reranker/decoder_only/layerwise/trainer.py @@ -35,8 +35,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): else: self.model.save(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/FlagEmbedding/finetune/reranker/encoder_only/base/runner.py b/FlagEmbedding/finetune/reranker/encoder_only/base/runner.py index 324c9e30..7842ee0b 100644 --- a/FlagEmbedding/finetune/reranker/encoder_only/base/runner.py +++ b/FlagEmbedding/finetune/reranker/encoder_only/base/runner.py @@ -71,6 +71,6 @@ def load_trainer(self) -> EncoderOnlyRerankerTrainer: args=self.training_args, train_dataset=self.train_dataset, data_collator=self.data_collator, - tokenizer=self.tokenizer + processing_class=self.tokenizer ) return trainer diff --git a/FlagEmbedding/finetune/reranker/encoder_only/base/trainer.py b/FlagEmbedding/finetune/reranker/encoder_only/base/trainer.py index 59120341..36ac40ce 100644 --- a/FlagEmbedding/finetune/reranker/encoder_only/base/trainer.py +++ b/FlagEmbedding/finetune/reranker/encoder_only/base/trainer.py @@ -30,8 +30,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): raise NotImplementedError(f'MODEL {self.model.__class__.__name__} ' f'does not support save_pretrained interface') else: self.model.save_pretrained(output_dir) - if self.tokenizer is not None and self.is_world_process_zero(): - self.tokenizer.save_pretrained(output_dir) + self._save_processing_class(output_dir) # Good practice: save your training arguments together with the trained model torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/examples/README.md b/examples/README.md index 3faea48d..abb6d8c8 100644 --- a/examples/README.md +++ b/examples/README.md @@ -104,13 +104,11 @@ torchrun --nproc_per_node 2 \ --query_instruction_format '{}{}' \ --knowledge_distillation False \ --output_dir ./test_encoder_only_base_bge-large-en-v1.5 \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ./finetune/ds_stage0.json \ --logging_steps 1 \ @@ -120,6 +118,7 @@ torchrun --nproc_per_node 2 \ --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. ``` ### 2. Reranker @@ -137,19 +136,18 @@ torchrun --nproc_per_node 2 \ --pad_to_multiple_of 8 \ --knowledge_distillation False \ --output_dir ./test_encoder_only_base_bge-reranker-large \ - --overwrite_output_dir \ --learning_rate 6e-5 \ --fp16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 1 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ./finetune/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. ``` ## 5. Evaluation @@ -187,4 +185,3 @@ python -m FlagEmbedding.evaluation.msmarco \ --devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \ --cache_dir ./cache/model ``` - diff --git a/examples/finetune/embedder/README.md b/examples/finetune/embedder/README.md index c9e141aa..b1197ee5 100644 --- a/examples/finetune/embedder/README.md +++ b/examples/finetune/embedder/README.md @@ -184,13 +184,11 @@ torchrun --nproc_per_node 2 \ --query_instruction_format '{}{}' \ --knowledge_distillation False \ --output_dir ./test_encoder_only_base_bge-large-en-v1.5 \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs 2 \ --per_device_train_batch_size 2 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../ds_stage0.json \ --logging_steps 1 \ @@ -200,6 +198,7 @@ torchrun --nproc_per_node 2 \ --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. ``` ### (2) bge-m3 @@ -223,13 +222,11 @@ torchrun --nproc_per_node 2 \ --small_threshold 0 \ --drop_threshold 0 \ --output_dir ./test_encoder_only_m3_bge-m3_sd \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs 2 \ --per_device_train_batch_size 2 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../ds_stage0.json \ --logging_steps 1 \ @@ -243,6 +240,7 @@ torchrun --nproc_per_node 2 \ --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. ``` Here are some new arguments: @@ -282,13 +280,11 @@ torchrun --nproc_per_node 2 \ --small_threshold 0 \ --drop_threshold 0 \ --output_dir ./test_decoder_only_base_bge-multilingual-gemma2_sd \ - --overwrite_output_dir \ --learning_rate 1e-4 \ --fp16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../ds_stage1.json \ --logging_steps 1 \ @@ -298,6 +294,7 @@ torchrun --nproc_per_node 2 \ --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. ``` Here are some new arguments: @@ -346,13 +343,11 @@ torchrun --nproc_per_node 2 \ --retrieval_use_examples True \ --icl_suffix_str '\n' \ --output_dir ./test_decoder_only_base_bge-en-icl_sd \ - --overwrite_output_dir \ --learning_rate 1e-4 \ --fp16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../ds_stage1.json \ --logging_steps 1 \ @@ -362,6 +357,7 @@ torchrun --nproc_per_node 2 \ --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. ``` Here are some new arguments: @@ -383,4 +379,3 @@ Here are some new arguments: - **`example_passage_max_len`**: The max length of example passage. - **`retrieval_use_examples`**: If passed, will use examples for retrieval. - **`icl_suffix_str`**: The suffix string for ICL dataset. - diff --git a/examples/finetune/embedder/decoder_only/base.sh b/examples/finetune/embedder/decoder_only/base.sh index 2b91f683..8268a222 100644 --- a/examples/finetune/embedder/decoder_only/base.sh +++ b/examples/finetune/embedder/decoder_only/base.sh @@ -42,13 +42,11 @@ data_args="\ training_args="\ --output_dir ./test_decoder_only_base_bge-multilingual-gemma2 \ - --overwrite_output_dir \ --learning_rate 1e-4 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage1.json \ --logging_steps 1 \ @@ -60,6 +58,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.decoder_only.base \ $model_args \ diff --git a/examples/finetune/embedder/decoder_only/base_same_dataset.sh b/examples/finetune/embedder/decoder_only/base_same_dataset.sh index d002264b..197e5288 100644 --- a/examples/finetune/embedder/decoder_only/base_same_dataset.sh +++ b/examples/finetune/embedder/decoder_only/base_same_dataset.sh @@ -45,13 +45,11 @@ data_args="\ training_args="\ --output_dir ./test_decoder_only_base_bge-multilingual-gemma2_sd \ - --overwrite_output_dir \ --learning_rate 1e-4 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage1.json \ --logging_steps 1 \ @@ -63,6 +61,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.decoder_only.base \ $model_args \ diff --git a/examples/finetune/embedder/decoder_only/icl_same_dataset.sh b/examples/finetune/embedder/decoder_only/icl_same_dataset.sh index 025ece9e..2e080f04 100644 --- a/examples/finetune/embedder/decoder_only/icl_same_dataset.sh +++ b/examples/finetune/embedder/decoder_only/icl_same_dataset.sh @@ -49,13 +49,11 @@ data_args="\ training_args="\ --output_dir ./test_decoder_only_base_bge-en-icl_sd \ - --overwrite_output_dir \ --learning_rate 1e-4 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage1.json \ --logging_steps 1 \ @@ -67,6 +65,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.decoder_only.icl \ $model_args \ diff --git a/examples/finetune/embedder/encoder_only/base.sh b/examples/finetune/embedder/encoder_only/base.sh index aae9e682..c345fa8d 100644 --- a/examples/finetune/embedder/encoder_only/base.sh +++ b/examples/finetune/embedder/encoder_only/base.sh @@ -36,13 +36,11 @@ data_args="\ training_args="\ --output_dir ./test_encoder_only_base_bge-large-en-v1.5 \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ @@ -54,6 +52,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.base \ $model_args \ diff --git a/examples/finetune/embedder/encoder_only/base_same_dataset.sh b/examples/finetune/embedder/encoder_only/base_same_dataset.sh index fe9a6b31..944e463e 100644 --- a/examples/finetune/embedder/encoder_only/base_same_dataset.sh +++ b/examples/finetune/embedder/encoder_only/base_same_dataset.sh @@ -39,13 +39,11 @@ data_args="\ training_args="\ --output_dir ./test_encoder_only_base_bge-large-en-v1.5_sd \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ @@ -57,6 +55,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.base \ $model_args \ diff --git a/examples/finetune/embedder/encoder_only/m3.sh b/examples/finetune/embedder/encoder_only/m3.sh index 49642a40..201cebb8 100644 --- a/examples/finetune/embedder/encoder_only/m3.sh +++ b/examples/finetune/embedder/encoder_only/m3.sh @@ -34,13 +34,11 @@ data_args="\ training_args="\ --output_dir ./test_encoder_only_m3_bge-m3 \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ @@ -56,6 +54,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.m3 \ $model_args \ diff --git a/examples/finetune/embedder/encoder_only/m3_same_dataset.sh b/examples/finetune/embedder/encoder_only/m3_same_dataset.sh index e0b7c8d8..3b7d1296 100644 --- a/examples/finetune/embedder/encoder_only/m3_same_dataset.sh +++ b/examples/finetune/embedder/encoder_only/m3_same_dataset.sh @@ -37,13 +37,11 @@ data_args="\ training_args="\ --output_dir ./test_encoder_only_m3_bge-m3_sd \ - --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ @@ -59,6 +57,8 @@ training_args="\ --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. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.m3 \ $model_args \ diff --git a/examples/finetune/reranker/README.md b/examples/finetune/reranker/README.md index 1e53ddf2..4fde383d 100644 --- a/examples/finetune/reranker/README.md +++ b/examples/finetune/reranker/README.md @@ -158,19 +158,18 @@ torchrun --nproc_per_node 2 \ --pad_to_multiple_of 8 \ --knowledge_distillation False \ --output_dir ./test_encoder_only_base_bge-reranker-base \ - --overwrite_output_dir \ --learning_rate 6e-5 \ --fp16 \ --num_train_epochs 2 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 1 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --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. ``` ### (2) bge-reranker-v2-gemma @@ -199,19 +198,18 @@ torchrun --nproc_per_node 2 \ --passage_instruction_for_rerank 'B: ' \ --passage_instruction_format '{}{}' \ --output_dir ./test_decoder_only_base_bge-reranker-v2-minicpm-layerwise \ - --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 1 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --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. ``` Here are some new arguments: @@ -258,19 +256,18 @@ torchrun --nproc_per_node 2 \ --passage_instruction_for_rerank 'B: ' \ --passage_instruction_format '{}{}' \ --output_dir ./test_decoder_only_base_bge-reranker-v2-minicpm-layerwise \ - --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 1 \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --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. ``` Here are some new arguments: diff --git a/examples/finetune/reranker/decoder_only/base.sh b/examples/finetune/reranker/decoder_only/base.sh index 1f862a88..2230edfa 100644 --- a/examples/finetune/reranker/decoder_only/base.sh +++ b/examples/finetune/reranker/decoder_only/base.sh @@ -44,14 +44,12 @@ data_args="\ training_args="\ --output_dir ./test_decoder_only_base_bge-reranker-v2-gemma \ - --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --gradient_accumulation_steps $gradient_accumulation_steps \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ../../ds_stage0.json \ @@ -59,6 +57,8 @@ training_args="\ --save_steps 1000 \ " +# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.reranker.decoder_only.base \ $model_args \ @@ -67,4 +67,4 @@ cmd="torchrun --nproc_per_node $num_gpus \ " echo $cmd -eval $cmd \ No newline at end of file +eval $cmd diff --git a/examples/finetune/reranker/decoder_only/layerwise.sh b/examples/finetune/reranker/decoder_only/layerwise.sh index f5d44c8b..8d2b6848 100644 --- a/examples/finetune/reranker/decoder_only/layerwise.sh +++ b/examples/finetune/reranker/decoder_only/layerwise.sh @@ -49,14 +49,12 @@ data_args="\ training_args="\ --output_dir ./test_decoder_only_base_bge-reranker-v2-minicpm-layerwise \ - --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --gradient_accumulation_steps $gradient_accumulation_steps \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ../../ds_stage0.json \ @@ -64,6 +62,8 @@ training_args="\ --save_steps 1000 \ " +# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.reranker.decoder_only.layerwise \ $model_args \ @@ -73,4 +73,3 @@ cmd="torchrun --nproc_per_node $num_gpus \ echo $cmd eval $cmd - diff --git a/examples/finetune/reranker/encoder_only/base.sh b/examples/finetune/reranker/encoder_only/base.sh index c8ca09fc..d0cbea6e 100644 --- a/examples/finetune/reranker/encoder_only/base.sh +++ b/examples/finetune/reranker/encoder_only/base.sh @@ -33,14 +33,12 @@ data_args="\ training_args="\ --output_dir ./test_encoder_only_base_bge-reranker-base \ - --overwrite_output_dir \ --learning_rate 6e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --gradient_accumulation_steps $gradient_accumulation_steps \ --dataloader_drop_last True \ - --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ../../ds_stage0.json \ @@ -48,6 +46,8 @@ training_args="\ --save_steps 1000 \ " +# For transformers<=4.57.3, add --warmup_ratio 0.1; for transformers>=5.0.0, add --warmup_steps 0.1. + cmd="torchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.reranker.encoder_only.base \ $model_args \ @@ -56,4 +56,4 @@ cmd="torchrun --nproc_per_node $num_gpus \ " echo $cmd -eval $cmd \ No newline at end of file +eval $cmd diff --git a/setup.py b/setup.py index 9e4feacb..cf067d92 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='FlagEmbedding', - version='1.4.0', + version='1.4.1', description='FlagEmbedding', long_description=readme, long_description_content_type="text/markdown", diff --git a/tests/README.md b/tests/README.md index 36e447c7..79939ae9 100644 --- a/tests/README.md +++ b/tests/README.md @@ -1,10 +1,11 @@ # FlagEmbedding Tests -This directory contains tests for the FlagEmbedding library, including compatibility tests for Transformers 5.0. +This directory contains tests for the FlagEmbedding library, including compatibility tests for Transformers 5.x. ## Test Files - `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_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. @@ -24,13 +25,16 @@ pytest tests/ # Run a specific test file pytest tests/test_imports_v5.py +# Run the fine-tuning Trainer compatibility tests +pytest tests/test_finetune_trainer_compat.py + # Run with verbose output pytest -v tests/ ``` -## Transformers 5.0 Compatibility +## Transformers 5.x Compatibility -The tests verify that FlagEmbedding works with Transformers 5.0, which removed the `is_torch_fx_available` function. +The tests verify that FlagEmbedding works with Transformers 5.x, which removed the `is_torch_fx_available` function and renamed the Trainer `tokenizer` argument to `processing_class`. The compatibility layer in `FlagEmbedding/utils/transformers_compat.py` provides this function for backward compatibility. **Note:** Transformers 5.0 requires Python 3.10 or higher. If you're using Python 3.9 or lower, you'll need to upgrade your Python version to test with Transformers 5.0. @@ -39,4 +43,4 @@ To test with a specific version of transformers (with Python 3.10+): ```bash pip install transformers==5.0.0 -pytest tests/ \ No newline at end of file +pytest tests/ diff --git a/tests/test_finetune_trainer_compat.py b/tests/test_finetune_trainer_compat.py new file mode 100644 index 00000000..b564590b --- /dev/null +++ b/tests/test_finetune_trainer_compat.py @@ -0,0 +1,153 @@ +"""Regression tests for the Transformers Trainer API migration.""" + +from pathlib import Path + +import torch +from transformers import Trainer, TrainingArguments + +from FlagEmbedding.finetune.embedder.decoder_only.base.trainer import ( + DecoderOnlyEmbedderTrainer, +) +from FlagEmbedding.finetune.embedder.decoder_only.icl.trainer import ( + DecoderOnlyEmbedderICLTrainer, +) +from FlagEmbedding.finetune.embedder.encoder_only.base.trainer import ( + EncoderOnlyEmbedderTrainer, +) +from FlagEmbedding.finetune.embedder.encoder_only.m3.trainer import ( + EncoderOnlyEmbedderM3Trainer, +) +from FlagEmbedding.finetune.reranker.decoder_only.base.trainer import ( + DecoderOnlyRerankerTrainer, +) +from FlagEmbedding.finetune.reranker.decoder_only.base.runner import ( + DecoderOnlyRerankerRunner, +) +from FlagEmbedding.finetune.reranker.decoder_only.layerwise.trainer import ( + DecoderOnlyRerankerTrainer as DecoderOnlyRerankerLayerwiseTrainer, +) +from FlagEmbedding.finetune.reranker.decoder_only.layerwise.runner import ( + DecoderOnlyRerankerRunner as DecoderOnlyRerankerLayerwiseRunner, +) +from FlagEmbedding.finetune.reranker.encoder_only.base.trainer import ( + EncoderOnlyRerankerTrainer, +) +from FlagEmbedding.finetune.reranker.encoder_only.base.runner import ( + EncoderOnlyRerankerRunner, +) + + +class RecordingProcessor: + def __init__(self, marker: str): + self.marker = marker + + def save_pretrained(self, output_dir): + Path(output_dir, "processor.marker").write_text(self.marker, encoding="utf-8") + + +class SaveableModel(torch.nn.Module): + def save(self, output_dir): + Path(output_dir, "model.marker").write_text("save", encoding="utf-8") + + def save_pretrained(self, output_dir): + Path(output_dir, "model.marker").write_text("save_pretrained", encoding="utf-8") + + +def training_args(tmp_path): + return TrainingArguments( + output_dir=str(tmp_path / "output"), + report_to=[], + ) + + +def test_processing_class_and_legacy_tokenizer_are_accepted(tmp_path): + processor = RecordingProcessor("processing") + trainer = EncoderOnlyEmbedderTrainer( + model=SaveableModel(), + args=training_args(tmp_path), + processing_class=processor, + ) + assert trainer.processing_class is processor + + legacy_processor = RecordingProcessor("legacy") + legacy_trainer = EncoderOnlyEmbedderTrainer( + model=SaveableModel(), + args=training_args(tmp_path / "legacy"), + tokenizer=legacy_processor, + ) + assert legacy_trainer.processing_class is legacy_processor + + +def test_all_flagembedding_finetune_trainers_save_the_processing_class(tmp_path): + trainer_types = [ + DecoderOnlyEmbedderTrainer, + DecoderOnlyEmbedderICLTrainer, + EncoderOnlyEmbedderTrainer, + EncoderOnlyEmbedderM3Trainer, + DecoderOnlyRerankerTrainer, + DecoderOnlyRerankerLayerwiseTrainer, + EncoderOnlyRerankerTrainer, + ] + + for trainer_type in trainer_types: + output_dir = tmp_path / trainer_type.__name__ + processor = RecordingProcessor(trainer_type.__name__) + trainer = trainer_type( + model=SaveableModel(), + args=training_args(output_dir), + processing_class=processor, + ) + + trainer._save(str(output_dir)) + + assert Path(output_dir, "processor.marker").read_text(encoding="utf-8") == trainer_type.__name__ + assert Path(output_dir, "training_args.bin").exists() + + +def test_reranker_runners_pass_processing_class(tmp_path): + runner_types = [ + EncoderOnlyRerankerRunner, + DecoderOnlyRerankerRunner, + DecoderOnlyRerankerLayerwiseRunner, + ] + + for runner_type in runner_types: + runner = runner_type.__new__(runner_type) + processor = RecordingProcessor(runner_type.__name__) + runner.model = SaveableModel() + runner.training_args = training_args(tmp_path / runner_type.__name__) + runner.train_dataset = None + runner.data_collator = None + runner.tokenizer = processor + + trainer = runner.load_trainer() + + assert trainer.processing_class is processor + + +def test_legacy_transformers_trainer_constructor_path(monkeypatch, tmp_path): + """The compatibility shim uses tokenizer= with pre-v5 Trainer APIs.""" + captured = {} + + def legacy_init(self, *args, tokenizer=None, **kwargs): + captured["tokenizer"] = tokenizer + self.tokenizer = tokenizer + self.model = kwargs.get("model") + self.args = kwargs.get("args") + + monkeypatch.setattr(Trainer, "__init__", legacy_init) + + from FlagEmbedding.abc.finetune.embedder.AbsTrainer import AbsEmbedderTrainer + + class LegacyTrainer(AbsEmbedderTrainer): + def _save(self, output_dir=None, state_dict=None): + pass + + processor = RecordingProcessor("legacy-transformers") + LegacyTrainer( + model=SaveableModel(), + args=training_args(tmp_path), + processing_class=processor, + ) + + assert captured["tokenizer"] is processor