app.py

import gradio as gr
import os
from huggingface_hub import InferenceClient
from datasets import load_dataset
import random
import re

# Load math datasets for sample problems
print("Loading datasets...")
fw = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
ds = load_dataset("HuggingFaceH4/ultrachat_200k", streaming=True)

def load_sample_problems():
    """Load sample problems from math datasets"""
    try:
        gsm8k = load_dataset("openai/gsm8k", "main", streaming=True)
        samples = []
        for i, item in enumerate(gsm8k["train"]):
            samples.append(item["question"])
            if i >= 50:
                break
        print(f"✅ Loaded {len(samples)} GSM8K samples")
        return samples
    except Exception as e:
        print(f"⚠️ Dataset error: {e}, using fallback")
        return [
            "What is the derivative of f(x) = 3x² + 2x - 1?",
            "A triangle has sides of length 5, 12, and 13. What is its area?",
            "If log₂(x) + log₂(x+6) = 4, find the value of x.",
            "Find the limit: lim(x->0) (sin(x)/x)",
            "Solve the system: x + 2y = 7, 3x - y = 4"
        ]

math_samples = load_sample_problems()

def create_math_system_message():
    """Create specialized system prompt for mathematics with LaTeX"""
    return """You are Mathetics AI, an advanced mathematics tutor and problem solver.

🧮 **Your Expertise:**
- Step-by-step problem solving with clear explanations
- Multiple solution approaches when applicable
- Proper mathematical notation and terminology using LaTeX
- Verification of answers through different methods

📐 **Problem Domains:**
- Arithmetic, Algebra, and Number Theory
- Geometry, Trigonometry, and Coordinate Geometry
- Calculus (Limits, Derivatives, Integrals)
- Statistics, Probability, and Data Analysis
- Competition Mathematics (AMC, AIME level)

💡 **Teaching Style:**
1. **Understand the Problem** - Identify what's being asked
2. **Plan the Solution** - Choose the appropriate method
3. **Execute Step-by-Step** - Show all work clearly with LaTeX formatting
4. **Verify the Answer** - Check if the result makes sense
5. **Alternative Methods** - Mention other possible approaches

**LaTeX Guidelines:**
- Use $...$ for inline math: $x^2 + y^2 = z^2$
- Use $$...$$ or \[...\] for display math
- Box final answers: \boxed{answer}
- Fractions: \frac{numerator}{denominator}
- Limits: \lim_{x \to 0}
- Derivatives: \frac{d}{dx} or f'(x)
- Integrals: \int f(x) \, dx

Always be precise, educational, and encourage mathematical thinking."""

