diff --git a/modules/chat.py b/modules/chat.py index 1140b5f..bc7623a 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -91,7 +91,7 @@ def extract_message_from_reply(reply, name1, name2, stop_at_newline): reply = fix_newlines(reply) return reply, next_character_found -def chatbot_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, regenerate=False, mode="cai-chat", end_of_turn=""): +def chatbot_wrapper(text, max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, regenerate=False, mode="cai-chat", end_of_turn=""): just_started = True eos_token = '\n' if stop_at_newline else None name1_original = name1 @@ -126,7 +126,7 @@ def chatbot_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical cumulative_reply = '' for i in range(chat_generation_attempts): reply = None - for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): + for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", max_new_tokens, generation_params, seed, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): reply = cumulative_reply + reply # Extracting the reply @@ -155,7 +155,7 @@ def chatbot_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical yield shared.history['visible'] -def impersonate_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): +def impersonate_wrapper(text, max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): eos_token = '\n' if stop_at_newline else None if 'pygmalion' in shared.model_name.lower(): @@ -169,7 +169,7 @@ def impersonate_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typ cumulative_reply = '' for i in range(chat_generation_attempts): reply = None - for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): + for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", max_new_tokens, generation_params, seed, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): reply = cumulative_reply + reply reply, next_character_found = extract_message_from_reply(reply, name1, name2, stop_at_newline) yield reply @@ -181,11 +181,11 @@ def impersonate_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typ yield reply -def cai_chatbot_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): - for history in chatbot_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts, regenerate=False, mode=mode, end_of_turn=end_of_turn): +def cai_chatbot_wrapper(text, max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): + for history in chatbot_wrapper(text, max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts, regenerate=False, mode=mode, end_of_turn=end_of_turn): yield chat_html_wrapper(history, name1, name2, mode) -def regenerate_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): +def regenerate_wrapper(text, max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts=1, mode="cai-chat", end_of_turn=""): if (shared.character != 'None' and len(shared.history['visible']) == 1) or len(shared.history['internal']) == 0: yield chat_html_wrapper(shared.history['visible'], name1, name2, mode) else: @@ -193,7 +193,7 @@ def regenerate_wrapper(text, max_new_tokens, do_sample, temperature, top_p, typi last_internal = shared.history['internal'].pop() # Yield '*Is typing...*' yield chat_html_wrapper(shared.history['visible']+[[last_visible[0], shared.processing_message]], name1, name2, mode) - for history in chatbot_wrapper(last_internal[0], max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts, regenerate=True, mode=mode, end_of_turn=end_of_turn): + for history in chatbot_wrapper(last_internal[0], max_new_tokens, generation_params, seed, name1, name2, context, stop_at_newline, chat_prompt_size, chat_generation_attempts, regenerate=True, mode=mode, end_of_turn=end_of_turn): shared.history['visible'][-1] = [last_visible[0], history[-1][1]] yield chat_html_wrapper(shared.history['visible'], name1, name2, mode) diff --git a/modules/text_generation.py b/modules/text_generation.py index 406c454..e6ccab6 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -102,10 +102,13 @@ def set_manual_seed(seed): def stop_everything_event(): shared.stop_everything = True -def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, encoder_repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, seed, eos_token=None, stopping_strings=[]): +def generate_reply(question, max_new_tokens, generation_params, seed, eos_token=None, stopping_strings=[]): + print(generation_params) + print('---------------') clear_torch_cache() set_manual_seed(seed) shared.stop_everything = False + updated_params = {} t0 = time.time() original_question = question @@ -117,9 +120,14 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi # These models are not part of Hugging Face, so we handle them # separately and terminate the function call earlier if any((shared.is_RWKV, shared.is_llamacpp)): + + for k in ['temperature', 'top_p', 'top_k', 'repetition_penalty']: + updated_params[k] = generation_params[k] + updated_params["token_count"] = generation_params["max_new_tokens"] + try: if shared.args.no_stream: - reply = shared.model.generate(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty) + reply = shared.model.generate(context=question, **updated_params) output = original_question+reply if not shared.is_chat(): reply = original_question + apply_extensions(reply, "output") @@ -130,7 +138,7 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi # RWKV has proper streaming, which is very nice. # No need to generate 8 tokens at a time. - for reply in shared.model.generate_with_streaming(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty): + for reply in shared.model.generate_with_streaming(context=question, **updated_params): output = original_question+reply if not shared.is_chat(): reply = original_question + apply_extensions(reply, "output") @@ -158,49 +166,39 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi t = [encode(string, 0, add_special_tokens=False) for string in stopping_strings] stopping_criteria_list.append(_SentinelTokenStoppingCriteria(sentinel_token_ids=t, starting_idx=len(input_ids[0]))) - generate_params = {} + updated_params["max_new_tokens"] = max_new_tokens if not shared.args.flexgen: - generate_params.update({ - "max_new_tokens": max_new_tokens, - "eos_token_id": eos_token_ids, - "stopping_criteria": stopping_criteria_list, - "do_sample": do_sample, - "temperature": temperature, - "top_p": top_p, - "typical_p": typical_p, - "repetition_penalty": repetition_penalty, - "encoder_repetition_penalty": encoder_repetition_penalty, - "top_k": top_k, - "min_length": min_length if shared.args.no_stream else 0, - "no_repeat_ngram_size": no_repeat_ngram_size, - "num_beams": num_beams, - "penalty_alpha": penalty_alpha, - "length_penalty": length_penalty, - "early_stopping": early_stopping, - }) + updated_params["eos_token_id"] = eos_token_ids + updated_params["stopping_criteria"] = stopping_criteria_list + for k in ["do_sample", "temperature", "top_p", "typical_p", "repetition_penalty", "encoder_repetition_penalty", "top_k", "min_length", "no_repeat_ngram_size", "num_beams", "penalty_alpha", "length_penalty", "early_stopping"]: + updated_params[k] = generation_params[k] + + if shared.args.no_stream: + updated_params["min_length"] = 0 else: - generate_params.update({ - "max_new_tokens": max_new_tokens if shared.args.no_stream else 8, - "do_sample": do_sample, - "temperature": temperature, - "stop": eos_token_ids[-1], - }) + for k in ["do_sample", "temperature"]: + updated_params[k] = generation_params[k] + updated_params["stop"] = generation_params["eos_token_ids"][-1] + if not shared.args.no_stream: + updated_params["max_new_tokens"] = 8 + print(updated_params) + if shared.args.no_cache: - generate_params.update({"use_cache": False}) + updated_params.update({"use_cache": False}) if shared.args.deepspeed: - generate_params.update({"synced_gpus": True}) + updated_params.update({"synced_gpus": True}) if shared.soft_prompt: inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) - generate_params.update({"inputs_embeds": inputs_embeds}) - generate_params.update({"inputs": filler_input_ids}) + updated_params.update({"inputs_embeds": inputs_embeds}) + updated_params.update({"inputs": filler_input_ids}) else: - generate_params.update({"inputs": input_ids}) + updated_params.update({"inputs": input_ids}) try: # Generate the entire reply at once. if shared.args.no_stream: with torch.no_grad(): - output = shared.model.generate(**generate_params)[0] + output = shared.model.generate(**updated_params)[0] if cuda: output = output.cuda() if shared.soft_prompt: @@ -228,7 +226,7 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi if not shared.is_chat(): yield formatted_outputs(original_question, shared.model_name) - with generate_with_streaming(**generate_params) as generator: + with generate_with_streaming(**updated_params) as generator: for output in generator: if shared.soft_prompt: output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) @@ -247,7 +245,7 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi for i in range(max_new_tokens//8+1): clear_torch_cache() with torch.no_grad(): - output = shared.model.generate(**generate_params)[0] + output = shared.model.generate(**updated_params)[0] if shared.soft_prompt: output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) @@ -263,10 +261,10 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi input_ids = np.reshape(output, (1, output.shape[0])) if shared.soft_prompt: inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) - generate_params.update({"inputs_embeds": inputs_embeds}) - generate_params.update({"inputs": filler_input_ids}) + updated_params.update({"inputs_embeds": inputs_embeds}) + updated_params.update({"inputs": filler_input_ids}) else: - generate_params.update({"inputs": input_ids}) + updated_params.update({"inputs": input_ids}) yield formatted_outputs(reply, shared.model_name) diff --git a/server.py b/server.py index 8bcb650..3897c1c 100644 --- a/server.py +++ b/server.py @@ -85,7 +85,7 @@ def load_lora_wrapper(selected_lora): add_lora_to_model(selected_lora) return selected_lora -def load_preset_values(preset_menu, return_dict=False): +def load_preset_values(preset_menu): generate_params = { 'do_sample': True, 'temperature': 1, @@ -110,10 +110,7 @@ def load_preset_values(preset_menu, return_dict=False): generate_params['temperature'] = min(1.99, generate_params['temperature']) - if return_dict: - return generate_params - else: - return generate_params['do_sample'], generate_params['temperature'], generate_params['top_p'], generate_params['typical_p'], generate_params['repetition_penalty'], generate_params['encoder_repetition_penalty'], generate_params['top_k'], generate_params['min_length'], generate_params['no_repeat_ngram_size'], generate_params['num_beams'], generate_params['penalty_alpha'], generate_params['length_penalty'], generate_params['early_stopping'] + return generate_params def upload_soft_prompt(file): with zipfile.ZipFile(io.BytesIO(file)) as zf: @@ -170,7 +167,8 @@ def create_prompt_menus(): shared.gradio['save_prompt'].click(save_prompt, [shared.gradio['textbox']], [shared.gradio['status']], show_progress=False) def create_settings_menus(default_preset): - generate_params = load_preset_values(default_preset if not shared.args.flexgen else 'Naive', return_dict=True) + generate_params = load_preset_values(default_preset if not shared.args.flexgen else 'Naive') + shared.gradio['generation_state'] = gr.State(generate_params) with gr.Row(): with gr.Column(): @@ -221,8 +219,26 @@ def create_settings_menus(default_preset): with gr.Row(): shared.gradio['upload_softprompt'] = gr.File(type='binary', file_types=['.zip']) + def update_dict(_dict, k, v): + _dict[k] = v + return _dict + + for k in ['do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping']: + if type(shared.gradio[k]) is gr.Checkbox: + shared.gradio[k].change( + lambda state, value, copy=k: update_dict(state, copy, value), + inputs=[shared.gradio['generation_state'], shared.gradio[k]], + outputs=shared.gradio['generation_state'], + ) + else: + shared.gradio[k].release( + lambda state, value, copy=k: update_dict(state, copy, value), + inputs=[shared.gradio['generation_state'], shared.gradio[k]], + outputs=shared.gradio['generation_state'], + ) + shared.gradio['model_menu'].change(load_model_wrapper, [shared.gradio['model_menu']], [shared.gradio['model_menu']], show_progress=True) - shared.gradio['preset_menu'].change(load_preset_values, [shared.gradio['preset_menu']], [shared.gradio[k] for k in ['do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping']]) + shared.gradio['preset_menu'].change(load_preset_values, [shared.gradio['preset_menu']], [shared.gradio[k] for k in ['generation_state']]) shared.gradio['lora_menu'].change(load_lora_wrapper, [shared.gradio['lora_menu']], [shared.gradio['lora_menu']], show_progress=True) shared.gradio['softprompts_menu'].change(load_soft_prompt, [shared.gradio['softprompts_menu']], [shared.gradio['softprompts_menu']], show_progress=True) shared.gradio['upload_softprompt'].upload(upload_soft_prompt, [shared.gradio['upload_softprompt']], [shared.gradio['softprompts_menu']]) @@ -376,7 +392,7 @@ def create_interface(): create_settings_menus(default_preset) - shared.input_params = [shared.gradio[k] for k in ['Chat input', 'max_new_tokens', 'do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping', 'seed', 'name1', 'name2', 'context', 'check', 'chat_prompt_size_slider', 'chat_generation_attempts', 'Chat mode', 'end_of_turn']] + shared.input_params = [shared.gradio[k] for k in ['Chat input', 'max_new_tokens', 'generation_state', 'seed', 'name1', 'name2', 'context', 'check', 'chat_prompt_size_slider', 'chat_generation_attempts', 'Chat mode', 'end_of_turn']] def set_chat_input(textbox): return textbox, "" @@ -456,7 +472,7 @@ def create_interface(): with gr.Tab("Parameters", elem_id="parameters"): create_settings_menus(default_preset) - shared.input_params = [shared.gradio[k] for k in ['textbox', 'max_new_tokens', 'do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping', 'seed']] + shared.input_params = [shared.gradio[k] for k in ['textbox', 'max_new_tokens', 'generation_state', 'seed']] output_params = [shared.gradio[k] for k in ['textbox', 'markdown', 'html']] gen_events.append(shared.gradio['Generate'].click(generate_reply, shared.input_params, output_params, show_progress=shared.args.no_stream, api_name='textgen')) gen_events.append(shared.gradio['textbox'].submit(generate_reply, shared.input_params, output_params, show_progress=shared.args.no_stream)) @@ -489,7 +505,7 @@ def create_interface(): with gr.Tab("Parameters", elem_id="parameters"): create_settings_menus(default_preset) - shared.input_params = [shared.gradio[k] for k in ['textbox', 'max_new_tokens', 'do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping', 'seed']] + shared.input_params = [shared.gradio[k] for k in ['textbox', 'max_new_tokens', 'generation_state', 'seed']] output_params = [shared.gradio[k] for k in ['output_textbox', 'markdown', 'html']] gen_events.append(shared.gradio['Generate'].click(generate_reply, shared.input_params, output_params, show_progress=shared.args.no_stream, api_name='textgen')) gen_events.append(shared.gradio['textbox'].submit(generate_reply, shared.input_params, output_params, show_progress=shared.args.no_stream))