from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
 
"""
将lora权重合并到大模型中
"""
 
def merge_lora_to_LLM():
    model_name_or_path = "/root/paddlejob/workspace/env_run/huitingfeng/models/llama-2-7b-hf"
    adapter_name_or_path = "/root/paddlejob/workspace/env_run/huitingfeng/model_outputs/lora-500k/checkpoint-15624"
    save_path = "/root/paddlejob/workspace/env_run/huitingfeng/model_outputs/lora-merged-500k"
    
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        trust_remote_code=True
    )
    tokenizer.pad_token = tokenizer.unk_token
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.bfloat16
    )
    model = PeftModel.from_pretrained(model, adapter_name_or_path)
    model = model.merge_and_unload()
 
    tokenizer.save_pretrained(save_path)
    model.save_pretrained(save_path)
 
if __name__ == "__main__":
    merge_lora_to_LLM()