From ebd698905c7b38aa9de387bd5bc89636bc6a2560 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 28 Feb 2023 00:04:04 -0300 Subject: [PATCH] Add streaming to RWKV --- modules/text_generation.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/modules/text_generation.py b/modules/text_generation.py index d879e14e..4e0056c6 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -87,9 +87,17 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi alpha_presence = 0.25, # Presence Penalty (as in GPT-3) token_ban = [0], # ban the generation of some tokens token_stop = []) # stop generation whenever you see any token here - reply = question + shared.model.generate(question, token_count=max_new_tokens, args=args, callback=None) - yield formatted_outputs(reply, None) - return formatted_outputs(reply, None) + + if shared.args.no_stream: + reply = question + shared.model.generate(question, token_count=max_new_tokens, args=args, callback=None) + yield formatted_outputs(reply, None) + return formatted_outputs(reply, None) + else: + for i in range(max_new_tokens//8): + reply = question + shared.model.generate(question, token_count=8, args=args, callback=None) + yield formatted_outputs(reply, None) + question = reply + return formatted_outputs(reply, None) original_question = question if not (shared.args.chat or shared.args.cai_chat):