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

# Load Hugging Face token from environment variables
hf_token = os.getenv("HUGGINGFACE_TOKEN")
if hf_token:
    login(hf_token)
    print("✅ Successfully logged in to Hugging Face Hub")
else:
    print("❌ Hugging Face token not found. Make sure it's set in 'Secrets'.")

# Model paths
BASE_MODEL = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
FINETUNED_MODEL_DIR = "./finetuned_model"  # Path to your fine-tuned adapter

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# Load base model efficiently
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch_dtype,
    device_map="auto" if torch.cuda.is_available() else None  # Use GPU if available
)

# Load and merge LoRA adapter
model = PeftModel.from_pretrained(base_model, FINETUNED_MODEL_DIR)
model = model.merge_and_unload()

# Move model to appropriate device
model.to(device)

def chat(message):
    inputs = tokenizer(message, return_tensors="pt").to(device)
    output = model.generate(**inputs, max_new_tokens=100)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# Gradio UI
interface = gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="Chat with Mistral (Fine-Tuned)",
    description="Talk to a fine-tuned Mistral-7B model."
)

if __name__ == "__main__":
    interface.launch()