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

# Load Model & Tokenizer
MODEL_NAME = "Qwen/Qwen2.5-3B"  # Change to your HF model repo if fine-tuned

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

def tinyzero_chat(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_length=512, temperature=0.7, top_p=0.9)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio Interface
demo = gr.Interface(
    fn=tinyzero_chat,
    inputs=gr.Textbox(placeholder="Ask TinyZero anything..."),
    outputs=gr.Textbox(),
    title="TinyZero Chatbot",
    description="An interactive demo of TinyZero trained on reasoning tasks."
)

# Launch
demo.launch()