Update app.py

app.py

@@ -1,46 +1,46 @@
 from transformers import pipeline
-from huggingface_hub import login
 import gradio as gr
+import torch
 import os
 
-# 1. First verify bitsandbytes is installed
-try:
-    import bitsandbytes  # noqa
-except ImportError:
-    raise ImportError("bitsandbytes not installed! Add it to requirements.txt")
+# Verify GPU support
+if not torch.cuda.is_available():
+    raise RuntimeError("GPU not available - enable GPU in Space settings")
 
-# 2. Authenticate
-login(token=os.environ.get("HF_TOKEN"))
-
-# 3. Load model WITHOUT 4-bit (for compatibility)
+# Load model with GPU acceleration
 math_pipeline = pipeline(
     "text-generation",
     model="google/gemma-2b-it",
-    device_map="auto",
-    torch_dtype="auto",  # Let transformers choose dtype
+    device=0,  # Force GPU usage
+    torch_dtype=torch.float16,
     model_kwargs={
-        "low_cpu_mem_usage": True  # Reduces memory spikes
+        "low_cpu_mem_usage": True,
+        "trust_remote_code": True
     }
 )
 
 def solve_math(question):
-    prompt = f"Solve concisely:\nQ: {question}\nA:"
+    prompt = f"Solve step by step:\nQ: {question}\nA:"
     try:
         result = math_pipeline(
             prompt,
-            max_new_tokens=100,  # Shorter = faster
+            max_new_tokens=150,
             temperature=0.3,
-            do_sample=False  # Faster generation
+            do_sample=False,
+            pad_token_id=math_pipeline.tokenizer.eos_token_id
         )
         return result[0]['generated_text'].split("A:")[-1].strip()
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Minimal UI
-with gr.Blocks() as demo:
-    gr.Markdown("## Math Solver")
-    question = gr.Textbox(label="Problem")
-    answer = gr.Textbox(label="Solution")
+# Preload model
+solve_math("2+2=")
+
+# Optimized UI
+with gr.Blocks(title="🚀 Math Solver") as demo:
+    gr.Markdown("## Enter a math problem:")
+    question = gr.Textbox(label="", placeholder="What is 2^8?")
+    answer = gr.Textbox(label="Solution", lines=3)
     question.submit(solve_math, question, answer)
 
-demo.launch()
+demo.launch(server_name="0.0.0.0")