From 00707a0b3bb922ce90af5ed7445c1122102fd517 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Sun, 29 Jan 2023 22:56:23 -0300 Subject: [PATCH] Add "Impersonate" button --- server.py | 76 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 44 insertions(+), 32 deletions(-) diff --git a/server.py b/server.py index b17a2f1a..5f6b0247 100644 --- a/server.py +++ b/server.py @@ -339,7 +339,7 @@ if args.chat or args.cai_chat: text = text.strip() return text - def generate_chat_prompt(text, tokens, name1, name2, context, history_size): + def generate_chat_prompt(text, tokens, name1, name2, context, history_size, impersonate=False): text = clean_chat_message(text) rows = [f"{context.strip()}\n"] @@ -354,16 +354,44 @@ if args.chat or args.cai_chat: i -= 1 if history_size != 0 and count >= history_size: break - rows.append(f"{name1}: {text}\n") - rows.append(apply_extensions(f"{name2}:", "bot_prefix")) - while len(rows) > 3 and len(encode(''.join(rows), tokens)[0]) >= 2048-tokens: + if not impersonate: + rows.append(f"{name1}: {text}\n") + rows.append(apply_extensions(f"{name2}:", "bot_prefix")) + limit = 3 + else: + rows.append(f"{name1}:") + limit = 2 + + while len(rows) > limit and len(encode(''.join(rows), tokens)[0]) >= 2048-tokens: rows.pop(1) rows.pop(1) question = ''.join(rows) return question + def extract_message_from_reply(question, reply, current, other, check, extensions=False): + next_character_found = False + previous_idx = [m.start() for m in re.finditer(f"(^|\n){current}:", question)] + idx = [m.start() for m in re.finditer(f"(^|\n){current}:", reply)] + idx = idx[len(previous_idx)-1] + + if extensions: + reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):] + else: + reply = reply[idx + 1 + len(f"{current}:"):] + + if check: + reply = reply.split('\n')[0].strip() + else: + idx = reply.find(f"\n{other}:") + if idx != -1: + reply = reply[:idx] + next_character_found = True + reply = clean_chat_message(reply) + + return reply, next_character_found + def chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size): original_text = text text = apply_extensions(text, "input") @@ -372,39 +400,21 @@ if args.chat or args.cai_chat: history['visible'].append(['', '']) eos_token = '\n' if check else None for reply in generate_reply(question, tokens, inference_settings, selected_model, eos_token=eos_token, stopping_string=f"\n{name1}:"): - next_character_found = False - - previous_idx = [m.start() for m in re.finditer(f"(^|\n){name2}:", question)] - idx = [m.start() for m in re.finditer(f"(^|\n){name2}:", reply)] - idx = idx[len(previous_idx)-1] - - reply = reply[idx + 1 + len(apply_extensions(f"{name2}:", "bot_prefix")):] - if check: - reply = reply.split('\n')[0].strip() - else: - idx = reply.find(f"\n{name1}:") - if idx != -1: - reply = reply[:idx] - next_character_found = True - reply = clean_chat_message(reply) - + reply, next_character_found = extract_message_from_reply(question, reply, name2, name1, check, extensions=True) history['internal'][-1] = [text, reply] history['visible'][-1] = [original_text, apply_extensions(reply, "output")] + yield history['visible'] if next_character_found: break - # Prevent the chat log from flashing if something like "\nYo" is generated just - # before "\nYou:" is completed - tmp = f"\n{name1}:" - next_character_substring_found = False - for j in range(1, len(tmp)): - if reply[-j:] == tmp[:j]: - next_character_substring_found = True - - if not next_character_substring_found: - yield history['visible'] - - yield history['visible'] + def impersonate_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size): + question = generate_chat_prompt(text, tokens, name1, name2, context, history_size, impersonate=True) + eos_token = '\n' if check else None + for reply in generate_reply(question, tokens, inference_settings, selected_model, eos_token=eos_token, stopping_string=f"\n{name2}:"): + reply, next_character_found = extract_message_from_reply(question, reply, name1, name2, check, extensions=False) + yield apply_extensions(reply, "output") + if next_character_found: + break def cai_chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size): for _history in chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size): @@ -614,6 +624,7 @@ if args.chat or args.cai_chat: with gr.Row(): buttons["Send last reply to input"] = gr.Button("Send last reply to input") buttons["Replace last reply"] = gr.Button("Replace last reply") + buttons["Impersonate"] = gr.Button("Impersonate") with gr.Row(): with gr.Column(): @@ -674,6 +685,7 @@ if args.chat or args.cai_chat: buttons["Send last reply to input"].click(send_last_reply_to_input, [], textbox, show_progress=args.no_stream) buttons["Replace last reply"].click(replace_last_reply, [textbox, name1, name2], display, show_progress=args.no_stream) + buttons["Impersonate"].click(impersonate_wrapper, input_params, textbox, show_progress=args.no_stream) buttons["Clear"].click(clear_chat_log, [character_menu, name1, name2], display) buttons["Remove last"].click(remove_last_message, [name1, name2], [display, textbox], show_progress=False) buttons["Stop"].click(None, None, None, cancels=gen_events)