Update app.py

app.py

@@ -1,85 +1,83 @@
 import gradio as gr
+from huggingface_hub import InferenceClient
 import re
-from typing import List, Tuple
-import html
 
-def create_collapsible_html(thinking: str) -> str:
-    """Create HTML for a collapsible thinking section"""
-    escaped_thinking = html.escape(thinking)
-    return f'''
-    <div class="thinking-container" style="margin: 8px 0;">
-        <details class="thinking-details" style="
-            border: 1px solid #ddd;
-            border-radius: 4px;
-            padding: 8px;
-            background-color: #f8f9fa;
-            cursor: pointer;
-        ">
-            <summary style="
-                font-weight: 500;
-                display: flex;
-                align-items: center;
-                gap: 8px;
-            ">
-                <span style="transform: rotate(-90deg); display: inline-block;">➤</span>
-                Show thoughts
-            </summary>
-            <div style="
-                margin-top: 8px;
-                padding: 8px;
-                border-top: 1px solid #eee;
-                white-space: pre-wrap;
-            ">{escaped_thinking}</div>
-        </details>
-    </div>
-    '''
+client = InferenceClient("Pinkstack/Superthoughts-lite-v1")
 
-def process_message(message: str) -> str:
-    """Process a message to convert thinking sections into collapsible elements"""
-    def replace_thinking(match):
-        thinking_content = match.group(1).strip()
-        return create_collapsible_html(thinking_content)
+def process_think(text):
+    """
+    Searches for text enclosed in <think>...</think> and replaces it with a
+    collapsible HTML details block.
+    """
+    pattern = re.compile(r'<think>(.*?)</think>', re.DOTALL)
     
-    # Replace thinking sections with collapsible elements
-    processed = re.sub(r'<think>(.*?)</think>', replace_thinking, message, flags=re.DOTALL)
+    def replacer(match):
+        content = match.group(1).strip()
+        # You can adjust the inline styles or classes as needed.
+        return (
+            '<details class="think-details">'
+            '<summary class="think-summary">Show thoughts</summary>'
+            f'<div class="think-content">{content}</div>'
+            '</details>'
+        )
     
-    # Remove output tags if present
-    processed = re.sub(r'<output>(.*?)</output>', r'\1', processed, flags=re.DOTALL)
+    return pattern.sub(replacer, text)
+
+def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
+    # Build the conversation history with the system prompt first.
+    messages = [{"role": "system", "content": system_message}]
+    for user_text, bot_text in history:
+        if user_text:
+            messages.append({"role": "user", "content": user_text})
+        if bot_text:
+            messages.append({"role": "assistant", "content": bot_text})
+    messages.append({"role": "user", "content": message})
     
-    return processed
+    response = ""
+    # Stream the response from the client.
+    for chunk in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = chunk.choices[0].delta.content
+        response += token
+        # Process the response to swap <think> sections with collapsible blocks.
+        yield process_think(response)
 
-class CustomChatInterface(gr.ChatInterface):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        
-    def render_message(self, message: str, is_user: bool) -> str:
-        if not is_user:
-            return process_message(message)
-        return message
+# Custom CSS to style the collapsible block
+css = """
+.think-details {
+    border: 1px solid #ccc;
+    border-radius: 5px;
+    margin: 10px 0;
+    padding: 5px;
+}
+.think-summary {
+    cursor: pointer;
+    font-weight: bold;
+    background-color: #f1f1f1;
+    padding: 5px;
+    border-radius: 3px;
+    user-select: none;
+}
+"""
 
-demo = CustomChatInterface(
+# Note: The parameter for allowing HTML rendering may vary.
+# In many cases, the chat output component will render HTML if you pass allow_html=True.
+demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="You must always include <think> ... </think> <output> </output> tokens.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
-    css="""
-    .thinking-details[open] summary span {
-        transform: rotate(0deg) !important;
-    }
-    .thinking-details summary::-webkit-details-marker {
-        display: none;
-    }
-    """
+    css=css,
+    allow_html=True  # Ensures the HTML in the output is rendered.
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()