Update app.py

app.py

@@ -78,7 +78,7 @@ def chat_llama3_8b(message: str,
     print(tokenizer.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False))
     print(input_ids)
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=False, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
         input_ids= input_ids,
@@ -97,10 +97,11 @@ def chat_llama3_8b(message: str,
 
     outputs = []
     for text in streamer:
-        # if "<think>" in text:
-            # text = text.replace("<think>", "").strip()
-        # if "</think>" in text:
-            # text = text.replace("</think>", "").strip()
+        # Remove thinking tags to prevent Gradio display issues
+        if "<think>" in text:
+            text = text.replace("[think]", "").strip()
+        if "</think>" in text:
+            text = text.replace("[/think]", "").strip()
         outputs.append(text)
         print("".join(outputs))
         yield "".join(outputs)