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


login(token=os.environ["HF_TOKEN"])
model_name = "mistralai/Mistral-7B-Instruct-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16,  
    trust_remote_code=True
)

def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=500) 
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(placeholder="💬 Type your message here...", label="📝 Your Input"),
    outputs=gr.Textbox(label="🤖 AI Response", lines=8),
    title="🤖 Ben AI Chatbot",
    description="✨ Experience a sleek and modern AI chatbot interface. 🚀",
    theme="soft"
)

demo.launch()