Update app.py

app.py

@@ -15,102 +15,20 @@ AVAILABLE_MODELS = {
 }
 
 def parse_thinking_tags(text):
-    """Parse text and extract thinking sections"""
+    """Parse text and extract thinking sections, return clean text and thinking content"""
     think_pattern = r'<think>(.*?)</think>'
-    parts = []
-    last_end = 0
+    thinking_blocks = []
     
+    # Extract all thinking blocks
     for match in re.finditer(think_pattern, text, re.DOTALL):
-        # Add text before thinking block
-        if match.start() > last_end:
-            before_text = text[last_end:match.start()].strip()
-            if before_text:
-                parts.append({"type": "text", "content": before_text})
-        
-        # Add thinking block
         thinking_content = match.group(1).strip()
         if thinking_content:
-            parts.append({"type": "thinking", "content": thinking_content})
-        
-        last_end = match.end()
+            thinking_blocks.append(thinking_content)
     
-    # Add remaining text
-    if last_end < len(text):
-        remaining_text = text[last_end:].strip()
-        if remaining_text:
-            parts.append({"type": "text", "content": remaining_text})
+    # Remove thinking tags from the main text
+    clean_text = re.sub(think_pattern, '', text, flags=re.DOTALL).strip()
     
-    return parts
-
-def create_thinking_html(thinking_content, token_count=None):
-    """Create HTML for thinking block"""
-    if token_count is None:
-        token_count = f"{len(thinking_content.split()) * 4:,}"  # Rough token estimate
-    
-    thinking_id = f"thinking_{hash(thinking_content) % 10000}"
-    
-    html = f"""
-    <div class="thinking-container" style="margin: 16px 0; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;">
-        <div class="thinking-header" onclick="toggleThinking('{thinking_id}')" 
-             style="background: #3a3a3a; border-radius: 8px; padding: 12px 16px; cursor: pointer; 
-                    display: flex; align-items: center; gap: 10px; transition: all 0.2s ease;
-                    border: 1px solid #4a4a4a;">
-            <div style="width: 16px; height: 16px; background: #6366f1; border-radius: 50%; position: relative;">
-                <div style="position: absolute; top: 50%; left: 50%; width: 6px; height: 6px; 
-                           background: white; border-radius: 50%; transform: translate(-50%, -50%);"></div>
-            </div>
-            <span style="font-weight: 500; color: #f0f0f0;">Thinking completed</span>
-            <span style="color: #888; font-size: 14px; margin-left: auto;">{token_count} tokens budget</span>
-            <div class="chevron" id="chevron_{thinking_id}" 
-                 style="width: 16px; height: 16px; border: 2px solid #888; border-left: none; 
-                        border-bottom: none; transform: rotate(45deg); transition: transform 0.2s ease;"></div>
-        </div>
-        <div class="thinking-content" id="{thinking_id}" 
-             style="background: #333; border-radius: 0 0 8px 8px; padding: 20px; 
-                    border: 1px solid #4a4a4a; border-top: none; display: none; 
-                    max-height: 400px; overflow-y: auto;">
-            <div style="white-space: pre-wrap; font-family: 'SF Mono', 'Monaco', 'Inconsolata', monospace; 
-                        font-size: 14px; line-height: 1.5; color: #d0d0d0;">{thinking_content}</div>
-        </div>
-    </div>
-    
-    <script>
-    function toggleThinking(id) {{
-        const content = document.getElementById(id);
-        const chevron = document.getElementById('chevron_' + id);
-        const header = content.previousElementSibling;
-        
-        if (content.style.display === 'none' || !content.style.display) {{
-            content.style.display = 'block';
-            chevron.style.transform = 'rotate(135deg)';
-            header.style.background = '#404040';
-        }} else {{
-            content.style.display = 'none';
-            chevron.style.transform = 'rotate(45deg)';
-            header.style.background = '#3a3a3a';
-        }}
-    }}
-    </script>
-    """
-    return html
-
-def format_response_with_thinking(response_text):
-    """Format response to include thinking blocks"""
-    parts = parse_thinking_tags(response_text)
-    
-    if not parts:
-        return response_text
-    
-    formatted_html = ""
-    
-    for part in parts:
-        if part["type"] == "thinking":
-            formatted_html += create_thinking_html(part["content"])
-        else:
-            # Regular text content
-            formatted_html += f'<div style="margin: 16px 0; line-height: 1.6; color: #e0e0e0;">{part["content"]}</div>'
-    
-    return formatted_html
+    return clean_text, thinking_blocks
 
 @spaces.GPU
 def initialize_model(model_name):
