mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-11-26 01:30:20 +01:00
179 lines
7.5 KiB
Python
179 lines
7.5 KiB
Python
import re
|
|
import time
|
|
import glob
|
|
from sys import exit
|
|
import torch
|
|
import argparse
|
|
from pathlib import Path
|
|
import gradio as gr
|
|
import transformers
|
|
from html_generator import *
|
|
from transformers import AutoTokenizer, T5Tokenizer
|
|
from transformers import AutoModelForCausalLM, T5ForConditionalGeneration
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
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 webui in notebook mode, where the output is written to the same text box as the input.')
|
|
args = parser.parse_args()
|
|
loaded_preset = None
|
|
available_models = sorted(set(map(lambda x : str(x.name).replace('.pt', ''), list(Path('models/').glob('*'))+list(Path('torch-dumps/').glob('*')))))
|
|
available_models = [item for item in available_models if not item.endswith('.txt')]
|
|
available_presets = sorted(set(map(lambda x : str(x.name).split('.')[0], list(Path('presets').glob('*.txt')))))
|
|
|
|
def load_model(model_name):
|
|
print(f"Loading {model_name}...")
|
|
t0 = time.time()
|
|
|
|
# Loading the model
|
|
if Path(f"torch-dumps/{model_name}.pt").exists():
|
|
print("Loading in .pt format...")
|
|
model = torch.load(Path(f"torch-dumps/{model_name}.pt")).cuda()
|
|
elif model_name.lower().startswith(('gpt-neo', 'opt-', 'galactica')):
|
|
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)
|
|
else:
|
|
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
|
|
elif model_name in ['gpt-j-6B']:
|
|
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
|
|
elif model_name in ['flan-t5', 't5-large']:
|
|
model = T5ForConditionalGeneration.from_pretrained(Path(f"models/{model_name}")).cuda()
|
|
else:
|
|
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
|
|
|
|
# Loading the tokenizer
|
|
if model_name.lower().startswith('gpt4chan'):
|
|
tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/"))
|
|
elif model_name in ['flan-t5']:
|
|
tokenizer = T5Tokenizer.from_pretrained(Path(f"models/{model_name}/"))
|
|
else:
|
|
tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{model_name}/"))
|
|
|
|
print(f"Loaded the model in {(time.time()-t0):.2f} seconds.")
|
|
return model, tokenizer
|
|
|
|
# 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
|
|
|
|
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'$')
|
|
return s
|
|
|
|
def generate_reply(question, temperature, max_length, inference_settings, selected_model):
|
|
global model, tokenizer, model_name, loaded_preset, preset
|
|
|
|
if selected_model != model_name:
|
|
model_name = selected_model
|
|
model = None
|
|
tokenier = None
|
|
torch.cuda.empty_cache()
|
|
model, tokenizer = load_model(model_name)
|
|
if inference_settings != loaded_preset:
|
|
with open(Path(f'presets/{inference_settings}.txt'), 'r') as infile:
|
|
preset = infile.read()
|
|
loaded_preset = inference_settings
|
|
|
|
torch.cuda.empty_cache()
|
|
input_text = question
|
|
input_ids = tokenizer.encode(str(input_text), return_tensors='pt').cuda()
|
|
|
|
output = eval(f"model.generate(input_ids, {preset}).cuda()")
|
|
reply = tokenizer.decode(output[0], skip_special_tokens=True)
|
|
|
|
if model_name.lower().startswith('galactica'):
|
|
reply = fix_galactica(reply)
|
|
return reply, reply, 'Only applicable for gpt4chan.'
|
|
elif model_name.lower().startswith('gpt4chan'):
|
|
reply = fix_gpt4chan(reply)
|
|
return reply, 'Only applicable for galactica models.', generate_html(reply)
|
|
else:
|
|
return reply, 'Only applicable for galactica models.', 'Only applicable for gpt4chan.'
|
|
|
|
# Choosing the default model
|
|
if args.model is not None:
|
|
model_name = args.model
|
|
else:
|
|
if len(available_models) == 0:
|
|
print("No models are available! Please download at least one.")
|
|
exit(0)
|
|
elif len(available_models) == 1:
|
|
i = 0
|
|
else:
|
|
print("The following models are available:\n")
|
|
for i,model in enumerate(available_models):
|
|
print(f"{i+1}. {model}")
|
|
print(f"\nWhich one do you want to load? 1-{len(available_models)}\n")
|
|
i = int(input())-1
|
|
model_name = available_models[i]
|
|
model, tokenizer = load_model(model_name)
|
|
|
|
if model_name.lower().startswith('gpt4chan'):
|
|
default_text = "-----\n--- 865467536\nInput text\n--- 865467537\n"
|
|
else:
|
|
default_text = "Common sense questions and answers\n\nQuestion: \nFactual answer:"
|
|
|
|
if args.notebook:
|
|
with gr.Blocks(css=".my-4 {margin-top: 0} .py-6 {padding-top: 2.5rem}") as interface:
|
|
gr.Markdown(
|
|
f"""
|
|
|
|
# Text generation lab
|
|
Generate text using Large Language Models.
|
|
"""
|
|
)
|
|
|
|
with gr.Tab('Raw'):
|
|
textbox = gr.Textbox(value=default_text, lines=23)
|
|
with gr.Tab('Markdown'):
|
|
markdown = gr.Markdown()
|
|
with gr.Tab('HTML'):
|
|
html = gr.HTML()
|
|
btn = gr.Button("Generate")
|
|
|
|
with gr.Row():
|
|
with gr.Column():
|
|
temp_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Temperature', value=0.7)
|
|
length_slider = gr.Slider(minimum=1, maximum=2000, step=1, label='max_length', value=200)
|
|
with gr.Column():
|
|
preset_menu = gr.Dropdown(choices=available_presets, value="NovelAI-Sphinx Moth", label='Preset')
|
|
model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model')
|
|
|
|
btn.click(generate_reply, [textbox, temp_slider, length_slider, preset_menu, model_menu], [textbox, markdown, html], show_progress=False)
|
|
else:
|
|
with gr.Blocks(css=".my-4 {margin-top: 0} .py-6 {padding-top: 2.5rem}") as interface:
|
|
gr.Markdown(
|
|
f"""
|
|
# Text generation lab
|
|
Generate text using Large Language Models.
|
|
"""
|
|
)
|
|
|
|
with gr.Row():
|
|
with gr.Column():
|
|
textbox = gr.Textbox(value=default_text, lines=15, label='Input')
|
|
temp_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Temperature', value=0.7)
|
|
length_slider = gr.Slider(minimum=1, maximum=2000, step=1, label='max_length', value=200)
|
|
preset_menu = gr.Dropdown(choices=available_presets, value="NovelAI-Sphinx Moth", label='Preset')
|
|
model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model')
|
|
btn = gr.Button("Generate")
|
|
with gr.Column():
|
|
with gr.Tab('Raw'):
|
|
output_textbox = gr.Textbox(value=default_text, lines=15, label='Output')
|
|
with gr.Tab('Markdown'):
|
|
markdown = gr.Markdown()
|
|
with gr.Tab('HTML'):
|
|
html = gr.HTML()
|
|
|
|
btn.click(generate_reply, [textbox, temp_slider, length_slider, preset_menu, model_menu], [output_textbox, markdown, html], show_progress=True)
|
|
|
|
interface.launch(share=False, server_name="0.0.0.0")
|