Refactor the code to make it more modular

This commit is contained in:
oobabooga 2023-02-23 12:05:25 -03:00
parent 18e0ec955e
commit 98af4bfb0d
10 changed files with 737 additions and 713 deletions

View File

@ -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. allowing you to use the API remotely.
''' '''
import requests import requests
# Server address # Server address

View File

@ -3,13 +3,12 @@
Converts a transformers model to a format compatible with flexgen. Converts a transformers model to a format compatible with flexgen.
''' '''
import argparse import argparse
import os import os
import numpy as np
from pathlib import Path from pathlib import Path
from sys import argv from sys import argv
import numpy as np
import torch import torch
from tqdm import tqdm from tqdm import tqdm
from transformers import AutoModelForCausalLM from transformers import AutoModelForCausalLM

View File

@ -10,7 +10,6 @@ Based on the original script by 81300:
https://gist.github.com/81300/fe5b08bff1cba45296a829b9d6b0f303 https://gist.github.com/81300/fe5b08bff1cba45296a829b9d6b0f303
''' '''
import argparse import argparse
from pathlib import Path from pathlib import Path
from sys import argv from sys import argv

369
modules/chat.py Normal file
View File

@ -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'<img src="data:image/jpeg;base64,{img_str}">'
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', '<br>')
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', '<br>')
# 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('<START>', '', dialogue)
dialogue = re.sub('<start>', '', 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<START>\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"')

41
modules/extensions.py Normal file
View File

@ -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")

View File

@ -3,9 +3,7 @@
This is a library for formatting GPT-4chan and chat outputs as nice HTML. This is a library for formatting GPT-4chan and chat outputs as nice HTML.
''' '''
import base64 import base64
import copy
import os import os
import re import re
from io import BytesIO from io import BytesIO

174
modules/prompt.py Normal file
View File

@ -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

39
modules/shared.py Normal file
View File

@ -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()

View File

@ -4,7 +4,6 @@ This code was copied from
https://github.com/PygmalionAI/gradio-ui/ https://github.com/PygmalionAI/gradio-ui/
''' '''
import torch import torch
import transformers import transformers

819
server.py

File diff suppressed because it is too large Load Diff