clean up

app.py

@@ -126,7 +126,7 @@ def chat_with_model(messages):
 
         # Early stopping if user reappears
         if "\nUser:" in output_text:
-            output_text = output_text.split("\nUser:")[0].rstrip()
+            output_text = output_text.split("\nUser")[0].rstrip()
             messages[-1]["content"] = output_text
             break
 
@@ -151,14 +151,13 @@ def chat_with_model(messages):
         messages[-1]["content"] = output_text
     
     # Wait for thread to finish
-    thread.join(timeout=1.0)
     current_model.to("cpu")
     torch.cuda.empty_cache()
 
     messages[-1]["content"] = output_text
     print(f'Step 3: {messages}')
 
-    yield messages
+    return messages