Update demo.py

demo.py

@@ -14,7 +14,7 @@ DEFAULT_INSTRUCTIONS = f"""The following is a conversation between a highly know
 """
 RETRY_COMMAND = "/retry"
 STOP_STR = f"\n{USER_NAME}:"
-STOP_SUSPECT_LIST = [":", "\n", "User"]
+STOP_SUSPECT_LIST = [":", "\n", "User", "\nUser"]
 
 
 def chat_accordion():
@@ -98,12 +98,12 @@ def chat():
         if not message or (message == RETRY_COMMAND and len(chat_history) == 0):
             yield chat_history
             return
-    
+
         if message == RETRY_COMMAND and chat_history:
             prev_turn = chat_history.pop(-1)
             user_message, _ = prev_turn
             message = user_message
-    
+
         prompt = format_chat_prompt(message, chat_history, instructions)
         chat_history = chat_history + [[message, ""]]
         stream = llm(
@@ -117,26 +117,20 @@ def chat():
         acc_text = ""
         for idx, response in enumerate(stream):
             text_token = response
-    
-            if acc_text.endswith(STOP_STR):
-                last_turn = list(chat_history.pop(-1))
-                last_turn[-1] += acc_text[:-len(STOP_STR)]
-                chat_history = chat_history + [last_turn]
-                yield chat_history
-                acc_text = text_token
+
+            if text_token in STOP_SUSPECT_LIST:
+                acc_text += text_token
                 continue
-    
+
             if idx == 0 and text_token.startswith(" "):
                 text_token = text_token[1:]
-    
+
             acc_text += text_token
-    
-        # If there's any remaining text after the loop
-        if acc_text:
             last_turn = list(chat_history.pop(-1))
             last_turn[-1] += acc_text
             chat_history = chat_history + [last_turn]
             yield chat_history
+            acc_text = ""
 
 
     def delete_last_turn(chat_history):