Update app.py

app.py

@@ -18,6 +18,7 @@ model_id = "deepgrove/maple-preview"
 print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
+# Explicitly load chat template
 try:
     template_path = hf_hub_download(repo_id=model_id, filename="chat_template.jinja")
     with open(template_path, "r", encoding="utf-8") as f:
@@ -34,6 +35,7 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
 )
 
+# 2. Aggressively replace the buggy fa3 wrapper in EVERY loaded module namespace
 def custom_flash_attention_forward(
     query_states, key_states, value_states, attention_mask=None, query_length=None,
     is_causal=True, dropout=0.0, position_ids=None, softmax_scale=None, **kwargs
@@ -63,44 +65,20 @@ for name, module in list(sys.modules.items()):
         
 print(f"✅ Replaced _flash_attention_forward in {patched_count} modules.")
 
-
-def format_for_display(raw_text: str) -> str:
-    """Turn <think>...</think> into a collapsible <details> block for on-screen rendering.
-    This is display-only — never fed back into the model."""
-    display_text = raw_text
-    if "<think>" in display_text:
-        if "</think>" in display_text:
-            display_text = display_text.replace(
-                "<think>", "<details><summary>🧠 Thought Process</summary>\n\n"
-            ).replace(
-                "</think>", "\n\n</details>\n\n---"
-            )
-        else:
-            display_text = display_text.replace(
-                "<think>", "<details><summary>🧠 Thought Process (Thinking...)</summary>\n\n"
-            ) + "\n\n</details>"
-    return display_text
-
-
-def user_submit(message, history):
-    """Add the user's message to history and clear the textbox."""
-    history = history + [{"role": "user", "content": message}]
-    return "", history
-
-
 @spaces.GPU
-def bot_respond(history):
-    """Generate a reply for the latest user turn. `history` here is the RAW
-    conversation (real <think> tags, no HTML) — this is what gets templated."""
+def predict(message, history):
+    messages = history + [{"role": "user", "content": message}]
+    
     prompt = tokenizer.apply_chat_template(
-        history,
-        tokenize=False,
+        messages, 
+        tokenize=False, 
         add_generation_prompt=True
     )
+    
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-
+    
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
+    
     generation_kwargs = dict(
         **inputs,
         streamer=streamer,
@@ -110,44 +88,40 @@ def bot_respond(history):
         top_p=0.9,
         pad_token_id=tokenizer.pad_token_id,
     )
-
+    
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
-    raw_text = ""
-    history = history + [{"role": "assistant", "content": ""}]
-
+    
+    generated_text = ""
     for new_text in streamer:
-        raw_text += new_text
-        history[-1]["content"] = format_for_display(raw_text)  # HTML for the UI only
-        yield history
-
-    # Once streaming finishes, overwrite with the RAW text (real <think> tags)
-    # so the next turn's apply_chat_template call parses reasoning correctly.
-    history[-1]["content"] = raw_text
-    yield history
-
+        generated_text += new_text
+        
+        # Format thinking tags into a collapsible HTML dropdown element
+        formatted_text = generated_text
+        if "<think>" in formatted_text:
+            if "</think>" in formatted_text:
+                formatted_text = formatted_text.replace("<think>", "<details><summary>🧠 Thought Process</summary>\n\n").replace("</think>", "\n\n</details>\n\n---")
+            else:
+                formatted_text = formatted_text.replace("<think>", "<details><summary>🧠 Thought Process (Thinking...)</summary>\n\n") + "\n\n</details>"
+                
+        yield formatted_text
 
 # 3. Custom Chatbot component with Copy & Copy All features enabled
 chatbot_ui = gr.Chatbot(
     type="messages",
-    show_copy_button=True,
-    show_copy_all_button=True,
+    show_copy_button=True,        # Adds a copy button to individual messages
+    show_copy_all_button=True,    # Adds a button to copy the entire conversation history
     allow_file_downloads=True,
 )
 
-with gr.Blocks(theme="soft") as demo:
-    gr.Markdown("# 🍁 Maple-Preview")
-    gr.Markdown("A 20B-A1B ternary-weight reasoning LLM by DeepGrove. Powered by ZeroGPU.")
-
-    chatbot_ui.render()
-    msg = gr.Textbox(placeholder="Type a message...", show_label=False)
-
-    msg.submit(
-        user_submit, [msg, chatbot_ui], [msg, chatbot_ui]
-    ).then(
-        bot_respond, chatbot_ui, chatbot_ui
-    )
+demo = gr.ChatInterface(
+    fn=predict,
+    type="messages",
+    chatbot=chatbot_ui,
+    title="🍁 Maple-Preview",
+    description="A 20B-A1B ternary-weight reasoning LLM by DeepGrove. Powered by ZeroGPU.",
+    theme="soft",
+)
 
 if __name__ == "__main__":
     demo.queue().launch()