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

Add support for custom chat styles (#1917)

Browse source
This commit is contained in:
oobabooga 2023-05-08 12:35:03 -03:00 • committed by GitHub
parent b040b4110d
commit b5260b24f1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 234 additions and 74 deletions
Download patch file Download diff file Expand all files Collapse all files

44
modules/chat.py
View file

@ -251,33 +251,33 @@ def impersonate_wrapper(text, state):
def cai_chatbot_wrapper(text, state):
for history in chatbot_wrapper(text, state):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'])
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def regenerate_wrapper(text, state):
if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'])
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'], state['chat_style'])
else:
for history in chatbot_wrapper('', state, regenerate=True):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'])
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def continue_wrapper(text, state):
if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'])
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'], state['chat_style'])
else:
for history in chatbot_wrapper('', state, _continue=True):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'])
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def remove_last_message(name1, name2, mode):
def remove_last_message(name1, name2, mode, style):
if len(shared.history['visible']) > 0 and shared.history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>':
last = shared.history['visible'].pop()
shared.history['internal'].pop()
else:
last = ['', '']
return chat_html_wrapper(shared.history['visible'], name1, name2, mode), last[0]
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style), last[0]
def send_last_reply_to_input():
@ -287,35 +287,35 @@ def send_last_reply_to_input():
return ''
def replace_last_reply(text, name1, name2, mode):
def replace_last_reply(text, name1, name2, mode, style):
if len(shared.history['visible']) > 0:
shared.history['visible'][-1][1] = text
shared.history['internal'][-1][1] = apply_extensions("input", text)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def send_dummy_message(text, name1, name2, mode):
def send_dummy_message(text, name1, name2, mode, style):
shared.history['visible'].append([text, ''])
shared.history['internal'].append([apply_extensions("input", text), ''])
return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def send_dummy_reply(text, name1, name2, mode):
def send_dummy_reply(text, name1, name2, mode, style):
if len(shared.history['visible']) > 0 and not shared.history['visible'][-1][1] == '':
shared.history['visible'].append(['', ''])
shared.history['internal'].append(['', ''])
shared.history['visible'][-1][1] = text
shared.history['internal'][-1][1] = apply_extensions("input", text)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def clear_html():
return chat_html_wrapper([], "", "")
def clear_chat_log(name1, name2, greeting, mode):
def clear_chat_log(name1, name2, greeting, mode, style):
shared.history['visible'] = []
shared.history['internal'] = []
@ -325,14 +325,14 @@ def clear_chat_log(name1, name2, greeting, mode):
# Save cleared logs
save_history(mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def redraw_html(name1, name2, mode):
return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
def redraw_html(name1, name2, mode, style):
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def tokenize_dialogue(dialogue, name1, name2, mode):
def tokenize_dialogue(dialogue, name1, name2, mode, style):
history = []
messages = []
dialogue = re.sub('<START>', '', dialogue)
@ -440,7 +440,7 @@ def generate_pfp_cache(character):
return None
def load_character(character, name1, name2, mode):
def load_character(character, name1, name2, mode, style):
shared.character = character
context = greeting = turn_template = ""
greeting_field = 'greeting'
@ -514,7 +514,7 @@ def load_character(character, name1, name2, mode):
# Create .json log files since they don't already exist
save_history(mode)
return name1, name2, picture, greeting, context, repr(turn_template)[1:-1], chat_html_wrapper(shared.history['visible'], name1, name2, mode)
return name1, name2, picture, greeting, context, repr(turn_template)[1:-1], chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def upload_character(json_file, img, tavern=False):
@ -549,7 +549,7 @@ def upload_tavern_character(img, name1, name2):
return upload_character(json.dumps(_json), img, tavern=True)
def upload_your_profile_picture(img, name1, name2, mode):
def upload_your_profile_picture(img, name1, name2, mode, style):
cache_folder = Path("cache")
if not cache_folder.exists():
cache_folder.mkdir()
@ -562,4 +562,4 @@ def upload_your_profile_picture(img, name1, name2, mode):
img.save(Path('cache/pfp_me.png'))
logging.info('Profile picture saved to "cache/pfp_me.png"')
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, reset_cache=True)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style, reset_cache=True)

29
modules/html_generator.py
View file

