From 21512e2790a88c3508884a497eb8aa208137c524 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 16 Feb 2023 21:21:45 -0300 Subject: [PATCH] Make the Stop button work more reliably --- server.py | 50 +++++++++++++++++++++++++++++++++++--------------- 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/server.py b/server.py index a540ca45..13d67b41 100644 --- a/server.py +++ b/server.py @@ -494,7 +494,7 @@ def create_settings_menus(): model_menu.change(load_model_wrapper, [model_menu], [model_menu], show_progress=True) preset_menu.change(load_preset_values, [preset_menu], [do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping]) softprompts_menu.change(load_soft_prompt, [softprompts_menu], [softprompts_menu], show_progress=True) - upload_softprompt.upload(upload_soft_prompt, [upload_softprompt], [softprompts_menu]) + upload_softprompt.change(upload_soft_prompt, [upload_softprompt], [softprompts_menu]) return preset_menu, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping # This gets the new line characters right. @@ -572,7 +572,14 @@ def generate_chat_picture(picture, name1, name2): visible_text = f'' return text, visible_text +def stop_everything_event(): + global stop_everything + stop_everything = True + def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): + global stop_everything + stop_everything = False + if 'pygmalion' in model_name.lower(): name1 = "You" @@ -583,13 +590,24 @@ def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, text = apply_extensions(text, "input") question = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size) - history['internal'].append(['', '']) - history['visible'].append(['', '']) eos_token = '\n' if check else None + first = True for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name1}:"): reply, next_character_found, substring_found = extract_message_from_reply(question, reply, name2, name1, check, extensions=True) + visible_reply = apply_extensions(reply, "output") + + # We need this global variable to handle the Stop event, + # otherwise gradio gets confused + if stop_everything: + return history['visible'] + + if first: + first = False + history['internal'].append(['', '']) + history['visible'].append(['', '']) + history['internal'][-1] = [text, reply] - history['visible'][-1] = [visible_text, apply_extensions(reply, "output")] + history['visible'][-1] = [visible_text, visible_reply] if not substring_found: yield history['visible'] if next_character_found: @@ -859,8 +877,10 @@ else: print() model_name = available_models[i] model, tokenizer = load_model(model_name) -loaded_preset = soft_prompt_tensor = None +loaded_preset = None +soft_prompt_tensor = None soft_prompt = False +stop_everything = False # UI settings if model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')): @@ -960,10 +980,10 @@ if args.chat or args.cai_chat: gen_events.append(buttons["Generate"].click(eval(function_call), input_params, display, show_progress=args.no_stream, api_name="textgen")) gen_events.append(textbox.submit(eval(function_call), input_params, display, show_progress=args.no_stream)) if args.picture: - picture_select.upload(eval(function_call), input_params, display, show_progress=args.no_stream) + gen_events.append(picture_select.change(eval(function_call), input_params, display, show_progress=args.no_stream)) gen_events.append(buttons["Regenerate"].click(regenerate_wrapper, input_params, display, show_progress=args.no_stream)) gen_events.append(buttons["Impersonate"].click(impersonate_wrapper, input_params, textbox, show_progress=args.no_stream)) - buttons["Stop"].click(None, None, None, cancels=gen_events) + buttons["Stop"].click(stop_everything_event, [], [], cancels=gen_events) 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) @@ -981,17 +1001,17 @@ if args.chat or args.cai_chat: textbox.submit(lambda : save_history(timestamp=False), [], [], show_progress=False) character_menu.change(load_character, [character_menu, name1, name2], [name2, context, display]) - upload_img_tavern.upload(upload_tavern_character, [upload_img_tavern, name1, name2], [character_menu]) - upload.upload(load_history, [upload, name1, name2], []) - upload_img_me.upload(upload_your_profile_picture, [upload_img_me], []) + upload_img_tavern.change(upload_tavern_character, [upload_img_tavern, name1, name2], [character_menu]) + upload.change(load_history, [upload, name1, name2], []) + upload_img_me.change(upload_your_profile_picture, [upload_img_me], []) if args.picture: - picture_select.upload(lambda : None, [], [picture_select], show_progress=False) + picture_select.change(lambda : None, [], [picture_select], show_progress=False) if args.cai_chat: - upload.upload(redraw_html, [name1, name2], [display]) - upload_img_me.upload(redraw_html, [name1, name2], [display]) + upload.change(redraw_html, [name1, name2], [display]) + upload_img_me.change(redraw_html, [name1, name2], [display]) else: - upload.upload(lambda : history['visible'], [], [display]) - upload_img_me.upload(lambda : history['visible'], [], [display]) + upload.change(lambda : history['visible'], [], [display]) + upload_img_me.change(lambda : history['visible'], [], [display]) elif args.notebook: with gr.Blocks(css=css, analytics_enabled=False) as interface: