import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def merge_weights():
    base_model_id = "Qwen/Qwen2.5-3B-Instruct"
    adapter_dir = "./digital_twin_adapters"
    merged_dir = "./digital_twin_merged"

    print("Loading base model in FP16 on CPU...")
    model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.float16,
        device_map="cpu",
        trust_remote_code=True
    )
    
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

    print("Loading fine-tuned LoRA adapters...")
    model = PeftModel.from_pretrained(model, adapter_dir)
    
    print("Permanently merging adapters into base model weights...")
    model = model.merge_and_unload()

    print(f"Saving merged FP16 model to '{merged_dir}'...")
    model.save_pretrained(merged_dir)
    tokenizer.save_pretrained(merged_dir)
    print("Weights successfully merged and saved.")

if __name__ == "__main__":
    merge_weights()