def render_latex(text):
    """Enhanced LaTeX rendering - fixes raw code output"""
    if not text or len(text) < 5:
        return text
    
    try:
        # Fix common LaTeX patterns from Qwen
        # Ensure proper inline math delimiters
        text = re.sub(r'(?<!\\)\$([^\$]+)\$(?!\$)', r'$\1$', text)
        
        # Convert display math
        text = re.sub(r'\$\$([^\$]+)\$\$', r'$$\1$$', text)
        text = re.sub(r'\\\[([^\\]+)\\\]', r'$$\1$$', text)
        text = re.sub(r'\\\(([^\\]+)\\\)', r'$\1$', text)
        
        # Fix escaped LaTeX commands
        text = re.sub(r'\\(lim|frac|sqrt|int|sum|prod|partial|nabla|infty|to|le|ge|neq|approx|cdot|times|div|deg|prime|log|ln|log|sin|cos|tan|cot|sec|csc|arcsin|arccos|arctan|sinh|cosh)', r'\1', text)
        
        # Ensure boxed answers render
        text = re.sub(r'\\boxed\{([^}]+)\}', r'$$\boxed{\1}$$', text)
        
        # Fix fractions
        text = re.sub(r'\\frac\{([^}]+)\}\{([^}]+)\}', r'$\frac{\1}{\2}$', text)
        
        # Clean up extra spaces
        text = re.sub(r'\s*([\$\\])\s*', r'\1', text)
        
    except Exception as e:
        print(f"⚠️ LaTeX formatting error: {e}")
        # Return original if processing fails
    
    return text

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """Enhanced response with proper LaTeX streaming"""
    # Show initial thinking
    yield "🤔 Thinking step-by-step..."
    
    # Use Qwen Math model
    client = InferenceClient(model="Qwen/Qwen2.5-Math-7B-Instruct")
    
    # Build messages properly
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})
    
    # Add conversation history
    for msg in history:
        if msg.get("role") == "user":
            messages.append({"role": "user", "content": msg["content"]})
        elif msg.get("role") == "assistant":
            messages.append({"role": "assistant", "content": msg["content"]})
    
    messages.append({"role": "user", "content": message})
    
    response = ""
    max_tokens = max(max_tokens, 1536)  # Ensure no truncation
    
    try:
        for message_chunk in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
            timeout=60  # Add timeout for stability
        ):
            choices = message_chunk.choices
            if len(choices) and choices[0].delta.content:
                token = choices[0].delta.content
                response += token
                
                # Stream with formatting every few tokens
                if len(response) % 50 == 0 or token.strip() in ['.', '!', '?', '\n']:
                    formatted = render_latex(response)
                    yield formatted
                else:
                    yield response
                    
        # Final formatted response
        final_formatted = render_latex(response.strip())
        yield final_formatted
            
    except Exception as e:
        error_msg = f"❌ **Error**: {str(e)[:100]}...\n\n💡 **Troubleshooting**:\n• Try a simpler problem\n• Check your connection\n• Wait a moment and retry"
        yield error_msg

def get_random_sample():
    """Get a random sample problem"""
    if math_samples:
        return random.choice(math_samples)
    return "Solve for x: 2x² + 5x - 3 = 0"

def insert_sample_to_chat(difficulty):
    """Insert random sample into chat input"""
    sample = get_random_sample()
    return sample, ""  # Return sample for input, clear status

def show_help():
    return """**🧮 Math Help Tips:**

1. **Be Specific**: "Find the derivative of x² + 3x" instead of "help with calculus"
2. **Request Steps**: "Show me step-by-step how to solve..."
3. **Ask for Verification**: "Check if my answer x=5 is correct"
4. **Alternative Methods**: "What's another way to solve this integral?"
5. **Use Clear Notation**: "lim(x->0)" instead of arrows for mobile

**LaTeX Examples:**
- Inline: $x^2 + y^2 = z^2$
- Display: $$\int x^2 \, dx = \frac{x^3}{3} + C$$
- Boxed: $$\boxed{x = 5}$$

**Pro Tip**: Crank tokens to 1500+ for competition problems!"""

# Enhanced ChatInterface with better LaTeX support
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    title="🧮 **Mathetics AI** - Advanced Mathematics Solver",
    description="""
    **Powered by Qwen 2.5-Math** | **Specialized for Mathematical Problem Solving**
    
    ✨ **Capabilities**: Algebra • Geometry • Calculus • Statistics • Competition Math
    📚 **Features**: Step-by-step solutions • Multiple approaches • Beautiful LaTeX rendering
    """,
    additional_inputs=[
        gr.Textbox(
            value=create_math_system_message(), 
            label="🧠 System Message (Math Tutor Personality)",
            lines=4,
            max_lines=12
        ),
        gr.Slider(minimum=256, maximum=2048, value=1024, step=128, label="📝 Max Tokens (Longer = More Detail)"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, label="🎯 Temperature (Creativity vs Precision)"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.85, step=0.05, label="🔍 Top-p (Response Diversity)"),
    ],
    examples=[
        ["Find the derivative of f(x) = 3x² + 2x - 1"],
        ["A triangle has sides 5, 12, and 13. What is its area?"],
        ["Solve: lim(x->0) (sin(x)/x)"],
        ["What is ∫(2x³ - 5x + 3) dx?"],
        ["Solve the system: x + 2y = 7, 3x - y = 4"]
    ],
    cache_examples=False,
    concurrency_limit=5,  # Reduce for stability
    retry_btn="🔄 Retry",
    undo_btn="↶ Undo",
    clear_btn="🗑️ Clear"
)