@@ -151,7 +69,7 @@ def generate_response(message, history, model_name, max_length=512, temperature=
     try:
         model_pipe = initialize_model(model_name)
     except Exception as e:
-        return f"Error loading model {model_name}: {str(e)}"
+        return f"Error loading model {model_name}: {str(e)}", []
     
     # Format the conversation history
     messages = []
@@ -160,9 +78,7 @@ def generate_response(message, history, model_name, max_length=512, temperature=
     for user_msg, assistant_msg in history:
         messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
-            # Parse HTML back to get original content if needed
-            clean_assistant_msg = re.sub(r'<[^>]+>', '', assistant_msg) if assistant_msg else ""
-            messages.append({"role": "assistant", "content": clean_assistant_msg})
+            messages.append({"role": "assistant", "content": assistant_msg})
     
     # Add current message
     messages.append({"role": "user", "content": message})
@@ -216,13 +132,13 @@ def generate_response(message, history, model_name, max_length=512, temperature=
             if "Assistant:" in assistant_response:
                 assistant_response = assistant_response.split("Assistant:")[-1].strip()
         
-        # Format the response with thinking blocks
-        formatted_response = format_response_with_thinking(assistant_response)
+        # Parse thinking tags
+        clean_response, thinking_blocks = parse_thinking_tags(assistant_response)
         
-        return formatted_response
+        return clean_response, thinking_blocks
         
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        return f"Error generating response: {str(e)}", []
 
 @spaces.GPU
 def generate(
@@ -284,54 +200,13 @@ def generate(
     if system_prompt and system_prompt.strip():
         final_user_input = f"System: {system_prompt}\n\nUser: {user_input}"
     
-    # Use the original generate_response function
-    return generate_response(final_user_input, gradio_history, model, max_tokens, temperature, 0.9)
+    # Use the generate_response function and return only the clean response
+    clean_response, thinking_blocks = generate_response(final_user_input, gradio_history, model, max_tokens, temperature, 0.9)
+    return clean_response
 
 # Create the Gradio interface
 def create_interface():
-    # Custom CSS for dark theme and thinking blocks
-    custom_css = """
-    .dark {
-        background: #1a1a1a !important;
-    }
-    
-    .chatbot .message-wrap .message {
-        background: #2a2a2a !important;
-        border: 1px solid #3a3a3a !important;
-    }
-    
-    .chatbot .message-wrap .message.user {
-        background: #2d3748 !important;
-    }
-    
-    .chatbot .message-wrap .message.bot {
-        background: #2a2a2a !important;
-    }
-    
-    .thinking-container {
-        margin: 16px 0;
-        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
-    }
-    
-    .thinking-header {
-        background: #3a3a3a;
-        border-radius: 8px;
-        padding: 12px 16px;
-        cursor: pointer;
-        display: flex;
-        align-items: center;
-        gap: 10px;
-        transition: all 0.2s ease;
-        border: 1px solid #4a4a4a;
-    }
-    
-    .thinking-header:hover {
-        background: #404040 !important;
-        border-color: #5a5a5a !important;
-    }
-    """
-    
-    with gr.Blocks(title="Multi-Model Chat", theme=gr.themes.Soft(), css=custom_css) as demo:
+    with gr.Blocks(title="Multi-Model Chat", theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # 🚀 Daedalus-1 Model Chat Interface
         
@@ -349,13 +224,24 @@ def create_interface():
                 info="Choose which model to use for generation"
             )
         
-        chatbot = gr.Chatbot(
-            height=500,
-            placeholder="Select a model and start chatting...",
-            label="Chat",
-            render_markdown=False,  # Disable markdown to allow custom HTML
-            elem_classes=["chatbot"]
-        )
+        with gr.Row():
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot(
+                    height=500,
+                    placeholder="Select a model and start chatting...",
+                    label="Chat"
+                )
+            
+            with gr.Column(scale=1):
+                thinking_display = gr.Accordion("💭 Thinking Process", open=True, visible=False)
+                with thinking_display:
+                    thinking_content = gr.Textbox(
+                        label="Model's Thinking",
+                        lines=15,
+                        interactive=False,
+                        show_label=False,
+                        container=False
+                    )
         
         msg = gr.Textbox(
             placeholder="Type your message here...",
@@ -393,29 +279,6 @@ def create_interface():
                 info="Controls diversity via nucleus sampling"
             )
         
-        # Add JavaScript for thinking block interactions
-        gr.HTML("""
-        <script>
-        function toggleThinking(id) {
-            const content = document.getElementById(id);
-            if (!content) return;
-            
-            const chevron = document.getElementById('chevron_' + id);
-            const header = content.previousElementSibling;
-            
-            if (content.style.display === 'none' || !content.style.display) {
-                content.style.display = 'block';
-                if (chevron) chevron.style.transform = 'rotate(135deg)';
-                if (header) header.style.background = '#404040';
-            } else {
-                content.style.display = 'none';
-                if (chevron) chevron.style.transform = 'rotate(45deg)';
-                if (header) header.style.background = '#3a3a3a';
-            }
-        }
-        </script>
-        """)
-        
         # Event handlers
         def user_message(message, history):
             return "", history + [[message, None]]
@@ -423,7 +286,7 @@ def create_interface():
         def bot_response(history, model_name, max_len, temp, top_p):
             if history:
                 user_message = history[-1][0]
-                bot_message = generate_response(
+                clean_response, thinking_blocks = generate_response(
                     user_message, 
                     history[:-1], 
                     model_name,
@@ -431,27 +294,41 @@ def create_interface():
                     temp, 
                     top_p
                 )
-                history[-1][1] = bot_message
-            return history
+                
+                # Update chat with clean response
+                history[-1][1] = clean_response
+                
+                # Format thinking content for display
+                thinking_text = ""
+                if thinking_blocks:
+                    for i, thinking in enumerate(thinking_blocks, 1):
+                        thinking_text += f"=== Thinking Block {i} ===\n\n{thinking}\n\n"
+                
+                return history, thinking_text, gr.update(visible=bool(thinking_blocks))
+            
+            return history, "", gr.update(visible=False)
+        
+        def clear_chat():
+            return None, "", gr.update(visible=False)
         
         def model_changed(model_name):
             return gr.update(placeholder=f"Chat with {model_name}...")
         
         # Wire up the events
-        msg.submit(user_message, [msg, chatbot], [msg, chatbot]).then(
-            bot_response, [chatbot, model_selector, max_length, temperature, top_p], chatbot
+        msg_submit = msg.submit(user_message, [msg, chatbot], [msg, chatbot]).then(
+            bot_response, [chatbot, model_selector, max_length, temperature, top_p], 
+            [chatbot, thinking_content, thinking_display]
         )
         
-        submit_btn.click(user_message, [msg, chatbot], [msg, chatbot]).then(
-            bot_response, [chatbot, model_selector, max_length, temperature, top_p], chatbot
+        btn_submit = submit_btn.click(user_message, [msg, chatbot], [msg, chatbot]).then(
+            bot_response, [chatbot, model_selector, max_length, temperature, top_p], 
+            [chatbot, thinking_content, thinking_display]
         )
         
-        clear_btn.click(lambda: None, None, chatbot, queue=False)
+        clear_btn.click(clear_chat, None, [chatbot, thinking_content, thinking_display], queue=False)
         
         model_selector.change(model_changed, model_selector, chatbot)
         
-
-        
         gr.Markdown("""
         ---
         
@@ -463,7 +340,7 @@ def create_interface():
 
         
         The model is designed for conversational AI and supports various text generation tasks. When the model uses thinking tags (`<think></think>`), 
-        this interface will render them as expandable sections similar to Claude's thinking process visualization.
+        this interface will show the thinking process in a separate panel on the right.
         
         This Space uses ZeroGPU for efficient GPU allocation.
         """)