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

# Load model and tokenizer from Hugging Face
model_name = "meta-llama/Llama-3.2-1b-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# Chat history
chat_history = []

def generate_response(message, history):
    # Combine message with history
    prompt = ""
    for user, bot in history:
        prompt += f"<|user|>{user}<|end|><|assistant|>{bot}<|end|>"
    prompt += f"<|user|>{message}<|end|><|assistant|>"

    # Tokenize and generate
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id
    )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract only the assistant's latest message
    response = result.split("<|assistant|>")[-1].strip()
    history.append((message, response))
    return response, history

# Gradio UI
chatbot = gr.ChatInterface(fn=generate_response, 
                           title="Llama 3.2 Chatbot",
                           chatbot=gr.Chatbot(),
                           textbox=gr.Textbox(placeholder="Ask me anything...", lines=2),
                           clear_btn="Clear")

chatbot.launch()