# Main interface with enhanced LaTeX CSS
with gr.Blocks(
    title="🧮 Mathetics AI", 
    theme=gr.themes.Soft(),
    css="""
    /* Enhanced math rendering */
    .markdown-body { 
        font-family: 'Times New Roman', Georgia, serif; 
        line-height: 1.6;
    }
    
    /* Katex math rendering */
    .katex { 
        font-size: 1.1em !important; 
        color: #2c3e50;
    }
    .katex-display { 
        font-size: 1.3em !important; 
        text-align: center; 
        margin: 1em 0;
        padding: 10px;
        background: #f8f9fa;
        border-radius: 8px;
    }
    
    /* Boxed answers */
    .katex .mord.text { font-weight: bold; }
    
    /* Chat message styling */
    .message { 
        margin: 10px 0; 
        padding: 12px; 
        border-radius: 8px;
    }
    .user { 
        background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
        border-left: 4px solid #2196f3;
    }
    .assistant { 
        background: linear-gradient(135deg, #f5f5f5 0%, #eeeeee 100%);
        border-left: 4px solid #4caf50;
    }
    
    /* Sidebar styling */
    .difficulty-selector { 
        background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%);
        padding: 15px; 
        border-radius: 10px; 
        margin: 10px 0;
        border: 1px solid #ffcc80;
    }
    
    .math-highlight { 
        background: linear-gradient(135deg, #e8f5e8 0%, #c8e6c9 100%);
        padding: 12px; 
        border-left: 4px solid #4caf50; 
        margin: 10px 0; 
        border-radius: 8px;
    }
    
    /* Responsive design */
    @media (max-width: 768px) {
        .katex { font-size: 1em !important; }
        .katex-display { font-size: 1.1em !important; }
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🧮 **Mathetics AI** - Advanced Mathematics Solver
    
    **Your Personal AI Math Tutor** | Specialized in step-by-step problem solving with beautiful LaTeX rendering
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot.render()
            
        with gr.Column(scale=1):
            with gr.Accordion("🎲 **Quick Actions**", open=True):
                difficulty_preset = gr.Dropdown(
                    choices=["Elementary", "High School", "College", "Competition"],
                    value="High School",
                    label="🎯 Problem Difficulty",
                    elem_classes=["difficulty-selector"]
                )
                
                sample_btn = gr.Button("🎯 Get Sample Problem", variant="secondary", size="sm")
                help_btn = gr.Button("❓ Math Help Tips", variant="secondary", size="sm")
                
                gr.Markdown("### 🔧 **Quick Tools**")
                gr.Markdown("""
                - **Algebra**: Equations, inequalities, factoring
                - **Geometry**: Area, volume, trigonometry  
                - **Calculus**: Derivatives, integrals, limits
                - **Statistics**: Probability, distributions
                - **Number Theory**: Prime factorization, GCD/LCM
                """)
    
    # Footer with enhanced info
    gr.Markdown("""
    ---
    **🔧 Technical Details:**
    - **Model**: Qwen/Qwen2.5-Math-7B-Instruct (Specialized for Mathematics)
    - **Features**: Real-time streaming • LaTeX rendering • Step-by-step solutions
    
    **💡 Usage Tips:**
    - Be specific about what you want to solve
    - Request step-by-step solutions explicitly  
    - Use "lim(x->0)" for limits (mobile-friendly)
    - Crank tokens to 1500+ for complex problems
    """)
    
    # Event handlers
    def handle_sample(difficulty):
        sample = get_random_sample()
        return sample
    
    sample_btn.click(
        handle_sample,
        inputs=[difficulty_preset],
        outputs=gr.Textbox(visible=False)  # This will trigger the chat input
    )
    
    help_btn.click(
        show_help,
        outputs=gr.Markdown(label="Help", visible=True)
    )

if __name__ == "__main__":
    demo.launch()