update embedder finetune code

FlagEmbedding/finetune/embedder/decoder_only/base/arguments.py

@@ -69,3 +69,8 @@ class DecoderOnlyEmbedderModelArguments(AbsEmbedderModelArguments):
         default=False,
         metadata={"help": "If passed, will merge the lora modules and save the entire model."}
     )
+
+    only_merge_lora_model: bool = field(
+        default=False,
+        metadata={"help": "If passed, will only merge the lora modules and save the entire model."}
+    )

FlagEmbedding/finetune/embedder/decoder_only/base/load_model.py

@@ -177,7 +177,9 @@ def save_merged_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir:
         model = PeftModel.from_pretrained(model, find_largest_checkpoint(output_dir))
         model = model.merge_and_unload()
 
-    model.save_pretrained(os.path.join(output_dir, 'merged_model'))
-
     tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=model_args.trust_remote_code)
     tokenizer.save_pretrained(os.path.join(output_dir, 'merged_model'))
+
+    # modify the vocab size in the model configuration
+    model.config.vocab_size = len(tokenizer)
+    model.save_pretrained(os.path.join(output_dir, 'merged_model'))

FlagEmbedding/finetune/embedder/decoder_only/base/runner.py

@@ -29,6 +29,9 @@ def __init__(
         training_args: AbsEmbedderTrainingArguments
     ):
         super().__init__(model_args, data_args, training_args)
+        self.model_args: DecoderOnlyEmbedderModelArguments
+        self.data_args: AbsEmbedderDataArguments
+        self.training_args: AbsEmbedderTrainingArguments
 
     def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
         """Load tokenizer and model.
@@ -117,11 +120,12 @@ def run(self):
         """
         Run the finetune.
         """
-        Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
+        if not self.model_args.only_merge_lora_model:
+            Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
 
-        # Training
-        self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
-        self.trainer.save_model()
+            # Training
+            self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
+            self.trainer.save_model()
 
         # save merged model
         if self.model_args.save_merged_lora_model and self.training_args.process_index == 0:

FlagEmbedding/finetune/embedder/decoder_only/icl/arguments.py

@@ -73,6 +73,11 @@ class DecoderOnlyEmbedderICLModelArguments(AbsEmbedderModelArguments):
         metadata={"help": "If passed, will merge the lora modules and save the entire model."}
     )
 
+    only_merge_lora_model: bool = field(
+        default=False,
+        metadata={"help": "If passed, will only merge the lora modules and save the entire model."}
+    )
+
 
 @dataclass
 class DecoderOnlyEmbedderICLDataArguments(AbsEmbedderDataArguments):

FlagEmbedding/finetune/embedder/decoder_only/icl/load_model.py

@@ -175,7 +175,9 @@ def save_merged_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_d
         model = PeftModel.from_pretrained(model, find_largest_checkpoint(output_dir))
         model = model.merge_and_unload()
 
-    model.save_pretrained(os.path.join(output_dir, 'merged_model'))
-
     tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=model_args.trust_remote_code)
     tokenizer.save_pretrained(os.path.join(output_dir, 'merged_model'))
+
+    # modify the vocab size in the model configuration
+    model.config.vocab_size = len(tokenizer)
+    model.save_pretrained(os.path.join(output_dir, 'merged_model'))

FlagEmbedding/finetune/embedder/decoder_only/icl/runner.py

@@ -145,11 +145,12 @@ def run(self):
         """
         Run the finetune.
         """
-        Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
-
-        # Training
-        self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
-        self.trainer.save_model()
+        if not self.model_args.only_merge_lora_model:
+            Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
+
+            # Training
+            self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
+            self.trainer.save_model()
 
         # save merged model
         if self.model_args.save_merged_lora_model and self.training_args.process_index == 0: