mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add chat-instruct mode (#2049)

Browse source
This commit is contained in:
oobabooga 2023-05-14 10:43:55 -03:00 • committed by GitHub
parent 5f6cf39f36
commit 3b886f9c9f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 118 additions and 82 deletions
Download patch file Download diff file Expand all files Collapse all files

126
modules/chat.py
View file

@ -19,25 +19,8 @@ from modules.text_generation import (generate_reply, get_encoded_length,
from modules.utils import replace_all
def generate_chat_prompt(user_input, state, **kwargs):
impersonate = kwargs.get('impersonate', False)
_continue = kwargs.get('_continue', False)
also_return_rows = kwargs.get('also_return_rows', False)
history = state.get('history', shared.history['internal'])
is_instruct = state['mode'] == 'instruct'
rows = [state['context_instruct'] if is_instruct else f"{state['context'].strip()}\n"]
min_rows = 3
# Finding the maximum prompt size
chat_prompt_size = state['chat_prompt_size']
if shared.soft_prompt:
chat_prompt_size -= shared.soft_prompt_tensor.shape[1]
max_length = min(get_max_prompt_length(state), chat_prompt_size)
# Building the turn templates
if is_instruct:
def get_turn_substrings(state, instruct=False):
if instruct:
if 'turn_template' not in state or state['turn_template'] == '':
template = '<|user|>\n<|user-message|>\n<|bot|>\n<|bot-message|>\n'
else:
@ -46,44 +29,92 @@ def generate_chat_prompt(user_input, state, **kwargs):
template = '<|user|>: <|user-message|>\n<|bot|>: <|bot-message|>\n'
replacements = {
'<|user|>': state['name1_instruct' if is_instruct else 'name1'].strip(),
'<|bot|>': state['name2_instruct' if is_instruct else 'name2'].strip(),
'<|user|>': state['name1_instruct' if instruct else 'name1'].strip(),
'<|bot|>': state['name2_instruct' if instruct else 'name2'].strip(),
}
user_turn = replace_all(template.split('<|bot|>')[0], replacements)
bot_turn = replace_all('<|bot|>' + template.split('<|bot|>')[1], replacements)
user_turn_stripped = replace_all(user_turn.split('<|user-message|>')[0], replacements)
bot_turn_stripped = replace_all(bot_turn.split('<|bot-message|>')[0], replacements)
output = {
'user_turn': template.split('<|bot|>')[0],
'bot_turn': '<|bot|>' + template.split('<|bot|>')[1],
'user_turn_stripped': template.split('<|bot|>')[0].split('<|user-message|>')[0],
'bot_turn_stripped': '<|bot|>' + template.split('<|bot|>')[1].split('<|bot-message|>')[0],
}
for k in output:
output[k] = replace_all(output[k], replacements)
return output
def generate_chat_prompt(user_input, state, **kwargs):
impersonate = kwargs.get('impersonate', False)
_continue = kwargs.get('_continue', False)
also_return_rows = kwargs.get('also_return_rows', False)
history = state.get('history', shared.history['internal'])
is_instruct = state['mode'] == 'instruct'
# Finding the maximum prompt size
chat_prompt_size = state['chat_prompt_size']
if shared.soft_prompt:
chat_prompt_size -= shared.soft_prompt_tensor.shape[1]
max_length = min(get_max_prompt_length(state), chat_prompt_size)
all_substrings = {
'chat': get_turn_substrings(state, instruct=False),
'instruct': get_turn_substrings(state, instruct=True)
}
substrings = all_substrings['instruct' if is_instruct else 'chat']
# Creating the template for "chat-instruct" mode
if state['mode'] == 'chat-instruct':
wrapper = ''
command = state['chat-instruct_command'].replace('<|character|>', state['name2'] if not impersonate else state['name1'])
wrapper += state['context_instruct']
wrapper += all_substrings['instruct']['user_turn'].replace('<|user-message|>', command)
wrapper += all_substrings['instruct']['bot_turn_stripped']
if impersonate:
wrapper += substrings['user_turn_stripped'].rstrip(' ')
else:
wrapper += apply_extensions("bot_prefix", substrings['bot_turn_stripped'].rstrip(' '))
else:
wrapper = '<|prompt|>'
# Building the prompt
min_rows = 3
i = len(history) - 1
while i >= 0 and get_encoded_length(''.join(rows)) < max_length:
rows = [state['context_instruct'] if is_instruct else f"{state['context'].strip()}\n"]
while i >= 0 and get_encoded_length(wrapper.replace('<|prompt|>', ''.join(rows))) < max_length:
if _continue and i == len(history) - 1:
rows.insert(1, bot_turn_stripped + history[i][1].strip())
rows.insert(1, substrings['bot_turn_stripped'] + history[i][1].strip())
else:
rows.insert(1, bot_turn.replace('<|bot-message|>', history[i][1].strip()))
rows.insert(1, substrings['bot_turn'].replace('<|bot-message|>', history[i][1].strip()))
string = history[i][0]
if string not in ['', '<|BEGIN-VISIBLE-CHAT|>']:
rows.insert(1, replace_all(user_turn, {'<|user-message|>': string.strip(), '<|round|>': str(i)}))
rows.insert(1, replace_all(substrings['user_turn'], {'<|user-message|>': string.strip(), '<|round|>': str(i)}))
i -= 1
if impersonate:
min_rows = 2
rows.append(user_turn_stripped.rstrip(' '))
if state['mode'] == 'chat-instruct':
min_rows = 1
else:
min_rows = 2
rows.append(substrings['user_turn_stripped'].rstrip(' '))
elif not _continue:
# Adding the user message
if len(user_input) > 0:
rows.append(replace_all(user_turn, {'<|user-message|>': user_input.strip(), '<|round|>': str(len(history))}))
rows.append(replace_all(substrings['user_turn'], {'<|user-message|>': user_input.strip(), '<|round|>': str(len(history))}))
# Adding the Character prefix
rows.append(apply_extensions("bot_prefix", bot_turn_stripped.rstrip(' ')))
if state['mode'] != 'chat-instruct':
rows.append(apply_extensions("bot_prefix", substrings['bot_turn_stripped'].rstrip(' ')))
while len(rows) > min_rows and get_encoded_length(''.join(rows)) >= max_length:
while len(rows) > min_rows and get_encoded_length(wrapper.replace('<|prompt|>', ''.join(rows))) >= max_length:
rows.pop(1)
prompt = ''.join(rows)
prompt = wrapper.replace('<|prompt|>', ''.join(rows))
if also_return_rows:
return prompt, rows
else:
@ -91,8 +122,9 @@ def generate_chat_prompt(user_input, state, **kwargs):
def get_stopping_strings(state):
if state['mode'] == 'instruct':
stopping_strings = [
stopping_strings = []
if state['mode'] in ['instruct', 'chat-instruct']:
stopping_strings += [
state['turn_template'].split('<|user-message|>')[1].split('<|bot|>')[0] + '<|bot|>',
state['turn_template'].split('<|bot-message|>')[1] + '<|user|>'
]
@ -104,8 +136,12 @@ def get_stopping_strings(state):
for i in range(len(stopping_strings)):
stopping_strings[i] = replace_all(stopping_strings[i], replacements).rstrip(' ').replace(r'\n', '\n')
else:
stopping_strings = [f"\n{state['name1']}:", f"\n{state['name2']}:"]
if state['mode'] in ['chat', 'chat-instruct']:
stopping_strings += [
f"\n{state['name1']}:",
f"\n{state['name2']}:"
]
stopping_strings += ast.literal_eval(f"[{state['custom_stopping_strings']}]")
return stopping_strings
@ -433,7 +469,7 @@ def generate_pfp_cache(character):
return None
def load_character(character, name1, name2, mode):
def load_character(character, name1, name2, instruct=False):
shared.character = character
context = greeting = turn_template = ""
greeting_field = 'greeting'
@ -444,7 +480,7 @@ def load_character(character, name1, name2, mode):
Path("cache/pfp_character.png").unlink()
if character != 'None':
folder = 'characters' if not mode == 'instruct' else 'characters/instruction-following'
folder = 'characters' if not instruct else 'characters/instruction-following'
picture = generate_pfp_cache(character)
for extension in ["yml", "yaml", "json"]:
filepath = Path(f'{folder}/{character}.{extension}')
@ -472,8 +508,8 @@ def load_character(character, name1, name2, mode):
if 'context' in data:
context = data['context']
if mode != 'instruct':
context = context.strip() + '\n\n'
if not instruct:
context = context.strip() + '\n'
elif "char_persona" in data:
context = build_pygmalion_style_context(data)
greeting_field = 'char_greeting'
@ -493,7 +529,7 @@ def load_character(character, name1, name2, mode):
greeting = shared.settings['greeting']
turn_template = shared.settings['turn_template']
if mode != 'instruct':
if not instruct:
shared.history['internal'] = []
shared.history['visible'] = []
if Path(f'logs/{shared.character}_persistent.json').exists():
@ -505,7 +541,7 @@ def load_character(character, name1, name2, mode):
shared.history['visible'] += [['', apply_extensions("output", greeting)]]
# Create .json log files since they don't already exist
save_history(mode)
save_history('instruct' if instruct else 'chat')
return name1, name2, picture, greeting, context, repr(turn_template)[1:-1]

1
modules/shared.py
View file

@ -54,6 +54,7 @@ settings = {
'mode': 'chat',
'chat_style': 'cai-chat',
'instruction_template': 'None',
'chat-instruct_command': 'Continue the chat dialogue below. Write a single reply for the character "<|character|>".\n\n<|prompt|>',
'chat_prompt_size': 2048,
'chat_prompt_size_min': 0,
'chat_prompt_size_max': 2048,

10
modules/ui.py
View file

@ -36,7 +36,7 @@ def list_model_elements():
def list_interface_input_elements(chat=False):
elements = ['max_new_tokens', 'seed', 'temperature', 'top_p', 'top_k', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'no_repeat_ngram_size', 'min_length', 'do_sample', 'penalty_alpha', 'num_beams', 'length_penalty', 'early_stopping', 'add_bos_token', 'ban_eos_token', 'truncation_length', 'custom_stopping_strings', 'skip_special_tokens', 'preset_menu', 'stream']
if chat:
elements += ['name1', 'name2', 'greeting', 'context', 'chat_prompt_size', 'chat_generation_attempts', 'stop_at_newline', 'mode', 'instruction_template', 'character_menu', 'name1_instruct', 'name2_instruct', 'context_instruct', 'turn_template', 'chat_style']
elements += ['name1', 'name2', 'greeting', 'context', 'chat_prompt_size', 'chat_generation_attempts', 'stop_at_newline', 'mode', 'instruction_template', 'character_menu', 'name1_instruct', 'name2_instruct', 'context_instruct', 'turn_template', 'chat_style', 'chat-instruct_command']
elements += list_model_elements()
return elements
@ -59,13 +59,7 @@ def apply_interface_values(state, use_persistent=False):
if len(state) == 0:
return [gr.update() for k in elements] # Dummy, do nothing
else:
if use_persistent and 'mode' in state:
if state['mode'] == 'instruct':
return [state[k] if (k not in ['character_menu'] and k in state) else gr.update() for k in elements]
else:
return [state[k] if (k not in ['instruction_template'] and k in state) else gr.update() for k in elements]
else:
return [state[k] if k in state else gr.update() for k in elements]
return [state[k] if k in state else gr.update() for k in elements]
class ToolButton(gr.Button, gr.components.FormComponent):