Update app.py

fixed column issue with gradio 5

app.py

@@ -2,18 +2,16 @@ import gradio as gr
 import re
 from vllm import LLM, SamplingParams
 import spaces
-import torch
 
 class VibeThinkerVLLM:
     def __init__(self, model_path="WeiboAI/VibeThinker-1.5B"):
         self.model_path = model_path
         print("Loading model with vLLM... This may take a minute.")
         
-        # T4 GPU has compute capability 7.5, which doesn't support bfloat16
-        # We need to use float16 instead
+        # T4 GPU compatible - using float16
         self.model = LLM(
             model=self.model_path,
-            dtype="float16",  # Changed from bfloat16 to float16 for T4 compatibility
+            dtype="float16",
             gpu_memory_utilization=0.9,
             max_model_len=40960,
             trust_remote_code=True
@@ -46,9 +44,7 @@ class VibeThinkerVLLM:
 
 
 def parse_model_output(text):
-    """
-    Parse model output into structured components
-    """
+    """Parse model output into structured components"""
     sections = []
     
     # Patterns
@@ -124,21 +120,129 @@ def parse_text_with_code(text):
     return sections
 
 
+def format_sections_to_html(sections):
+    """
+    Convert parsed sections to rich HTML with collapsible elements
+    This approach works reliably with Gradio 5's HTML component
+    """
+    html_parts = []
+    
+    # Add JavaScript for interactivity
+    html_parts.append("""
+    <script>
+    function copyCode(elementId) {
+        const codeElement = document.getElementById(elementId);
+        const code = codeElement.textContent;
+        navigator.clipboard.writeText(code).then(() => {
+            // Show temporary success message
+            const btn = event.target;
+            const originalText = btn.textContent;
+            btn.textContent = '✅ Copied!';
+            setTimeout(() => { btn.textContent = originalText; }, 2000);
+        }).catch(err => {
+            console.error('Failed to copy:', err);
+            alert('Failed to copy code');
+        });
+    }
+    
+    function downloadCode(elementId, language) {
+        const codeElement = document.getElementById(elementId);
+        const code = codeElement.textContent;
+        
+        const extensions = {
+            'python': 'py', 'javascript': 'js', 'typescript': 'ts',
+            'html': 'html', 'css': 'css', 'java': 'java',
+            'cpp': 'cpp', 'c': 'c', 'ruby': 'rb',
+            'go': 'go', 'rust': 'rs', 'swift': 'swift',
+            'kotlin': 'kt', 'plaintext': 'txt'
+        };
+        
+        const ext = extensions[language.toLowerCase()] || 'txt';
+        const filename = `code_snippet.${ext}`;
+        
+        const blob = new Blob([code], { type: 'text/plain' });
+        const url = window.URL.createObjectURL(blob);
+        const a = document.createElement('a');
+        a.href = url;
+        a.download = filename;
+        document.body.appendChild(a);
+        a.click();
+        document.body.removeChild(a);
+        window.URL.revokeObjectURL(url);
+    }
+    </script>
+    """)
+    
+    for i, section in enumerate(sections):
+        if section['type'] == 'thinking':
+            # Collapsible thinking section
+            html_parts.append(f"""
+            <details class="thinking-section" style="margin: 15px 0; border: 2px solid #f39c12; border-radius: 8px; background-color: #fff9e6;">
+                <summary style="padding: 12px; cursor: pointer; font-weight: bold; color: #d68910; user-select: none;">
+                    🤔 Thinking Process (Click to expand)
+                </summary>
+                <div style="padding: 15px; border-top: 1px solid #f39c12; background-color: #fffef7; white-space: pre-wrap; font-family: 'Courier New', monospace; font-size: 13px; color: #333; line-height: 1.6; max-height: 500px; overflow-y: auto;">
+{section['content']}
+                </div>
+            </details>
+            """)
+        
+        elif section['type'] == 'code':
+            # Code block with copy/download buttons
+            code_id = f"code-{i}"
+            # Escape HTML in code
+            escaped_code = section['content'].replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
+            
+            html_parts.append(f"""
+            <details class="code-section" open style="margin: 15px 0; border: 2px solid #3498db; border-radius: 8px; background-color: #e8f4fd;">
+                <summary style="padding: 12px; cursor: pointer; font-weight: bold; color: #2874a6; user-select: none;">
+                    💻 Code ({section['language']}) - Click to collapse
+                </summary>
+                <div style="position: relative; padding: 0;">
+                    <div style="position: absolute; top: 10px; right: 10px; z-index: 10;">
+                        <button onclick="copyCode('{code_id}')" style="padding: 6px 12px; margin-right: 5px; background-color: #3498db; color: white; border: none; border-radius: 4px; cursor: pointer; font-size: 12px;">
+                            📋 Copy
+                        </button>
+                        <button onclick="downloadCode('{code_id}', '{section['language']}')" style="padding: 6px 12px; background-color: #27ae60; color: white; border: none; border-radius: 4px; cursor: pointer; font-size: 12px;">
+                            ⬇️ Download
+                        </button>
+                    </div>
+                    <pre id="{code_id}" style="margin: 0; padding: 40px 15px 15px 15px; background-color: #f8f9fa; border-top: 1px solid #3498db; overflow-x: auto; font-family: 'Courier New', monospace; font-size: 13px; line-height: 1.5;"><code class="language-{section['language']}">{escaped_code}</code></pre>
+                </div>
+            </details>
+            """)
+        
+        else:  # text
+            # Regular text output with markdown-style rendering
+            # Convert markdown to HTML
+            text_html = section['content']
+            # Basic markdown conversions
+            text_html = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', text_html)
+            text_html = re.sub(r'\*(.*?)\*', r'<em>\1</em>', text_html)
+            text_html = re.sub(r'`(.*?)`', r'<code style="background-color: #f4f4f4; padding: 2px 5px; border-radius: 3px;">\1</code>', text_html)
+            
+            html_parts.append(f"""
+            <div class="text-section" style="margin: 15px 0; padding: 15px; border: 1px solid #bdc3c7; border-radius: 8px; background-color: #ffffff; white-space: pre-wrap; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif; font-size: 14px; line-height: 1.8; color: #2c3e50;">
+{text_html}
+            </div>
+            """)
+    
+    return "\n".join(html_parts)
+
+
 # Initialize model
 print("Initializing VibeThinker-1.5B with vLLM...")
 model = VibeThinkerVLLM()
 
 
 def generate_response(prompt, temperature, max_tokens, top_p):
-    """
-    Generate and parse response - yields components for streaming display
-    """
+    """Generate and return formatted HTML response"""
     if not prompt.strip():
-        return [gr.Markdown("⚠️ Please enter a question.")]
+        return "<div style='color: #e74c3c; padding: 20px; text-align: center;'>⚠️ Please enter a question.</div>"
     
     try:
         # Show generating message
-        yield [gr.Markdown("🤖 Generating response...")]
+        yield "<div style='text-align: center; padding: 40px; color: #3498db;'><h3>🤖 Generating response...</h3><p>This may take a moment...</p></div>"
         
         # Generate raw response
         raw_response = model.infer_text(
@@ -151,52 +255,20 @@ def generate_response(prompt, temperature, max_tokens, top_p):
         # Parse the response
         sections = parse_model_output(raw_response)
         
-        # Build component list
-        output_components = []
+        # Convert to HTML
+        html_output = format_sections_to_html(sections)
         
-        for i, section in enumerate(sections):
-            if section['type'] == 'thinking':
-                # Thinking section with Accordion
-                output_components.append(
-                    gr.Accordion("🤔 Thinking Process", open=False)
-                )
-                output_components.append(
-                    gr.Textbox(
-                        value=section['content'],
-                        lines=20,
-                        max_lines=50,
-                        show_label=False,
-                        container=False,
-                        interactive=False,
-                        elem_classes=["thinking-box"]
-                    )
-                )
-            
-            elif section['type'] == 'code':
-                # Code section with native Code component
-                output_components.append(
-                    gr.Accordion(f"💻 Code ({section['language']})", open=True)
-                )
-                output_components.append(
-                    gr.Code(
-                        value=section['content'],
-                        language=section['language'],
-                        lines=20,
-                        show_label=False,
-                        interactive=False,
-                        elem_classes=["code-box"]
-                    )
-                )
-            
-            else:  # text
-                output_components.append(
-                    gr.Markdown(section['content'])
-                )
-        
-        yield output_components
+        yield html_output
     
     except Exception as e:
-        yield [gr.Markdown(f"❌ **Error:** {str(e)}\n\nPlease try again or adjust the parameters.")]
+        error_html = f"""
+        <div style='color: #e74c3c; padding: 20px; border: 2px solid #e74c3c; border-radius: 8px; background-color: #fadbd8; margin: 15px 0;'>
+            <h3>❌ Error</h3>
+            <p><strong>{str(e)}</strong></p>
+            <p>Please try again or adjust the parameters.</p>
+        </div>
+        """
+        yield error_html
 
 
 # Custom theme for Gradio 5
@@ -218,16 +290,6 @@ with gr.Blocks(
     title="VibeThinker-1.5B Advanced",
     theme=theme,
     fill_height=False,
-    css="""
-    .thinking-box {
-        font-family: 'Courier New', monospace;
-        font-size: 13px;
-        line-height: 1.6;
-    }
-    .code-box {
-        font-family: 'Fira Code', 'Courier New', monospace;
-    }
-    """,
 ) as demo:
     
     gr.Markdown("""
@@ -237,8 +299,8 @@ with gr.Blocks(
     
     ### ✨ Features:
     - 🤔 **Collapsible Thinking Sections** - Explore the model's reasoning process
-    - 💻 **Syntax-Highlighted Code** - Native code display with copy functionality
-    - 📝 **Clean Markdown Output** - Beautiful formatting for text responses
+    - 💻 **Interactive Code Blocks** - Copy or download code with one click
+    - 📝 **Clean Formatted Output** - Beautiful rendering for all content types
     
     **Best for:** Competitive math problems and algorithm coding challenges
     
@@ -295,18 +357,15 @@ with gr.Blocks(
                 )
         
         with gr.Column(scale=1):
-            # Output area
-            output_area = gr.Column()
-            
-            with output_area:
-                initial_message = gr.Markdown(
-                    """
-                    <div style='text-align: center; padding: 60px; color: #7f8c8d;'>
-                        <h3>👋 Ready to solve problems!</h3>
-                        <p>Enter your question and click Generate Solution</p>
-                    </div>
-                    """
-                )
+            # Output area using HTML component
+            output_html = gr.HTML(
+                value="""
+                <div style='text-align: center; padding: 60px; color: #7f8c8d;'>
+                    <h3>👋 Ready to solve problems!</h3>
+                    <p>Enter your question and click Generate Solution</p>
+                </div>
+                """
+            )
     
     # Example problems
     gr.Examples(
@@ -337,26 +396,24 @@ with gr.Blocks(
     
     # Event handlers
     def clear_interface():
-        return "", [gr.Markdown(
-            """
-            <div style='text-align: center; padding: 60px; color: #7f8c8d;'>
-                <h3>👋 Ready to solve problems!</h3>
-                <p>Enter your question and click Generate Solution</p>
-            </div>
-            """
-        )]
+        return "", """
+        <div style='text-align: center; padding: 60px; color: #7f8c8d;'>
+            <h3>👋 Ready to solve problems!</h3>
+            <p>Enter your question and click Generate Solution</p>
+        </div>
+        """
     
     submit_btn.click(
         fn=generate_response,
         inputs=[prompt_input, temperature_slider, max_tokens_slider, top_p_slider],
-        outputs=[output_area],
+        outputs=output_html,
         show_progress="full"
     )
     
     clear_btn.click(
         fn=clear_interface,
         inputs=[],
-        outputs=[prompt_input, output_area]
+        outputs=[prompt_input, output_html]
     )