Update app.py

app.py

@@ -12,6 +12,35 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 
+# Helper function to format the streaming tokens with custom HTML/Markdown styling
+def format_reasoning(text):
+    """
+    Parses <think> tags and applies custom styles.
+    Uses HTML <details> for a clean, distinct 'accordion' look for reasoning.
+    """
+    if "<think>" in text:
+        if "</think>" in text:
+            # Thinking phase is complete
+            parts = text.split("</think>")
+            thinking_content = parts[0].replace("<think>", "").strip()
+            answer_content = parts[1].strip()
+            
+            return f"""<details open><summary style="color: #ff7a00; font-weight: bold; cursor: pointer;">🤔 Thinking Process (Click to collapse)</summary>
+<div style="color: #666; font-style: italic; background-color: #f9f9f9; padding: 10px; border-left: 3px solid #ff7a00; margin: 5px 0 15px 0;">
+{thinking_content}
+</div></details>
+
+{answer_content}"""
+        else:
+            # Currently thinking
+            thinking_content = text.replace("<think>", "").strip()
+            return f"""<details open><summary style="color: #ff7a00; font-weight: bold;">🤔 Thinking...</summary>
+<div style="color: #666; font-style: italic; background-color: #f9f9f9; padding: 10px; border-left: 3px solid #ff7a00; margin: 5px 0 15px 0;">
+{thinking_content}
+</div></details>"""
+    
+    return text
+
 # 2. Main Generation Logic Handler
 def generate_response(message, history, system_prompt, temperature, top_p, top_k, max_tokens):
     # Setup standard base template structures
@@ -24,7 +53,17 @@ def generate_response(message, history, system_prompt, temperature, top_p, top_k
         if role == "user":
             prompt += f"User: {content}\n"
         elif role == "assistant":
-            prompt += f"Assistant: {content}\n"
+            # Strip out our custom HTML formatting back to raw tokens for context history
+            # (Simple fallback: if the history already has html, you might want to keep a raw text history, 
+            # but for this generation block we clean basic tags if present)
+            clean_content = content
+            if "🤔" in content:
+                # Reconstruct rough raw tags for the model's history context
+                import re
+                clean_content = re.sub(r'<details.*?>.*?<\/summary>', '<think>', clean_content)
+                clean_content = clean_content.replace('</div></details>', '</think>')
+                clean_content = re.sub(r'<div.*?>', '', clean_content)
+            prompt += f"Assistant: {clean_content}\n"
     
     # Inject active raw message with Quasar's default target sequence patterns
     prompt += f"User: {message}\nAssistant: <think>"
@@ -48,7 +87,8 @@ def generate_response(message, history, system_prompt, temperature, top_p, top_k
     partial_text = "<think>"
     for new_text in streamer:
         partial_text += new_text
-        yield partial_text
+        # Yield the beautifully formatted text instead of raw markdown tags
+        yield format_reasoning(partial_text)
 
 # 3. Clean Interface Assembly 
 with gr.Blocks() as demo:
@@ -57,7 +97,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         # Left Workspace Column: Chat Layout
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=500)
+            # Line wrapping and markdown parsing enabled by default
+            chatbot = gr.Chatbot(height=500, type="messages")
             msg_input = gr.Textbox(placeholder="Enter reasoning prompt...", label="Your Message")
             with gr.Row():
                 submit_btn = gr.Button("Send", variant="primary")
@@ -84,8 +125,8 @@ with gr.Blocks() as demo:
     def bot_side_inference(history, system_p, temp, t_p, t_k, max_t):
         user_message = history[-1]["content"]
         
-        # Initialize assistant message dictionary element
-        history.append({"role": "assistant", "content": "<think>"})
+        # Initialize assistant message element with a loading state style
+        history.append({"role": "assistant", "content": "🤔 *Thinking...*"})
         
         for partial_reply in generate_response(user_message, history[:-2], system_p, temp, t_p, t_k, max_t):
             history[-1]["content"] = partial_reply