app.py

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

# Load the model and tokenizer
model_name = "Qwen/Qwen2-Math-1.5B-Instruct"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
).to(device)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define a function for Gradio to handle user input
def solve_math(prompt):
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generation_config = GenerationConfig(
        do_sample=False,  # For greedy decoding
        max_new_tokens=512
    )

    generated_ids = model.generate(
        **model_inputs,
        generation_config=generation_config
    )

    # Remove the input tokens from the output
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    # Decode the generated output and return the result
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response

# Create the Gradio interface
iface = gr.Interface(
    fn=solve_math,  # Function to call
    inputs="text",  # Text input for the user prompt
    outputs="text", # Text output for the model's response
    title="Math Solver",  # App title
    description="Provide a math problem and the model will solve it."
)

# Launch the app
if __name__ == "__main__":
    iface.launch()