From ad2970374adeb58aec1d7748b02a8c82cc524c0a Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 8 Mar 2023 03:00:06 -0300 Subject: [PATCH] Readability improvements --- modules/text_generation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/text_generation.py b/modules/text_generation.py index a8157a7..9477fe4 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -195,8 +195,8 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi for output in eval(f"generate_with_streaming({', '.join(generate_params)})"): if shared.soft_prompt: output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) - reply = decode(output) + if not (shared.args.chat or shared.args.cai_chat): reply = original_question + apply_extensions(reply[len(question):], "output") yield formatted_outputs(reply, shared.model_name) @@ -213,16 +213,16 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi output = eval(f"shared.model.generate({', '.join(generate_params)})")[0] if shared.soft_prompt: output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) - reply = decode(output) + if not (shared.args.chat or shared.args.cai_chat): reply = original_question + apply_extensions(reply[len(question):], "output") yield formatted_outputs(reply, shared.model_name) if np.count_nonzero(input_ids[0] == n) < np.count_nonzero(output == n): break - input_ids = np.reshape(output, (1, output.shape[0])) + input_ids = np.reshape(output, (1, output.shape[0])) if shared.soft_prompt: inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)