Update app.py

app.py

@@ -1,82 +1,91 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-import re
 
 client = InferenceClient("Pinkstack/Superthoughts-lite-v1")
 
-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)
-    
-    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>'
-        )
-    
-    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.
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
     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})
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
     messages.append({"role": "user", "content": message})
-    
     response = ""
-    # Stream the response from the client.
-    for chunk in client.chat_completion(
+    for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        token = chunk.choices[0].delta.content
+        token = message.choices[0].delta.content
         response += token
-        # Process the response to swap <think> sections with collapsible blocks.
-        yield process_think(response)
+        yield response
+
+def format_response(response):
+    # Replace <think>...</think> with a collapsible section
+    response = response.replace("<think>", '<details><summary>Show thoughts</summary><div class="thoughts">')
+    response = response.replace("</think>", "</div></details>")
+    return response
 
-# Custom CSS to style the collapsible block
 css = """
-.think-details {
+.thoughts {
     border: 1px solid #ccc;
+    padding: 10px;
+    background-color: #f9f9f9;
     border-radius: 5px;
-    margin: 10px 0;
-    padding: 5px;
 }
-.think-summary {
+details summary {
     cursor: pointer;
-    font-weight: bold;
-    background-color: #f1f1f1;
     padding: 5px;
-    border-radius: 3px;
-    user-select: none;
+    background-color: #e0e0e0;
+    border-radius: 5px;
+    font-weight: bold;
+}
+details summary::-webkit-details-marker {
+    display: none;
+}
+details summary:after {
+    content: " ▶";
+}
+details[open] summary:after {
+    content: " ▼";
 }
 """
 
-# 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)"),
-    ],
-    css=css,
-)
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("## Chat with Superthoughts")
+    gr.Markdown("**Warning:** The first output from the AI may take a few moments. After the first message, it should work quickly.")
+
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    system_message = gr.Textbox(value="You must always include <think> ... </think> <output> </output> tokens.", label="System message")
+    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
+    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
+
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history, system_message, max_tokens, temperature, top_p):
+        user_message, _ = history[-1]
+        response = ""
+        for partial_response in respond(user_message, history[:-1], system_message, max_tokens, temperature, top_p):
+            response = partial_response
+        formatted_response = format_response(response)
+        history[-1][1] = formatted_response
+        return history
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, system_message, max_tokens, temperature, top_p], chatbot
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()