diff --git a/extensions/openai/typing.py b/extensions/openai/typing.py index f4923967..9c4a04f0 100644 --- a/extensions/openai/typing.py +++ b/extensions/openai/typing.py @@ -37,6 +37,7 @@ class GenerationOptions(BaseModel): early_stopping: bool = False truncation_length: int = 0 max_tokens_second: int = 0 + prompt_lookup_num_tokens: int = 0 custom_token_bans: str = "" auto_max_new_tokens: bool = False ban_eos_token: bool = False diff --git a/modules/loaders.py b/modules/loaders.py index 9423b231..5620c95a 100644 --- a/modules/loaders.py +++ b/modules/loaders.py @@ -192,6 +192,7 @@ def transformers_samplers(): 'add_bos_token', 'skip_special_tokens', 'auto_max_new_tokens', + 'prompt_lookup_num_tokens' } diff --git a/modules/shared.py b/modules/shared.py index 1c53ba38..e4bdacaa 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -47,6 +47,7 @@ settings = { 'truncation_length_max': 200000, 'max_tokens_second': 0, 'max_updates_second': 0, + 'prompt_lookup_num_tokens': 0, 'custom_stopping_strings': '', 'custom_token_bans': '', 'auto_max_new_tokens': False, diff --git a/modules/text_generation.py b/modules/text_generation.py index f17d2c8e..c8562450 100644 --- a/modules/text_generation.py +++ b/modules/text_generation.py @@ -291,6 +291,9 @@ def generate_reply_HF(question, original_question, seed, state, stopping_strings if state['negative_prompt'] != '': generate_params['negative_prompt_ids'] = encode(state['negative_prompt']) + if state['prompt_lookup_num_tokens'] > 0: + generate_params['prompt_lookup_num_tokens'] = state['prompt_lookup_num_tokens'] + for k in ['epsilon_cutoff', 'eta_cutoff']: if state[k] > 0: generate_params[k] = state[k] * 1e-4 diff --git a/modules/ui.py b/modules/ui.py index 3c978225..b639c4df 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -112,6 +112,7 @@ def list_interface_input_elements(): 'auto_max_new_tokens', 'max_tokens_second', 'max_updates_second', + 'prompt_lookup_num_tokens', 'seed', 'temperature', 'temperature_last', diff --git a/modules/ui_parameters.py b/modules/ui_parameters.py index 07108553..fafdcacc 100644 --- a/modules/ui_parameters.py +++ b/modules/ui_parameters.py @@ -72,6 +72,7 @@ def create_ui(default_preset): shared.gradio['truncation_length'] = gr.Slider(value=get_truncation_length(), minimum=shared.settings['truncation_length_min'], maximum=shared.settings['truncation_length_max'], step=256, label='Truncate the prompt up to this length', info='The leftmost tokens are removed if the prompt exceeds this length. Most models require this to be at most 2048.') shared.gradio['max_tokens_second'] = gr.Slider(value=shared.settings['max_tokens_second'], minimum=0, maximum=20, step=1, label='Maximum tokens/second', info='To make text readable in real time.') shared.gradio['max_updates_second'] = gr.Slider(value=shared.settings['max_updates_second'], minimum=0, maximum=24, step=1, label='Maximum UI updates/second', info='Set this if you experience lag in the UI during streaming.') + shared.gradio['prompt_lookup_num_tokens'] = gr.Slider(value=shared.settings['prompt_lookup_num_tokens'], minimum=0, maximum=10, step=1, label='prompt_lookup_num_tokens', info='Activates Prompt Lookup Decoding.') shared.gradio['custom_stopping_strings'] = gr.Textbox(lines=1, value=shared.settings["custom_stopping_strings"] or None, label='Custom stopping strings', info='In addition to the defaults. Written between "" and separated by commas.', placeholder='"\\n", "\\nYou:"') shared.gradio['custom_token_bans'] = gr.Textbox(value=shared.settings['custom_token_bans'] or None, label='Custom token bans', info='Specific token IDs to ban from generating, comma-separated. The IDs can be found in the Default or Notebook tab.') diff --git a/settings-template.yaml b/settings-template.yaml index 8ebbc1ab..87101116 100644 --- a/settings-template.yaml +++ b/settings-template.yaml @@ -16,6 +16,7 @@ truncation_length_min: 0 truncation_length_max: 200000 max_tokens_second: 0 max_updates_second: 0 +prompt_lookup_num_tokens: 0 custom_stopping_strings: '' custom_token_bans: '' auto_max_new_tokens: false