from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from common import BASE

base = AutoModelForCausalLM.from_pretrained(BASE)
model = PeftModel.from_pretrained(base, "lora-out").merge_and_unload()
model.save_pretrained("merged")
AutoTokenizer.from_pretrained(BASE).save_pretrained("merged")
print("Merged model saved to merged/")