Update app.py

app.py

@@ -1,47 +1,75 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import gradio as gr
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Math-1.5B")
-model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Math-1.5B", torch_dtype=torch.bfloat16, device_map="cpu")
+# Lightweight CPU-friendly model
+model_name = "microsoft/phi-1_5"
 
-
-from transformers import GenerationConfig
-model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen2.5-Math-1.5B")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+model.generation_config = GenerationConfig.from_pretrained(model_name)
 model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
 
-def solve_math_problem(questions):
-    if isinstance(questions, str):  # If input is a single string
-        questions = [questions]
-
-    results = []
-    for question in questions:
-        messages = [{"role": "user", "content": question}]
-        input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
-        outputs = model.generate(input_tensor.to(model.device), max_new_tokens=100)
-        result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
-        results.append(result)
-    return results
-
-
-interface = gr.Interface(
-    fn=solve_math_problem,
-    inputs="text",
-    outputs="text",
-    title="Math Wizard",
-    description="""
-        Welcome to the Math Wizard!
-        Ask any math question, and let the wizard guide you through the solution step-by-step.
-    """,
-    allow_flagging=False,
-    examples=[
-        ["What is the integral of x^2?"],
-        ["How do I solve a quadratic equation?"],
-        ["Tell me about Ramanujan"]
-    ]
-)
+def solve_math_problem(question):
+    messages = [{"role": "user", "content": question}]
+    input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+    input_tensor = input_tensor.to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(input_tensor, max_new_tokens=150)
+
+    response = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
+    return response.strip()
+
+
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
+    gr.Markdown("# 🧙‍♂️ Math Wizard AI")
+    gr.Markdown("""
+        <div style="font-size: 16px; line-height: 1.5">
+        Welcome to the <b>Math Wizard</b> – your intelligent assistant for solving math problems of all kinds! <br>
+        Ask anything from algebra, calculus, or even about famous mathematicians.<br><br>
+        </div>
+    """)
+
+    with gr.Tabs():
+        with gr.Tab("🧮 General Math"):
+            with gr.Row():
+                with gr.Column():
+                    question_box = gr.Textbox(
+                        label="Ask your question here:",
+                        placeholder="E.g. What is the derivative of x^2 + 3x + 2?",
+                        lines=3
+                    )
+                    submit_btn = gr.Button("🔍 Solve Now")
+                    clear_btn = gr.Button("❌ Clear")
+
+                with gr.Column():
+                    answer_box = gr.Textbox(label="📘 Answer from the Wizard", lines=8, interactive=False)
+                    copy_btn = gr.Button("📋 Copy Answer")
+
+            submit_btn.click(fn=solve_math_problem, inputs=question_box, outputs=answer_box)
+            clear_btn.click(lambda: ("", ""), outputs=[question_box, answer_box])
+            copy_btn.click(lambda x: x, inputs=answer_box, outputs=answer_box, show_progress=False)
+
+        with gr.Tab("🧠 Examples & Inspiration"):
+            gr.Markdown("""
+            <h4>Try asking things like:</h4>
+            <ul>
+            <li>🧩 Solve the equation x² + 2x - 3 = 0</li>
+            <li>📐 What is the Pythagorean theorem?</li>
+            <li>📊 Integrate sin(x) from 0 to π</li>
+            <li>🧮 Tell me about Euclid</li>
+            </ul>
+            """)
+
+        with gr.Tab("📚 About"):
+            gr.Markdown("""
+            <h4>About Math Wizard</h4>
+            <p>This assistant is powered by a lightweight AI model that runs smoothly even on CPUs.</p>
+            <p>Built with ❤️ using Gradio + HuggingFace Transformers</p>
+            <p>Model: <code>microsoft/phi-1_5</code> optimized for reasoning and small footprint.</p>
+            """)
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()