@ -12,6 +12,8 @@ from pathlib import Path
import markdown
from PIL import Image, ImageOps
from modules.utils import get_available_chat_styles
# This is to store the paths to the thumbnails of the profile pictures
image_cache = {}
@ -19,13 +21,14 @@ with open(Path(__file__).resolve().parent / '../css/html_readable_style.css', 'r
readable_css = f.read()
with open(Path(__file__).resolve().parent / '../css/html_4chan_style.css', 'r') as css_f:
_4chan_css = css_f.read()
with open(Path(__file__).resolve().parent / '../css/html_cai_style.css', 'r') as f:
cai_css = f.read()
with open(Path(__file__).resolve().parent / '../css/html_bubble_chat_style.css', 'r') as f:
bubble_chat_css = f.read()
with open(Path(__file__).resolve().parent / '../css/html_instruct_style.css', 'r') as f:
instruct_css = f.read()
# Custom chat styles
chat_styles = {}
for k in get_available_chat_styles():
chat_styles[k] = open(Path(f'css/chat_style-{k}.css'), 'r').read()
def fix_newlines(string):
string = string.replace('\n', '\n\n')
@ -185,8 +188,8 @@ def generate_instruct_html(history):
return output
def generate_cai_chat_html(history, name1, name2, reset_cache=False):
output = f'<style>{cai_css}</style><div class="chat" id="chat">'
def generate_cai_chat_html(history, name1, name2, style, reset_cache=False):
output = f'<style>{chat_styles[style]}</style><div class="chat" id="chat">'
# We use ?name2 and ?time.time() to force the browser to reset caches
img_bot = f'<img src="file/cache/pfp_character.png?{name2}">' if Path("cache/pfp_character.png").exists() else ''
@ -235,7 +238,7 @@ def generate_cai_chat_html(history, name1, name2, reset_cache=False):
def generate_chat_html(history, name1, name2, reset_cache=False):
output = f'<style>{bubble_chat_css}</style><div class="chat" id="chat">'
output = f'<style>{chat_styles["wpp"]}</style><div class="chat" id="chat">'
for i, _row in enumerate(history[::-1]):
row = [convert_to_markdown(entry) for entry in _row]
@ -267,12 +270,10 @@ def generate_chat_html(history, name1, name2, reset_cache=False):
return output
def chat_html_wrapper(history, name1, name2, mode, reset_cache=False):
if mode == "cai-chat":
return generate_cai_chat_html(history, name1, name2, reset_cache)
elif mode == "chat":
return generate_chat_html(history, name1, name2)
elif mode == "instruct":
def chat_html_wrapper(history, name1, name2, mode, style, reset_cache=False):
if mode == 'instruct':
return generate_instruct_html(history)
elif style == 'wpp':
return generate_chat_html(history, name1, name2)
else:
return ''
return generate_cai_chat_html(history, name1, name2, style, reset_cache)

9
modules/shared.py
View file

@ -49,7 +49,8 @@ settings = {
'truncation_length': 2048,
'truncation_length_min': 0,
'truncation_length_max': 8192,
'mode': 'cai-chat',
'mode': 'chat',
'chat_style': 'cai-chat',
'instruction_template': 'None',
'chat_prompt_size': 2048,
'chat_prompt_size_min': 0,
@ -95,7 +96,6 @@ parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpForma
# Basic settings
parser.add_argument('--notebook', action='store_true', help='Launch the web UI in notebook mode, where the output is written to the same text box as the input.')
parser.add_argument('--chat', action='store_true', help='Launch the web UI in chat mode with a style similar to the Character.AI website.')
parser.add_argument('--cai-chat', action='store_true', help='DEPRECATED: use --chat instead.')
parser.add_argument('--character', type=str, help='The name of the character to load in chat mode by default.')
parser.add_argument('--model', type=str, help='Name of the model to load by default.')
parser.add_argument('--lora', type=str, nargs="+", help='The list of LoRAs to load. If you want to load more than one LoRA, write the names separated by spaces.')
@ -176,11 +176,6 @@ for k in deprecated_dict:
logging.warning(f"--{k} is deprecated and will be removed. Use --{deprecated_dict[k][0]} instead.")
setattr(args, deprecated_dict[k][0], getattr(args, k))
# Deprecation warnings for parameters that have been removed
if args.cai_chat:
logging.warning("--cai-chat is deprecated. Use --chat instead.")
args.chat = True
# Security warnings
if args.trust_remote_code:
logging.warning("trust_remote_code is enabled. This is dangerous.")

2
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']
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 += list_model_elements()
return elements

4
modules/utils.py
View file

@ -59,3 +59,7 @@ def get_available_loras():
def get_datasets(path: str, ext: str):
return ['None'] + sorted(set([k.stem for k in Path(path).glob(f'*.{ext}') if k.stem != 'put-trainer-datasets-here']), key=natural_keys)
def get_available_chat_styles():
return sorted(set(('-'.join(k.stem.split('-')[1:]) for k in Path('css').glob('chat_style*.css'))), key=natural_keys)