# load_and_merge.py
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import os

base_model_path   = "YOUR_PATH_TO_LLAMA-3.1-8B"  
lora_adapter_path = "YOUR_PATH_TO_LORA"
save_dir          = "YOUR_SAVE_DIR"      

os.makedirs(save_dir, exist_ok=True)

print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
)

tokenizer = AutoTokenizer.from_pretrained(lora_adapter_path, legacy=True)


base_model.resize_token_embeddings(len(tokenizer))

print("Loading LoRA adapter...")
peft_model = PeftModel.from_pretrained(base_model, lora_adapter_path)


print("Merging...")
merged_model = peft_model.merge_and_unload()
breakpoint()


print("Saving merged model...")
merged_model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)

print(f"✅ 合并完成")