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

BASE_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

ADAPTER_CHECKPOINT_PATH = "./model_output/phi2_finetuned_logs/checkpoint-575"

# D:\phi2_tuning\model_output\phi2_finetuned_logs\checkpoint-500

MERGED_MODEL_PATH = "./updated_logger"

print(f"loading base model from: {BASE_MODEL_NAME}")

try:
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_NAME,
        low_cpu_mem_usage=True,
        return_dict=True,
        torch_dtype = torch.float16,
        trust_remote_code=True,
        device_map="auto"
    )
except Exception as e:
    print(f"error loading model: {e}")
    exit()

tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL_NAME,
    trust_remote_code=True
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

try:
    model = PeftModel.from_pretrained(base_model, ADAPTER_CHECKPOINT_PATH)
except Exception as e:
    print(f"error loading the adapter checkpoint")
    print("ensure the adapter checkpoint is correct and retry again")

merged_model = model.merge_and_unload()
print("adapters merged successfully!!")

print("saving the merged model...")

os.makedirs(MERGED_MODEL_PATH, exist_ok=True)
merged_model.save_pretrained(MERGED_MODEL_PATH)
tokenizer.save_pretrained(MERGED_MODEL_PATH)

print(f"model merged and saved to {MERGED_MODEL_PATH}")