import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
from huggingface_hub import login

def merge_model(hf_token):
    try:
        login(token=hf_token)
        yield "✓ Logged in successfully\n"
        
        yield "Loading base model (this takes 2-3 min)...\n"
        base_model = AutoModelForCausalLM.from_pretrained(
            "HuggingFaceH4/zephyr-7b-beta",
            torch_dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True
        )
        yield "✓ Base model loaded\n"
        
        yield "Loading tokenizer...\n"
        tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
        yield "✓ Tokenizer loaded\n"
        
        yield "Loading your adapter...\n"
        model = PeftModel.from_pretrained(base_model, "youtiva-123/autotrain-r6w6w-7pv2i")
        yield "✓ Adapter loaded\n"
        
        yield "Merging adapter with base model...\n"
        merged_model = model.merge_and_unload()
        yield "✓ Models merged\n"
        
        yield "Uploading to youtiva-123/zephyr-youtiva-merged...\n"
        merged_model.push_to_hub("youtiva-123/zephyr-youtiva-merged")
        tokenizer.push_to_hub("youtiva-123/zephyr-youtiva-merged")
        
        yield "✅ SUCCESS! Your merged model is ready at:\nhttps://huggingface.co/youtiva-123/zephyr-youtiva-merged\n\nNow go deploy this model as an endpoint!"
    except Exception as e:
        yield f"❌ Error: {str(e)}"

interface = gr.Interface(
    fn=merge_model,
    inputs=gr.Textbox(label="Hugging Face Token (needs write access)", type="password"),
    outputs=gr.Textbox(label="Status", lines=15),
    title="Merge LoRA Adapter with Base Model"
)

interface.launch()