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

# Load model and tokenizer
model_name = "Qwen/Qwen2.5-3B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Chat history
chat_history = []

# System prompt
SYSTEM_PROMPT = "You are Qwen/Qwen2.5-3B-Instruct, created by Alibaba Cloud. You are a helpful assistant."

def generate_response(user_input, history):
    # Build message list
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": user_input})

    # Apply chat template
    prompt_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Tokenize
    model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device)

    # Generate response
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )

    # Only return new tokens
    new_tokens = generated_ids[0][model_inputs.input_ids.shape[-1]:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    # Update chat history
    history.append((user_input, response))
    return history, history

# Launch Gradio Chatbot UI
chatbot_ui = gr.ChatInterface(
    fn=generate_response,
    title="🧠 Qwen 2.5 3B - Chatbot",
    description="A simple chatbot interface powered by Qwen2.5-3B-Instruct (Alibaba Cloud).",
    theme="soft",
    examples = [
    "How can virtual reality (VR) influence consumer behavior towards sustainability?",
    "What impact does sustainable packaging have on consumer purchasing decisions?",
    "In what ways can education promote more sustainable consumer behaviors?"
],

)

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