import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig

# Model name
MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,              # quantize to 4-bit
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

# Load tokenizer and model for CPU
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",      
    torch_dtype=torch.float16  
)

# Build pipeline
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="cpu"
)

system_prompt = """
You are Gompei the Goat, the beloved mascot of Worcester Polytechnic Institute (WPI).
Speak in a friendly, school-spirited, and enthusiastic tone. Respond in 1-2 sentences for greetings. 
Provide detailed WPI facts only when asked a question.
"""

def chatbot(message, history):
    # Build context with system prompt
    context = system_prompt
    for user, bot in history:
        context += f"\nUser: {user}\nGompei: {bot}"
    context += f"\nUser: {message}\nGompei:"

    # Generate response
    response = generator(
        context,
        max_new_tokens=140,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    reply = response[0]["generated_text"].split("Gompei:")[-1].strip()
    return reply

demo = gr.ChatInterface(
    fn=chatbot,
    title="Chat with Gompei the Goat 🐐",
    description="Ask Gompei questions about WPI!",
    type="messages"  # ✅ avoids the deprecation warning
)

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