From e055967974936c05035e470f24506a9218fe4b98 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 17 Jan 2024 17:09:36 -0300
Subject: [PATCH] Add prompt_lookup_num_tokens parameter (#5296)

---
 extensions/openai/typing.py | 1 +
 modules/loaders.py          | 1 +
 modules/shared.py           | 1 +
 modules/text_generation.py  | 3 +++
 modules/ui.py               | 1 +
 modules/ui_parameters.py    | 1 +
 settings-template.yaml      | 1 +
 7 files changed, 9 insertions(+)

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