text-generation-webui/modules/extensions.py

120 lines
4.2 KiB
Python
Raw Normal View History

2023-03-24 20:51:27 +01:00
import traceback
2023-04-24 01:32:22 +02:00
from functools import partial
2023-03-24 20:51:27 +01:00
2023-03-15 19:11:16 +01:00
import gradio as gr
import extensions
2023-02-23 18:41:42 +01:00
import modules.shared as shared
2023-02-24 14:01:21 +01:00
state = {}
available_extensions = []
setup_called = set()
2023-02-23 18:49:02 +01:00
def load_extensions():
2023-04-07 17:20:57 +02:00
global state, setup_called
2023-02-24 14:01:21 +01:00
for i, name in enumerate(shared.args.extensions):
if name in available_extensions:
if name != 'api':
print(f'Loading the extension "{name}"... ', end='')
2023-03-16 03:29:56 +01:00
try:
exec(f"import extensions.{name}.script")
extension = getattr(extensions, name).script
2023-04-07 17:20:57 +02:00
if extension not in setup_called and hasattr(extension, "setup"):
setup_called.add(extension)
extension.setup()
2023-03-16 03:29:56 +01:00
state[name] = [True, i]
if name != 'api':
print('Ok.')
2023-03-16 03:29:56 +01:00
except:
if name != 'api':
print('Fail.')
2023-03-24 20:51:27 +01:00
traceback.print_exc()
2023-02-23 18:49:02 +01:00
2023-04-07 17:20:57 +02:00
# This iterator returns the extensions in the order specified in the command-line
2023-02-24 14:01:21 +01:00
def iterator():
for name in sorted(state, key=lambda x: state[x][1]):
2023-04-07 05:52:02 +02:00
if state[name][0]:
yield getattr(extensions, name).script, name
2023-02-24 14:01:21 +01:00
2023-04-07 17:20:57 +02:00
# Extension functions that map string -> string
2023-04-24 01:32:22 +02:00
def _apply_string_extensions(function_name, text):
2023-02-24 14:01:21 +01:00
for extension, _ in iterator():
2023-04-24 01:32:22 +02:00
if hasattr(extension, function_name):
text = getattr(extension, function_name)(text)
return text
2023-04-24 01:32:22 +02:00
# Input hijack of extensions
def _apply_input_hijack(text, visible_text):
for extension, _ in iterator():
if hasattr(extension, 'input_hijack') and extension.input_hijack['state']:
extension.input_hijack['state'] = False
if callable(extension.input_hijack['value']):
text, visible_text = extension.input_hijack['value'](text, visible_text)
else:
text, visible_text = extension.input_hijack['value']
return text, visible_text
# custom_generate_chat_prompt handling
def _apply_custom_generate_chat_prompt(text, state, **kwargs):
custom_generate_chat_prompt = None
for extension, _ in iterator():
if custom_generate_chat_prompt is None and hasattr(extension, 'custom_generate_chat_prompt'):
custom_generate_chat_prompt = extension.custom_generate_chat_prompt
if custom_generate_chat_prompt is not None:
return custom_generate_chat_prompt(text, state, **kwargs)
return None
# Extension functions that override the default tokenizer output
def _apply_tokenizer_extensions(function_name, state, prompt, input_ids, input_embeds):
for extension, _ in iterator():
if hasattr(extension, function_name):
prompt, input_ids, input_embeds = getattr(extension, function_name)(state, prompt, input_ids, input_embeds)
return prompt, input_ids, input_embeds
EXTENSION_MAP = {
"input": partial(_apply_string_extensions, "input_modifier"),
"output": partial(_apply_string_extensions, "output_modifier"),
"bot_prefix": partial(_apply_string_extensions, "bot_prefix_modifier"),
"tokenizer": partial(_apply_tokenizer_extensions, "tokenizer_modifier"),
"input_hijack": _apply_input_hijack,
"custom_generate_chat_prompt": _apply_custom_generate_chat_prompt
}
def apply_extensions(typ, *args, **kwargs):
if typ not in EXTENSION_MAP:
raise ValueError(f"Invalid extension type {typ}")
return EXTENSION_MAP[typ](*args, **kwargs)
2023-02-23 18:55:21 +01:00
def create_extensions_block():
2023-03-19 14:25:49 +01:00
global setup_called
2023-03-24 20:51:27 +01:00
2023-02-24 23:00:11 +01:00
# Updating the default values
2023-02-24 14:01:21 +01:00
for extension, name in iterator():
if hasattr(extension, 'params'):
for param in extension.params:
_id = f"{name}-{param}"
if _id in shared.settings:
extension.params[param] = shared.settings[_id]
2023-02-23 18:55:21 +01:00
should_display_ui = False
for extension, name in iterator():
if hasattr(extension, "ui"):
should_display_ui = True
2023-03-19 14:22:24 +01:00
2023-02-24 23:00:11 +01:00
# Creating the extension ui elements
if should_display_ui:
2023-03-27 03:20:30 +02:00
with gr.Column(elem_id="extensions"):
2023-03-16 03:29:56 +01:00
for extension, name in iterator():
2023-03-27 03:20:30 +02:00
gr.Markdown(f"\n### {name}")
2023-03-16 03:29:56 +01:00
if hasattr(extension, "ui"):
extension.ui()