Break on "User:"

app.py

@@ -92,6 +92,9 @@ def chat_with_model(messages):
             streamer.end_of_generation.set()  # signal to stop generation thread
             break
 
+        if "\nUser:" in token_str:
+            break
+
         if "<think>" in token_str:
             in_think = True
             token_str = token_str.replace("<think>", "")
@@ -118,6 +121,8 @@ def chat_with_model(messages):
     current_model.to("cpu")
     torch.cuda.empty_cache()
 
+    return messages
+
 
 # Globals
 current_model = None