2023-05-04 02:43:17 +02:00
|
|
|
import logging
|
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
|
|
|
|
|
2023-02-23 16:05:25 +01:00
|
|
|
import extensions
|
2023-02-23 18:41:42 +01:00
|
|
|
import modules.shared as shared
|
2023-02-23 16:05:25 +01:00
|
|
|
|
2023-02-24 14:01:21 +01:00
|
|
|
state = {}
|
2023-02-23 16:05:25 +01:00
|
|
|
available_extensions = []
|
2023-03-29 04:27:02 +02:00
|
|
|
setup_called = set()
|
2023-02-23 16:05:25 +01:00
|
|
|
|
2023-04-07 05:15:45 +02:00
|
|
|
|
2023-04-25 05:23:11 +02:00
|
|
|
def apply_settings(extension, name):
|
|
|
|
if not hasattr(extension, 'params'):
|
|
|
|
return
|
|
|
|
|
|
|
|
for param in extension.params:
|
|
|
|
_id = f"{name}-{param}"
|
|
|
|
if _id not in shared.settings:
|
|
|
|
continue
|
|
|
|
|
|
|
|
extension.params[param] = shared.settings[_id]
|
|
|
|
|
|
|
|
|
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:
|
2023-04-23 20:52:43 +02:00
|
|
|
if name != 'api':
|
2023-05-04 02:43:17 +02:00
|
|
|
logging.info(f'Loading the extension "{name}"...')
|
2023-03-16 03:29:56 +01:00
|
|
|
try:
|
|
|
|
exec(f"import extensions.{name}.script")
|
2023-04-16 06:36:50 +02:00
|
|
|
extension = getattr(extensions, name).script
|
2023-04-25 06:16:23 +02:00
|
|
|
apply_settings(extension, name)
|
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-04-25 05:10:21 +02:00
|
|
|
|
2023-03-16 03:29:56 +01:00
|
|
|
state[name] = [True, i]
|
|
|
|
except:
|
2023-05-07 22:44:05 +02:00
|
|
|
logging.error(f'Failed to load the extension "{name}".')
|
2023-03-24 20:51:27 +01:00
|
|
|
traceback.print_exc()
|
2023-02-23 18:49:02 +01:00
|
|
|
|
2023-04-07 05:15:45 +02: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():
|
2023-04-07 05:15:45 +02:00
|
|
|
for name in sorted(state, key=lambda x: state[x][1]):
|
2023-04-07 05:52:02 +02:00
|
|
|
if state[name][0]:
|
2023-04-16 06:36:50 +02:00
|
|
|
yield getattr(extensions, name).script, name
|
2023-02-24 14:01:21 +01:00
|
|
|
|
2023-04-07 05:15:45 +02: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)
|
2023-04-25 05:10:21 +02:00
|
|
|
|
2023-02-23 16:05:25 +01:00
|
|
|
return text
|
|
|
|
|
2023-04-07 05:15:45 +02:00
|
|
|
|
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']
|
2023-04-25 05:10:21 +02:00
|
|
|
|
2023-04-24 01:32:22 +02:00
|
|
|
return text, visible_text
|
|
|
|
|
|
|
|
|
2023-05-10 01:18:02 +02:00
|
|
|
# custom_generate_chat_prompt handling - currently only the first one will work
|
2023-04-24 01:32:22 +02:00
|
|
|
def _apply_custom_generate_chat_prompt(text, state, **kwargs):
|
|
|
|
custom_generate_chat_prompt = None
|
|
|
|
for extension, _ in iterator():
|
2023-05-10 01:18:02 +02:00
|
|
|
if hasattr(extension, 'custom_generate_chat_prompt'):
|
|
|
|
return custom_generate_chat_prompt(text, state, **kwargs)
|
2023-04-25 05:10:21 +02:00
|
|
|
|
2023-04-24 01:32:22 +02:00
|
|
|
return None
|
|
|
|
|
|
|
|
|
2023-05-05 23:53:03 +02:00
|
|
|
# Extension that modifies the input parameters before they are used
|
|
|
|
def _apply_state_modifier_extensions(state):
|
|
|
|
for extension, _ in iterator():
|
|
|
|
if hasattr(extension, "state_modifier"):
|
|
|
|
state = getattr(extension, "state_modifier")(state)
|
|
|
|
|
|
|
|
return state
|
2023-05-10 03:49:39 +02:00
|
|
|
|
2023-05-05 23:53:03 +02:00
|
|
|
|
2023-05-10 01:18:02 +02:00
|
|
|
# Extension functions that override the default tokenizer output - currently only the first one will work
|
2023-04-24 01:32:22 +02:00
|
|
|
def _apply_tokenizer_extensions(function_name, state, prompt, input_ids, input_embeds):
|
|
|
|
for extension, _ in iterator():
|
|
|
|
if hasattr(extension, function_name):
|
2023-05-10 01:18:02 +02:00
|
|
|
return getattr(extension, function_name)(state, prompt, input_ids, input_embeds)
|
2023-04-25 05:10:21 +02:00
|
|
|
|
2023-04-24 01:32:22 +02:00
|
|
|
return prompt, input_ids, input_embeds
|
|
|
|
|
|
|
|
|
2023-05-10 01:18:02 +02:00
|
|
|
# Get prompt length in tokens after applying extension functions which override the default tokenizer output
|
|
|
|
# currently only the first one will work
|
|
|
|
def _apply_custom_tokenized_length(prompt):
|
|
|
|
for extension, _ in iterator():
|
|
|
|
if hasattr(extension, 'custom_tokenized_length'):
|
|
|
|
return getattr(extension, 'custom_tokenized_length')(prompt)
|
2023-05-10 03:49:39 +02:00
|
|
|
|
2023-05-10 01:18:02 +02:00
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
# Custom generate reply handling - currently only the first one will work
|
2023-05-05 23:53:03 +02:00
|
|
|
def _apply_custom_generate_reply():
|
|
|
|
for extension, _ in iterator():
|
|
|
|
if hasattr(extension, 'custom_generate_reply'):
|
|
|
|
return getattr(extension, 'custom_generate_reply')
|
|
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
2023-04-24 01:32:22 +02:00
|
|
|
EXTENSION_MAP = {
|
|
|
|
"input": partial(_apply_string_extensions, "input_modifier"),
|
|
|
|
"output": partial(_apply_string_extensions, "output_modifier"),
|
2023-05-05 23:53:03 +02:00
|
|
|
"state": _apply_state_modifier_extensions,
|
2023-04-24 01:32:22 +02:00
|
|
|
"bot_prefix": partial(_apply_string_extensions, "bot_prefix_modifier"),
|
|
|
|
"tokenizer": partial(_apply_tokenizer_extensions, "tokenizer_modifier"),
|
|
|
|
"input_hijack": _apply_input_hijack,
|
2023-05-05 23:53:03 +02:00
|
|
|
"custom_generate_chat_prompt": _apply_custom_generate_chat_prompt,
|
2023-05-10 01:18:02 +02:00
|
|
|
"custom_generate_reply": _apply_custom_generate_reply,
|
|
|
|
"tokenized_length": _apply_custom_tokenized_length
|
2023-04-24 01:32:22 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
def apply_extensions(typ, *args, **kwargs):
|
|
|
|
if typ not in EXTENSION_MAP:
|
|
|
|
raise ValueError(f"Invalid extension type {typ}")
|
2023-04-25 05:10:21 +02:00
|
|
|
|
2023-04-24 01:32:22 +02:00
|
|
|
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-03-19 14:31:21 +01:00
|
|
|
should_display_ui = False
|
2023-03-29 04:27:02 +02:00
|
|
|
for extension, name in iterator():
|
|
|
|
if hasattr(extension, "ui"):
|
|
|
|
should_display_ui = True
|
2023-04-25 05:23:11 +02:00
|
|
|
break
|
2023-03-19 14:22:24 +01:00
|
|
|
|
2023-02-24 23:00:11 +01:00
|
|
|
# Creating the extension ui elements
|
2023-03-19 14:31:21 +01:00
|
|
|
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():
|
|
|
|
if hasattr(extension, "ui"):
|
2023-05-03 00:20:02 +02:00
|
|
|
gr.Markdown(f"\n### {name}")
|
2023-03-16 03:29:56 +01:00
|
|
|
extension.ui()
|