Update app.py

app.py

@@ -170,6 +170,16 @@ def chat_with_model(messages, pid):
             updated_messages[-1]["content"] = output_text
             break
 
+        if "\nSystem" in output_text:
+            output_text = output_text.split("\nSystem")[0].rstrip()
+            updated_messages[-1]["content"] = output_text
+            break
+
+        if "\nAssistant" in output_text:
+            output_text = output_text.split("\nAssistant")[0].rstrip()
+            updated_messages[-1]["content"] = output_text
+            break
+
         generated_tokens += 1
         if generated_tokens >= max_new_tokens:
             break