Update app.py

app.py

@@ -12,16 +12,6 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 phi4_model = AutoModelForCausalLM.from_pretrained(phi4_model_path, device_map="auto", torch_dtype="auto")
 phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
 
-# Function to process LaTeX in the output
-def process_latex(text):
-    """
-    Process LaTeX equations in text:
-    - Inline equations: $...$
-    - Block equations: $$...$$
-    """
-    # No special processing needed as Gradio's markdown component supports LaTeX
-    return text
-
 @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():
@@ -34,8 +24,19 @@ def generate_response(user_message, max_tokens, temperature, top_k, top_p, repet
     sep_tag = "<|im_sep|>"
     end_tag = "<|im_end|>"
 
-    # Recommended prompt settings by Microsoft
-    system_message = "Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <think> {Thought section} </think> {Solution section}. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion. You can use LaTeX for mathematical expressions, enclosed in $ for inline equations and $$ for block equations. Now, try to solve the following question through the above guidelines:"
+    # Add a prompt to encourage LaTeX usage for mathematical expressions
+    system_message = """Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. 
+
+Please structure your response into two main sections: Thought and Solution using the specified format: <think> {Thought section} </think> {Solution section}. 
+
+In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. 
+
+In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion.
+
+IMPORTANT: When expressing mathematical formulas or equations, always use LaTeX format. Use single dollar signs for inline equations (e.g., $x^2$) and double dollar signs for block equations (e.g., $$\\frac{a}{b}$$). Ensure all mathematical symbols, fractions, square roots, and complex expressions are properly formatted in LaTeX.
+
+Now, try to solve the following question through the above guidelines:"""
+    
     prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
     for message in history_state:
         if message["role"] == "user":
@@ -46,8 +47,6 @@ def generate_response(user_message, max_tokens, temperature, top_k, top_p, repet
 
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-    do_sample = not (temperature == 1.0 and top_k >= 100 and top_p == 1.0)
-
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
 
     # sampling techniques
@@ -56,10 +55,10 @@ def generate_response(user_message, max_tokens, temperature, top_k, top_p, repet
         "attention_mask": inputs["attention_mask"],
         "max_new_tokens": int(max_tokens),
         "do_sample": True,
-        "temperature": 0.8,
+        "temperature": float(temperature),
         "top_k": int(top_k),
-        "top_p": 0.95,
-        "repetition_penalty": repetition_penalty,
+        "top_p": float(top_p),
+        "repetition_penalty": float(repetition_penalty),
         "streamer": streamer,
     }
 
@@ -72,39 +71,28 @@ def generate_response(user_message, max_tokens, temperature, top_k, top_p, repet
         {"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
-        # Process the response to handle LaTeX content
-        processed_response = process_latex(assistant_response.strip())
-        new_history[-1]["content"] = processed_response
+        new_history[-1]["content"] = assistant_response.strip()
         yield new_history, new_history
 
     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.",
-    "Math with LaTeX": "Solve the quadratic equation $ax^2 + bx + c = 0$ for $x$ using the quadratic formula. Then, solve the specific equation $3x^2 - 6x + 2 = 0$."
-}
-
-# Customize the Gradio CSS to enhance the LaTeX rendering
-css = """
-.math-container .katex { 
-    font-size: 1.2em;
-}
-.math-block .katex {
-    font-size: 1.5em;
-    display: block;
-    margin: 1em 0;
+    "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."
 }
-"""
 
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # Try the example problems below to see how the model breaks down complex reasoning problems with LaTeX support.
+        # Problem Solving with LaTeX Math Support
+        This application uses advanced reasoning to solve complex problems with LaTeX formatting for mathematical expressions.
         """
     )
     
@@ -148,14 +136,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 )
         
         with gr.Column(scale=4):
-            # Use Chatbot with latex_delimiters to enable LaTeX rendering
+            # Use the markdown flag to ensure proper rendering of LaTeX
             chatbot = gr.Chatbot(
                 label="Chat", 
-                type="messages",
-                latex_delimiters=[
-                    {"left": "$$", "right": "$$", "display": True},
-                    {"left": "$", "right": "$", "display": False}
-                ]
+                render_markdown=True,
+                elem_id="chatbot",
+                show_copy_button=True,
+                bubble_full_width=False,
+                avatar_images=(None, None)
             )
             with gr.Row():
                 user_input = gr.Textbox(
@@ -170,7 +158,49 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 example1_button = gr.Button("Math reasoning")
                 example2_button = gr.Button("Logic puzzle")
                 example3_button = gr.Button("Physics problem")
-                example4_button = gr.Button("Math with LaTeX")
+                example4_button = gr.Button("LaTeX example")
+
+    # Add custom JavaScript to ensure LaTeX rendering
+    demo.load(None, None, None, _js="""
+    function() {
+        // 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'
+            };
+        }
+        
+        // Function to render LaTeX in the chatbot
+        function renderMathInChatbot() {
+            if (window.MathJax && window.MathJax.Hub) {
+                window.MathJax.Hub.Queue(['Typeset', window.MathJax.Hub, document.getElementById('chatbot')]);
+            }
+        }
+        
+        // Set up a mutation observer to detect changes in the chatbot
+        const observer = new MutationObserver(renderMathInChatbot);
+        const chatbot = document.getElementById('chatbot');
+        if (chatbot) {
+            observer.observe(chatbot, { childList: true, subtree: true });
+        }
+        
+        // Initial render
+        renderMathInChatbot();
+    }
+    """)
 
     submit_button.click(
         fn=generate_response,
@@ -204,7 +234,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=user_input
     )
     example4_button.click(
-        fn=lambda: gr.update(value=example_messages["Math with LaTeX"]),
+        fn=lambda: gr.update(value=example_messages["LaTeX example"]),
         inputs=None,
         outputs=user_input
     )