import gradio as gr
import torch
from transformers import pipeline

# Model configuration
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

# Load the model once when the application starts
device = 0 if torch.cuda.is_available() else -1

generator = pipeline(
    task="text-generation",
    model=MODEL_ID,
    device=device,
    torch_dtype="auto",
)

def chat(message, history):
    """Generate an answer from the user's message."""

    if not message or not message.strip():
        return "Please enter a message."

    messages = [
        {
            "role": "system",
            "content": "You are a helpful, concise AI assistant."
        }
    ]

    # Include previous conversation messages
    for user_message, assistant_message in history or []:
        messages.append({
            "role": "user",
            "content": user_message
        })
        messages.append({
            "role": "assistant",
            "content": assistant_message
        })

    messages.append({
        "role": "user",
        "content": message.strip()
    })

    try:
        result = generator(
            messages,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )

        generated = result[0]["generated_text"]

        # Chat-template pipelines may return the full
        # conversation, so extract the last assistant reply.
        if isinstance(generated, list):
            answer = generated[-1]["content"]
        else:
            answer = generated

        return answer.strip()

    except Exception as error:
        print(f"Generation error: {error}")
        return "Sorry, I couldn't generate a response. Please try again."

# Build the user interface
demo = gr.ChatInterface(
    fn=chat,
    title="My Small AI Assistant",
    description=(
        "A lightweight AI chatbot powered by a "
        "pretrained Hugging Face language model."
    ),
    examples=[
        "Explain Python in simple words.",
        "Write a short story about a robot.",
        "What is machine learning?",
    ],
)

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