diff --git a/modules/chat.py b/modules/chat.py index a297d84c..0e6318f6 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -24,16 +24,16 @@ def clean_chat_message(text): text = text.strip() return text -def generate_chat_prompt(user_input, tokens, name1, name2, context, chat_prompt_size, impersonate=False): +def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat_prompt_size, impersonate=False): user_input = clean_chat_message(user_input) rows = [f"{context.strip()}\n"] if shared.soft_prompt: chat_prompt_size -= shared.soft_prompt_tensor.shape[1] - max_length = min(get_max_prompt_length(tokens), chat_prompt_size) + max_length = min(get_max_prompt_length(max_new_tokens), chat_prompt_size) i = len(shared.history['internal'])-1 - while i >= 0 and len(encode(''.join(rows), tokens)[0]) < max_length: + while i >= 0 and len(encode(''.join(rows), max_new_tokens)[0]) < max_length: rows.insert(1, f"{name2}: {shared.history['internal'][i][1].strip()}\n") if not (shared.history['internal'][i][0] == '<|BEGIN-VISIBLE-CHAT|>'): rows.insert(1, f"{name1}: {shared.history['internal'][i][0].strip()}\n") @@ -47,7 +47,7 @@ def generate_chat_prompt(user_input, tokens, name1, name2, context, chat_prompt_ rows.append(f"{name1}:") limit = 2 - while len(rows) > limit and len(encode(''.join(rows), tokens)[0]) >= max_length: + while len(rows) > limit and len(encode(''.join(rows), max_new_tokens)[0]) >= max_length: rows.pop(1) prompt = ''.join(rows) @@ -95,7 +95,7 @@ def generate_chat_picture(picture, name1, name2): def stop_everything_event(): shared.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): +def chatbot_wrapper(text, max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture=None): shared.stop_everything = False just_started = True eos_token = '\n' if check else None @@ -110,10 +110,10 @@ def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, if shared.args.chat: visible_text = visible_text.replace('\n', '
') text = apply_extensions(text, "input") - prompt = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size) + prompt = generate_chat_prompt(text, max_new_tokens, name1, name2, context, chat_prompt_size) # Generate - for reply in generate_reply(prompt, 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}:"): + for reply in generate_reply(prompt, max_new_tokens, 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, eos_token=eos_token, stopping_string=f"\n{name1}:"): # Extracting the reply reply, next_character_found, substring_found = extract_message_from_reply(prompt, reply, name2, name1, check, extensions=True) @@ -138,15 +138,15 @@ def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, break yield shared.history['visible'] -def impersonate_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): +def impersonate_wrapper(text, max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture=None): eos_token = '\n' if check else None if 'pygmalion' in shared.model_name.lower(): name1 = "You" - prompt = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size, impersonate=True) + prompt = generate_chat_prompt(text, max_new_tokens, name1, name2, context, chat_prompt_size, impersonate=True) - for reply in generate_reply(prompt, 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{name2}:"): + for reply in generate_reply(prompt, max_new_tokens, 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, eos_token=eos_token, stopping_string=f"\n{name2}:"): reply, next_character_found, substring_found = extract_message_from_reply(prompt, reply, name1, name2, check, extensions=False) if not substring_found: yield reply @@ -154,11 +154,11 @@ def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, to break yield reply -def cai_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): - for _history in 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): +def cai_chatbot_wrapper(text, max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture=None): + for _history in chatbot_wrapper(text, max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture): yield generate_chat_html(_history, name1, name2, shared.character) -def regenerate_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): +def regenerate_wrapper(text, max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture=None): if shared.character != 'None' and len(shared.history['visible']) == 1: if shared.args.cai_chat: yield generate_chat_html(shared.history['visible'], name1, name2, shared.character) @@ -168,7 +168,7 @@ def regenerate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top last_visible = shared.history['visible'].pop() last_internal = shared.history['internal'].pop() - for _history in chatbot_wrapper(last_internal[0], 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): + for _history in chatbot_wrapper(last_internal[0], max_new_tokens, 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, name1, name2, context, check, chat_prompt_size, picture): if shared.args.cai_chat: shared.history['visible'][-1] = [last_visible[0], _history[-1][1]] yield generate_chat_html(shared.history['visible'], name1, name2, shared.character) diff --git a/modules/text_generation.py b/modules/text_generation.py index 81e8da9f..834101ff 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -72,14 +72,14 @@ def formatted_outputs(reply, model_name): else: return reply -def 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=None, stopping_string=None): +def generate_reply(question, max_new_tokens, 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, eos_token=None, stopping_string=None): original_question = question if not (shared.args.chat or shared.args.cai_chat): question = apply_extensions(question, "input") if shared.args.verbose: print(f"\n\n{question}\n--------------------\n") - input_ids = encode(question, tokens) + input_ids = encode(question, max_new_tokens) cuda = "" if (shared.args.cpu or shared.args.deepspeed or shared.args.flexgen) else ".cuda()" if not shared.args.flexgen: n = shared.tokenizer.eos_token_id if eos_token is None else shared.tokenizer.encode(eos_token, return_tensors='pt')[0][-1] @@ -126,7 +126,7 @@ def generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top if shared.args.deepspeed: generate_params.append("synced_gpus=True") if shared.args.no_stream: - generate_params.append("max_new_tokens=tokens") + generate_params.append("max_new_tokens=max_new_tokens") else: generate_params.append("max_new_tokens=8") @@ -156,7 +156,7 @@ def generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top # Generate the reply 8 tokens at a time else: yield formatted_outputs(original_question, shared.model_name) - for i in tqdm(range(tokens//8+1)): + for i in tqdm(range(max_new_tokens//8+1)): with torch.no_grad(): output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0] if shared.soft_prompt: diff --git a/server.py b/server.py index c63127df..74ecf76f 100644 --- a/server.py +++ b/server.py @@ -252,7 +252,7 @@ if shared.args.chat or shared.args.cai_chat: with gr.Tab("Extensions"): extensions_module.create_extensions_block() - input_params = [shared.gradio[i] for i in ['textbox', 'max_new_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_slider']] + input_params = [shared.gradio[i] for i in ['textbox', 'max_new_tokens', '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', 'name1', 'name2', 'context', 'check', 'chat_prompt_size_slider']] if shared.args.picture: input_params.append(shared.gradio['picture_select']) function_call = "chat.cai_chatbot_wrapper" if shared.args.cai_chat else "chat.chatbot_wrapper" @@ -312,7 +312,7 @@ elif shared.args.notebook: if shared.args.extensions is not None: extensions_module.create_extensions_block() - input_params = [shared.gradio[k] for k in ('textbox', 'max_new_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')] + input_params = [shared.gradio[k] for k in ('textbox', 'max_new_tokens', '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')] output_params = [shared.gradio[k] for k in ["textbox", "markdown", "html"]] gen_events.append(shared.gradio["Generate"].click(generate_reply, input_params, output_params, show_progress=shared.args.no_stream, api_name="textgen")) gen_events.append(shared.gradio['textbox'].submit(generate_reply, input_params, output_params, show_progress=shared.args.no_stream)) @@ -344,7 +344,7 @@ else: with gr.Tab('HTML'): shared.gradio['html'] = gr.HTML() - input_params = [shared.gradio[k] for k in ['textbox', 'max_new_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']] + input_params = [shared.gradio[k] for k in ['textbox', 'max_new_tokens', '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']] output_params = [shared.gradio[k] for k in ['output_textbox', 'markdown', 'html']] gen_events.append(shared.gradio['Generate'].click(generate_reply, input_params, output_params, show_progress=shared.args.no_stream, api_name="textgen")) gen_events.append(shared.gradio['textbox'].submit(generate_reply, input_params, output_params, show_progress=shared.args.no_stream))