Update app.py

app.py

@@ -1,47 +1,88 @@
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    pipeline
+)
 import torch
+import gradio as gr
 
-# Load WizardMath (adjust model size if needed)
-model_name = "WizardLM/WizardMath-7B-V1.1"  # Smaller: "WizardLM/WizardMath-70B-V1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Configuration
+MODEL_NAME = "WizardLM/WizardMath-7B-V1.1"  # Use 7B model for Spaces
+CACHE_DIR = "/tmp"  # For Spaces limited storage
+
+# 4-bit quantization setup
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True
+)
+
+# Load model with optimizations
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,  # Reduce memory usage
-    device_map="auto"  # Auto-select GPU/CPU
+    MODEL_NAME,
+    quantization_config=quant_config,
+    device_map="auto",
+    cache_dir=CACHE_DIR,
+    trust_remote_code=True
+)
+
+# Create a text generation pipeline
+math_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-def solve_math_problem(question):
-    # Format input for WizardMath
-    prompt = f"USER: Solve this math problem: {question}\nASSISTANT:"
+def solve_math(question):
+    prompt = f"USER: {question}\nASSISTANT:"
     
-    # Generate response
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        inputs.input_ids,
+    # Generate response with adjusted parameters
+    outputs = math_pipeline(
+        prompt,
         max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # Decode and clean output
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = answer.split("ASSISTANT:")[-1].strip()  # Extract the answer part
-    
+    # Extract and clean the answer
+    full_response = outputs[0]["generated_text"]
+    answer = full_response.split("ASSISTANT:")[-1].strip()
     return answer
 
 # Gradio Interface
 demo = gr.Interface(
-    fn=solve_math_problem,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your math problem here..."),
-    outputs=gr.Textbox(label="Solution"),
-    title="🧙 WizardMath Problem Solver",
+    fn=solve_math,
+    inputs=gr.Textbox(
+        label="Math Problem",
+        placeholder="Enter your math question here...",
+        lines=3
+    ),
+    outputs=gr.Textbox(
+        label="Solution",
+        lines=5
+    ),
+    title="🧮 WizardMath Solver",
+    description="Solves math problems using WizardMath-7B (4-bit quantized)",
     examples=[
-        ["What is the integral of x^2 from 0 to 3?"],
-        ["Solve for x: 2x + 5 = 15"],
-        ["Calculate the area of a circle with radius 4."]
+        ["What is 2^10 + 5*3?"],
+        ["Solve for x: 3x + 7 = 22"],
+        ["Calculate the area of a circle with radius 5"]
     ],
-    theme="soft"  # Try "default" or "huggingface"
+    allow_flagging="never"
 )
 
+# Launch with Space-optimized settings
 if __name__ == "__main__":
-    demo.launch(server_port=7860, share=False)  # Set share=True for public link
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False  # Set to True for public link during testing
+    )