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

# Load the Hugging Face token from Secrets (recommended for public Spaces)
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 details
BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.1"
FINETUNED_MODEL_DIR = "./finetuned_model"  # Ensure this matches your uploaded folder

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# Load model with quantization for efficiency
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,  # Use 4-bit quantization for lower memory usage
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)

model = AutoModelForCausalLM.from_pretrained(
    FINETUNED_MODEL_DIR,  # Load your fine-tuned model
    quantization_config=quantization_config,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Inference function
def chat(message):
    inputs = tokenizer(message, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    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()