app.py

import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
import torch
from threading import Thread
import re

phi4_model_path = "Daemontatox/Manticore-32B"

device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Configure 4-bit quantization
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# Load model with 4-bit quantization
phi4_model = AutoModelForCausalLM.from_pretrained(
    phi4_model_path, 
    quantization_config=quantization_config,
    device_map="auto", 
    torch_dtype=torch.float16,
    trust_remote_code=True
)

phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path, trust_remote_code=True)

# Ensure pad token is set
if phi4_tokenizer.pad_token is None:
    phi4_tokenizer.pad_token = phi4_tokenizer.eos_token

@spaces.GPU(duration=120)
def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
    if not user_message.strip():
        return history_state, history_state
        
    # Phi-4 model settings
    model = phi4_model
    tokenizer = phi4_tokenizer
    start_tag = "<|im_start|>"
    sep_tag = "<|im_sep|>"
    end_tag = "<|im_end|>"

    # Add a prompt to encourage LaTeX usage for mathematical expressions
    system_message = """
   # Analytical Reasoning Assistant

You are an analytical reasoning assistant designed to solve problems through a rigorous, transparent thinking process. Your goal is to provide not just accurate solutions, but also to demonstrate the comprehensive reasoning path that led to those solutions.

## Response Structure

Structure your responses in two distinct sections:

<think>
Your detailed reasoning process, including all exploratory paths, considerations, and evaluations.
</think>

Your final, refined solution based on the analysis in the thinking section.

## Thinking Section Guidelines

In the <think> section:

1. **Question Analysis**: Carefully deconstruct the problem, identifying key variables, constraints, and objectives.
2. **Knowledge Activation**: Recall relevant principles, formulas, or methods needed for this type of problem.
3. **Strategy Formation**: Outline possible approaches to solving the problem.
4. **Exploration**: Work through each promising approach step-by-step, showing all calculations and logical steps.
5. **Critical Evaluation**: Assess the validity of each step and approach, identifying potential errors or limitations.
6. **Refinement**: Correct any mistakes and optimize your solution method.
7. **Verification**: Check your answer against the original constraints and through alternative methods when possible.

Document your entire thinking journey, including false starts and corrections. This transparency helps demonstrate the full problem-solving process.

## Solution Section Guidelines

In the solution section:

1. Present your final answer in a clear, structured format.
2. Include a concise explanation of the correct approach.
3. Provide the step-by-step solution using the most effective method identified.
4. Ensure the solution is complete, accurate, and directly addresses the original question.

## Mathematical Notation

When expressing mathematical content:

- Use LaTeX formatting for all mathematical expressions.
- For inline equations, use single dollar signs: $x^2 + y^2 = z^2$
- For displayed equations, use double dollar signs:
$$\int_{a}^{b} f(x) dx = F(b) - F(a)$$
- Format all mathematical symbols appropriately, including:
  - Fractions: $\frac{numerator}{denominator}$
  - Square roots: $\sqrt{expression}$
  - Exponents: $x^{power}$
  - Subscripts: $x_{index}$
  - Vectors: $\vec{v}$ or $\mathbf{v}$
  - Matrices: $\begin{bmatrix} a & b \\ c & d \end{bmatrix}$
- For complex mathematical expressions, prefer display mode for clarity.

## Multiple Solution Paths

When appropriate, explore multiple valid solution methods to demonstrate different perspectives on the problem. Compare their efficiency, elegance, and insights provided.

Remember that your purpose is not just to provide answers, but to model exemplary analytical thinking that helps users understand both the solution and the reasoning process.
Ensure the final answer is in LATEX format. 
{Latex Answer}

"""
    
    prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
    for message in history_state:
        if message["role"] == "user":
            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
        elif message["role"] == "assistant" and message["content"]:
            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
    prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"

    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=4096)
    
    # Move inputs to the same device as the model
    if torch.cuda.is_available():
        inputs = {k: v.to(device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    # sampling techniques
    generation_kwargs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "max_new_tokens": int(max_tokens),
        "do_sample": True,
        "temperature": float(temperature),
        "top_k": int(top_k),
        "top_p": float(top_p),
        "repetition_penalty": float(repetition_penalty),
        "streamer": streamer,
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Stream the response
    assistant_response = ""
    new_history = history_state + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": ""}
    ]
    
    for new_token in streamer:
        cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
        assistant_response += cleaned_token
        new_history[-1]["content"] = assistant_response.strip()
        yield new_history, new_history

    thread.join()  # Ensure thread completion
    yield new_history, new_history

# Add an example that explicitly shows LaTeX formatting
example_messages = {
    "Math reasoning": "If a rectangular prism has a length of 6 cm, a width of 4 cm, and a height of 5 cm, what is the length of the longest line segment that can be drawn from one vertex to another?",
    "Logic puzzle": "Four people (Alex, Blake, Casey, and Dana) each have a different favorite color (red, blue, green, yellow) and a different favorite fruit (apple, banana, cherry, date). Given the following clues: 1) The person who likes red doesn't like dates. 2) Alex likes yellow. 3) The person who likes blue likes cherries. 4) Blake doesn't like apples or bananas. 5) Casey doesn't like yellow or green. Who likes what color and what fruit?",
    "Physics problem": "A ball is thrown upward with an initial velocity of 15 m/s from a height of 2 meters above the ground. Assuming the acceleration due to gravity is 9.8 m/s², determine: 1) The maximum height the ball reaches. 2) The total time the ball is in the air before hitting the ground. 3) The velocity with which the ball hits the ground.",
    "LaTeX example": "Solve the quadratic equation ax^2 + bx + c = 0 and explain the solution. Then calculate the roots of 2x^2 - 5x + 3 = 0."
}

# Custom CSS for better LaTeX display
css = """
.markdown-body .katex { 
    font-size: 1.2em; 
}
.markdown-body .katex-display { 
    margin: 1em 0; 
    overflow-x: auto;
    overflow-y: hidden;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.Markdown(
        """
        # Problem Solving with LaTeX Math Support (4-bit Quantized)
        This application uses advanced reasoning to solve complex problems with LaTeX formatting for mathematical expressions.
        The model is loaded with 4-bit quantization to reduce memory usage.
        """
    )
    
    # Add JavaScript for MathJax loading
    gr.HTML("""
    <script>
    // Check if MathJax is available
    if (typeof window.MathJax === 'undefined') {
        // Load MathJax if not available
        const script = document.createElement('script');
        script.src = 'https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/MathJax.js?config=TeX-MML-AM_CHTML';
        script.async = true;
        document.head.appendChild(script);
        
        // Configure MathJax
        window.MathJax = {
            tex2jax: {
                inlineMath: [['$', '$']],
                displayMath: [['$$', '$$']],
                processEscapes: true
            },
            showProcessingMessages: false,
            messageStyle: 'none'
        };
    }
    
    // Set up a rerender function
    function rerender() {
        if (window.MathJax && window.MathJax.Hub) {
            window.MathJax.Hub.Queue(['Typeset', window.MathJax.Hub]);
        }
    }
    
    // Call rerender periodically
    setInterval(rerender, 1000);
    </script>
    """)
    
    history_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Settings")
            max_tokens_slider = gr.Slider(
                minimum=64,
                maximum=32768,
                step=1024,
                value=4096,
                label="Max Tokens"
            )
            with gr.Accordion("Advanced Settings", open=False):
                temperature_slider = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.8,
                    label="Temperature"
                )
                top_k_slider = gr.Slider(
                    minimum=1,
                    maximum=100,
                    step=1,
                    value=50,
                    label="Top-k"
                )
                top_p_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    label="Top-p"
                )
                repetition_penalty_slider = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    value=1.0,
                    label="Repetition Penalty"
                )
        
        with gr.Column(scale=4):
            # Use the markdown flag and type='messages' to ensure proper rendering of LaTeX
            chatbot = gr.Chatbot(
                label="Chat", 
                render_markdown=True,
                type="messages",
                elem_id="chatbot",
                show_copy_button=True,
                avatar_images=(None, None)
            )
            with gr.Row():
                user_input = gr.Textbox(
                    label="Your message",
                    placeholder="Type your message here...",
                    scale=3
                )
                submit_button = gr.Button("Send", variant="primary", scale=1)
                clear_button = gr.Button("Clear", scale=1)
            gr.Markdown("**Try these examples:**")
            with gr.Row():
                example1_button = gr.Button("Math reasoning")
                example2_button = gr.Button("Logic puzzle")
                example3_button = gr.Button("Physics problem")
                example4_button = gr.Button("LaTeX example")

    submit_button.click(
        fn=generate_response,
        inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],
        outputs=[chatbot, history_state]
    ).then(
        fn=lambda: gr.update(value=""),
        inputs=None,
        outputs=user_input
    )

    clear_button.click(
        fn=lambda: ([], []),
        inputs=None,
        outputs=[chatbot, history_state]
    )

    example1_button.click(
        fn=lambda: gr.update(value=example_messages["Math reasoning"]),
        inputs=None,
        outputs=user_input
    )
    example2_button.click(
        fn=lambda: gr.update(value=example_messages["Logic puzzle"]),
        inputs=None,
        outputs=user_input
    )
    example3_button.click(
        fn=lambda: gr.update(value=example_messages["Physics problem"]),
        inputs=None,
        outputs=user_input
    )
    example4_button.click(
        fn=lambda: gr.update(value=example_messages["LaTeX example"]),
        inputs=None,
        outputs=user_input
    )

demo.launch(ssr_mode=False)