1 files changed, 5 insertions, 2 deletions
diff --git a/g4f/locals/provider.py b/g4f/locals/provider.py
index 45041539..d9d73455 100644
--- a/g4f/locals/provider.py
+++ b/g4f/locals/provider.py
@@ -66,9 +66,12 @@ class LocalProvider:
             if message["role"] != "system"
         ) + "\nASSISTANT: "
 
+        def should_not_stop(token_id: int, token: str):
+            return "USER" not in token
+
         with model.chat_session(system_message, prompt_template):
             if stream:
-                for token in model.generate(conversation, streaming=True):
+                for token in model.generate(conversation, streaming=True, callback=should_not_stop):
                     yield token
             else:
-                yield model.generate(conversation)
-\ No newline at end of file
+                yield model.generate(conversation, callback=should_not_stop)
+\ No newline at end of file