import gradio as gr
from transformers import pipeline

# טעינת המודל
print("Loading model...")
pipe = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-Math-1.5B-Instruct",
    device_map="auto"
)

def generate_problems(topic, num_problems):
    prompt = f"""Generate {num_problems} math problems about {topic}.

Format:
Problem 1: [clear problem statement]
Problem 2: [clear problem statement]

Make them educational and challenging."""
    
    result = pipe(
        prompt,
        max_new_tokens=1000,
        temperature=0.7,
        do_sample=True
    )
    return result[0]['generated_text']

# ממשק פשוט
demo = gr.Interface(
    fn=generate_problems,
    inputs=[
        gr.Textbox(label="Topic", placeholder="algebra, calculus, geometry..."),
        gr.Slider(1, 10, value=3, step=1, label="Number of problems")
    ],
    outputs=gr.Textbox(label="Generated Problems", lines=20),
    title="Math Problem Generator",
    description="Generate math problems using AI"
)

demo.launch()