From 98af4bfb0db142f72749c91ce4dadec7bf08e2f8 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 23 Feb 2023 12:05:25 -0300 Subject: [PATCH] Refactor the code to make it more modular --- api-example.py | 1 - convert-to-flexgen.py | 3 +- convert-to-safetensors.py | 1 - modules/chat.py | 369 ++++++++++++++++ modules/extensions.py | 41 ++ modules/html_generator.py | 2 - modules/prompt.py | 174 ++++++++ modules/shared.py | 39 ++ modules/stopping_criteria.py | 1 - server.py | 819 +++++------------------------------ 10 files changed, 737 insertions(+), 713 deletions(-) create mode 100644 modules/chat.py create mode 100644 modules/extensions.py create mode 100644 modules/prompt.py create mode 100644 modules/shared.py diff --git a/api-example.py b/api-example.py index a5967a9d..f2f6c51e 100644 --- a/api-example.py +++ b/api-example.py @@ -10,7 +10,6 @@ Optionally, you can also add the --share flag to generate a public gradio URL, allowing you to use the API remotely. ''' - import requests # Server address diff --git a/convert-to-flexgen.py b/convert-to-flexgen.py index 18afa9bd..043e1941 100644 --- a/convert-to-flexgen.py +++ b/convert-to-flexgen.py @@ -3,13 +3,12 @@ Converts a transformers model to a format compatible with flexgen. ''' - import argparse import os -import numpy as np from pathlib import Path from sys import argv +import numpy as np import torch from tqdm import tqdm from transformers import AutoModelForCausalLM diff --git a/convert-to-safetensors.py b/convert-to-safetensors.py index 60770843..f1506646 100644 --- a/convert-to-safetensors.py +++ b/convert-to-safetensors.py @@ -10,7 +10,6 @@ Based on the original script by 81300: https://gist.github.com/81300/fe5b08bff1cba45296a829b9d6b0f303 ''' - import argparse from pathlib import Path from sys import argv diff --git a/modules/chat.py b/modules/chat.py new file mode 100644 index 00000000..d02acaae --- /dev/null +++ b/modules/chat.py @@ -0,0 +1,369 @@ +import io +import json +import re +from datetime import datetime +from pathlib import Path + +import modules.shared as shared +from modules.extensions import apply_extensions +from modules.html_generator import * +from modules.prompt import encode +from modules.prompt import generate_reply +from modules.prompt import get_max_prompt_length + +history = {'internal': [], 'visible': []} +character = None + +# This gets the new line characters right. +def clean_chat_message(text): + text = text.replace('\n', '\n\n') + text = re.sub(r"\n{3,}", "\n\n", text) + text = text.strip() + return text + +def generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size, impersonate=False): + text = clean_chat_message(text) + rows = [f"{context.strip()}\n"] + i = len(history['internal'])-1 + count = 0 + + if shared.soft_prompt: + chat_prompt_size -= shared.soft_prompt_tensor.shape[1] + max_length = min(get_max_prompt_length(tokens), chat_prompt_size) + + while i >= 0 and len(encode(''.join(rows), tokens)[0]) < max_length: + rows.insert(1, f"{name2}: {history['internal'][i][1].strip()}\n") + count += 1 + if not (history['internal'][i][0] == '<|BEGIN-VISIBLE-CHAT|>'): + rows.insert(1, f"{name1}: {history['internal'][i][0].strip()}\n") + count += 1 + i -= 1 + + if not impersonate: + rows.append(f"{name1}: {text}\n") + rows.append(apply_extensions(f"{name2}:", "bot_prefix")) + limit = 3 + else: + rows.append(f"{name1}:") + limit = 2 + + while len(rows) > limit and len(encode(''.join(rows), tokens)[0]) >= max_length: + rows.pop(1) + rows.pop(1) + + question = ''.join(rows) + return question + +def extract_message_from_reply(question, reply, current, other, check, extensions=False): + next_character_found = False + substring_found = False + + previous_idx = [m.start() for m in re.finditer(f"(^|\n){re.escape(current)}:", question)] + idx = [m.start() for m in re.finditer(f"(^|\n){re.escape(current)}:", reply)] + idx = idx[len(previous_idx)-1] + + if extensions: + reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):] + else: + reply = reply[idx + 1 + len(f"{current}:"):] + + if check: + reply = reply.split('\n')[0].strip() + else: + idx = reply.find(f"\n{other}:") + if idx != -1: + reply = reply[:idx] + next_character_found = True + reply = clean_chat_message(reply) + + # Detect if something like "\nYo" is generated just before + # "\nYou:" is completed + tmp = f"\n{other}:" + for j in range(1, len(tmp)): + if reply[-j:] == tmp[:j]: + substring_found = True + + return reply, next_character_found, substring_found + +def generate_chat_picture(picture, name1, name2): + text = f'*{name1} sends {name2} a picture that contains the following: "{bot_picture.caption_image(picture)}"*' + buffer = BytesIO() + picture.save(buffer, format="JPEG") + img_str = base64.b64encode(buffer.getvalue()).decode('utf-8') + visible_text = f'' + return text, visible_text + +def stop_everything_event(): + global stop_everything + stop_everything = True + +def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): + global stop_everything + stop_everything = False + + if 'pygmalion' in shared.model_name.lower(): + name1 = "You" + + if shared.args.picture and picture is not None: + text, visible_text = generate_chat_picture(picture, name1, name2) + else: + visible_text = text + if shared.args.chat: + visible_text = visible_text.replace('\n', '
') + + text = apply_extensions(text, "input") + question = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size) + eos_token = '\n' if check else None + first = True + for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name1}:"): + reply, next_character_found, substring_found = extract_message_from_reply(question, reply, name2, name1, check, extensions=True) + visible_reply = apply_extensions(reply, "output") + if shared.args.chat: + visible_reply = visible_reply.replace('\n', '
') + + # We need this global variable to handle the Stop event, + # otherwise gradio gets confused + if stop_everything: + return history['visible'] + + if first: + first = False + history['internal'].append(['', '']) + history['visible'].append(['', '']) + + history['internal'][-1] = [text, reply] + history['visible'][-1] = [visible_text, visible_reply] + if not substring_found: + yield history['visible'] + if next_character_found: + break + yield history['visible'] + +def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): + if 'pygmalion' in shared.model_name.lower(): + name1 = "You" + + question = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size, impersonate=True) + eos_token = '\n' if check else None + for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name2}:"): + reply, next_character_found, substring_found = extract_message_from_reply(question, reply, name1, name2, check, extensions=False) + if not substring_found: + yield reply + if next_character_found: + break + yield reply + +def cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): + for _history in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture): + yield generate_chat_html(_history, name1, name2, character) + +def regenerate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): + if character is not None and len(history['visible']) == 1: + if shared.args.cai_chat: + yield generate_chat_html(history['visible'], name1, name2, character) + else: + yield history['visible'] + else: + last_visible = history['visible'].pop() + last_internal = history['internal'].pop() + + for _history in chatbot_wrapper(last_internal[0], tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture): + if shared.args.cai_chat: + history['visible'][-1] = [last_visible[0], _history[-1][1]] + yield generate_chat_html(history['visible'], name1, name2, character) + else: + history['visible'][-1] = (last_visible[0], _history[-1][1]) + yield history['visible'] + +def remove_last_message(name1, name2): + if not history['internal'][-1][0] == '<|BEGIN-VISIBLE-CHAT|>': + last = history['visible'].pop() + history['internal'].pop() + else: + last = ['', ''] + if shared.args.cai_chat: + return generate_chat_html(history['visible'], name1, name2, character), last[0] + else: + return history['visible'], last[0] + +def send_last_reply_to_input(): + if len(history['internal']) > 0: + return history['internal'][-1][1] + else: + return '' + +def replace_last_reply(text, name1, name2): + if len(history['visible']) > 0: + if shared.args.cai_chat: + history['visible'][-1][1] = text + else: + history['visible'][-1] = (history['visible'][-1][0], text) + history['internal'][-1][1] = apply_extensions(text, "input") + + if shared.args.cai_chat: + return generate_chat_html(history['visible'], name1, name2, character) + else: + return history['visible'] + +def clear_html(): + return generate_chat_html([], "", "", character) + +def clear_chat_log(_character, name1, name2): + global history + if _character != 'None': + for i in range(len(history['internal'])): + if '<|BEGIN-VISIBLE-CHAT|>' in history['internal'][i][0]: + history['visible'] = [['', history['internal'][i][1]]] + history['internal'] = history['internal'][:i+1] + break + else: + history['internal'] = [] + history['visible'] = [] + if shared.args.cai_chat: + return generate_chat_html(history['visible'], name1, name2, character) + else: + return history['visible'] + +def redraw_html(name1, name2): + global history + return generate_chat_html(history['visible'], name1, name2, character) + +def tokenize_dialogue(dialogue, name1, name2): + _history = [] + + dialogue = re.sub('', '', dialogue) + dialogue = re.sub('', '', dialogue) + dialogue = re.sub('(\n|^)[Aa]non:', '\\1You:', dialogue) + dialogue = re.sub('(\n|^)\[CHARACTER\]:', f'\\g<1>{name2}:', dialogue) + idx = [m.start() for m in re.finditer(f"(^|\n)({re.escape(name1)}|{re.escape(name2)}):", dialogue)] + if len(idx) == 0: + return _history + + messages = [] + for i in range(len(idx)-1): + messages.append(dialogue[idx[i]:idx[i+1]].strip()) + messages.append(dialogue[idx[-1]:].strip()) + + entry = ['', ''] + for i in messages: + if i.startswith(f'{name1}:'): + entry[0] = i[len(f'{name1}:'):].strip() + elif i.startswith(f'{name2}:'): + entry[1] = i[len(f'{name2}:'):].strip() + if not (len(entry[0]) == 0 and len(entry[1]) == 0): + _history.append(entry) + entry = ['', ''] + + print(f"\033[1;32;1m\nDialogue tokenized to:\033[0;37;0m\n", end='') + for row in _history: + for column in row: + print("\n") + for line in column.strip().split('\n'): + print("| "+line+"\n") + print("|\n") + print("------------------------------") + + return _history + +def save_history(timestamp=True): + if timestamp: + fname = f"{character or ''}{'_' if character else ''}{datetime.now().strftime('%Y%m%d-%H%M%S')}.json" + else: + fname = f"{character or ''}{'_' if character else ''}persistent.json" + if not Path('logs').exists(): + Path('logs').mkdir() + with open(Path(f'logs/{fname}'), 'w') as f: + f.write(json.dumps({'data': history['internal'], 'data_visible': history['visible']}, indent=2)) + return Path(f'logs/{fname}') + +def load_history(file, name1, name2): + global history + file = file.decode('utf-8') + try: + j = json.loads(file) + if 'data' in j: + history['internal'] = j['data'] + if 'data_visible' in j: + history['visible'] = j['data_visible'] + else: + history['visible'] = copy.deepcopy(history['internal']) + # Compatibility with Pygmalion AI's official web UI + elif 'chat' in j: + history['internal'] = [':'.join(x.split(':')[1:]).strip() for x in j['chat']] + if len(j['chat']) > 0 and j['chat'][0].startswith(f'{name2}:'): + history['internal'] = [['<|BEGIN-VISIBLE-CHAT|>', history['internal'][0]]] + [[history['internal'][i], history['internal'][i+1]] for i in range(1, len(history['internal'])-1, 2)] + history['visible'] = copy.deepcopy(history['internal']) + history['visible'][0][0] = '' + else: + history['internal'] = [[history['internal'][i], history['internal'][i+1]] for i in range(0, len(history['internal'])-1, 2)] + history['visible'] = copy.deepcopy(history['internal']) + except: + history['internal'] = tokenize_dialogue(file, name1, name2) + history['visible'] = copy.deepcopy(history['internal']) + +def load_character(_character, name1, name2): + global history, character + context = "" + history['internal'] = [] + history['visible'] = [] + if _character != 'None': + character = _character + data = json.loads(open(Path(f'characters/{_character}.json'), 'r').read()) + name2 = data['char_name'] + if 'char_persona' in data and data['char_persona'] != '': + context += f"{data['char_name']}'s Persona: {data['char_persona']}\n" + if 'world_scenario' in data and data['world_scenario'] != '': + context += f"Scenario: {data['world_scenario']}\n" + context = f"{context.strip()}\n\n" + if 'example_dialogue' in data and data['example_dialogue'] != '': + history['internal'] = tokenize_dialogue(data['example_dialogue'], name1, name2) + if 'char_greeting' in data and len(data['char_greeting'].strip()) > 0: + history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', data['char_greeting']]] + history['visible'] += [['', apply_extensions(data['char_greeting'], "output")]] + else: + history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', "Hello there!"]] + history['visible'] += [['', "Hello there!"]] + else: + character = None + context = settings['context_pygmalion'] + name2 = settings['name2_pygmalion'] + + if Path(f'logs/{character}_persistent.json').exists(): + load_history(open(Path(f'logs/{character}_persistent.json'), 'rb').read(), name1, name2) + + if shared.args.cai_chat: + return name2, context, generate_chat_html(history['visible'], name1, name2, character) + else: + return name2, context, history['visible'] + +def upload_character(json_file, img, tavern=False): + json_file = json_file if type(json_file) == str else json_file.decode('utf-8') + data = json.loads(json_file) + outfile_name = data["char_name"] + i = 1 + while Path(f'characters/{outfile_name}.json').exists(): + outfile_name = f'{data["char_name"]}_{i:03d}' + i += 1 + if tavern: + outfile_name = f'TavernAI-{outfile_name}' + with open(Path(f'characters/{outfile_name}.json'), 'w') as f: + f.write(json_file) + if img is not None: + img = Image.open(io.BytesIO(img)) + img.save(Path(f'characters/{outfile_name}.png')) + print(f'New character saved to "characters/{outfile_name}.json".') + return outfile_name + +def upload_tavern_character(img, name1, name2): + _img = Image.open(io.BytesIO(img)) + _img.getexif() + decoded_string = base64.b64decode(_img.info['chara']) + _json = json.loads(decoded_string) + _json = {"char_name": _json['name'], "char_persona": _json['description'], "char_greeting": _json["first_mes"], "example_dialogue": _json['mes_example'], "world_scenario": _json['scenario']} + _json['example_dialogue'] = _json['example_dialogue'].replace('{{user}}', name1).replace('{{char}}', _json['char_name']) + return upload_character(json.dumps(_json), img, tavern=True) + +def upload_your_profile_picture(img): + img = Image.open(io.BytesIO(img)) + img.save(Path(f'img_me.png')) + print(f'Profile picture saved to "img_me.png"') diff --git a/modules/extensions.py b/modules/extensions.py new file mode 100644 index 00000000..eeb10bcd --- /dev/null +++ b/modules/extensions.py @@ -0,0 +1,41 @@ +import modules.shared as shared + +import extensions + +extension_state = {} +available_extensions = [] + +def apply_extensions(text, typ): + for ext in sorted(extension_state, key=lambda x : extension_state[x][1]): + if extension_state[ext][0] == True: + ext_string = f"extensions.{ext}.script" + if typ == "input" and hasattr(eval(ext_string), "input_modifier"): + text = eval(f"{ext_string}.input_modifier(text)") + elif typ == "output" and hasattr(eval(ext_string), "output_modifier"): + text = eval(f"{ext_string}.output_modifier(text)") + elif typ == "bot_prefix" and hasattr(eval(ext_string), "bot_prefix_modifier"): + text = eval(f"{ext_string}.bot_prefix_modifier(text)") + return text + +def update_extensions_parameters(*kwargs): + i = 0 + for ext in sorted(extension_state, key=lambda x : extension_state[x][1]): + if extension_state[ext][0] == True: + params = eval(f"extensions.{ext}.script.params") + for param in params: + if len(kwargs) >= i+1: + params[param] = eval(f"kwargs[{i}]") + i += 1 + +def load_extensions(): + global extension_state + for i,ext in enumerate(shared.args.extensions.split(',')): + if ext in available_extensions: + print(f'Loading the extension "{ext}"... ', end='') + ext_string = f"extensions.{ext}.script" + exec(f"import {ext_string}") + extension_state[ext] = [True, i] + print(f'Ok.') + +def get_params(name): + return eval(f"extensions.{name}.script.params") diff --git a/modules/html_generator.py b/modules/html_generator.py index f0e26392..ed9996fc 100644 --- a/modules/html_generator.py +++ b/modules/html_generator.py @@ -3,9 +3,7 @@ This is a library for formatting GPT-4chan and chat outputs as nice HTML. ''' - import base64 -import copy import os import re from io import BytesIO diff --git a/modules/prompt.py b/modules/prompt.py new file mode 100644 index 00000000..b95897aa --- /dev/null +++ b/modules/prompt.py @@ -0,0 +1,174 @@ +import time + +import modules.shared as shared +import torch +import transformers +from modules.extensions import apply_extensions +from modules.html_generator import * +from modules.stopping_criteria import _SentinelTokenStoppingCriteria +from tqdm import tqdm + + +def get_max_prompt_length(tokens): + max_length = 2048-tokens + if shared.soft_prompt: + max_length -= shared.soft_prompt_tensor.shape[1] + return max_length + +def encode(prompt, tokens_to_generate=0, add_special_tokens=True): + input_ids = shared.tokenizer.encode(str(prompt), return_tensors='pt', truncation=True, max_length=get_max_prompt_length(tokens_to_generate), add_special_tokens=add_special_tokens) + if shared.args.cpu or shared.args.flexgen: + return input_ids + elif shared.args.deepspeed: + return input_ids.to(device=local_rank) + else: + return input_ids.cuda() + +def decode(output_ids): + reply = shared.tokenizer.decode(output_ids, skip_special_tokens=True) + reply = reply.replace(r'<|endoftext|>', '') + return reply + +def generate_softprompt_input_tensors(input_ids): + inputs_embeds = shared.model.transformer.wte(input_ids) + inputs_embeds = torch.cat((shared.soft_prompt_tensor, inputs_embeds), dim=1) + filler_input_ids = torch.zeros((1, inputs_embeds.shape[1]), dtype=input_ids.dtype).to(shared.model.device) + filler_input_ids += shared.model.config.bos_token_id # setting dummy input_ids to bos tokens + return inputs_embeds, filler_input_ids + +# Removes empty replies from gpt4chan outputs +def fix_gpt4chan(s): + for i in range(10): + s = re.sub("--- [0-9]*\n>>[0-9]*\n---", "---", s) + s = re.sub("--- [0-9]*\n *\n---", "---", s) + s = re.sub("--- [0-9]*\n\n\n---", "---", s) + return s + +# Fix the LaTeX equations in galactica +def fix_galactica(s): + s = s.replace(r'\[', r'$') + s = s.replace(r'\]', r'$') + s = s.replace(r'\(', r'$') + s = s.replace(r'\)', r'$') + s = s.replace(r'$$', r'$') + s = re.sub(r'\n', r'\n\n', s) + s = re.sub(r"\n{3,}", "\n\n", s) + return s + +def formatted_outputs(reply, model_name): + if not (shared.args.chat or shared.args.cai_chat): + if shared.model_name.lower().startswith('galactica'): + reply = fix_galactica(reply) + return reply, reply, generate_basic_html(reply) + elif shared.model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')): + reply = fix_gpt4chan(reply) + return reply, 'Only applicable for GALACTICA models.', generate_4chan_html(reply) + else: + return reply, 'Only applicable for GALACTICA models.', generate_basic_html(reply) + else: + return reply + +def generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=None, stopping_string=None): + original_question = question + if not (shared.args.chat or shared.args.cai_chat): + question = apply_extensions(question, "input") + if shared.args.verbose: + print(f"\n\n{question}\n--------------------\n") + + input_ids = encode(question, tokens) + cuda = "" if (shared.args.cpu or shared.args.deepspeed or shared.args.flexgen) else ".cuda()" + if not shared.args.flexgen: + n = shared.tokenizer.eos_token_id if eos_token is None else shared.tokenizer.encode(eos_token, return_tensors='pt')[0][-1] + else: + n = shared.tokenizer(eos_token).input_ids[0] if eos_token else None + + if stopping_string is not None: + # The stopping_criteria code below was copied from + # https://github.com/PygmalionAI/gradio-ui/blob/master/src/model.py + t = encode(stopping_string, 0, add_special_tokens=False) + stopping_criteria_list = transformers.StoppingCriteriaList([ + _SentinelTokenStoppingCriteria( + sentinel_token_ids=t, + starting_idx=len(input_ids[0]) + ) + ]) + else: + stopping_criteria_list = None + + if not shared.args.flexgen: + generate_params = [ + f"eos_token_id={n}", + f"stopping_criteria=stopping_criteria_list", + f"do_sample={do_sample}", + f"temperature={temperature}", + f"top_p={top_p}", + f"typical_p={typical_p}", + f"repetition_penalty={repetition_penalty}", + f"top_k={top_k}", + f"min_length={min_length if shared.args.no_stream else 0}", + f"no_repeat_ngram_size={no_repeat_ngram_size}", + f"num_beams={num_beams}", + f"penalty_alpha={penalty_alpha}", + f"length_penalty={length_penalty}", + f"early_stopping={early_stopping}", + ] + else: + generate_params = [ + f"do_sample={do_sample}", + f"temperature={temperature}", + f"stop={n}", + ] + + if shared.args.deepspeed: + generate_params.append("synced_gpus=True") + if shared.args.no_stream: + generate_params.append(f"max_new_tokens=tokens") + else: + generate_params.append(f"max_new_tokens=8") + + if shared.soft_prompt: + inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) + generate_params.insert(0, "inputs_embeds=inputs_embeds") + generate_params.insert(0, "filler_input_ids") + else: + generate_params.insert(0, "input_ids") + + # Generate the entire reply at once + if shared.args.no_stream: + t0 = time.time() + with torch.no_grad(): + output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0] + if shared.soft_prompt: + output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) + + reply = decode(output) + if not (shared.args.chat or shared.args.cai_chat): + reply = original_question + apply_extensions(reply[len(question):], "output") + yield formatted_outputs(reply, shared.model_name) + + t1 = time.time() + print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output)-len(input_ids[0]))/(t1-t0)/8:.2f} it/s, {len(output)-len(input_ids[0])} tokens)") + + # Generate the reply 8 tokens at a time + else: + yield formatted_outputs(original_question, shared.model_name) + for i in tqdm(range(tokens//8+1)): + with torch.no_grad(): + output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0] + if shared.soft_prompt: + output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) + + reply = decode(output) + if not (shared.args.chat or shared.args.cai_chat): + reply = original_question + apply_extensions(reply[len(question):], "output") + yield formatted_outputs(reply, shared.model_name) + + if not shared.args.flexgen: + input_ids = torch.reshape(output, (1, output.shape[0])) + else: + input_ids = np.reshape(output, (1, output.shape[0])) + if shared.soft_prompt: + inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) + + if output[-1] == n: + break diff --git a/modules/shared.py b/modules/shared.py new file mode 100644 index 00000000..72622894 --- /dev/null +++ b/modules/shared.py @@ -0,0 +1,39 @@ +import argparse + +global tokenizer + +model = None +tokenizer = None +model_name = "" +soft_prompt_tensor = None +soft_prompt = False +stop_everything = False + +parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog,max_help_position=54)) +parser.add_argument('--model', type=str, help='Name of the model to load by default.') +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.') +parser.add_argument('--cai-chat', action='store_true', help='Launch the web UI in chat mode with a style similar to Character.AI\'s. If the file img_bot.png or img_bot.jpg exists in the same folder as server.py, this image will be used as the bot\'s profile picture. Similarly, img_me.png or img_me.jpg will be used as your profile picture.') +parser.add_argument('--picture', action='store_true', help='Adds an ability to send pictures in chat UI modes. Captions are generated by BLIP.') +parser.add_argument('--cpu', action='store_true', help='Use the CPU to generate text.') +parser.add_argument('--load-in-8bit', action='store_true', help='Load the model with 8-bit precision.') +parser.add_argument('--bf16', action='store_true', help='Load the model with bfloat16 precision. Requires NVIDIA Ampere GPU.') +parser.add_argument('--auto-devices', action='store_true', help='Automatically split the model across the available GPU(s) and CPU.') +parser.add_argument('--disk', action='store_true', help='If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk.') +parser.add_argument('--disk-cache-dir', type=str, default="cache", help='Directory to save the disk cache to. Defaults to "cache".') +parser.add_argument('--gpu-memory', type=int, help='Maximum GPU memory in GiB to allocate. This is useful if you get out of memory errors while trying to generate text. Must be an integer number.') +parser.add_argument('--cpu-memory', type=int, help='Maximum CPU memory in GiB to allocate for offloaded weights. Must be an integer number. Defaults to 99.') +parser.add_argument('--flexgen', action='store_true', help='Enable the use of FlexGen offloading.') +parser.add_argument('--percent', nargs="+", type=int, default=[0, 100, 100, 0, 100, 0], help='FlexGen: allocation percentages. Must be 6 numbers separated by spaces (default: 0, 100, 100, 0, 100, 0).') +parser.add_argument("--compress-weight", action="store_true", help="FlexGen: activate weight compression.") +parser.add_argument('--deepspeed', action='store_true', help='Enable the use of DeepSpeed ZeRO-3 for inference via the Transformers integration.') +parser.add_argument('--nvme-offload-dir', type=str, help='DeepSpeed: Directory to use for ZeRO-3 NVME offloading.') +parser.add_argument('--local_rank', type=int, default=0, help='DeepSpeed: Optional argument for distributed setups.') +parser.add_argument('--no-stream', action='store_true', help='Don\'t stream the text output in real time. This improves the text generation performance.') +parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.') +parser.add_argument('--extensions', type=str, help='The list of extensions to load. If you want to load more than one extension, write the names separated by commas and between quotation marks, "like,this".') +parser.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.') +parser.add_argument('--listen-port', type=int, help='The listening port that the server will use.') +parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.') +parser.add_argument('--verbose', action='store_true', help='Print the prompts to the terminal.') +args = parser.parse_args() diff --git a/modules/stopping_criteria.py b/modules/stopping_criteria.py index 3baadf6c..3e403ffe 100644 --- a/modules/stopping_criteria.py +++ b/modules/stopping_criteria.py @@ -4,7 +4,6 @@ This code was copied from https://github.com/PygmalionAI/gradio-ui/ ''' - import torch import transformers diff --git a/server.py b/server.py index d1066917..d1a0a01f 100644 --- a/server.py +++ b/server.py @@ -1,17 +1,11 @@ -import argparse -import base64 -import copy import gc -import glob import io import json import os import re import sys import time -import warnings import zipfile -from datetime import datetime from pathlib import Path import gradio as gr @@ -19,48 +13,23 @@ import numpy as np import torch import transformers from PIL import Image -from tqdm import tqdm from transformers import AutoConfig from transformers import AutoModelForCausalLM from transformers import AutoTokenizer -from io import BytesIO +import modules.chat as chat +import modules.extensions as extensions_module +import modules.shared as shared +from modules.extensions import extension_state +from modules.extensions import load_extensions +from modules.extensions import update_extensions_parameters from modules.html_generator import * -from modules.stopping_criteria import _SentinelTokenStoppingCriteria +from modules.prompt import generate_reply from modules.ui import * transformers.logging.set_verbosity_error() -parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog,max_help_position=54)) -parser.add_argument('--model', type=str, help='Name of the model to load by default.') -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.') -parser.add_argument('--cai-chat', action='store_true', help='Launch the web UI in chat mode with a style similar to Character.AI\'s. If the file img_bot.png or img_bot.jpg exists in the same folder as server.py, this image will be used as the bot\'s profile picture. Similarly, img_me.png or img_me.jpg will be used as your profile picture.') -parser.add_argument('--picture', action='store_true', help='Adds an ability to send pictures in chat UI modes. Captions are generated by BLIP.') -parser.add_argument('--cpu', action='store_true', help='Use the CPU to generate text.') -parser.add_argument('--load-in-8bit', action='store_true', help='Load the model with 8-bit precision.') -parser.add_argument('--bf16', action='store_true', help='Load the model with bfloat16 precision. Requires NVIDIA Ampere GPU.') -parser.add_argument('--auto-devices', action='store_true', help='Automatically split the model across the available GPU(s) and CPU.') -parser.add_argument('--disk', action='store_true', help='If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk.') -parser.add_argument('--disk-cache-dir', type=str, default="cache", help='Directory to save the disk cache to. Defaults to "cache".') -parser.add_argument('--gpu-memory', type=int, help='Maximum GPU memory in GiB to allocate. This is useful if you get out of memory errors while trying to generate text. Must be an integer number.') -parser.add_argument('--cpu-memory', type=int, help='Maximum CPU memory in GiB to allocate for offloaded weights. Must be an integer number. Defaults to 99.') -parser.add_argument('--flexgen', action='store_true', help='Enable the use of FlexGen offloading.') -parser.add_argument('--percent', nargs="+", type=int, default=[0, 100, 100, 0, 100, 0], help='FlexGen: allocation percentages. Must be 6 numbers separated by spaces (default: 0, 100, 100, 0, 100, 0).') -parser.add_argument("--compress-weight", action="store_true", help="FlexGen: activate weight compression.") -parser.add_argument('--deepspeed', action='store_true', help='Enable the use of DeepSpeed ZeRO-3 for inference via the Transformers integration.') -parser.add_argument('--nvme-offload-dir', type=str, help='DeepSpeed: Directory to use for ZeRO-3 NVME offloading.') -parser.add_argument('--local_rank', type=int, default=0, help='DeepSpeed: Optional argument for distributed setups.') -parser.add_argument('--no-stream', action='store_true', help='Don\'t stream the text output in real time. This improves the text generation performance.') -parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.') -parser.add_argument('--extensions', type=str, help='The list of extensions to load. If you want to load more than one extension, write the names separated by commas and between quotation marks, "like,this".') -parser.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.') -parser.add_argument('--listen-port', type=int, help='The listening port that the server will use.') -parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.') -parser.add_argument('--verbose', action='store_true', help='Print the prompts to the terminal.') -args = parser.parse_args() - -if (args.chat or args.cai_chat) and not args.no_stream: +if (shared.args.chat or shared.args.cai_chat) and not shared.args.no_stream: print("Warning: chat mode currently becomes somewhat slower with text streaming on.\nConsider starting the web UI with the --no-stream option.\n") settings = { @@ -84,28 +53,28 @@ settings = { 'stop_at_newline_pygmalion': False, } -if args.settings is not None and Path(args.settings).exists(): - new_settings = json.loads(open(Path(args.settings), 'r').read()) +if shared.args.settings is not None and Path(shared.args.settings).exists(): + new_settings = json.loads(open(Path(shared.args.settings), 'r').read()) for item in new_settings: settings[item] = new_settings[item] -if args.flexgen: +if shared.args.flexgen: from flexgen.flex_opt import (Policy, OptLM, TorchDevice, TorchDisk, TorchMixedDevice, CompressionConfig, Env, Task, get_opt_config) -if args.deepspeed: +if shared.args.deepspeed: import deepspeed from transformers.deepspeed import HfDeepSpeedConfig, is_deepspeed_zero3_enabled from modules.deepspeed_parameters import generate_ds_config # Distributed setup - local_rank = args.local_rank if args.local_rank is not None else int(os.getenv("LOCAL_RANK", "0")) + local_rank = shared.args.local_rank if shared.args.local_rank is not None else int(os.getenv("LOCAL_RANK", "0")) world_size = int(os.getenv("WORLD_SIZE", "1")) torch.cuda.set_device(local_rank) deepspeed.init_distributed() - ds_config = generate_ds_config(args.bf16, 1 * world_size, args.nvme_offload_dir) + ds_config = generate_ds_config(shared.args.bf16, 1 * world_size, shared.args.nvme_offload_dir) dschf = HfDeepSpeedConfig(ds_config) # Keep this object alive for the Transformers integration -if args.picture and (args.cai_chat or args.chat): +if shared.args.picture and (shared.args.cai_chat or shared.args.chat): import modules.bot_picture as bot_picture def load_model(model_name): @@ -113,27 +82,27 @@ def load_model(model_name): t0 = time.time() # Default settings - if not (args.cpu or args.load_in_8bit or args.auto_devices or args.disk or args.gpu_memory is not None or args.cpu_memory is not None or args.deepspeed or args.flexgen): - if any(size in model_name.lower() for size in ('13b', '20b', '30b')): - model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), device_map='auto', load_in_8bit=True) + if not (shared.args.cpu or shared.args.load_in_8bit or shared.args.auto_devices or shared.args.disk or shared.args.gpu_memory is not None or shared.args.cpu_memory is not None or shared.args.deepspeed or shared.args.flexgen): + if any(size in shared.model_name.lower() for size in ('13b', '20b', '30b')): + model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), device_map='auto', load_in_8bit=True) else: - model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 if args.bf16 else torch.float16).cuda() + model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16).cuda() # FlexGen - elif args.flexgen: + elif shared.args.flexgen: gpu = TorchDevice("cuda:0") cpu = TorchDevice("cpu") - disk = TorchDisk(args.disk_cache_dir) + disk = TorchDisk(shared.args.disk_cache_dir) env = Env(gpu=gpu, cpu=cpu, disk=disk, mixed=TorchMixedDevice([gpu, cpu, disk])) # Offloading policy policy = Policy(1, 1, - args.percent[0], args.percent[1], - args.percent[2], args.percent[3], - args.percent[4], args.percent[5], + shared.args.percent[0], shared.args.percent[1], + shared.args.percent[2], shared.args.percent[3], + shared.args.percent[4], shared.args.percent[5], overlap=True, sep_layer=True, pin_weight=True, cpu_cache_compute=False, attn_sparsity=1.0, - compress_weight=args.compress_weight, + compress_weight=shared.args.compress_weight, comp_weight_config=CompressionConfig( num_bits=4, group_size=64, group_dim=0, symmetric=False), @@ -142,13 +111,13 @@ def load_model(model_name): num_bits=4, group_size=64, group_dim=2, symmetric=False)) - opt_config = get_opt_config(f"facebook/{model_name}") + opt_config = get_opt_config(f"facebook/{shared.model_name}") model = OptLM(opt_config, env, "models", policy) model.init_all_weights() # DeepSpeed ZeRO-3 - elif args.deepspeed: - model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), torch_dtype=torch.bfloat16 if args.bf16 else torch.float16) + elif shared.args.deepspeed: + model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16) model = deepspeed.initialize(model=model, config_params=ds_config, model_parameters=None, optimizer=None, lr_scheduler=None)[0] model.module.eval() # Inference print(f"DeepSpeed ZeRO-3 is enabled: {is_deepspeed_zero3_enabled()}") @@ -157,49 +126,47 @@ def load_model(model_name): else: command = "AutoModelForCausalLM.from_pretrained" params = ["low_cpu_mem_usage=True"] - if not args.cpu and not torch.cuda.is_available(): + if not shared.args.cpu and not torch.cuda.is_available(): print("Warning: no GPU has been detected.\nFalling back to CPU mode.\n") - args.cpu = True + shared.args.cpu = True - if args.cpu: + if shared.args.cpu: params.append("low_cpu_mem_usage=True") params.append("torch_dtype=torch.float32") else: params.append("device_map='auto'") - params.append("load_in_8bit=True" if args.load_in_8bit else "torch_dtype=torch.bfloat16" if args.bf16 else "torch_dtype=torch.float16") + params.append("load_in_8bit=True" if shared.args.load_in_8bit else "torch_dtype=torch.bfloat16" if shared.args.bf16 else "torch_dtype=torch.float16") - if args.gpu_memory: - params.append(f"max_memory={{0: '{args.gpu_memory or '99'}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}") - elif not args.load_in_8bit: + if shared.args.gpu_memory: + params.append(f"max_memory={{0: '{shared.args.gpu_memory or '99'}GiB', 'cpu': '{shared.args.cpu_memory or '99'}GiB'}}") + elif not shared.args.load_in_8bit: total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024)) suggestion = round((total_mem-1000)/1000)*1000 if total_mem-suggestion < 800: suggestion -= 1000 suggestion = int(round(suggestion/1000)) print(f"\033[1;32;1mAuto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors.\nYou can manually set other values.\033[0;37;0m") - params.append(f"max_memory={{0: '{suggestion}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}") - if args.disk: - params.append(f"offload_folder='{args.disk_cache_dir}'") + params.append(f"max_memory={{0: '{suggestion}GiB', 'cpu': '{shared.args.cpu_memory or '99'}GiB'}}") + if shared.args.disk: + params.append(f"offload_folder='{shared.args.disk_cache_dir}'") - command = f"{command}(Path(f'models/{model_name}'), {', '.join(set(params))})" + command = f"{command}(Path(f'models/{shared.model_name}'), {', '.join(set(params))})" model = eval(command) # Loading the tokenizer - if model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) and Path(f"models/gpt-j-6B/").exists(): + if shared.model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) and Path(f"models/gpt-j-6B/").exists(): tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/")) else: - tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{model_name}/")) + tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{shared.model_name}/")) tokenizer.truncation_side = 'left' print(f"Loaded the model in {(time.time()-t0):.2f} seconds.") return model, tokenizer def load_soft_prompt(name): - global soft_prompt, soft_prompt_tensor - if name == 'None': - soft_prompt = False - soft_prompt_tensor = None + shared.soft_prompt = False + shared.soft_prompt_tensor = None else: with zipfile.ZipFile(Path(f'softprompts/{name}.zip')) as zf: zf.extract('tensor.npy') @@ -216,11 +183,11 @@ def load_soft_prompt(name): tensor = np.load('tensor.npy') Path('tensor.npy').unlink() Path('meta.json').unlink() - tensor = torch.Tensor(tensor).to(device=model.device, dtype=model.dtype) + tensor = torch.Tensor(tensor).to(device=shared.model.device, dtype=shared.model.dtype) tensor = torch.reshape(tensor, (1, tensor.shape[0], tensor.shape[1])) - soft_prompt = True - soft_prompt_tensor = tensor + shared.soft_prompt = True + shared.soft_prompt_tensor = tensor return name @@ -237,15 +204,13 @@ def upload_soft_prompt(file): return name def load_model_wrapper(selected_model): - global model_name, model, tokenizer - - if selected_model != model_name: - model_name = selected_model - model = tokenizer = None - if not args.cpu: + if selected_model != shared.model_name: + shared.model_name = selected_model + model = shared.tokenizer = None + if not shared.args.cpu: gc.collect() torch.cuda.empty_cache() - model, tokenizer = load_model(model_name) + shared.model, shared.tokenizer = load_model(shared.model_name) return selected_model @@ -278,196 +243,6 @@ def load_preset_values(preset_menu, return_dict=False): else: return generate_params['do_sample'], generate_params['temperature'], generate_params['top_p'], generate_params['typical_p'], generate_params['repetition_penalty'], generate_params['top_k'], generate_params['min_length'], generate_params['no_repeat_ngram_size'], generate_params['num_beams'], generate_params['penalty_alpha'], generate_params['length_penalty'], generate_params['early_stopping'] -# Removes empty replies from gpt4chan outputs -def fix_gpt4chan(s): - for i in range(10): - s = re.sub("--- [0-9]*\n>>[0-9]*\n---", "---", s) - s = re.sub("--- [0-9]*\n *\n---", "---", s) - s = re.sub("--- [0-9]*\n\n\n---", "---", s) - return s - -# Fix the LaTeX equations in galactica -def fix_galactica(s): - s = s.replace(r'\[', r'$') - s = s.replace(r'\]', r'$') - s = s.replace(r'\(', r'$') - s = s.replace(r'\)', r'$') - s = s.replace(r'$$', r'$') - s = re.sub(r'\n', r'\n\n', s) - s = re.sub(r"\n{3,}", "\n\n", s) - return s - -def get_max_prompt_length(tokens): - global soft_prompt, soft_prompt_tensor - max_length = 2048-tokens - if soft_prompt: - max_length -= soft_prompt_tensor.shape[1] - return max_length - -def encode(prompt, tokens_to_generate=0, add_special_tokens=True): - input_ids = tokenizer.encode(str(prompt), return_tensors='pt', truncation=True, max_length=get_max_prompt_length(tokens_to_generate), add_special_tokens=add_special_tokens) - if args.cpu or args.flexgen: - return input_ids - elif args.deepspeed: - return input_ids.to(device=local_rank) - else: - return input_ids.cuda() - -def decode(output_ids): - reply = tokenizer.decode(output_ids, skip_special_tokens=True) - reply = reply.replace(r'<|endoftext|>', '') - return reply - -def formatted_outputs(reply, model_name): - if not (args.chat or args.cai_chat): - if model_name.lower().startswith('galactica'): - reply = fix_galactica(reply) - return reply, reply, generate_basic_html(reply) - elif model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')): - reply = fix_gpt4chan(reply) - return reply, 'Only applicable for GALACTICA models.', generate_4chan_html(reply) - else: - return reply, 'Only applicable for GALACTICA models.', generate_basic_html(reply) - else: - return reply - -def generate_softprompt_input_tensors(input_ids): - inputs_embeds = model.transformer.wte(input_ids) - inputs_embeds = torch.cat((soft_prompt_tensor, inputs_embeds), dim=1) - filler_input_ids = torch.zeros((1, inputs_embeds.shape[1]), dtype=input_ids.dtype).to(model.device) - filler_input_ids += model.config.bos_token_id # setting dummy input_ids to bos tokens - return inputs_embeds, filler_input_ids - -def generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=None, stopping_string=None): - global model_name, model, tokenizer, soft_prompt, soft_prompt_tensor - - original_question = question - if not (args.chat or args.cai_chat): - question = apply_extensions(question, "input") - if args.verbose: - print(f"\n\n{question}\n--------------------\n") - - input_ids = encode(question, tokens) - cuda = "" if (args.cpu or args.deepspeed or args.flexgen) else ".cuda()" - if not args.flexgen: - n = tokenizer.eos_token_id if eos_token is None else tokenizer.encode(eos_token, return_tensors='pt')[0][-1] - else: - n = tokenizer(eos_token).input_ids[0] if eos_token else None - - if stopping_string is not None: - # The stopping_criteria code below was copied from - # https://github.com/PygmalionAI/gradio-ui/blob/master/src/model.py - t = encode(stopping_string, 0, add_special_tokens=False) - stopping_criteria_list = transformers.StoppingCriteriaList([ - _SentinelTokenStoppingCriteria( - sentinel_token_ids=t, - starting_idx=len(input_ids[0]) - ) - ]) - else: - stopping_criteria_list = None - - if not args.flexgen: - generate_params = [ - f"eos_token_id={n}", - f"stopping_criteria=stopping_criteria_list", - f"do_sample={do_sample}", - f"temperature={temperature}", - f"top_p={top_p}", - f"typical_p={typical_p}", - f"repetition_penalty={repetition_penalty}", - f"top_k={top_k}", - f"min_length={min_length if args.no_stream else 0}", - f"no_repeat_ngram_size={no_repeat_ngram_size}", - f"num_beams={num_beams}", - f"penalty_alpha={penalty_alpha}", - f"length_penalty={length_penalty}", - f"early_stopping={early_stopping}", - ] - else: - generate_params = [ - f"do_sample={do_sample}", - f"temperature={temperature}", - f"stop={n}", - ] - - if args.deepspeed: - generate_params.append("synced_gpus=True") - if args.no_stream: - generate_params.append(f"max_new_tokens=tokens") - else: - generate_params.append(f"max_new_tokens=8") - - if soft_prompt: - inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) - generate_params.insert(0, "inputs_embeds=inputs_embeds") - generate_params.insert(0, "filler_input_ids") - else: - generate_params.insert(0, "input_ids") - - # Generate the entire reply at once - if args.no_stream: - t0 = time.time() - with torch.no_grad(): - output = eval(f"model.generate({', '.join(generate_params)}){cuda}")[0] - if soft_prompt: - output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) - - reply = decode(output) - if not (args.chat or args.cai_chat): - reply = original_question + apply_extensions(reply[len(question):], "output") - yield formatted_outputs(reply, model_name) - - t1 = time.time() - print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output)-len(input_ids[0]))/(t1-t0)/8:.2f} it/s, {len(output)-len(input_ids[0])} tokens)") - - # Generate the reply 8 tokens at a time - else: - yield formatted_outputs(original_question, model_name) - for i in tqdm(range(tokens//8+1)): - with torch.no_grad(): - output = eval(f"model.generate({', '.join(generate_params)}){cuda}")[0] - if soft_prompt: - output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:])) - - reply = decode(output) - if not (args.chat or args.cai_chat): - reply = original_question + apply_extensions(reply[len(question):], "output") - yield formatted_outputs(reply, model_name) - - if not args.flexgen: - input_ids = torch.reshape(output, (1, output.shape[0])) - else: - input_ids = np.reshape(output, (1, output.shape[0])) - if soft_prompt: - inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids) - - if output[-1] == n: - break - -def apply_extensions(text, typ): - global available_extensions, extension_state - for ext in sorted(extension_state, key=lambda x : extension_state[x][1]): - if extension_state[ext][0] == True: - ext_string = f"extensions.{ext}.script" - if typ == "input" and hasattr(eval(ext_string), "input_modifier"): - text = eval(f"{ext_string}.input_modifier(text)") - elif typ == "output" and hasattr(eval(ext_string), "output_modifier"): - text = eval(f"{ext_string}.output_modifier(text)") - elif typ == "bot_prefix" and hasattr(eval(ext_string), "bot_prefix_modifier"): - text = eval(f"{ext_string}.bot_prefix_modifier(text)") - return text - -def update_extensions_parameters(*kwargs): - i = 0 - for ext in sorted(extension_state, key=lambda x : extension_state[x][1]): - if extension_state[ext][0] == True: - params = eval(f"extensions.{ext}.script.params") - for param in params: - if len(kwargs) >= i+1: - params[param] = eval(f"kwargs[{i}]") - i += 1 - def get_available_models(): return sorted([item.name for item in list(Path('models/').glob('*')) if not item.name.endswith(('.txt', '-np'))], key=str.lower) @@ -486,11 +261,11 @@ def get_available_softprompts(): def create_extensions_block(): extensions_ui_elements = [] default_values = [] - if not (args.chat or args.cai_chat): + if not (shared.args.chat or shared.args.cai_chat): gr.Markdown('## Extensions parameters') for ext in sorted(extension_state, key=lambda x : extension_state[x][1]): if extension_state[ext][0] == True: - params = eval(f"extensions.{ext}.script.params") + params = extensions_module.get_params(ext) for param in params: _id = f"{ext}-{param}" default_value = settings[_id] if _id in settings else params[param] @@ -507,16 +282,16 @@ def create_extensions_block(): btn_extensions.click(update_extensions_parameters, [*extensions_ui_elements], []) def create_settings_menus(): - generate_params = load_preset_values(settings[f'preset{suffix}'] if not args.flexgen else 'Naive', return_dict=True) + generate_params = load_preset_values(settings[f'preset{suffix}'] if not shared.args.flexgen else 'Naive', return_dict=True) with gr.Row(): with gr.Column(): with gr.Row(): - model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model') + model_menu = gr.Dropdown(choices=available_models, value=shared.model_name, label='Model') create_refresh_button(model_menu, lambda : None, lambda : {"choices": get_available_models()}, "refresh-button") with gr.Column(): with gr.Row(): - preset_menu = gr.Dropdown(choices=available_presets, value=settings[f'preset{suffix}'] if not args.flexgen else 'Naive', label='Generation parameters preset') + preset_menu = gr.Dropdown(choices=available_presets, value=settings[f'preset{suffix}'] if not shared.args.flexgen else 'Naive', label='Generation parameters preset') create_refresh_button(preset_menu, lambda : None, lambda : {"choices": get_available_presets()}, "refresh-button") with gr.Accordion("Custom generation parameters", open=False, elem_id="accordion"): @@ -531,7 +306,7 @@ def create_settings_menus(): no_repeat_ngram_size = gr.Slider(0, 20, step=1, value=generate_params["no_repeat_ngram_size"], label="no_repeat_ngram_size") with gr.Row(): typical_p = gr.Slider(0.0,1.0,value=generate_params['typical_p'],step=0.01,label="typical_p") - min_length = gr.Slider(0, 2000, step=1, value=generate_params["min_length"] if args.no_stream else 0, label="min_length", interactive=args.no_stream) + min_length = gr.Slider(0, 2000, step=1, value=generate_params["min_length"] if shared.args.no_stream else 0, label="min_length", interactive=shared.args.no_stream) gr.Markdown("Contrastive search:") penalty_alpha = gr.Slider(0, 5, value=generate_params["penalty_alpha"], label="penalty_alpha") @@ -557,381 +332,18 @@ def create_settings_menus(): upload_softprompt.upload(upload_soft_prompt, [upload_softprompt], [softprompts_menu]) return preset_menu, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping -# This gets the new line characters right. -def clean_chat_message(text): - text = text.replace('\n', '\n\n') - text = re.sub(r"\n{3,}", "\n\n", text) - text = text.strip() - return text - -def generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size, impersonate=False): - global soft_prompt, soft_prompt_tensor - - text = clean_chat_message(text) - rows = [f"{context.strip()}\n"] - i = len(history['internal'])-1 - count = 0 - - if soft_prompt: - chat_prompt_size -= soft_prompt_tensor.shape[1] - max_length = min(get_max_prompt_length(tokens), chat_prompt_size) - - while i >= 0 and len(encode(''.join(rows), tokens)[0]) < max_length: - rows.insert(1, f"{name2}: {history['internal'][i][1].strip()}\n") - count += 1 - if not (history['internal'][i][0] == '<|BEGIN-VISIBLE-CHAT|>'): - rows.insert(1, f"{name1}: {history['internal'][i][0].strip()}\n") - count += 1 - i -= 1 - - if not impersonate: - rows.append(f"{name1}: {text}\n") - rows.append(apply_extensions(f"{name2}:", "bot_prefix")) - limit = 3 - else: - rows.append(f"{name1}:") - limit = 2 - - while len(rows) > limit and len(encode(''.join(rows), tokens)[0]) >= max_length: - rows.pop(1) - rows.pop(1) - - question = ''.join(rows) - return question - -def extract_message_from_reply(question, reply, current, other, check, extensions=False): - next_character_found = False - substring_found = False - - previous_idx = [m.start() for m in re.finditer(f"(^|\n){re.escape(current)}:", question)] - idx = [m.start() for m in re.finditer(f"(^|\n){re.escape(current)}:", reply)] - idx = idx[len(previous_idx)-1] - - if extensions: - reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):] - else: - reply = reply[idx + 1 + len(f"{current}:"):] - - if check: - reply = reply.split('\n')[0].strip() - else: - idx = reply.find(f"\n{other}:") - if idx != -1: - reply = reply[:idx] - next_character_found = True - reply = clean_chat_message(reply) - - # Detect if something like "\nYo" is generated just before - # "\nYou:" is completed - tmp = f"\n{other}:" - for j in range(1, len(tmp)): - if reply[-j:] == tmp[:j]: - substring_found = True - - return reply, next_character_found, substring_found - -def generate_chat_picture(picture, name1, name2): - text = f'*{name1} sends {name2} a picture that contains the following: "{bot_picture.caption_image(picture)}"*' - buffer = BytesIO() - picture.save(buffer, format="JPEG") - img_str = base64.b64encode(buffer.getvalue()).decode('utf-8') - visible_text = f'' - return text, visible_text - -def stop_everything_event(): - global stop_everything - stop_everything = True - -def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): - global stop_everything - stop_everything = False - - if 'pygmalion' in model_name.lower(): - name1 = "You" - - if args.picture and picture is not None: - text, visible_text = generate_chat_picture(picture, name1, name2) - else: - visible_text = text - if args.chat: - visible_text = visible_text.replace('\n', '
') - - text = apply_extensions(text, "input") - question = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size) - eos_token = '\n' if check else None - first = True - for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name1}:"): - reply, next_character_found, substring_found = extract_message_from_reply(question, reply, name2, name1, check, extensions=True) - visible_reply = apply_extensions(reply, "output") - if args.chat: - visible_reply = visible_reply.replace('\n', '
') - - # We need this global variable to handle the Stop event, - # otherwise gradio gets confused - if stop_everything: - return history['visible'] - - if first: - first = False - history['internal'].append(['', '']) - history['visible'].append(['', '']) - - history['internal'][-1] = [text, reply] - history['visible'][-1] = [visible_text, visible_reply] - if not substring_found: - yield history['visible'] - if next_character_found: - break - yield history['visible'] - -def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): - if 'pygmalion' in model_name.lower(): - name1 = "You" - - question = generate_chat_prompt(text, tokens, name1, name2, context, chat_prompt_size, impersonate=True) - eos_token = '\n' if check else None - for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name2}:"): - reply, next_character_found, substring_found = extract_message_from_reply(question, reply, name1, name2, check, extensions=False) - if not substring_found: - yield reply - if next_character_found: - break - yield reply - -def cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): - for _history in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture): - yield generate_chat_html(_history, name1, name2, character) - -def regenerate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture=None): - if character is not None and len(history['visible']) == 1: - if args.cai_chat: - yield generate_chat_html(history['visible'], name1, name2, character) - else: - yield history['visible'] - else: - last_visible = history['visible'].pop() - last_internal = history['internal'].pop() - - for _history in chatbot_wrapper(last_internal[0], tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size, picture): - if args.cai_chat: - history['visible'][-1] = [last_visible[0], _history[-1][1]] - yield generate_chat_html(history['visible'], name1, name2, character) - else: - history['visible'][-1] = (last_visible[0], _history[-1][1]) - yield history['visible'] - -def remove_last_message(name1, name2): - if not history['internal'][-1][0] == '<|BEGIN-VISIBLE-CHAT|>': - last = history['visible'].pop() - history['internal'].pop() - else: - last = ['', ''] - if args.cai_chat: - return generate_chat_html(history['visible'], name1, name2, character), last[0] - else: - return history['visible'], last[0] - -def send_last_reply_to_input(): - if len(history['internal']) > 0: - return history['internal'][-1][1] - else: - return '' - -def replace_last_reply(text, name1, name2): - if len(history['visible']) > 0: - if args.cai_chat: - history['visible'][-1][1] = text - else: - history['visible'][-1] = (history['visible'][-1][0], text) - history['internal'][-1][1] = apply_extensions(text, "input") - - if args.cai_chat: - return generate_chat_html(history['visible'], name1, name2, character) - else: - return history['visible'] - -def clear_html(): - return generate_chat_html([], "", "", character) - -def clear_chat_log(_character, name1, name2): - global history - if _character != 'None': - for i in range(len(history['internal'])): - if '<|BEGIN-VISIBLE-CHAT|>' in history['internal'][i][0]: - history['visible'] = [['', history['internal'][i][1]]] - history['internal'] = history['internal'][:i+1] - break - else: - history['internal'] = [] - history['visible'] = [] - if args.cai_chat: - return generate_chat_html(history['visible'], name1, name2, character) - else: - return history['visible'] - -def redraw_html(name1, name2): - global history - return generate_chat_html(history['visible'], name1, name2, character) - -def tokenize_dialogue(dialogue, name1, name2): - _history = [] - - dialogue = re.sub('', '', dialogue) - dialogue = re.sub('', '', dialogue) - dialogue = re.sub('(\n|^)[Aa]non:', '\\1You:', dialogue) - dialogue = re.sub('(\n|^)\[CHARACTER\]:', f'\\g<1>{name2}:', dialogue) - idx = [m.start() for m in re.finditer(f"(^|\n)({re.escape(name1)}|{re.escape(name2)}):", dialogue)] - if len(idx) == 0: - return _history - - messages = [] - for i in range(len(idx)-1): - messages.append(dialogue[idx[i]:idx[i+1]].strip()) - messages.append(dialogue[idx[-1]:].strip()) - - entry = ['', ''] - for i in messages: - if i.startswith(f'{name1}:'): - entry[0] = i[len(f'{name1}:'):].strip() - elif i.startswith(f'{name2}:'): - entry[1] = i[len(f'{name2}:'):].strip() - if not (len(entry[0]) == 0 and len(entry[1]) == 0): - _history.append(entry) - entry = ['', ''] - - print(f"\033[1;32;1m\nDialogue tokenized to:\033[0;37;0m\n", end='') - for row in _history: - for column in row: - print("\n") - for line in column.strip().split('\n'): - print("| "+line+"\n") - print("|\n") - print("------------------------------") - - return _history - -def save_history(timestamp=True): - if timestamp: - fname = f"{character or ''}{'_' if character else ''}{datetime.now().strftime('%Y%m%d-%H%M%S')}.json" - else: - fname = f"{character or ''}{'_' if character else ''}persistent.json" - if not Path('logs').exists(): - Path('logs').mkdir() - with open(Path(f'logs/{fname}'), 'w') as f: - f.write(json.dumps({'data': history['internal'], 'data_visible': history['visible']}, indent=2)) - return Path(f'logs/{fname}') - -def load_history(file, name1, name2): - global history - file = file.decode('utf-8') - try: - j = json.loads(file) - if 'data' in j: - history['internal'] = j['data'] - if 'data_visible' in j: - history['visible'] = j['data_visible'] - else: - history['visible'] = copy.deepcopy(history['internal']) - # Compatibility with Pygmalion AI's official web UI - elif 'chat' in j: - history['internal'] = [':'.join(x.split(':')[1:]).strip() for x in j['chat']] - if len(j['chat']) > 0 and j['chat'][0].startswith(f'{name2}:'): - history['internal'] = [['<|BEGIN-VISIBLE-CHAT|>', history['internal'][0]]] + [[history['internal'][i], history['internal'][i+1]] for i in range(1, len(history['internal'])-1, 2)] - history['visible'] = copy.deepcopy(history['internal']) - history['visible'][0][0] = '' - else: - history['internal'] = [[history['internal'][i], history['internal'][i+1]] for i in range(0, len(history['internal'])-1, 2)] - history['visible'] = copy.deepcopy(history['internal']) - except: - history['internal'] = tokenize_dialogue(file, name1, name2) - history['visible'] = copy.deepcopy(history['internal']) - -def load_character(_character, name1, name2): - global history, character - context = "" - history['internal'] = [] - history['visible'] = [] - if _character != 'None': - character = _character - data = json.loads(open(Path(f'characters/{_character}.json'), 'r').read()) - name2 = data['char_name'] - if 'char_persona' in data and data['char_persona'] != '': - context += f"{data['char_name']}'s Persona: {data['char_persona']}\n" - if 'world_scenario' in data and data['world_scenario'] != '': - context += f"Scenario: {data['world_scenario']}\n" - context = f"{context.strip()}\n\n" - if 'example_dialogue' in data and data['example_dialogue'] != '': - history['internal'] = tokenize_dialogue(data['example_dialogue'], name1, name2) - if 'char_greeting' in data and len(data['char_greeting'].strip()) > 0: - history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', data['char_greeting']]] - history['visible'] += [['', apply_extensions(data['char_greeting'], "output")]] - else: - history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', "Hello there!"]] - history['visible'] += [['', "Hello there!"]] - else: - character = None - context = settings['context_pygmalion'] - name2 = settings['name2_pygmalion'] - - if Path(f'logs/{character}_persistent.json').exists(): - load_history(open(Path(f'logs/{character}_persistent.json'), 'rb').read(), name1, name2) - - if args.cai_chat: - return name2, context, generate_chat_html(history['visible'], name1, name2, character) - else: - return name2, context, history['visible'] - -def upload_character(json_file, img, tavern=False): - json_file = json_file if type(json_file) == str else json_file.decode('utf-8') - data = json.loads(json_file) - outfile_name = data["char_name"] - i = 1 - while Path(f'characters/{outfile_name}.json').exists(): - outfile_name = f'{data["char_name"]}_{i:03d}' - i += 1 - if tavern: - outfile_name = f'TavernAI-{outfile_name}' - with open(Path(f'characters/{outfile_name}.json'), 'w') as f: - f.write(json_file) - if img is not None: - img = Image.open(io.BytesIO(img)) - img.save(Path(f'characters/{outfile_name}.png')) - print(f'New character saved to "characters/{outfile_name}.json".') - return outfile_name - -def upload_tavern_character(img, name1, name2): - _img = Image.open(io.BytesIO(img)) - _img.getexif() - decoded_string = base64.b64decode(_img.info['chara']) - _json = json.loads(decoded_string) - _json = {"char_name": _json['name'], "char_persona": _json['description'], "char_greeting": _json["first_mes"], "example_dialogue": _json['mes_example'], "world_scenario": _json['scenario']} - _json['example_dialogue'] = _json['example_dialogue'].replace('{{user}}', name1).replace('{{char}}', _json['char_name']) - return upload_character(json.dumps(_json), img, tavern=True) - -def upload_your_profile_picture(img): - img = Image.open(io.BytesIO(img)) - img.save(Path(f'img_me.png')) - print(f'Profile picture saved to "img_me.png"') - # Global variables available_models = get_available_models() available_presets = get_available_presets() available_characters = get_available_characters() -available_extensions = get_available_extensions() +extensions_module.available_extensions = get_available_extensions() available_softprompts = get_available_softprompts() -extension_state = {} -if args.extensions is not None: - for i,ext in enumerate(args.extensions.split(',')): - if ext in available_extensions: - print(f'Loading the extension "{ext}"... ', end='') - ext_string = f"extensions.{ext}.script" - exec(f"import {ext_string}") - extension_state[ext] = [True, i] - print(f'Ok.') +if shared.args.extensions is not None: + load_extensions() # Choosing the default model -if args.model is not None: - model_name = args.model +if shared.args.model is not None: + shared.model_name = shared.args.model else: if len(available_models) == 0: print("No models are available! Please download at least one.") @@ -945,38 +357,33 @@ else: print(f"\nWhich one do you want to load? 1-{len(available_models)}\n") i = int(input())-1 print() - model_name = available_models[i] -model, tokenizer = load_model(model_name) + shared.model_name = available_models[i] +shared.model, shared.tokenizer = load_model(shared.model_name) loaded_preset = None -soft_prompt_tensor = None -soft_prompt = False -stop_everything = False # UI settings -if model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')): +if shared.model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')): default_text = settings['prompt_gpt4chan'] -elif re.match('(rosey|chip|joi)_.*_instruct.*', model_name.lower()) is not None: +elif re.match('(rosey|chip|joi)_.*_instruct.*', shared.model_name.lower()) is not None: default_text = 'User: \n' else: default_text = settings['prompt'] description = f"\n\n# Text generation lab\nGenerate text using Large Language Models.\n" -suffix = '_pygmalion' if 'pygmalion' in model_name.lower() else '' +suffix = '_pygmalion' if 'pygmalion' in shared.model_name.lower() else '' buttons = {} gen_events = [] -history = {'internal': [], 'visible': []} -character = None -if args.chat or args.cai_chat: +if shared.args.chat or shared.args.cai_chat: if Path(f'logs/persistent.json').exists(): - load_history(open(Path(f'logs/persistent.json'), 'rb').read(), settings[f'name1{suffix}'], settings[f'name2{suffix}']) + chat.load_history(open(Path(f'logs/persistent.json'), 'rb').read(), settings[f'name1{suffix}'], settings[f'name2{suffix}']) with gr.Blocks(css=css+chat_css, analytics_enabled=False) as interface: - if args.cai_chat: - display = gr.HTML(value=generate_chat_html(history['visible'], settings[f'name1{suffix}'], settings[f'name2{suffix}'], character)) + if shared.args.cai_chat: + display = gr.HTML(value=generate_chat_html(chat.history['visible'], settings[f'name1{suffix}'], settings[f'name2{suffix}'], chat.character)) else: - display = gr.Chatbot(value=history['visible']) + display = gr.Chatbot(value=chat.history['visible']) textbox = gr.Textbox(label='Input') with gr.Row(): buttons["Stop"] = gr.Button("Stop") @@ -989,7 +396,7 @@ if args.chat or args.cai_chat: with gr.Row(): buttons["Send last reply to input"] = gr.Button("Send last reply to input") buttons["Replace last reply"] = gr.Button("Replace last reply") - if args.picture: + if shared.args.picture: with gr.Row(): picture_select = gr.Image(label="Send a picture", type='pil') @@ -1036,52 +443,52 @@ if args.chat or args.cai_chat: preset_menu, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping = create_settings_menus() - if args.extensions is not None: + if shared.args.extensions is not None: with gr.Tab("Extensions"): create_extensions_block() input_params = [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, chat_prompt_size_slider] - if args.picture: + if shared.args.picture: input_params.append(picture_select) - function_call = "cai_chatbot_wrapper" if args.cai_chat else "chatbot_wrapper" + function_call = "chat.cai_chatbot_wrapper" if shared.args.cai_chat else "chat.chatbot_wrapper" - gen_events.append(buttons["Generate"].click(eval(function_call), input_params, display, show_progress=args.no_stream, api_name="textgen")) - gen_events.append(textbox.submit(eval(function_call), input_params, display, show_progress=args.no_stream)) - if args.picture: - picture_select.upload(eval(function_call), input_params, display, show_progress=args.no_stream) - gen_events.append(buttons["Regenerate"].click(regenerate_wrapper, input_params, display, show_progress=args.no_stream)) - gen_events.append(buttons["Impersonate"].click(impersonate_wrapper, input_params, textbox, show_progress=args.no_stream)) - buttons["Stop"].click(stop_everything_event, [], [], cancels=gen_events) + gen_events.append(buttons["Generate"].click(eval(function_call), input_params, display, show_progress=shared.args.no_stream, api_name="textgen")) + gen_events.append(textbox.submit(eval(function_call), input_params, display, show_progress=shared.args.no_stream)) + if shared.args.picture: + picture_select.upload(eval(function_call), input_params, display, show_progress=shared.args.no_stream) + gen_events.append(buttons["Regenerate"].click(chat.regenerate_wrapper, input_params, display, show_progress=shared.args.no_stream)) + gen_events.append(buttons["Impersonate"].click(chat.impersonate_wrapper, input_params, textbox, show_progress=shared.args.no_stream)) + buttons["Stop"].click(chat.stop_everything_event, [], [], cancels=gen_events) - buttons["Send last reply to input"].click(send_last_reply_to_input, [], textbox, show_progress=args.no_stream) - buttons["Replace last reply"].click(replace_last_reply, [textbox, name1, name2], display, show_progress=args.no_stream) - buttons["Clear history"].click(clear_chat_log, [character_menu, name1, name2], display) - buttons["Remove last"].click(remove_last_message, [name1, name2], [display, textbox], show_progress=False) - buttons["Download"].click(save_history, inputs=[], outputs=[download]) - buttons["Upload character"].click(upload_character, [upload_char, upload_img], [character_menu]) + buttons["Send last reply to input"].click(chat.send_last_reply_to_input, [], textbox, show_progress=shared.args.no_stream) + buttons["Replace last reply"].click(chat.replace_last_reply, [textbox, name1, name2], display, show_progress=shared.args.no_stream) + buttons["Clear history"].click(chat.clear_chat_log, [character_menu, name1, name2], display) + buttons["Remove last"].click(chat.remove_last_message, [name1, name2], [display, textbox], show_progress=False) + buttons["Download"].click(chat.save_history, inputs=[], outputs=[download]) + buttons["Upload character"].click(chat.upload_character, [upload_char, upload_img], [character_menu]) # Clearing stuff and saving the history for i in ["Generate", "Regenerate", "Replace last reply"]: buttons[i].click(lambda x: "", textbox, textbox, show_progress=False) - buttons[i].click(lambda : save_history(timestamp=False), [], [], show_progress=False) - buttons["Clear history"].click(lambda : save_history(timestamp=False), [], [], show_progress=False) + buttons[i].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False) + buttons["Clear history"].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False) textbox.submit(lambda x: "", textbox, textbox, show_progress=False) - textbox.submit(lambda : save_history(timestamp=False), [], [], show_progress=False) + textbox.submit(lambda : chat.save_history(timestamp=False), [], [], show_progress=False) - character_menu.change(load_character, [character_menu, name1, name2], [name2, context, display]) - upload_chat_history.upload(load_history, [upload_chat_history, name1, name2], []) - upload_img_tavern.upload(upload_tavern_character, [upload_img_tavern, name1, name2], [character_menu]) - upload_img_me.upload(upload_your_profile_picture, [upload_img_me], []) - if args.picture: + character_menu.change(chat.load_character, [character_menu, name1, name2], [name2, context, display]) + upload_chat_history.upload(chat.load_history, [upload_chat_history, name1, name2], []) + upload_img_tavern.upload(chat.upload_tavern_character, [upload_img_tavern, name1, name2], [character_menu]) + upload_img_me.upload(chat.upload_your_profile_picture, [upload_img_me], []) + if shared.args.picture: picture_select.upload(lambda : None, [], [picture_select], show_progress=False) - if args.cai_chat: - upload_chat_history.upload(redraw_html, [name1, name2], [display]) - upload_img_me.upload(redraw_html, [name1, name2], [display]) + if shared.args.cai_chat: + upload_chat_history.upload(chat.redraw_html, [name1, name2], [display]) + upload_img_me.upload(chat.redraw_html, [name1, name2], [display]) else: - upload_chat_history.upload(lambda : history['visible'], [], [display]) - upload_img_me.upload(lambda : history['visible'], [], [display]) + upload_chat_history.upload(lambda : chat.history['visible'], [], [display]) + upload_img_me.upload(lambda : chat.history['visible'], [], [display]) -elif args.notebook: +elif shared.args.notebook: with gr.Blocks(css=css, analytics_enabled=False) as interface: gr.Markdown(description) with gr.Tab('Raw'): @@ -1098,11 +505,11 @@ elif args.notebook: preset_menu, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping = create_settings_menus() - if args.extensions is not None: + if shared.args.extensions is not None: create_extensions_block() - gen_events.append(buttons["Generate"].click(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [textbox, markdown, html], show_progress=args.no_stream, api_name="textgen")) - gen_events.append(textbox.submit(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [textbox, markdown, html], show_progress=args.no_stream)) + gen_events.append(buttons["Generate"].click(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [textbox, markdown, html], show_progress=shared.args.no_stream, api_name="textgen")) + gen_events.append(textbox.submit(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [textbox, markdown, html], show_progress=shared.args.no_stream)) buttons["Stop"].click(None, None, None, cancels=gen_events) else: @@ -1120,7 +527,7 @@ else: buttons["Stop"] = gr.Button("Stop") preset_menu, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping = create_settings_menus() - if args.extensions is not None: + if shared.args.extensions is not None: create_extensions_block() with gr.Column(): @@ -1131,16 +538,16 @@ else: with gr.Tab('HTML'): html = gr.HTML() - gen_events.append(buttons["Generate"].click(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=args.no_stream, api_name="textgen")) - gen_events.append(textbox.submit(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=args.no_stream)) - gen_events.append(buttons["Continue"].click(generate_reply, [output_textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=args.no_stream)) + gen_events.append(buttons["Generate"].click(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=shared.args.no_stream, api_name="textgen")) + gen_events.append(textbox.submit(generate_reply, [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=shared.args.no_stream)) + gen_events.append(buttons["Continue"].click(generate_reply, [output_textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping], [output_textbox, markdown, html], show_progress=shared.args.no_stream)) buttons["Stop"].click(None, None, None, cancels=gen_events) interface.queue() -if args.listen: - interface.launch(prevent_thread_lock=True, share=args.share, server_name="0.0.0.0", server_port=args.listen_port) +if shared.args.listen: + interface.launch(prevent_thread_lock=True, share=shared.args.share, server_name="0.0.0.0", server_port=shared.args.listen_port) else: - interface.launch(prevent_thread_lock=True, share=args.share, server_port=args.listen_port) + interface.launch(prevent_thread_lock=True, share=shared.args.share, server_port=shared.args.listen_port) # I think that I will need this later while True: