Update app.py

app.py

@@ -4,21 +4,22 @@ import torch
 import gradio as gr
 import os
 
-# Authenticate
-login(token=os.environ.get("HF_TOKEN"))
-
 # Configuration
 MODEL_NAME = "google/gemma-2b-it"
 CACHE_DIR = "/tmp"
+MAX_TOKENS = 200  # Reduced for faster responses
+
+# Authenticate (HF_TOKEN must be set in Space secrets)
+login(token=os.environ.get("HF_TOKEN"))
 
-# 4-bit quantization config
+# 4-bit quantization for memory efficiency
 quant_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_quant_type="nf4"
 )
 
-# Load model
+# Load model with error handling
 try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
     model = AutoModelForCausalLM.from_pretrained(
@@ -29,44 +30,67 @@ try:
         cache_dir=CACHE_DIR
     )
 except Exception as e:
-    raise gr.Error(f"Model loading failed: {str(e)}")
+    raise gr.Error(f"⚠️ Model loading failed. Please check your token and try again.\nError: {str(e)}")
 
 def solve_math(question):
+    """Generate step-by-step solutions with error handling"""
     try:
-        prompt = f"Solve step by step: {question}\nAnswer:"
+        prompt = f"Solve this step by step:\n\nQuestion: {question}\nAnswer:"
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
-            temperature=0.3,
-            do_sample=True
+            max_new_tokens=MAX_TOKENS,
+            temperature=0.3,  # Lower = more deterministic answers
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return answer.split("Answer:")[-1].strip()
+    
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Error generating answer: {str(e)}"
+
+# Preload model for faster first response
+solve_math("2+2=")  # Warm-up call
 
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""<h1><center>🧮 Gemma-2B Math Solver</center></h1>""")
+    gr.Markdown("""<h1><center>⚡ Gemma-2B Math Solver</center></h1>""")
+    
     with gr.Row():
         question = gr.Textbox(
-            label="Math Problem",
-            placeholder="What is 2^10 + 5*3?",
+            label="Enter your math problem",
+            placeholder="What is the integral of x^2 from 0 to 3?",
             lines=3
         )
+    
     with gr.Row():
         submit_btn = gr.Button("Solve", variant="primary")
+    
     with gr.Row():
         answer = gr.Textbox(
-            label="Solution",
-            lines=5,
+            label="Step-by-step solution",
+            lines=6,
             interactive=False
         )
     
+    # Examples for quick testing
+    gr.Examples(
+        examples=[
+            ["What is 2^10 + 5*3?"],
+            ["Solve for x: 3x + 5 = 20"],
+            ["Calculate the area of a circle with radius 4"]
+        ],
+        inputs=question
+    )
+
     submit_btn.click(
         fn=solve_math,
         inputs=question,
-        outputs=answer
+        outputs=answer,
+        api_name="solve"
     )
 
 if __name__ == "__main__":