Merge pull request #1108 from 545999961/master

update bge-en-icl

Author: 545999961

FlagEmbedding/llm_dense_retriever/README.md

@@ -208,7 +208,8 @@ run.py \
 --use_special_tokens \
 --symmetric_batch_size 256 \
 --symmetric_train_group_size 8 \
---max_class_neg 7
+--max_class_neg 7 \
+--save_merged_lora_model True
 ```
 
 ## Citation

FlagEmbedding/llm_dense_retriever/finetune/arguments.py

@@ -55,10 +55,10 @@ class ModelArguments:
         metadata={"help": "If passed, will use flash attention to train the model."}
     )
     token: str = field(
-        default="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"
+        default=".."
     )
     cache_dir: str = field(
-        default="/share/LMs"
+        default="../LMs"
     )
     from_peft: str = field(
         default=None

FlagEmbedding/llm_dense_retriever/finetune/load_model.py

@@ -1,8 +1,25 @@
 import os
+import re
+
 import torch
 from transformers import AutoConfig, AutoModel, AutoTokenizer
 from peft import LoraConfig, TaskType, get_peft_model, PeftModel
 
+def find_largest_checkpoint(checkpoint_dir):
+    checkpoint_pattern = re.compile(r'checkpoint-(\d+)')
+    max_number = -1
+    max_checkpoint_file = None
+    for file in os.listdir(checkpoint_dir):
+        match = checkpoint_pattern.search(file)
+        if match:
+            number = int(match.group(1))
+            if number > max_number:
+                max_number = number
+                max_checkpoint_file = file
+    if max_checkpoint_file:
+        return os.path.join(checkpoint_dir, max_checkpoint_file)
+    else:
+        return None
 
 def get_model(model_args, output_dir, resize, resize_tokens):
 
@@ -112,8 +129,12 @@ def save_merged_model(model_args, output_dir):
     if os.path.exists(os.path.join(output_dir, 'embedding', 'emb.pth')):
         model.set_input_embeddings(torch.load(os.path.join(output_dir, 'embedding', 'emb.pth')))
 
-    model = PeftModel.from_pretrained(model, output_dir)
-    model = model.merge_and_unload()
+    try:
+        model = PeftModel.from_pretrained(model, output_dir)
+        model = model.merge_and_unload()
+    except:
+        model = PeftModel.from_pretrained(model, find_largest_checkpoint(output_dir))
+        model = model.merge_and_unload()
 
     model.save_pretrained(os.path.join(output_dir, 'full_model'))
 

FlagEmbedding/llm_dense_retriever/finetune/run.py

@@ -14,7 +14,7 @@
 from data import SameDatasetTrainDataset, SameEmbedCollator
 from modeling import BiEncoderModel
 from trainer import BiTrainer
-from load_model import get_model
+from load_model import get_model, save_merged_model
 
 logger = logging.getLogger(__name__)
 
@@ -143,6 +143,16 @@ def main():
         # os.makedirs(os.path.join(training_args.output_dir, 'embedding'), exist_ok=True)
         # torch.save(base_model.model.model.embed_tokens, os.path.join(training_args.output_dir, 'embedding', 'emb.pth'))
 
+def save_model():
+    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    model_args: ModelArguments
+    data_args: DataArguments
+    training_args: TrainingArguments
+
+    if model_args.save_merged_lora_model and training_args.process_index == 0:
+        save_merged_model(model_args, training_args.output_dir)
 
 if __name__ == "__main__":
-    main()
+    main()
+    save_model()