Update app.py

app.py

@@ -1,100 +1,47 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import pipeline
 from huggingface_hub import login
-import torch
 import gradio as gr
 import os
 
-# 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)
+# Authenticate (set HF_TOKEN in Space secrets)
 login(token=os.environ.get("HF_TOKEN"))
 
-# 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 lightweight pipeline (faster than full model load)
+math_pipeline = pipeline(
+    "text-generation",
+    model="google/gemma-2b-it",
+    device_map="auto",
+    torch_dtype="auto",
+    model_kwargs={"load_in_4bit": True}
 )
 
-# Load model with error handling
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        quantization_config=quant_config,
-        device_map="auto",
-        torch_dtype=torch.float16,
-        cache_dir=CACHE_DIR
-    )
-except Exception as 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"""
+    """Super-fast response with optimized prompt"""
+    prompt = f"""Solve this math problem concisely:
+    
+Question: {question}
+Answer:"""
+    
     try:
-        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=MAX_TOKENS,
-            temperature=0.3,  # Lower = more deterministic answers
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+        result = math_pipeline(
+            prompt,
+            max_new_tokens=150,  # Shorter = faster
+            temperature=0.1,     # More deterministic
+            do_sample=False,     # Faster generation
+            num_return_sequences=1
         )
-        
-        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return answer.split("Answer:")[-1].strip()
-    
+        return result[0]['generated_text'].split("Answer:")[-1].strip()
     except Exception as e:
-        return f"❌ Error generating answer: {str(e)}"
+        return f"🚨 Error: {str(e)}"
 
-# Preload model for faster first response
+# Preload pipeline
 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>""")
-    
-    with gr.Row():
-        question = gr.Textbox(
-            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="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,
-        api_name="solve"
-    )
+# Minimal UI for fastest response
+with gr.Blocks(title="⚡ Instant Math Solver") as demo:
+    gr.Markdown("### Enter a math problem:")
+    question = gr.Textbox(lines=2)
+    answer = gr.Textbox(label="Answer", lines=3)
+    question.submit(solve_math, question, answer)
 
-if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+demo.launch(server_name="0.0.0.0")