Update app.py

app.py

@@ -88,7 +88,19 @@ example_messages = {
     "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()) as demo:
+# 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
@@ -96,6 +108,41 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         """
     )
     
+    # 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():
@@ -136,13 +183,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 )
         
         with gr.Column(scale=4):
-            # Use the markdown flag to ensure proper rendering of LaTeX
+            # 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,
-                bubble_full_width=False,
                 avatar_images=(None, None)
             )
             with gr.Row():
@@ -160,48 +207,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 example3_button = gr.Button("Physics problem")
                 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,
         inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],