From 98fa73a9741d5db7f7518adc0a92249f65383d90 Mon Sep 17 00:00:00 2001 From: Brian Dashore Date: Sat, 7 Oct 2023 18:46:42 -0400 Subject: [PATCH] Text Generation: stop if EOS token is reached (#4213) --- modules/text_generation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/text_generation.py b/modules/text_generation.py index a7f3509..0f24dc5 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -353,10 +353,11 @@ def generate_reply_HF(question, original_question, seed, state, stopping_strings with generate_with_streaming(**generate_params) as generator: for output in generator: - yield get_reply_from_output_ids(output, input_ids, original_question, state, is_chat=is_chat) if output[-1] in eos_token_ids: break + yield get_reply_from_output_ids(output, input_ids, original_question, state, is_chat=is_chat) + except Exception: traceback.print_exc() finally: