Spaces:
Runtime error
Runtime error
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() | |