Skip to content

Commit

Permalink
Prevent *Is typing* from disappearing instantly while streaming
Browse files Browse the repository at this point in the history
  • Loading branch information
oobabooga committed Mar 15, 2023
1 parent 4146ac4 commit cf2da86
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions modules/text_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,8 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
reply = shared.model.generate(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k)
yield formatted_outputs(reply, shared.model_name)
else:
yield formatted_outputs(question, shared.model_name)
if not (shared.args.chat or shared.args.cai_chat):
yield formatted_outputs(question, shared.model_name)
# RWKV has proper streaming, which is very nice.
# No need to generate 8 tokens at a time.
for reply in shared.model.generate_with_streaming(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k):
Expand Down Expand Up @@ -197,7 +198,8 @@ def generate_with_callback(callback=None, **kwargs):
def generate_with_streaming(**kwargs):
return Iteratorize(generate_with_callback, kwargs, callback=None)

yield formatted_outputs(original_question, shared.model_name)
if not (shared.args.chat or shared.args.cai_chat):
yield formatted_outputs(original_question, shared.model_name)
with generate_with_streaming(**generate_params) as generator:
for output in generator:
if shared.soft_prompt:
Expand Down

0 comments on commit cf2da86

Please sign in to comment.