summaryrefslogtreecommitdiffstats
path: root/g4f/locals/provider.py
diff options
context:
space:
mode:
authorHeiner Lohaus <hlohaus@users.noreply.github.com>2024-04-13 18:02:47 +0200
committerHeiner Lohaus <hlohaus@users.noreply.github.com>2024-04-13 18:02:47 +0200
commit0a3fe0b7df7b2779ae795c8590d2a707d6672345 (patch)
treeec033e0233d18846cee37196e2a2615377af7c49 /g4f/locals/provider.py
parentFix GPT4All import error (diff)
downloadgpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar.gz
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar.bz2
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar.lz
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar.xz
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.tar.zst
gpt4free-0a3fe0b7df7b2779ae795c8590d2a707d6672345.zip
Diffstat (limited to 'g4f/locals/provider.py')
-rw-r--r--g4f/locals/provider.py7
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