2022-12-21 17:27:31 +01:00
import re
2023-01-06 05:33:21 +01:00
import time
import glob
2023-01-06 23:56:44 +01:00
from sys import exit
2022-12-21 17:27:31 +01:00
import torch
2023-01-06 23:56:44 +01:00
import argparse
2023-01-15 19:23:41 +01:00
import json
2023-01-07 20:33:43 +01:00
from pathlib import Path
2022-12-21 17:27:31 +01:00
import gradio as gr
import transformers
2023-01-07 03:14:08 +01:00
from html_generator import *
2023-01-16 03:15:30 +01:00
from transformers import AutoTokenizer , AutoModelForCausalLM
2023-01-15 04:39:51 +01:00
import warnings
2023-01-19 16:01:58 +01:00
import gc
2023-01-19 16:20:57 +01:00
from tqdm import tqdm
2022-12-21 17:27:31 +01:00
2023-01-07 03:14:08 +01:00
2023-01-15 19:23:41 +01:00
transformers . logging . set_verbosity_error ( )
2023-01-06 23:56:44 +01:00
parser = argparse . ArgumentParser ( )
2023-01-07 00:22:26 +01:00
parser . add_argument ( ' --model ' , type = str , help = ' Name of the model to load by default. ' )
2023-01-16 14:10:09 +01:00
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 profile.png or profile.jpg exists in the same folder as server.py, this image will be used as the bot \' s profile picture. ' )
2023-01-09 14:58:46 +01:00
parser . add_argument ( ' --cpu ' , action = ' store_true ' , help = ' Use the CPU to generate text. ' )
2023-01-11 03:16:33 +01:00
parser . add_argument ( ' --load-in-8bit ' , action = ' store_true ' , help = ' Load the model with 8-bit precision. ' )
2023-01-19 15:09:24 +01:00
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. ' )
2023-01-16 14:10:09 +01:00
parser . add_argument ( ' --max-gpu-memory ' , type = int , help = ' Maximum memory in GiB to allocate to the GPU when loading the model. This is useful if you get out of memory errors while trying to generate text. Must be an integer number. ' )
2023-01-19 14:43:05 +01:00
parser . add_argument ( ' --no-stream ' , action = ' store_true ' , help = ' Don \' t stream the text output in real time. This slightly improves the text generation performance. ' )
2023-01-16 20:35:45 +01:00
parser . add_argument ( ' --settings ' , type = str , help = ' Load the default interface settings from this json file. See settings-template.json for an example. ' )
2023-01-19 21:31:29 +01:00
parser . add_argument ( ' --no-listen ' , action = ' store_true ' , help = ' Make the web UI unreachable from your local network. ' )
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. ' )
2023-01-06 23:56:44 +01:00
args = parser . parse_args ( )
2023-01-15 04:39:51 +01:00
2023-01-06 06:06:59 +01:00
loaded_preset = None
2023-01-16 20:35:45 +01:00
available_models = sorted ( set ( [ item . replace ( ' .pt ' , ' ' ) for item in map ( lambda x : str ( x . name ) , list ( Path ( ' models/ ' ) . glob ( ' * ' ) ) + list ( Path ( ' torch-dumps/ ' ) . glob ( ' * ' ) ) ) if not item . endswith ( ' .txt ' ) ] ) , key = str . lower )
2023-01-20 01:56:33 +01:00
available_presets = sorted ( set ( map ( lambda x : ' . ' . join ( str ( x . name ) . split ( ' . ' ) [ : - 1 ] ) , Path ( ' presets ' ) . glob ( ' *.txt ' ) ) ) , key = str . lower )
available_characters = sorted ( set ( map ( lambda x : ' . ' . join ( str ( x . name ) . split ( ' . ' ) [ : - 1 ] ) , Path ( ' characters ' ) . glob ( ' *.json ' ) ) ) , key = str . lower )
2023-01-06 05:33:21 +01:00
2023-01-15 19:23:41 +01:00
settings = {
' max_new_tokens ' : 200 ,
' max_new_tokens_min ' : 1 ,
' max_new_tokens_max ' : 2000 ,
' preset ' : ' NovelAI-Sphinx Moth ' ,
' name1 ' : ' Person 1 ' ,
' name2 ' : ' Person 2 ' ,
' context ' : ' This is a conversation between two people. ' ,
' prompt ' : ' Common sense questions and answers \n \n Question: \n Factual answer: ' ,
' prompt_gpt4chan ' : ' ----- \n --- 865467536 \n Input text \n --- 865467537 \n ' ,
' stop_at_newline ' : True ,
2023-01-19 22:58:45 +01:00
' preset_pygmalion ' : ' Pygmalion ' ,
' name1_pygmalion ' : ' You ' ,
' name2_pygmalion ' : ' Kawaii ' ,
' context_pygmalion ' : ' This is a conversation between two people. \n <START> ' ,
2023-01-19 20:46:46 +01:00
' stop_at_newline_pygmalion ' : False ,
2023-01-15 19:23:41 +01:00
}
2023-01-16 20:35:45 +01:00
if args . settings is not None and Path ( args . settings ) . exists ( ) :
with open ( Path ( args . settings ) , ' r ' ) as f :
2023-01-15 19:23:41 +01:00
new_settings = json . load ( f )
2023-01-16 20:35:45 +01:00
for item in new_settings :
if item in settings :
settings [ item ] = new_settings [ item ]
2023-01-15 04:39:51 +01:00
2022-12-21 17:27:31 +01:00
def load_model ( model_name ) :
2023-01-06 05:41:52 +01:00
print ( f " Loading { model_name } ... " )
2022-12-21 17:27:31 +01:00
t0 = time . time ( )
2023-01-06 05:41:52 +01:00
2023-01-11 03:16:33 +01:00
# Default settings
2023-01-19 16:11:44 +01:00
if not ( args . cpu or args . load_in_8bit or args . auto_devices or args . disk or args . max_gpu_memory is not None ) :
2023-01-11 03:16:33 +01:00
if Path ( f " torch-dumps/ { model_name } .pt " ) . exists ( ) :
print ( " Loading in .pt format... " )
model = torch . load ( Path ( f " torch-dumps/ { model_name } .pt " ) )
elif model_name . lower ( ) . startswith ( ( ' gpt-neo ' , ' opt- ' , ' galactica ' ) ) and 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 ( )
# Custom
2023-01-06 06:54:33 +01:00
else :
2023-01-11 03:16:33 +01:00
settings = [ " low_cpu_mem_usage=True " ]
2023-01-11 03:39:50 +01:00
command = " AutoModelForCausalLM.from_pretrained "
2023-01-11 03:16:33 +01:00
2023-01-09 20:28:04 +01:00
if args . cpu :
2023-01-11 03:16:33 +01:00
settings . append ( " torch_dtype=torch.float32 " )
2023-01-09 20:28:04 +01:00
else :
2023-01-16 03:01:51 +01:00
settings . append ( " device_map= ' auto ' " )
2023-01-16 02:33:35 +01:00
if args . max_gpu_memory is not None :
settings . append ( f " max_memory= {{ 0: ' { args . max_gpu_memory } GiB ' , ' cpu ' : ' 99GiB ' }} " )
2023-01-19 15:09:24 +01:00
if args . disk :
settings . append ( " offload_folder= ' cache ' " )
2023-01-16 03:01:51 +01:00
if args . load_in_8bit :
2023-01-11 03:16:33 +01:00
settings . append ( " load_in_8bit=True " )
else :
settings . append ( " torch_dtype=torch.float16 " )
2023-01-16 20:35:45 +01:00
settings = ' , ' . join ( set ( settings ) )
2023-01-16 03:01:51 +01:00
command = f " { command } (Path(f ' models/ { model_name } ' ), { settings } ) "
2023-01-11 03:16:33 +01:00
model = eval ( command )
2022-12-21 17:27:31 +01:00
2023-01-06 06:54:33 +01:00
# Loading the tokenizer
2023-01-11 05:10:11 +01:00
if model_name . lower ( ) . startswith ( ( ' gpt4chan ' , ' gpt-4chan ' , ' 4chan ' ) ) and Path ( f " models/gpt-j-6B/ " ) . exists ( ) :
2023-01-07 20:33:43 +01:00
tokenizer = AutoTokenizer . from_pretrained ( Path ( " models/gpt-j-6B/ " ) )
2022-12-21 17:27:31 +01:00
else :
2023-01-07 20:33:43 +01:00
tokenizer = AutoTokenizer . from_pretrained ( Path ( f " models/ { model_name } / " ) )
2023-01-16 17:43:23 +01:00
tokenizer . truncation_side = ' left '
2022-12-21 17:27:31 +01:00
2023-01-06 06:06:59 +01:00
print ( f " Loaded the model in { ( time . time ( ) - t0 ) : .2f } seconds. " )
2022-12-21 17:27:31 +01:00
return model , tokenizer
2023-01-06 06:26:33 +01:00
# Removes empty replies from gpt4chan outputs
2022-12-21 17:27:31 +01:00
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
2023-01-16 20:35:45 +01:00
# Fix the LaTeX equations in galactica
2023-01-07 05:56:21 +01:00
def fix_galactica ( s ) :
s = s . replace ( r ' \ [ ' , r ' $ ' )
s = s . replace ( r ' \ ] ' , r ' $ ' )
2023-01-07 16:13:09 +01:00
s = s . replace ( r ' \ ( ' , r ' $ ' )
s = s . replace ( r ' \ ) ' , r ' $ ' )
s = s . replace ( r ' $$ ' , r ' $ ' )
2023-01-07 05:56:21 +01:00
return s
2023-01-18 00:16:23 +01:00
def encode ( prompt , tokens ) :
if not args . cpu :
torch . cuda . empty_cache ( )
input_ids = tokenizer . encode ( str ( prompt ) , return_tensors = ' pt ' , truncation = True , max_length = 2048 - tokens ) . cuda ( )
else :
input_ids = tokenizer . encode ( str ( prompt ) , return_tensors = ' pt ' , truncation = True , max_length = 2048 - tokens )
return input_ids
2023-01-19 14:43:05 +01:00
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 ) :
2023-01-19 18:57:01 +01:00
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 ' ) :
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 )
2023-01-19 14:43:05 +01:00
else :
2023-01-19 18:57:01 +01:00
return reply
2023-01-19 14:43:05 +01:00
2023-01-13 18:28:53 +01:00
def generate_reply ( question , tokens , inference_settings , selected_model , eos_token = None ) :
2023-01-06 06:06:59 +01:00
global model , tokenizer , model_name , loaded_preset , preset
2022-12-21 17:27:31 +01:00
if selected_model != model_name :
model_name = selected_model
model = None
2023-01-08 18:37:43 +01:00
tokenizer = None
2023-01-09 14:58:46 +01:00
if not args . cpu :
2023-01-19 16:01:58 +01:00
gc . collect ( )
2023-01-09 14:58:46 +01:00
torch . cuda . empty_cache ( )
2022-12-21 17:27:31 +01:00
model , tokenizer = load_model ( model_name )
2023-01-06 06:06:59 +01:00
if inference_settings != loaded_preset :
2023-01-07 20:33:43 +01:00
with open ( Path ( f ' presets/ { inference_settings } .txt ' ) , ' r ' ) as infile :
2023-01-06 05:33:21 +01:00
preset = infile . read ( )
2023-01-06 06:06:59 +01:00
loaded_preset = inference_settings
2022-12-21 17:27:31 +01:00
2023-01-19 04:41:57 +01:00
cuda = " " if args . cpu else " .cuda() "
2023-01-19 14:43:05 +01:00
n = None if eos_token is None else tokenizer . encode ( eos_token , return_tensors = ' pt ' ) [ 0 ] [ - 1 ]
2023-01-20 04:45:02 +01:00
input_ids = encode ( question , tokens )
2023-01-19 14:43:05 +01:00
# Generate the entire reply at once
if args . no_stream :
output = eval ( f " model.generate(input_ids, eos_token_id= { n } , { preset } ) { cuda } " )
reply = decode ( output [ 0 ] )
yield formatted_outputs ( reply , model_name )
# Generate the reply 1 token at a time
else :
2023-01-19 15:09:24 +01:00
yield formatted_outputs ( question , model_name )
2023-01-19 03:56:42 +01:00
preset = preset . replace ( ' max_new_tokens=tokens ' , ' max_new_tokens=1 ' )
2023-01-19 16:20:57 +01:00
for i in tqdm ( range ( tokens ) ) :
2023-01-19 03:56:42 +01:00
output = eval ( f " model.generate(input_ids, { preset } ) { cuda } " )
2023-01-19 14:43:05 +01:00
reply = decode ( output [ 0 ] )
2023-01-19 03:56:42 +01:00
if eos_token is not None and reply [ - 1 ] == eos_token :
break
2023-01-19 14:43:05 +01:00
yield formatted_outputs ( reply , model_name )
2023-01-19 03:56:42 +01:00
input_ids = output
2023-01-19 01:37:21 +01:00
2023-01-06 23:56:44 +01:00
# Choosing the default model
if args . model is not None :
model_name = args . model
else :
2023-01-07 02:05:37 +01:00
if len ( available_models ) == 0 :
2023-01-06 23:56:44 +01:00
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 " \n Which one do you want to load? 1- { len ( available_models ) } \n " )
i = int ( input ( ) ) - 1
2023-01-09 16:56:54 +01:00
print ( )
2023-01-06 23:56:44 +01:00
model_name = available_models [ i ]
2022-12-21 17:27:31 +01:00
model , tokenizer = load_model ( model_name )
2023-01-06 23:56:44 +01:00
2023-01-09 00:10:31 +01:00
# UI settings
2023-01-07 23:11:21 +01:00
if model_name . lower ( ) . startswith ( ' gpt4chan ' ) :
2023-01-15 19:23:41 +01:00
default_text = settings [ ' prompt_gpt4chan ' ]
2022-12-21 17:27:31 +01:00
else :
2023-01-15 19:23:41 +01:00
default_text = settings [ ' prompt ' ]
2023-01-07 23:11:21 +01:00
2023-01-15 19:23:41 +01:00
description = f " \n \n # Text generation lab \n Generate text using Large Language Models. \n "
2023-01-15 22:16:46 +01:00
css = " .my-4 { margin-top: 0} .py-6 { padding-top: 2.5rem} "
2023-01-19 02:44:47 +01:00
if args . chat or args . cai_chat :
2023-01-08 02:52:46 +01:00
history = [ ]
2023-01-19 20:46:46 +01:00
character = None
2023-01-08 02:52:46 +01:00
2023-01-15 04:39:51 +01:00
# This gets the new line characters right.
2023-01-18 23:06:50 +01:00
def clean_chat_message ( text ) :
2023-01-15 03:50:34 +01:00
text = text . replace ( ' \n ' , ' \n \n ' )
text = re . sub ( r " \ n { 3,} " , " \n \n " , text )
text = text . strip ( )
2023-01-15 04:39:51 +01:00
return text
2023-01-18 23:06:50 +01:00
def generate_chat_prompt ( text , tokens , name1 , name2 , context ) :
text = clean_chat_message ( text )
2023-01-15 03:50:34 +01:00
2023-01-18 00:16:23 +01:00
rows = [ f " { context } \n \n " ]
i = len ( history ) - 1
while i > = 0 and len ( encode ( ' ' . join ( rows ) , tokens ) [ 0 ] ) < 2048 - tokens :
rows . insert ( 1 , f " { name2 } : { history [ i ] [ 1 ] . strip ( ) } \n " )
rows . insert ( 1 , f " { name1 } : { history [ i ] [ 0 ] . strip ( ) } \n " )
i - = 1
rows . append ( f " { name1 } : { text } \n " )
rows . append ( f " { name2 } : " )
while len ( rows ) > 3 and len ( encode ( ' ' . join ( rows ) , tokens ) [ 0 ] ) > = 2048 - tokens :
rows . pop ( 1 )
rows . pop ( 1 )
question = ' ' . join ( rows )
2023-01-18 23:06:50 +01:00
return question
2023-01-08 02:52:46 +01:00
2023-01-18 23:06:50 +01:00
def chatbot_wrapper ( text , tokens , inference_settings , selected_model , name1 , name2 , context , check ) :
question = generate_chat_prompt ( text , tokens , name1 , name2 , context )
2023-01-19 02:08:23 +01:00
history . append ( [ ' ' , ' ' ] )
2023-01-18 23:06:50 +01:00
eos_token = ' \n ' if check else None
2023-01-19 18:57:01 +01:00
for reply in generate_reply ( question , tokens , inference_settings , selected_model , eos_token = eos_token ) :
2023-01-19 01:51:18 +01:00
next_character_found = False
2023-01-18 23:06:50 +01:00
2023-01-19 23:59:34 +01:00
previous_idx = [ m . start ( ) for m in re . finditer ( f " (^| \n ) { name2 } : " , question ) ]
2023-01-19 18:57:01 +01:00
idx = [ m . start ( ) for m in re . finditer ( f " (^| \n ) { name2 } : " , reply ) ]
idx = idx [ len ( previous_idx ) - 1 ]
2023-01-19 23:59:34 +01:00
2023-01-19 18:57:01 +01:00
reply = reply [ idx + len ( f " \n { name2 } : " ) : ]
2023-01-18 23:06:50 +01:00
if check :
2023-01-19 18:57:01 +01:00
reply = reply . split ( ' \n ' ) [ 0 ] . strip ( )
2023-01-18 23:06:50 +01:00
else :
idx = reply . find ( f " \n { name1 } : " )
if idx != - 1 :
reply = reply [ : idx ]
2023-01-19 01:51:18 +01:00
next_character_found = True
2023-01-18 23:06:50 +01:00
reply = clean_chat_message ( reply )
history [ - 1 ] = [ text , reply ]
2023-01-19 14:43:05 +01:00
if next_character_found :
break
2023-01-18 23:06:50 +01:00
# Prevent the chat log from flashing if something like "\nYo" is generated just
# before "\nYou:" is completed
tmp = f " \n { name1 } : "
2023-01-19 01:51:18 +01:00
next_character_substring_found = False
2023-01-19 14:43:05 +01:00
for j in range ( 1 , len ( tmp ) ) :
2023-01-18 23:06:50 +01:00
if reply [ - j : ] == tmp [ : j ] :
2023-01-19 01:51:18 +01:00
next_character_substring_found = True
2023-01-18 23:06:50 +01:00
2023-01-19 01:51:18 +01:00
if not next_character_substring_found :
2023-01-18 23:06:50 +01:00
yield history
2023-01-08 02:52:46 +01:00
2023-01-19 14:43:05 +01:00
yield history
2023-01-19 01:51:18 +01:00
2023-01-15 16:20:04 +01:00
def cai_chatbot_wrapper ( text , tokens , inference_settings , selected_model , name1 , name2 , context , check ) :
2023-01-18 23:06:50 +01:00
for history in chatbot_wrapper ( text , tokens , inference_settings , selected_model , name1 , name2 , context , check ) :
2023-01-19 20:46:46 +01:00
yield generate_chat_html ( history , name1 , name2 , character )
2023-01-15 16:20:04 +01:00
def remove_last_message ( name1 , name2 ) :
2023-01-15 07:19:09 +01:00
history . pop ( )
2023-01-15 16:20:04 +01:00
if args . cai_chat :
2023-01-19 20:46:46 +01:00
return generate_chat_html ( history , name1 , name2 , character )
2023-01-15 16:20:04 +01:00
else :
return history
2023-01-15 07:19:09 +01:00
2023-01-08 02:52:46 +01:00
def clear ( ) :
global history
history = [ ]
2023-01-15 16:20:04 +01:00
def clear_html ( ) :
2023-01-19 20:46:46 +01:00
return generate_chat_html ( [ ] , " " , " " , character )
2023-01-15 16:20:04 +01:00
2023-01-19 18:03:47 +01:00
def redraw_html ( name1 , name2 ) :
global history
2023-01-19 20:46:46 +01:00
return generate_chat_html ( history , name1 , name2 , character )
2023-01-19 18:03:47 +01:00
def save_history ( ) :
if not Path ( ' logs ' ) . exists ( ) :
Path ( ' logs ' ) . mkdir ( )
with open ( Path ( ' logs/conversation.json ' ) , ' w ' ) as f :
f . write ( json . dumps ( { ' data ' : history } ) )
return Path ( ' logs/conversation.json ' )
def load_history ( file ) :
global history
history = json . loads ( file . decode ( ' utf-8 ' ) ) [ ' data ' ]
2023-01-19 20:46:46 +01:00
def load_character ( _character , name1 , name2 ) :
global history , character
context = " "
history = [ ]
if _character != ' None ' :
character = _character
with open ( Path ( f ' characters/ { _character } .json ' ) , ' r ' ) as f :
data = json . loads ( f . 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 "
2023-01-19 23:04:54 +01:00
context = f " { context . strip ( ) } \n <START> \n "
2023-01-19 20:46:46 +01:00
if ' example_dialogue ' in data and data [ ' example_dialogue ' ] != ' ' :
2023-01-19 23:04:54 +01:00
context + = f " { data [ ' example_dialogue ' ] . strip ( ) } \n "
2023-01-19 20:46:46 +01:00
if ' char_greeting ' in data :
history = [ [ ' ' , data [ ' char_greeting ' ] ] ]
else :
character = None
context = settings [ ' context_pygmalion ' ]
name2 = settings [ ' name2_pygmalion ' ]
if args . cai_chat :
return name2 , context , generate_chat_html ( history , name1 , name2 , character )
else :
return name2 , context , history
suffix = ' _pygmalion ' if ' pygmalion ' in model_name . lower ( ) else ' '
2023-01-15 22:16:46 +01:00
with gr . Blocks ( css = css + " .h- \ [40vh \ ] { height: 66.67vh} .gradio-container { max-width: 800px; margin-left: auto; margin-right: auto} " , analytics_enabled = False ) as interface :
if args . cai_chat :
2023-01-19 20:46:46 +01:00
display1 = gr . HTML ( value = generate_chat_html ( [ ] , " " , " " , character ) )
2023-01-15 22:16:46 +01:00
else :
display1 = gr . Chatbot ( )
textbox = gr . Textbox ( lines = 2 , label = ' Input ' )
btn = gr . Button ( " Generate " )
2023-01-09 21:23:43 +01:00
with gr . Row ( ) :
2023-01-19 02:44:47 +01:00
btn2 = gr . Button ( " Clear history " )
stop = gr . Button ( " Stop " )
btn3 = gr . Button ( " Remove last message " )
2023-01-13 19:02:17 +01:00
2023-01-15 22:16:46 +01:00
length_slider = gr . Slider ( minimum = settings [ ' max_new_tokens_min ' ] , maximum = settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = settings [ ' max_new_tokens ' ] )
with gr . Row ( ) :
2023-01-08 02:52:46 +01:00
with gr . Column ( ) :
2023-01-15 22:16:46 +01:00
model_menu = gr . Dropdown ( choices = available_models , value = model_name , label = ' Model ' )
with gr . Column ( ) :
2023-01-19 22:58:45 +01:00
preset_menu = gr . Dropdown ( choices = available_presets , value = settings [ f ' preset { suffix } ' ] , label = ' Settings preset ' )
2023-01-15 22:16:46 +01:00
2023-01-19 22:58:45 +01:00
name1 = gr . Textbox ( value = settings [ f ' name1 { suffix } ' ] , lines = 1 , label = ' Your name ' )
name2 = gr . Textbox ( value = settings [ f ' name2 { suffix } ' ] , lines = 1 , label = ' Bot \' s name ' )
context = gr . Textbox ( value = settings [ f ' context { suffix } ' ] , lines = 2 , label = ' Context ' )
2023-01-15 22:16:46 +01:00
with gr . Row ( ) :
2023-01-19 20:46:46 +01:00
character_menu = gr . Dropdown ( choices = [ " None " ] + available_characters , value = " None " , label = ' Character ' )
with gr . Row ( ) :
2023-01-19 22:58:45 +01:00
check = gr . Checkbox ( value = settings [ f ' stop_at_newline { suffix } ' ] , label = ' Stop generating at new line character? ' )
2023-01-19 18:03:47 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
gr . Markdown ( " Upload chat history " )
upload = gr . File ( type = ' binary ' )
with gr . Column ( ) :
gr . Markdown ( " Download chat history " )
save_btn = gr . Button ( value = " Click me " )
download = gr . File ( )
2023-01-15 22:16:46 +01:00
2023-01-15 16:20:04 +01:00
if args . cai_chat :
2023-01-19 03:56:42 +01:00
gen_event = btn . click ( cai_chatbot_wrapper , [ textbox , length_slider , preset_menu , model_menu , name1 , name2 , context , check ] , display1 , show_progress = args . no_stream , api_name = " textgen " )
gen_event2 = textbox . submit ( cai_chatbot_wrapper , [ textbox , length_slider , preset_menu , model_menu , name1 , name2 , context , check ] , display1 , show_progress = args . no_stream )
2023-01-19 02:44:47 +01:00
btn2 . click ( clear_html , [ ] , display1 , show_progress = False )
2023-01-15 16:20:04 +01:00
else :
2023-01-19 03:56:42 +01:00
gen_event = btn . click ( chatbot_wrapper , [ textbox , length_slider , preset_menu , model_menu , name1 , name2 , context , check ] , display1 , show_progress = args . no_stream , api_name = " textgen " )
gen_event2 = textbox . submit ( chatbot_wrapper , [ textbox , length_slider , preset_menu , model_menu , name1 , name2 , context , check ] , display1 , show_progress = args . no_stream )
2023-01-19 02:44:47 +01:00
btn2 . click ( lambda x : " " , display1 , display1 , show_progress = False )
2023-01-15 16:20:04 +01:00
2023-01-08 02:52:46 +01:00
btn2 . click ( clear )
2023-01-15 16:20:04 +01:00
btn3 . click ( remove_last_message , [ name1 , name2 ] , display1 , show_progress = False )
2023-01-08 05:10:02 +01:00
btn . click ( lambda x : " " , textbox , textbox , show_progress = False )
2023-01-08 05:33:45 +01:00
textbox . submit ( lambda x : " " , textbox , textbox , show_progress = False )
2023-01-19 02:44:47 +01:00
stop . click ( None , None , None , cancels = [ gen_event , gen_event2 ] )
2023-01-19 18:03:47 +01:00
save_btn . click ( save_history , inputs = [ ] , outputs = [ download ] )
upload . upload ( load_history , [ upload ] , [ ] )
2023-01-19 20:46:46 +01:00
character_menu . change ( load_character , [ character_menu , name1 , name2 ] , [ name2 , context , display1 ] )
2023-01-19 19:05:42 +01:00
if args . cai_chat :
upload . upload ( redraw_html , [ name1 , name2 ] , [ display1 ] )
else :
upload . upload ( lambda : history , [ ] , [ display1 ] )
2023-01-19 18:03:47 +01:00
2023-01-19 02:44:47 +01:00
elif args . notebook :
with gr . Blocks ( css = css , analytics_enabled = False ) as interface :
gr . Markdown ( description )
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 " )
stop = gr . Button ( " Stop " )
2023-01-11 05:33:57 +01:00
2023-01-19 02:44:47 +01:00
length_slider = gr . Slider ( minimum = settings [ ' max_new_tokens_min ' ] , maximum = settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = settings [ ' max_new_tokens ' ] )
with gr . Row ( ) :
with gr . Column ( ) :
model_menu = gr . Dropdown ( choices = available_models , value = model_name , label = ' Model ' )
with gr . Column ( ) :
preset_menu = gr . Dropdown ( choices = available_presets , value = settings [ ' preset ' ] , label = ' Settings preset ' )
2023-01-19 03:56:42 +01:00
gen_event = btn . click ( generate_reply , [ textbox , length_slider , preset_menu , model_menu ] , [ textbox , markdown , html ] , show_progress = args . no_stream , api_name = " textgen " )
gen_event2 = textbox . submit ( generate_reply , [ textbox , length_slider , preset_menu , model_menu ] , [ textbox , markdown , html ] , show_progress = args . no_stream )
2023-01-19 02:44:47 +01:00
stop . click ( None , None , None , cancels = [ gen_event , gen_event2 ] )
else :
2023-01-09 00:10:31 +01:00
with gr . Blocks ( css = css , analytics_enabled = False ) as interface :
gr . Markdown ( description )
2023-01-07 02:05:37 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
textbox = gr . Textbox ( value = default_text , lines = 15 , label = ' Input ' )
2023-01-15 19:23:41 +01:00
length_slider = gr . Slider ( minimum = settings [ ' max_new_tokens_min ' ] , maximum = settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = settings [ ' max_new_tokens ' ] )
preset_menu = gr . Dropdown ( choices = available_presets , value = settings [ ' preset ' ] , label = ' Settings preset ' )
2023-01-07 02:05:37 +01:00
model_menu = gr . Dropdown ( choices = available_models , value = model_name , label = ' Model ' )
btn = gr . Button ( " Generate " )
2023-01-19 02:44:47 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
cont = gr . Button ( " Continue " )
with gr . Column ( ) :
stop = gr . Button ( " Stop " )
2023-01-07 02:05:37 +01:00
with gr . Column ( ) :
with gr . Tab ( ' Raw ' ) :
2023-01-11 05:36:11 +01:00
output_textbox = gr . Textbox ( lines = 15 , label = ' Output ' )
2023-01-07 02:05:37 +01:00
with gr . Tab ( ' Markdown ' ) :
markdown = gr . Markdown ( )
2023-01-07 03:14:08 +01:00
with gr . Tab ( ' HTML ' ) :
html = gr . HTML ( )
2023-01-07 02:05:37 +01:00
2023-01-19 03:56:42 +01:00
gen_event = btn . click ( generate_reply , [ textbox , length_slider , preset_menu , model_menu ] , [ output_textbox , markdown , html ] , show_progress = args . no_stream , api_name = " textgen " )
gen_event2 = textbox . submit ( generate_reply , [ textbox , length_slider , preset_menu , model_menu ] , [ output_textbox , markdown , html ] , show_progress = args . no_stream )
cont_event = cont . click ( generate_reply , [ output_textbox , length_slider , preset_menu , model_menu ] , [ output_textbox , markdown , html ] , show_progress = args . no_stream )
2023-01-19 02:44:47 +01:00
stop . click ( None , None , None , cancels = [ gen_event , gen_event2 , cont_event ] )
2022-12-21 17:27:31 +01:00
2023-01-18 23:06:50 +01:00
interface . queue ( )
2023-01-11 05:10:11 +01:00
if args . no_listen :
2023-01-19 21:31:29 +01:00
interface . launch ( share = args . share )
2023-01-11 05:10:11 +01:00
else :
2023-01-19 21:31:29 +01:00
interface . launch ( share = args . share , server_name = " 0.0.0.0 " )