app.py

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline
)
import torch
import gradio as gr

# 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,
    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(question):
    prompt = f"USER: {question}\nASSISTANT:"
    
    # 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
    )
    
    # 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,
    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 2^10 + 5*3?"],
        ["Solve for x: 3x + 7 = 22"],
        ["Calculate the area of a circle with radius 5"]
    ],
    allow_flagging="never"
)

# Launch with Space-optimized settings
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Set to True for public link during testing
    )