File size: 2,661 Bytes
ea19bf1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import gradio as gr
from datasets import load_dataset
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration # Changed model class
# Check if GPU is available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
# Load dataset
ds = load_dataset("AI-MO/NuminaMath-CoT")
# Load model and tokenizer
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device) # Changed model class
def process_example(example):
"""Process a single example from the dataset"""
question = example['question']
solution = example['solution']
answer = example['answer']
return f"Question: {question}\nSolution: {solution}\nAnswer: {answer}"
def get_random_example():
"""Get a random example from the dataset"""
import random
idx = random.randint(0, len(ds['train']) - 1)
return process_example(ds['train'][idx])
def solve_math_problem(question):
"""Generate solution for a given math problem"""
# Add prefix for T5
input_text = "solve math: " + question
inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True).to(device)
# Generate response
outputs = model.generate(
inputs["input_ids"],
max_length=200,
num_return_sequences=1,
temperature=0.7,
do_sample=True,
top_p=0.9,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# Create Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Math Problem Solver")
gr.Markdown("Using FLAN-T5 model to solve mathematical problems with step-by-step solutions.")
with gr.Row():
with gr.Column():
input_text = gr.Textbox(
label="Enter your math problem",
placeholder="Type your math problem here...",
lines=3
)
with gr.Row():
submit_btn = gr.Button("Solve Problem", variant="primary")
example_btn = gr.Button("Show Random Example")
with gr.Column():
output_text = gr.Textbox(
label="Solution",
lines=8,
show_copy_button=True
)
# Set up event handlers
submit_btn.click(
fn=solve_math_problem,
inputs=input_text,
outputs=output_text
)
example_btn.click(
fn=get_random_example,
inputs=None,
outputs=input_text
)
# Launch the interface
demo.launch(share=True) |