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

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="cpu", torch_dtype="auto")

def generate(user_message, system_message, max_tokens, temperature):
    # 1. Construct the Conversation Structure properly
    # This ensures the model knows exactly who is talking
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message}
    ]
    
    # 2. Apply the specific Chat Template for Qwen 2.5
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    outputs = model.generate(
        **inputs, 
        max_new_tokens=max_tokens, 
        temperature=temperature,
        do_sample=True
    )
    
    # 3. Decode only the new response
    return tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)

# Define inputs: User Msg, System Msg, Max Tokens, Temperature
gr.Interface(
    fn=generate, 
    inputs=[
        gr.Textbox(label="User Message"), 
        gr.Textbox(label="System Prompt", value="You are a helpful assistant."),
        gr.Slider(10, 500, value=200, label="Max Tokens"),
        gr.Slider(0.0, 1.0, value=0.5, label="Temperature")
    ], 
    outputs="text"
).launch()