Update app.py

app.py

@@ -72,7 +72,7 @@ def generate_response(user_message, max_tokens, temperature, top_p, history_stat
         {"role": "user", "content": user_message},
         {"role": "assistant", "content": ""}
     ]
-
+'''
     try:
         for new_token in streamer:
             if "<|end" in new_token:
@@ -84,6 +84,23 @@ def generate_response(user_message, max_tokens, temperature, top_p, history_stat
     except Exception:
         pass
 
+    yield new_history, new_history'''
+
+    try:
+        for new_token in streamer:
+            if "<|end" in new_token:
+                continue
+            cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
+            assistant_response += cleaned_token
+            new_history[-1]["content"] = assistant_response.strip()
+            yield new_history, new_history
+    except Exception:
+        new_history[-1]["content"] += "\n⚠️ Incomplete response due to generation limit or stream error."
+    yield new_history, new_history
+    return
+
+    # Final yield to clean up in case of clean stop
+    new_history[-1]["content"] = assistant_response.strip()
     yield new_history, new_history
 
 example_messages = {