forestcalled/text-generation-webui
0
1import base642import copy3import functools4import html5import json6import re7from datetime import datetime8from functools import partial9from pathlib import Path10 11import gradio as gr12import yaml13from jinja2.sandbox import ImmutableSandboxedEnvironment14from PIL import Image15 16import modules.shared as shared17from modules.extensions import apply_extensions18from modules.html_generator import chat_html_wrapper, make_thumbnail19from modules.logging_colors import logger20from modules.text_generation import (21 generate_reply,22 get_encoded_length,23 get_max_prompt_length24)25from modules.utils import delete_file, get_available_characters, save_file26 27# Copied from the Transformers library28jinja_env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)29 30 31def str_presenter(dumper, data):32 """33 Copied from https://github.com/yaml/pyyaml/issues/24034 Makes pyyaml output prettier multiline strings.35 """36 37 if data.count('\n') > 0:38 return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')39 40 return dumper.represent_scalar('tag:yaml.org,2002:str', data)41 42 43yaml.add_representer(str, str_presenter)44yaml.representer.SafeRepresenter.add_representer(str, str_presenter)45 46 47def get_generation_prompt(renderer, impersonate=False, strip_trailing_spaces=True):48 '''49 Given a Jinja template, reverse-engineers the prefix and the suffix for50 an assistant message (if impersonate=False) or an user message51 (if impersonate=True)52 '''53 54 if impersonate:55 messages = [56 {"role": "user", "content": "<<|user-message-1|>>"},57 {"role": "user", "content": "<<|user-message-2|>>"},58 ]59 else:60 messages = [61 {"role": "assistant", "content": "<<|user-message-1|>>"},62 {"role": "assistant", "content": "<<|user-message-2|>>"},63 ]64 65 prompt = renderer(messages=messages)66 67 suffix_plus_prefix = prompt.split("<<|user-message-1|>>")[1].split("<<|user-message-2|>>")[0]68 suffix = prompt.split("<<|user-message-2|>>")[1]69 prefix = suffix_plus_prefix[len(suffix):]70 71 if strip_trailing_spaces:72 prefix = prefix.rstrip(' ')73 74 return prefix, suffix75 76 77def generate_chat_prompt(user_input, state, **kwargs):78 impersonate = kwargs.get('impersonate', False)79 _continue = kwargs.get('_continue', False)80 also_return_rows = kwargs.get('also_return_rows', False)81 history = kwargs.get('history', state['history'])['internal']82 83 # Templates84 chat_template = jinja_env.from_string(state['chat_template_str'])85 instruction_template = jinja_env.from_string(state['instruction_template_str'])86 chat_renderer = partial(chat_template.render, add_generation_prompt=False, name1=state['name1'], name2=state['name2'])87 instruct_renderer = partial(instruction_template.render, add_generation_prompt=False)88 89 messages = []90 91 if state['mode'] == 'instruct':92 renderer = instruct_renderer93 if state['custom_system_message'].strip() != '':94 messages.append({"role": "system", "content": state['custom_system_message']})95 else:96 renderer = chat_renderer97 if state['context'].strip() != '':98 context = replace_character_names(state['context'], state['name1'], state['name2'])99 messages.append({"role": "system", "content": context})100 101 insert_pos = len(messages)102 for user_msg, assistant_msg in reversed(history):103 user_msg = user_msg.strip()104 assistant_msg = assistant_msg.strip()105 106 if assistant_msg:107 messages.insert(insert_pos, {"role": "assistant", "content": assistant_msg})108 109 if user_msg not in ['', '<|BEGIN-VISIBLE-CHAT|>']:110 messages.insert(insert_pos, {"role": "user", "content": user_msg})111 112 user_input = user_input.strip()113 if user_input and not impersonate and not _continue:114 messages.append({"role": "user", "content": user_input})115 116 def remove_extra_bos(prompt):117 for bos_token in ['<s>', '<|startoftext|>']:118 while prompt.startswith(bos_token):119 prompt = prompt[len(bos_token):]120 121 return prompt122 123 def make_prompt(messages):124 if state['mode'] == 'chat-instruct' and _continue:125 prompt = renderer(messages=messages[:-1])126 else:127 prompt = renderer(messages=messages)128 129 if state['mode'] == 'chat-instruct':130 outer_messages = []131 if state['custom_system_message'].strip() != '':132 outer_messages.append({"role": "system", "content": state['custom_system_message']})133 134 prompt = remove_extra_bos(prompt)135 command = state['chat-instruct_command']136 command = command.replace('<|character|>', state['name2'] if not impersonate else state['name1'])137 command = command.replace('<|prompt|>', prompt)138 139 if _continue:140 prefix = get_generation_prompt(renderer, impersonate=impersonate, strip_trailing_spaces=False)[0]141 prefix += messages[-1]["content"]142 else:143 prefix = get_generation_prompt(renderer, impersonate=impersonate)[0]144 if not impersonate:145 prefix = apply_extensions('bot_prefix', prefix, state)146 147 outer_messages.append({"role": "user", "content": command})148 outer_messages.append({"role": "assistant", "content": prefix})149 150 prompt = instruction_template.render(messages=outer_messages)151 suffix = get_generation_prompt(instruct_renderer, impersonate=False)[1]152 prompt = prompt[:-len(suffix)]153 154 else:155 if _continue:156 suffix = get_generation_prompt(renderer, impersonate=impersonate)[1]157 prompt = prompt[:-len(suffix)]158 else:159 prefix = get_generation_prompt(renderer, impersonate=impersonate)[0]160 if state['mode'] == 'chat' and not impersonate:161 prefix = apply_extensions('bot_prefix', prefix, state)162 163 prompt += prefix164 165 prompt = remove_extra_bos(prompt)166 return prompt167 168 prompt = make_prompt(messages)169 170 # Handle truncation171 max_length = get_max_prompt_length(state)172 while len(messages) > 0 and get_encoded_length(prompt) > max_length:173 # Try to save the system message174 if len(messages) > 1 and messages[0]['role'] == 'system':175 messages.pop(1)176 else:177 messages.pop(0)178 179 prompt = make_prompt(messages)180 181 if also_return_rows:182 return prompt, [message['content'] for message in messages]183 else:184 return prompt185 186 187def get_stopping_strings(state):188 stopping_strings = []189 renderers = []190 191 if state['mode'] in ['instruct', 'chat-instruct']:192 template = jinja_env.from_string(state['instruction_template_str'])193 renderer = partial(template.render, add_generation_prompt=False)194 renderers.append(renderer)195 196 if state['mode'] in ['chat', 'chat-instruct']:197 template = jinja_env.from_string(state['chat_template_str'])198 renderer = partial(template.render, add_generation_prompt=False, name1=state['name1'], name2=state['name2'])199 renderers.append(renderer)200 201 for renderer in renderers:202 prefix_bot, suffix_bot = get_generation_prompt(renderer, impersonate=False)203 prefix_user, suffix_user = get_generation_prompt(renderer, impersonate=True)204 205 stopping_strings += [206 suffix_user + prefix_bot,207 suffix_user + prefix_user,208 suffix_bot + prefix_bot,209 suffix_bot + prefix_user,210 ]211 212 if 'stopping_strings' in state and isinstance(state['stopping_strings'], list):213 stopping_strings += state.pop('stopping_strings')214 215 return list(set(stopping_strings))216 217 218def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_message=True, for_ui=False):219 history = state['history']220 output = copy.deepcopy(history)221 output = apply_extensions('history', output)222 state = apply_extensions('state', state)223 224 visible_text = None225 stopping_strings = get_stopping_strings(state)226 is_stream = state['stream']227 228 # Prepare the input229 if not (regenerate or _continue):230 visible_text = html.escape(text)231 232 # Apply extensions233 text, visible_text = apply_extensions('chat_input', text, visible_text, state)234 text = apply_extensions('input', text, state, is_chat=True)235 236 output['internal'].append([text, ''])237 output['visible'].append([visible_text, ''])238 239 # *Is typing...*240 if loading_message:241 yield {242 'visible': output['visible'][:-1] + [[output['visible'][-1][0], shared.processing_message]],243 'internal': output['internal']244 }245 else:246 text, visible_text = output['internal'][-1][0], output['visible'][-1][0]247 if regenerate:248 if loading_message:249 yield {250 'visible': output['visible'][:-1] + [[visible_text, shared.processing_message]],251 'internal': output['internal'][:-1] + [[text, '']]252 }253 elif _continue:254 last_reply = [output['internal'][-1][1], output['visible'][-1][1]]255 if loading_message:256 yield {257 'visible': output['visible'][:-1] + [[visible_text, last_reply[1] + '...']],258 'internal': output['internal']259 }260 261 if shared.model_name == 'None' or shared.model is None:262 raise ValueError("No model is loaded! Select one in the Model tab.")263 264 # Generate the prompt265 kwargs = {266 '_continue': _continue,267 'history': output if _continue else {k: v[:-1] for k, v in output.items()}268 }269 prompt = apply_extensions('custom_generate_chat_prompt', text, state, **kwargs)270 if prompt is None:271 prompt = generate_chat_prompt(text, state, **kwargs)272 273 # Generate274 reply = None275 for j, reply in enumerate(generate_reply(prompt, state, stopping_strings=stopping_strings, is_chat=True, for_ui=for_ui)):276 277 # Extract the reply278 visible_reply = reply279 if state['mode'] in ['chat', 'chat-instruct']:280 visible_reply = re.sub("(<USER>|<user>|{{user}})", state['name1'], reply)281 282 visible_reply = html.escape(visible_reply)283 284 if shared.stop_everything:285 output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True)286 yield output287 return288 289 if _continue:290 output['internal'][-1] = [text, last_reply[0] + reply]291 output['visible'][-1] = [visible_text, last_reply[1] + visible_reply]292 if is_stream:293 yield output294 elif not (j == 0 and visible_reply.strip() == ''):295 output['internal'][-1] = [text, reply.lstrip(' ')]296 output['visible'][-1] = [visible_text, visible_reply.lstrip(' ')]297 if is_stream:298 yield output299 300 output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True)301 yield output302 303 304def impersonate_wrapper(text, state):305 306 static_output = chat_html_wrapper(state['history'], state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])307 308 if shared.model_name == 'None' or shared.model is None:309 logger.error("No model is loaded! Select one in the Model tab.")310 yield '', static_output311 return312 313 prompt = generate_chat_prompt('', state, impersonate=True)314 stopping_strings = get_stopping_strings(state)315 316 yield text + '...', static_output317 reply = None318 for reply in generate_reply(prompt + text, state, stopping_strings=stopping_strings, is_chat=True):319 yield (text + reply).lstrip(' '), static_output320 if shared.stop_everything:321 return322 323 324def generate_chat_reply(text, state, regenerate=False, _continue=False, loading_message=True, for_ui=False):325 history = state['history']326 if regenerate or _continue:327 text = ''328 if (len(history['visible']) == 1 and not history['visible'][0][0]) or len(history['internal']) == 0:329 yield history330 return331 332 for history in chatbot_wrapper(text, state, regenerate=regenerate, _continue=_continue, loading_message=loading_message, for_ui=for_ui):333 yield history334 335 336def character_is_loaded(state, raise_exception=False):337 if state['mode'] in ['chat', 'chat-instruct'] and state['name2'] == '':338 logger.error('It looks like no character is loaded. Please load one under Parameters > Character.')339 if raise_exception:340 raise ValueError341 342 return False343 else:344 return True345 346 347def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False):348 '''349 Same as above but returns HTML for the UI350 '''351 352 if not character_is_loaded(state):353 return354 355 if state['start_with'] != '' and not _continue:356 if regenerate:357 text, state['history'] = remove_last_message(state['history'])358 regenerate = False359 360 _continue = True361 send_dummy_message(text, state)362 send_dummy_reply(state['start_with'], state)363 364 for i, history in enumerate(generate_chat_reply(text, state, regenerate, _continue, loading_message=True, for_ui=True)):365 yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu']), history366 367 368def remove_last_message(history):369 if len(history['visible']) > 0 and history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>':370 last = history['visible'].pop()371 history['internal'].pop()372 else:373 last = ['', '']374 375 return html.unescape(last[0]), history376 377 378def send_last_reply_to_input(history):379 if len(history['visible']) > 0:380 return html.unescape(history['visible'][-1][1])381 else:382 return ''383 384 385def replace_last_reply(text, state):386 history = state['history']387 388 if len(text.strip()) == 0:389 return history390 elif len(history['visible']) > 0:391 history['visible'][-1][1] = html.escape(text)392 history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True)393 394 return history395 396 397def send_dummy_message(text, state):398 history = state['history']399 history['visible'].append([html.escape(text), ''])400 history['internal'].append([apply_extensions('input', text, state, is_chat=True), ''])401 return history402 403 404def send_dummy_reply(text, state):405 history = state['history']406 if len(history['visible']) > 0 and not history['visible'][-1][1] == '':407 history['visible'].append(['', ''])408 history['internal'].append(['', ''])409 410 history['visible'][-1][1] = html.escape(text)411 history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True)412 return history413 414 415def redraw_html(history, name1, name2, mode, style, character, reset_cache=False):416 return chat_html_wrapper(history, name1, name2, mode, style, character, reset_cache=reset_cache)417 418 419def start_new_chat(state):420 mode = state['mode']421 history = {'internal': [], 'visible': []}422 423 if mode != 'instruct':424 greeting = replace_character_names(state['greeting'], state['name1'], state['name2'])425 if greeting != '':426 history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', greeting]]427 history['visible'] += [['', apply_extensions('output', greeting, state, is_chat=True)]]428 429 unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')430 save_history(history, unique_id, state['character_menu'], state['mode'])431 432 return history433 434 435def get_history_file_path(unique_id, character, mode):436 if mode == 'instruct':437 p = Path(f'logs/instruct/{unique_id}.json')438 else:439 p = Path(f'logs/chat/{character}/{unique_id}.json')440 441 return p442 443 444def save_history(history, unique_id, character, mode):445 if shared.args.multi_user:446 return447 448 p = get_history_file_path(unique_id, character, mode)449 if not p.parent.is_dir():450 p.parent.mkdir(parents=True)451 452 with open(p, 'w', encoding='utf-8') as f:453 f.write(json.dumps(history, indent=4))454 455 456def rename_history(old_id, new_id, character, mode):457 if shared.args.multi_user:458 return459 460 old_p = get_history_file_path(old_id, character, mode)461 new_p = get_history_file_path(new_id, character, mode)462 if new_p.parent != old_p.parent:463 logger.error(f"The following path is not allowed: {new_p}.")464 elif new_p == old_p:465 logger.info("The provided path is identical to the old one.")466 else:467 logger.info(f"Renaming {old_p} to {new_p}")468 old_p.rename(new_p)469 470 471def find_all_histories(state):472 if shared.args.multi_user:473 return ['']474 475 if state['mode'] == 'instruct':476 paths = Path('logs/instruct').glob('*.json')477 else:478 character = state['character_menu']479 480 # Handle obsolete filenames and paths481 old_p = Path(f'logs/{character}_persistent.json')482 new_p = Path(f'logs/persistent_{character}.json')483 if old_p.exists():484 logger.warning(f"Renaming {old_p} to {new_p}")485 old_p.rename(new_p)486 if new_p.exists():487 unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')488 p = get_history_file_path(unique_id, character, state['mode'])489 logger.warning(f"Moving {new_p} to {p}")490 p.parent.mkdir(exist_ok=True)491 new_p.rename(p)492 493 paths = Path(f'logs/chat/{character}').glob('*.json')494 495 histories = sorted(paths, key=lambda x: x.stat().st_mtime, reverse=True)496 histories = [path.stem for path in histories]497 498 return histories499 500 501def load_latest_history(state):502 '''503 Loads the latest history for the given character in chat or chat-instruct504 mode, or the latest instruct history for instruct mode.505 '''506 507 if shared.args.multi_user:508 return start_new_chat(state)509 510 histories = find_all_histories(state)511 512 if len(histories) > 0:513 unique_id = Path(histories[0]).stem514 history = load_history(unique_id, state['character_menu'], state['mode'])515 else:516 history = start_new_chat(state)517 518 return history519 520 521def load_history(unique_id, character, mode):522 p = get_history_file_path(unique_id, character, mode)523 524 f = json.loads(open(p, 'rb').read())525 if 'internal' in f and 'visible' in f:526 history = f527 else:528 history = {529 'internal': f['data'],530 'visible': f['data_visible']531 }532 533 return history534 535 536def load_history_json(file, history):537 try:538 file = file.decode('utf-8')539 f = json.loads(file)540 if 'internal' in f and 'visible' in f:541 history = f542 else:543 history = {544 'internal': f['data'],545 'visible': f['data_visible']546 }547 548 return history549 except:550 return history551 552 553def delete_history(unique_id, character, mode):554 p = get_history_file_path(unique_id, character, mode)555 delete_file(p)556 557 558def replace_character_names(text, name1, name2):559 text = text.replace('{{user}}', name1).replace('{{char}}', name2)560 return text.replace('<USER>', name1).replace('<BOT>', name2)561 562 563def generate_pfp_cache(character):564 cache_folder = Path("cache")565 if not cache_folder.exists():566 cache_folder.mkdir()567 568 for path in [Path(f"characters/{character}.{extension}") for extension in ['png', 'jpg', 'jpeg']]:569 if path.exists():570 original_img = Image.open(path)571 original_img.save(Path('cache/pfp_character.png'), format='PNG')572 573 thumb = make_thumbnail(original_img)574 thumb.save(Path('cache/pfp_character_thumb.png'), format='PNG')575 576 return thumb577 578 return None579 580 581def load_character(character, name1, name2):582 context = greeting = ""583 greeting_field = 'greeting'584 picture = None585 586 filepath = None587 for extension in ["yml", "yaml", "json"]:588 filepath = Path(f'characters/{character}.{extension}')589 if filepath.exists():590 break591 592 if filepath is None or not filepath.exists():593 logger.error(f"Could not find the character \"{character}\" inside characters/. No character has been loaded.")594 raise ValueError595 596 file_contents = open(filepath, 'r', encoding='utf-8').read()597 data = json.loads(file_contents) if extension == "json" else yaml.safe_load(file_contents)598 599 for path in [Path("cache/pfp_character.png"), Path("cache/pfp_character_thumb.png")]:600 if path.exists():601 path.unlink()602 603 picture = generate_pfp_cache(character)604 605 # Finding the bot's name606 for k in ['name', 'bot', '<|bot|>', 'char_name']:607 if k in data and data[k] != '':608 name2 = data[k]609 break610 611 # Find the user name (if any)612 for k in ['your_name', 'user', '<|user|>']:613 if k in data and data[k] != '':614 name1 = data[k]615 break616 617 if 'context' in data:618 context = data['context'].strip()619 elif "char_persona" in data:620 context = build_pygmalion_style_context(data)621 greeting_field = 'char_greeting'622 623 greeting = data.get(greeting_field, greeting)624 return name1, name2, picture, greeting, context625 626 627def load_instruction_template(template):628 for filepath in [Path(f'instruction-templates/{template}.yaml'), Path('instruction-templates/Alpaca.yaml')]:629 if filepath.exists():630 break631 else:632 return ''633 634 file_contents = open(filepath, 'r', encoding='utf-8').read()635 data = yaml.safe_load(file_contents)636 if 'instruction_template' in data:637 return data['instruction_template']638 else:639 return jinja_template_from_old_format(data)640 641 642@functools.cache643def load_character_memoized(character, name1, name2):644 return load_character(character, name1, name2)645 646 647@functools.cache648def load_instruction_template_memoized(template):649 return load_instruction_template(template)650 651 652def upload_character(file, img, tavern=False):653 decoded_file = file if isinstance(file, str) else file.decode('utf-8')654 try:655 data = json.loads(decoded_file)656 except:657 data = yaml.safe_load(decoded_file)658 659 if 'char_name' in data:660 name = data['char_name']661 greeting = data['char_greeting']662 context = build_pygmalion_style_context(data)663 yaml_data = generate_character_yaml(name, greeting, context)664 else:665 name = data['name']666 yaml_data = generate_character_yaml(data['name'], data['greeting'], data['context'])667 668 outfile_name = name669 i = 1670 while Path(f'characters/{outfile_name}.yaml').exists():671 outfile_name = f'{name}_{i:03d}'672 i += 1673 674 with open(Path(f'characters/{outfile_name}.yaml'), 'w', encoding='utf-8') as f:675 f.write(yaml_data)676 677 if img is not None:678 img.save(Path(f'characters/{outfile_name}.png'))679 680 logger.info(f'New character saved to "characters/{outfile_name}.yaml".')681 return gr.update(value=outfile_name, choices=get_available_characters())682 683 684def build_pygmalion_style_context(data):685 context = ""686 if 'char_persona' in data and data['char_persona'] != '':687 context += f"{data['char_name']}'s Persona: {data['char_persona']}\n"688 689 if 'world_scenario' in data and data['world_scenario'] != '':690 context += f"Scenario: {data['world_scenario']}\n"691 692 if 'example_dialogue' in data and data['example_dialogue'] != '':693 context += f"{data['example_dialogue'].strip()}\n"694 695 context = f"{context.strip()}\n"696 return context697 698 699def upload_tavern_character(img, _json):700 _json = {'char_name': _json['name'], 'char_persona': _json['description'], 'char_greeting': _json['first_mes'], 'example_dialogue': _json['mes_example'], 'world_scenario': _json['scenario']}701 return upload_character(json.dumps(_json), img, tavern=True)702 703 704def check_tavern_character(img):705 if "chara" not in img.info:706 return "Not a TavernAI card", None, None, gr.update(interactive=False)707 708 decoded_string = base64.b64decode(img.info['chara']).replace(b'\\r\\n', b'\\n')709 _json = json.loads(decoded_string)710 if "data" in _json:711 _json = _json["data"]712 713 return _json['name'], _json['description'], _json, gr.update(interactive=True)714 715 716def upload_your_profile_picture(img):717 cache_folder = Path("cache")718 if not cache_folder.exists():719 cache_folder.mkdir()720 721 if img is None:722 if Path("cache/pfp_me.png").exists():723 Path("cache/pfp_me.png").unlink()724 else:725 img = make_thumbnail(img)726 img.save(Path('cache/pfp_me.png'))727 logger.info('Profile picture saved to "cache/pfp_me.png"')728 729 730def generate_character_yaml(name, greeting, context):731 data = {732 'name': name,733 'greeting': greeting,734 'context': context,735 }736 737 data = {k: v for k, v in data.items() if v} # Strip falsy738 return yaml.dump(data, sort_keys=False, width=float("inf"))739 740 741def generate_instruction_template_yaml(instruction_template):742 data = {743 'instruction_template': instruction_template744 }745 746 return my_yaml_output(data)747 748 749def save_character(name, greeting, context, picture, filename):750 if filename == "":751 logger.error("The filename is empty, so the character will not be saved.")752 return753 754 data = generate_character_yaml(name, greeting, context)755 filepath = Path(f'characters/{filename}.yaml')756 save_file(filepath, data)757 path_to_img = Path(f'characters/{filename}.png')758 if picture is not None:759 picture.save(path_to_img)760 logger.info(f'Saved {path_to_img}.')761 762 763def delete_character(name, instruct=False):764 for extension in ["yml", "yaml", "json"]:765 delete_file(Path(f'characters/{name}.{extension}'))766 767 delete_file(Path(f'characters/{name}.png'))768 769 770def jinja_template_from_old_format(params, verbose=False):771 MASTER_TEMPLATE = """772{%- set ns = namespace(found=false) -%}773{%- for message in messages -%}774 {%- if message['role'] == 'system' -%}775 {%- set ns.found = true -%}776 {%- endif -%}777{%- endfor -%}778{%- if not ns.found -%}779 {{- '<|PRE-SYSTEM|>' + '<|SYSTEM-MESSAGE|>' + '<|POST-SYSTEM|>' -}}780{%- endif %}781{%- for message in messages %}782 {%- if message['role'] == 'system' -%}783 {{- '<|PRE-SYSTEM|>' + message['content'] + '<|POST-SYSTEM|>' -}}784 {%- else -%}785 {%- if message['role'] == 'user' -%}786 {{-'<|PRE-USER|>' + message['content'] + '<|POST-USER|>'-}}787 {%- else -%}788 {{-'<|PRE-ASSISTANT|>' + message['content'] + '<|POST-ASSISTANT|>' -}}789 {%- endif -%}790 {%- endif -%}791{%- endfor -%}792{%- if add_generation_prompt -%}793 {{-'<|PRE-ASSISTANT-GENERATE|>'-}}794{%- endif -%}795"""796 797 if 'context' in params and '<|system-message|>' in params['context']:798 pre_system = params['context'].split('<|system-message|>')[0]799 post_system = params['context'].split('<|system-message|>')[1]800 else:801 pre_system = ''802 post_system = ''803 804 pre_user = params['turn_template'].split('<|user-message|>')[0].replace('<|user|>', params['user'])805 post_user = params['turn_template'].split('<|user-message|>')[1].split('<|bot|>')[0]806 807 pre_assistant = '<|bot|>' + params['turn_template'].split('<|bot-message|>')[0].split('<|bot|>')[1]808 pre_assistant = pre_assistant.replace('<|bot|>', params['bot'])809 post_assistant = params['turn_template'].split('<|bot-message|>')[1]810 811 def preprocess(string):812 return string.replace('\n', '\\n').replace('\'', '\\\'')813 814 pre_system = preprocess(pre_system)815 post_system = preprocess(post_system)816 pre_user = preprocess(pre_user)817 post_user = preprocess(post_user)818 pre_assistant = preprocess(pre_assistant)819 post_assistant = preprocess(post_assistant)820 821 if verbose:822 print(823 '\n',824 repr(pre_system) + '\n',825 repr(post_system) + '\n',826 repr(pre_user) + '\n',827 repr(post_user) + '\n',828 repr(pre_assistant) + '\n',829 repr(post_assistant) + '\n',830 )831 832 result = MASTER_TEMPLATE833 if 'system_message' in params:834 result = result.replace('<|SYSTEM-MESSAGE|>', preprocess(params['system_message']))835 else:836 result = result.replace('<|SYSTEM-MESSAGE|>', '')837 838 result = result.replace('<|PRE-SYSTEM|>', pre_system)839 result = result.replace('<|POST-SYSTEM|>', post_system)840 result = result.replace('<|PRE-USER|>', pre_user)841 result = result.replace('<|POST-USER|>', post_user)842 result = result.replace('<|PRE-ASSISTANT|>', pre_assistant)843 result = result.replace('<|PRE-ASSISTANT-GENERATE|>', pre_assistant.rstrip(' '))844 result = result.replace('<|POST-ASSISTANT|>', post_assistant)845 846 result = result.strip()847 848 return result849 850 851def my_yaml_output(data):852 '''853 pyyaml is very inconsistent with multiline strings.854 for simple instruction template outputs, this is enough.855 '''856 result = ""857 for k in data:858 result += k + ": |-\n"859 for line in data[k].splitlines():860 result += " " + line.rstrip(' ') + "\n"861 862 return result863