"""
Merge LoRA adapter with base model and push to Hugging Face Hub.
"""

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

# --- Config ---
BASE_MODEL = "nvidia/Llama-3.1-Nemotron-Nano-4B-v1.1"
ADAPTER_PATH = "training/output"
HUB_REPO = "aslanconfig/tech-advisor-nemotron-4b"

# --- Load base model ---
print(f"Loading base model: {BASE_MODEL}")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# --- Load and merge adapter ---
print(f"Loading adapter from: {ADAPTER_PATH}")
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
print("Merging adapter into base model...")
model = model.merge_and_unload()

# --- Push to Hub ---
print(f"Pushing merged model to: {HUB_REPO}")
model.push_to_hub(HUB_REPO)
tokenizer.push_to_hub(HUB_REPO)
print("Done! Model available at: https://huggingface.co/" + HUB_REPO)