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

MODEL_NAME = "Qwen/Qwen3-0.6B"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

def chat(message, history):
    prompt = f"<|user|>\n{message}\n<|assistant|>\n"

    inputs = tokenizer(prompt, return_tensors="pt")

    output = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        do_sample=True,
        top_p=0.95
    )

    response = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    answer = response.split("<|assistant|>")[-1].strip()

    return answer

demo = gr.ChatInterface(
    fn=chat,
    title="TinyLlama Chatbot",
    description="A chatbot powered by TinyLlama"
)

demo.launch()