2023-01-22 04:02:46 +01:00
import gc
2023-01-28 23:16:37 +01:00
import io
2023-02-10 19:40:03 +01:00
import json
import re
import sys
import time
2023-02-13 19:25:16 +01:00
import zipfile
2023-01-07 20:33:43 +01:00
from pathlib import Path
2023-02-10 19:40:03 +01:00
2022-12-21 17:27:31 +01:00
import gradio as gr
2023-02-10 19:40:03 +01:00
import torch
2023-02-23 16:05:25 +01:00
import modules . chat as chat
import modules . extensions as extensions_module
import modules . shared as shared
2023-02-23 17:03:52 +01:00
import modules . ui as ui
from modules . html_generator import generate_chat_html
2023-02-23 18:41:42 +01:00
from modules . models import load_model , load_soft_prompt
2023-02-23 17:03:52 +01:00
from modules . text_generation import generate_reply
2022-12-21 17:27:31 +01:00
2023-02-23 18:31:28 +01:00
# Loading custom settings
2023-03-06 14:57:45 +01:00
settings_file = None
2023-02-23 16:05:25 +01:00
if shared . args . settings is not None and Path ( shared . args . settings ) . exists ( ) :
2023-03-06 14:57:45 +01:00
settings_file = Path ( shared . args . settings )
elif Path ( ' settings.json ' ) . exists ( ) :
settings_file = Path ( ' settings.json ' )
if settings_file is not None :
print ( f " Loading settings from { settings_file } ... " )
new_settings = json . loads ( open ( settings_file , ' r ' ) . read ( ) )
2023-01-16 20:35:45 +01:00
for item in new_settings :
2023-02-23 17:28:30 +01:00
shared . settings [ item ] = new_settings [ item ]
2023-02-13 19:25:16 +01:00
2023-02-23 18:31:28 +01:00
def get_available_models ( ) :
2023-03-02 16:03:57 +01:00
if shared . args . flexgen :
return sorted ( [ re . sub ( ' -np$ ' , ' ' , item . name ) for item in list ( Path ( ' models/ ' ) . glob ( ' * ' ) ) if item . name . endswith ( ' -np ' ) ] , key = str . lower )
else :
2023-03-10 01:54:50 +01:00
return sorted ( [ item . name for item in list ( Path ( ' models/ ' ) . glob ( ' * ' ) ) if not item . name . endswith ( ( ' .txt ' , ' -np ' , ' .pt ' ) ) ] , key = str . lower )
2023-02-23 18:31:28 +01:00
def get_available_presets ( ) :
return sorted ( set ( map ( lambda x : ' . ' . join ( str ( x . name ) . split ( ' . ' ) [ : - 1 ] ) , Path ( ' presets ' ) . glob ( ' *.txt ' ) ) ) , key = str . lower )
def get_available_characters ( ) :
2023-02-24 21:31:23 +01:00
return [ ' None ' ] + sorted ( set ( map ( lambda x : ' . ' . join ( str ( x . name ) . split ( ' . ' ) [ : - 1 ] ) , Path ( ' characters ' ) . glob ( ' *.json ' ) ) ) , key = str . lower )
2023-02-23 18:31:28 +01:00
def get_available_extensions ( ) :
return sorted ( set ( map ( lambda x : x . parts [ 1 ] , Path ( ' extensions ' ) . glob ( ' */script.py ' ) ) ) , key = str . lower )
def get_available_softprompts ( ) :
2023-02-24 21:31:23 +01:00
return [ ' None ' ] + sorted ( set ( map ( lambda x : ' . ' . join ( str ( x . name ) . split ( ' . ' ) [ : - 1 ] ) , Path ( ' softprompts ' ) . glob ( ' *.zip ' ) ) ) , key = str . lower )
2023-02-23 18:31:28 +01:00
2023-02-08 02:08:21 +01:00
def load_model_wrapper ( selected_model ) :
2023-02-23 16:05:25 +01:00
if selected_model != shared . model_name :
shared . model_name = selected_model
2023-02-23 17:28:30 +01:00
shared . model = shared . tokenizer = None
2023-02-23 16:05:25 +01:00
if not shared . args . cpu :
2023-02-08 02:08:21 +01:00
gc . collect ( )
torch . cuda . empty_cache ( )
2023-02-23 16:05:25 +01:00
shared . model , shared . tokenizer = load_model ( shared . model_name )
2023-02-08 02:08:21 +01:00
2023-02-12 13:36:27 +01:00
return selected_model
2023-02-08 02:08:21 +01:00
def load_preset_values ( preset_menu , return_dict = False ) :
2023-02-08 04:19:20 +01:00
generate_params = {
2023-02-08 02:08:21 +01:00
' do_sample ' : True ,
' temperature ' : 1 ,
' top_p ' : 1 ,
' typical_p ' : 1 ,
' repetition_penalty ' : 1 ,
2023-03-15 15:04:30 +01:00
' encoder_repetition_penalty ' : 1 ,
2023-02-08 02:08:21 +01:00
' top_k ' : 50 ,
2023-02-08 03:11:04 +01:00
' num_beams ' : 1 ,
2023-02-11 18:48:12 +01:00
' penalty_alpha ' : 0 ,
2023-02-08 03:11:04 +01:00
' min_length ' : 0 ,
' length_penalty ' : 1 ,
' no_repeat_ngram_size ' : 0 ,
' early_stopping ' : False ,
2023-02-08 02:08:21 +01:00
}
with open ( Path ( f ' presets/ { preset_menu } .txt ' ) , ' r ' ) as infile :
preset = infile . read ( )
2023-02-11 18:54:29 +01:00
for i in preset . splitlines ( ) :
i = i . rstrip ( ' , ' ) . strip ( ) . split ( ' = ' )
2023-02-08 02:08:21 +01:00
if len ( i ) == 2 and i [ 0 ] . strip ( ) != ' tokens ' :
2023-02-08 04:19:20 +01:00
generate_params [ i [ 0 ] . strip ( ) ] = eval ( i [ 1 ] . strip ( ) )
2023-02-08 02:08:21 +01:00
2023-02-08 04:19:20 +01:00
generate_params [ ' temperature ' ] = min ( 1.99 , generate_params [ ' temperature ' ] )
2023-02-08 02:08:21 +01:00
if return_dict :
2023-02-08 04:19:20 +01:00
return generate_params
2023-02-08 02:08:21 +01:00
else :
2023-03-15 15:04:30 +01:00
return generate_params [ ' do_sample ' ] , generate_params [ ' temperature ' ] , generate_params [ ' top_p ' ] , generate_params [ ' typical_p ' ] , generate_params [ ' repetition_penalty ' ] , generate_params [ ' encoder_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 ' ]
2023-02-08 02:08:21 +01:00
2023-02-23 17:28:30 +01:00
def upload_soft_prompt ( file ) :
with zipfile . ZipFile ( io . BytesIO ( file ) ) as zf :
zf . extract ( ' meta.json ' )
j = json . loads ( open ( ' meta.json ' , ' r ' ) . read ( ) )
name = j [ ' name ' ]
Path ( ' meta.json ' ) . unlink ( )
with open ( Path ( f ' softprompts/ { name } .zip ' ) , ' wb ' ) as f :
f . write ( file )
return name
2023-03-02 15:25:04 +01:00
def create_settings_menus ( default_preset ) :
generate_params = load_preset_values ( default_preset if not shared . args . flexgen else ' Naive ' , return_dict = True )
2023-01-22 04:49:59 +01:00
2023-02-08 02:08:21 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
with gr . Row ( ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' model_menu ' ] = gr . Dropdown ( choices = available_models , value = shared . model_name , label = ' Model ' )
2023-02-24 21:31:23 +01:00
ui . create_refresh_button ( shared . gradio [ ' model_menu ' ] , lambda : None , lambda : { ' choices ' : get_available_models ( ) } , ' refresh-button ' )
2023-02-08 02:08:21 +01:00
with gr . Column ( ) :
with gr . Row ( ) :
2023-03-02 15:25:04 +01:00
shared . gradio [ ' preset_menu ' ] = gr . Dropdown ( choices = available_presets , value = default_preset if not shared . args . flexgen else ' Naive ' , label = ' Generation parameters preset ' )
2023-02-24 21:31:23 +01:00
ui . create_refresh_button ( shared . gradio [ ' preset_menu ' ] , lambda : None , lambda : { ' choices ' : get_available_presets ( ) } , ' refresh-button ' )
2023-01-22 04:49:59 +01:00
2023-02-24 21:31:23 +01:00
with gr . Accordion ( ' Custom generation parameters ' , open = False , elem_id = ' accordion ' ) :
2023-02-08 02:08:21 +01:00
with gr . Row ( ) :
2023-02-25 19:17:44 +01:00
with gr . Column ( ) :
shared . gradio [ ' temperature ' ] = gr . Slider ( 0.01 , 1.99 , value = generate_params [ ' temperature ' ] , step = 0.01 , label = ' temperature ' )
shared . gradio [ ' top_p ' ] = gr . Slider ( 0.0 , 1.0 , value = generate_params [ ' top_p ' ] , step = 0.01 , label = ' top_p ' )
2023-03-15 15:04:30 +01:00
shared . gradio [ ' top_k ' ] = gr . Slider ( 0 , 200 , value = generate_params [ ' top_k ' ] , step = 1 , label = ' top_k ' )
2023-02-25 19:21:40 +01:00
shared . gradio [ ' typical_p ' ] = gr . Slider ( 0.0 , 1.0 , value = generate_params [ ' typical_p ' ] , step = 0.01 , label = ' typical_p ' )
2023-03-15 15:04:30 +01:00
with gr . Column ( ) :
shared . gradio [ ' repetition_penalty ' ] = gr . Slider ( 1.0 , 1.5 , value = generate_params [ ' repetition_penalty ' ] , step = 0.01 , label = ' repetition_penalty ' )
shared . gradio [ ' encoder_repetition_penalty ' ] = gr . Slider ( 0.8 , 1.5 , value = generate_params [ ' encoder_repetition_penalty ' ] , step = 0.01 , label = ' encoder_repetition_penalty ' )
2023-02-25 19:17:44 +01:00
shared . gradio [ ' no_repeat_ngram_size ' ] = gr . Slider ( 0 , 20 , step = 1 , value = generate_params [ ' no_repeat_ngram_size ' ] , label = ' no_repeat_ngram_size ' )
shared . gradio [ ' 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 )
2023-03-15 15:04:30 +01:00
shared . gradio [ ' do_sample ' ] = gr . Checkbox ( value = generate_params [ ' do_sample ' ] , label = ' do_sample ' )
2023-02-17 20:18:01 +01:00
2023-02-24 21:31:23 +01:00
gr . Markdown ( ' Contrastive search: ' )
shared . gradio [ ' penalty_alpha ' ] = gr . Slider ( 0 , 5 , value = generate_params [ ' penalty_alpha ' ] , label = ' penalty_alpha ' )
2023-02-17 20:18:01 +01:00
2023-02-24 21:31:23 +01:00
gr . Markdown ( ' Beam search (uses a lot of VRAM): ' )
2023-02-17 20:33:27 +01:00
with gr . Row ( ) :
2023-02-25 19:17:44 +01:00
with gr . Column ( ) :
shared . gradio [ ' num_beams ' ] = gr . Slider ( 1 , 20 , step = 1 , value = generate_params [ ' num_beams ' ] , label = ' num_beams ' )
with gr . Column ( ) :
shared . gradio [ ' length_penalty ' ] = gr . Slider ( - 5 , 5 , value = generate_params [ ' length_penalty ' ] , label = ' length_penalty ' )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' early_stopping ' ] = gr . Checkbox ( value = generate_params [ ' early_stopping ' ] , label = ' early_stopping ' )
2023-02-08 02:08:21 +01:00
2023-02-24 21:31:23 +01:00
with gr . Accordion ( ' Soft prompt ' , open = False , elem_id = ' accordion ' ) :
2023-02-13 19:25:16 +01:00
with gr . Row ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' softprompts_menu ' ] = gr . Dropdown ( choices = available_softprompts , value = ' None ' , label = ' Soft prompt ' )
ui . create_refresh_button ( shared . gradio [ ' softprompts_menu ' ] , lambda : None , lambda : { ' choices ' : get_available_softprompts ( ) } , ' refresh-button ' )
2023-02-13 19:25:16 +01:00
2023-02-14 03:34:04 +01:00
gr . Markdown ( ' Upload a soft prompt (.zip format): ' )
2023-02-13 19:25:16 +01:00
with gr . Row ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_softprompt ' ] = gr . File ( type = ' binary ' , file_types = [ ' .zip ' ] )
2023-02-13 19:25:16 +01:00
2023-02-24 20:46:50 +01:00
shared . gradio [ ' model_menu ' ] . change ( load_model_wrapper , [ shared . gradio [ ' model_menu ' ] ] , [ shared . gradio [ ' model_menu ' ] ] , show_progress = True )
2023-03-15 15:04:30 +01:00
shared . gradio [ ' preset_menu ' ] . change ( load_preset_values , [ shared . gradio [ ' preset_menu ' ] ] , [ shared . gradio [ ' do_sample ' ] , shared . gradio [ ' temperature ' ] , shared . gradio [ ' top_p ' ] , shared . gradio [ ' typical_p ' ] , shared . gradio [ ' repetition_penalty ' ] , shared . gradio [ ' encoder_repetition_penalty ' ] , shared . gradio [ ' top_k ' ] , shared . gradio [ ' min_length ' ] , shared . gradio [ ' no_repeat_ngram_size ' ] , shared . gradio [ ' num_beams ' ] , shared . gradio [ ' penalty_alpha ' ] , shared . gradio [ ' length_penalty ' ] , shared . gradio [ ' early_stopping ' ] ] )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' softprompts_menu ' ] . change ( load_soft_prompt , [ shared . gradio [ ' softprompts_menu ' ] ] , [ shared . gradio [ ' softprompts_menu ' ] ] , show_progress = True )
shared . gradio [ ' upload_softprompt ' ] . upload ( upload_soft_prompt , [ shared . gradio [ ' upload_softprompt ' ] ] , [ shared . gradio [ ' softprompts_menu ' ] ] )
2023-02-08 02:08:21 +01:00
2023-01-22 04:49:59 +01:00
available_models = get_available_models ( )
available_presets = get_available_presets ( )
available_characters = get_available_characters ( )
2023-02-13 19:25:16 +01:00
available_softprompts = get_available_softprompts ( )
2023-02-23 17:03:52 +01:00
2023-03-02 15:25:04 +01:00
# Default extensions
2023-02-23 17:03:52 +01:00
extensions_module . available_extensions = get_available_extensions ( )
2023-02-28 06:20:11 +01:00
if shared . args . chat or shared . args . cai_chat :
for extension in shared . settings [ ' chat_default_extensions ' ] :
shared . args . extensions = shared . args . extensions or [ ]
if extension not in shared . args . extensions :
shared . args . extensions . append ( extension )
else :
for extension in shared . settings [ ' default_extensions ' ] :
shared . args . extensions = shared . args . extensions or [ ]
if extension not in shared . args . extensions :
shared . args . extensions . append ( extension )
if shared . args . extensions is not None and len ( shared . args . extensions ) > 0 :
2023-02-23 18:55:21 +01:00
extensions_module . load_extensions ( )
2023-01-22 04:49:59 +01:00
2023-03-02 15:25:04 +01:00
# Default model
2023-02-23 16:05:25 +01:00
if shared . args . model is not None :
shared . model_name = shared . args . model
2023-01-06 23:56:44 +01:00
else :
2023-01-07 02:05:37 +01:00
if len ( available_models ) == 0 :
2023-02-24 21:31:23 +01:00
print ( ' No models are available! Please download at least one. ' )
2023-01-30 18:17:12 +01:00
sys . exit ( 0 )
2023-01-06 23:56:44 +01:00
elif len ( available_models ) == 1 :
i = 0
else :
2023-02-24 21:31:23 +01:00
print ( ' The following models are available: \n ' )
2023-02-23 18:31:28 +01:00
for i , model in enumerate ( available_models ) :
2023-02-24 21:31:23 +01:00
print ( f ' { i + 1 } . { model } ' )
print ( f ' \n Which one do you want to load? 1- { len ( available_models ) } \n ' )
2023-01-06 23:56:44 +01:00
i = int ( input ( ) ) - 1
2023-01-09 16:56:54 +01:00
print ( )
2023-02-23 16:05:25 +01:00
shared . model_name = available_models [ i ]
shared . model , shared . tokenizer = load_model ( shared . model_name )
2023-01-06 23:56:44 +01:00
2023-03-02 15:25:04 +01:00
# Default UI settings
2023-02-23 18:31:28 +01:00
gen_events = [ ]
2023-03-02 15:25:04 +01:00
default_preset = shared . settings [ ' presets ' ] [ next ( ( k for k in shared . settings [ ' presets ' ] if re . match ( k . lower ( ) , shared . model_name . lower ( ) ) ) , ' default ' ) ]
default_text = shared . settings [ ' prompts ' ] [ next ( ( k for k in shared . settings [ ' prompts ' ] if re . match ( k . lower ( ) , shared . model_name . lower ( ) ) ) , ' default ' ) ]
2023-03-09 13:44:08 +01:00
title = ' Text generation web UI '
2023-02-24 21:31:23 +01:00
description = ' \n \n # Text generation lab \n Generate text using Large Language Models. \n '
2023-03-02 15:25:04 +01:00
suffix = ' _pygmalion ' if ' pygmalion ' in shared . model_name . lower ( ) else ' '
2023-01-19 18:03:47 +01:00
2023-02-23 16:05:25 +01:00
if shared . args . chat or shared . args . cai_chat :
2023-03-09 13:44:08 +01:00
with gr . Blocks ( css = ui . css + ui . chat_css , analytics_enabled = False , title = title ) as shared . gradio [ ' interface ' ] :
2023-02-23 16:05:25 +01:00
if shared . args . cai_chat :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' display ' ] = gr . HTML ( value = generate_chat_html ( shared . history [ ' visible ' ] , shared . settings [ f ' name1 { suffix } ' ] , shared . settings [ f ' name2 { suffix } ' ] , shared . character ) )
2023-01-15 22:16:46 +01:00
else :
2023-02-26 04:51:15 +01:00
shared . gradio [ ' display ' ] = gr . Chatbot ( value = shared . history [ ' visible ' ] ) . style ( color_map = ( " #326efd " , " #212528 " ) )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' textbox ' ] = gr . Textbox ( label = ' Input ' )
2023-01-09 21:23:43 +01:00
with gr . Row ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Stop ' ] = gr . Button ( ' Stop ' )
shared . gradio [ ' Generate ' ] = gr . Button ( ' Generate ' )
2023-02-05 02:53:42 +01:00
with gr . Row ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Impersonate ' ] = gr . Button ( ' Impersonate ' )
2023-02-25 19:35:43 +01:00
shared . gradio [ ' Regenerate ' ] = gr . Button ( ' Regenerate ' )
2023-01-29 16:02:44 +01:00
with gr . Row ( ) :
2023-02-25 20:33:46 +01:00
shared . gradio [ ' Copy last reply ' ] = gr . Button ( ' Copy last reply ' )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Replace last reply ' ] = gr . Button ( ' Replace last reply ' )
2023-02-25 19:35:43 +01:00
shared . gradio [ ' Remove last ' ] = gr . Button ( ' Remove last ' )
2023-02-27 15:41:21 +01:00
2023-02-25 19:35:43 +01:00
shared . gradio [ ' Clear history ' ] = gr . Button ( ' Clear history ' )
2023-02-27 15:41:21 +01:00
shared . gradio [ ' Clear history-confirm ' ] = gr . Button ( ' Confirm ' , variant = " stop " , visible = False )
2023-02-28 05:34:07 +01:00
shared . gradio [ ' Clear history-cancel ' ] = gr . Button ( ' Cancel ' , visible = False )
2023-02-24 21:31:23 +01:00
with gr . Tab ( ' Chat settings ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' name1 ' ] = gr . Textbox ( value = shared . settings [ f ' name1 { suffix } ' ] , lines = 1 , label = ' Your name ' )
shared . gradio [ ' name2 ' ] = gr . Textbox ( value = shared . settings [ f ' name2 { suffix } ' ] , lines = 1 , label = ' Bot \' s name ' )
2023-03-01 23:32:04 +01:00
shared . gradio [ ' context ' ] = gr . Textbox ( value = shared . settings [ f ' context { suffix } ' ] , lines = 5 , label = ' Context ' )
2023-02-16 00:55:32 +01:00
with gr . Row ( ) :
2023-02-26 21:19:36 +01:00
shared . gradio [ ' character_menu ' ] = gr . Dropdown ( choices = available_characters , value = ' None ' , label = ' Character ' , elem_id = ' character-menu ' )
2023-02-24 21:31:23 +01:00
ui . create_refresh_button ( shared . gradio [ ' character_menu ' ] , lambda : None , lambda : { ' choices ' : get_available_characters ( ) } , ' refresh-button ' )
2023-01-15 22:16:46 +01:00
2023-02-16 00:55:32 +01:00
with gr . Row ( ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' check ' ] = gr . Checkbox ( value = shared . settings [ f ' stop_at_newline { suffix } ' ] , label = ' Stop generating at new line character? ' )
2023-02-16 00:55:32 +01:00
with gr . Row ( ) :
with gr . Tab ( ' Chat history ' ) :
with gr . Row ( ) :
with gr . Column ( ) :
gr . Markdown ( ' Upload ' )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_chat_history ' ] = gr . File ( type = ' binary ' , file_types = [ ' .json ' , ' .txt ' ] )
2023-02-16 00:55:32 +01:00
with gr . Column ( ) :
gr . Markdown ( ' Download ' )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' download ' ] = gr . File ( )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' download_button ' ] = gr . Button ( value = ' Click me ' )
2023-02-16 00:55:32 +01:00
with gr . Tab ( ' Upload character ' ) :
with gr . Row ( ) :
with gr . Column ( ) :
gr . Markdown ( ' 1. Select the JSON file ' )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_json ' ] = gr . File ( type = ' binary ' , file_types = [ ' .json ' ] )
2023-02-16 00:55:32 +01:00
with gr . Column ( ) :
gr . Markdown ( ' 2. Select your character \' s profile picture (optional) ' )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_img_bot ' ] = gr . File ( type = ' binary ' , file_types = [ ' image ' ] )
shared . gradio [ ' Upload character ' ] = gr . Button ( value = ' Submit ' )
2023-02-16 00:55:32 +01:00
with gr . Tab ( ' Upload your profile picture ' ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_img_me ' ] = gr . File ( type = ' binary ' , file_types = [ ' image ' ] )
2023-02-16 00:55:32 +01:00
with gr . Tab ( ' Upload TavernAI Character Card ' ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' upload_img_tavern ' ] = gr . File ( type = ' binary ' , file_types = [ ' image ' ] )
2023-02-16 00:55:32 +01:00
2023-02-24 21:31:23 +01:00
with gr . Tab ( ' Generation settings ' ) :
2023-02-16 00:55:32 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' max_new_tokens ' ] = gr . Slider ( minimum = shared . settings [ ' max_new_tokens_min ' ] , maximum = shared . settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = shared . settings [ ' max_new_tokens ' ] )
2023-02-16 00:55:32 +01:00
with gr . Column ( ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' chat_prompt_size_slider ' ] = gr . Slider ( minimum = shared . settings [ ' chat_prompt_size_min ' ] , maximum = shared . settings [ ' chat_prompt_size_max ' ] , step = 1 , label = ' Maximum prompt size in tokens ' , value = shared . settings [ ' chat_prompt_size ' ] )
2023-02-26 05:51:59 +01:00
shared . gradio [ ' chat_generation_attempts ' ] = gr . Slider ( minimum = shared . settings [ ' chat_generation_attempts_min ' ] , maximum = shared . settings [ ' chat_generation_attempts_max ' ] , value = shared . settings [ ' chat_generation_attempts ' ] , step = 1 , label = ' Generation attempts (for longer replies) ' )
2023-03-02 15:25:04 +01:00
create_settings_menus ( default_preset )
2023-02-25 04:23:51 +01:00
2023-03-15 15:04:30 +01:00
shared . input_params = [ shared . gradio [ k ] for k in [ ' textbox ' , ' max_new_tokens ' , ' do_sample ' , ' temperature ' , ' top_p ' , ' typical_p ' , ' repetition_penalty ' , ' encoder_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 ' , ' chat_generation_attempts ' ] ]
2023-02-23 16:05:25 +01:00
if shared . args . extensions is not None :
2023-02-24 21:31:23 +01:00
with gr . Tab ( ' Extensions ' ) :
2023-02-23 18:55:21 +01:00
extensions_module . create_extensions_block ( )
2023-01-29 03:00:51 +01:00
2023-02-24 21:31:23 +01:00
function_call = ' chat.cai_chatbot_wrapper ' if shared . args . cai_chat else ' chat.chatbot_wrapper '
2023-02-23 16:05:25 +01:00
2023-03-14 20:31:27 +01:00
gen_events . append ( shared . gradio [ ' Generate ' ] . click ( eval ( function_call ) , shared . input_params , shared . gradio [ ' display ' ] , show_progress = shared . args . no_stream ) )
2023-03-13 00:38:40 +01:00
gen_events . append ( shared . gradio [ ' textbox ' ] . submit ( eval ( function_call ) , shared . input_params , shared . gradio [ ' display ' ] , show_progress = shared . args . no_stream ) )
gen_events . append ( shared . gradio [ ' Regenerate ' ] . click ( chat . regenerate_wrapper , shared . input_params , shared . gradio [ ' display ' ] , show_progress = shared . args . no_stream ) )
gen_events . append ( shared . gradio [ ' Impersonate ' ] . click ( chat . impersonate_wrapper , shared . input_params , shared . gradio [ ' textbox ' ] , show_progress = shared . args . no_stream ) )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Stop ' ] . click ( chat . stop_everything_event , [ ] , [ ] , cancels = gen_events )
2023-02-24 20:46:50 +01:00
2023-02-25 20:33:46 +01:00
shared . gradio [ ' Copy last reply ' ] . click ( chat . send_last_reply_to_input , [ ] , shared . gradio [ ' textbox ' ] , show_progress = shared . args . no_stream )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Replace last reply ' ] . click ( chat . replace_last_reply , [ shared . gradio [ ' textbox ' ] , shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , shared . gradio [ ' display ' ] , show_progress = shared . args . no_stream )
2023-02-27 15:41:21 +01:00
# Clear history with confirmation
clear_arr = [ shared . gradio [ k ] for k in [ ' Clear history-confirm ' , ' Clear history ' , ' Clear history-cancel ' ] ]
shared . gradio [ ' Clear history ' ] . click ( lambda : [ gr . update ( visible = True ) , gr . update ( visible = False ) , gr . update ( visible = True ) ] , None , clear_arr )
shared . gradio [ ' Clear history-confirm ' ] . click ( lambda : [ gr . update ( visible = False ) , gr . update ( visible = True ) , gr . update ( visible = False ) ] , None , clear_arr )
2023-02-27 17:41:00 +01:00
shared . gradio [ ' Clear history-confirm ' ] . click ( chat . clear_chat_log , [ shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , shared . gradio [ ' display ' ] )
2023-02-27 15:41:21 +01:00
shared . gradio [ ' Clear history-cancel ' ] . click ( lambda : [ gr . update ( visible = False ) , gr . update ( visible = True ) , gr . update ( visible = False ) ] , None , clear_arr )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Remove last ' ] . click ( chat . remove_last_message , [ shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , [ shared . gradio [ ' display ' ] , shared . gradio [ ' textbox ' ] ] , show_progress = False )
shared . gradio [ ' download_button ' ] . click ( chat . save_history , inputs = [ ] , outputs = [ shared . gradio [ ' download ' ] ] )
shared . gradio [ ' Upload character ' ] . click ( chat . upload_character , [ shared . gradio [ ' upload_json ' ] , shared . gradio [ ' upload_img_bot ' ] ] , [ shared . gradio [ ' character_menu ' ] ] )
2023-02-15 15:38:44 +01:00
2023-02-15 16:46:11 +01:00
# Clearing stuff and saving the history
2023-02-24 21:31:23 +01:00
for i in [ ' Generate ' , ' Regenerate ' , ' Replace last reply ' ] :
shared . gradio [ i ] . click ( lambda x : ' ' , shared . gradio [ ' textbox ' ] , shared . gradio [ ' textbox ' ] , show_progress = False )
2023-02-24 20:46:50 +01:00
shared . gradio [ i ] . click ( lambda : chat . save_history ( timestamp = False ) , [ ] , [ ] , show_progress = False )
2023-02-27 17:41:00 +01:00
shared . gradio [ ' Clear history-confirm ' ] . click ( lambda : chat . save_history ( timestamp = False ) , [ ] , [ ] , show_progress = False )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' textbox ' ] . submit ( lambda x : ' ' , shared . gradio [ ' textbox ' ] , shared . gradio [ ' textbox ' ] , show_progress = False )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' textbox ' ] . submit ( lambda : chat . save_history ( timestamp = False ) , [ ] , [ ] , show_progress = False )
shared . gradio [ ' character_menu ' ] . change ( chat . load_character , [ shared . gradio [ ' character_menu ' ] , shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , [ shared . gradio [ ' name2 ' ] , shared . gradio [ ' context ' ] , shared . gradio [ ' display ' ] ] )
shared . gradio [ ' upload_chat_history ' ] . upload ( chat . load_history , [ shared . gradio [ ' upload_chat_history ' ] , shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , [ ] )
shared . gradio [ ' upload_img_tavern ' ] . upload ( chat . upload_tavern_character , [ shared . gradio [ ' upload_img_tavern ' ] , shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] , [ shared . gradio [ ' character_menu ' ] ] )
shared . gradio [ ' upload_img_me ' ] . upload ( chat . upload_your_profile_picture , [ shared . gradio [ ' upload_img_me ' ] ] , [ ] )
2023-02-23 23:50:23 +01:00
reload_func = chat . redraw_html if shared . args . cai_chat else lambda : shared . history [ ' visible ' ]
2023-02-24 20:46:50 +01:00
reload_inputs = [ shared . gradio [ ' name1 ' ] , shared . gradio [ ' name2 ' ] ] if shared . args . cai_chat else [ ]
shared . gradio [ ' upload_chat_history ' ] . upload ( reload_func , reload_inputs , [ shared . gradio [ ' display ' ] ] )
shared . gradio [ ' upload_img_me ' ] . upload ( reload_func , reload_inputs , [ shared . gradio [ ' display ' ] ] )
2023-03-11 14:50:18 +01:00
shared . gradio [ ' Stop ' ] . click ( reload_func , reload_inputs , [ shared . gradio [ ' display ' ] ] )
2023-02-25 04:23:51 +01:00
shared . gradio [ ' interface ' ] . load ( lambda : chat . load_default_history ( shared . settings [ f ' name1 { suffix } ' ] , shared . settings [ f ' name2 { suffix } ' ] ) , None , None )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' interface ' ] . load ( reload_func , reload_inputs , [ shared . gradio [ ' display ' ] ] , show_progress = True )
2023-01-19 18:03:47 +01:00
2023-02-23 16:05:25 +01:00
elif shared . args . notebook :
2023-03-09 13:44:08 +01:00
with gr . Blocks ( css = ui . css , analytics_enabled = False , title = title ) as shared . gradio [ ' interface ' ] :
2023-01-19 02:44:47 +01:00
gr . Markdown ( description )
with gr . Tab ( ' Raw ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' textbox ' ] = gr . Textbox ( value = default_text , lines = 23 )
2023-01-19 02:44:47 +01:00
with gr . Tab ( ' Markdown ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' markdown ' ] = gr . Markdown ( )
2023-01-19 02:44:47 +01:00
with gr . Tab ( ' HTML ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' html ' ] = gr . HTML ( )
2023-01-11 05:33:57 +01:00
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Generate ' ] = gr . Button ( ' Generate ' )
shared . gradio [ ' Stop ' ] = gr . Button ( ' Stop ' )
2023-02-24 20:46:50 +01:00
shared . gradio [ ' max_new_tokens ' ] = gr . Slider ( minimum = shared . settings [ ' max_new_tokens_min ' ] , maximum = shared . settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = shared . settings [ ' max_new_tokens ' ] )
2023-01-19 02:44:47 +01:00
2023-03-02 15:25:04 +01:00
create_settings_menus ( default_preset )
2023-02-23 16:05:25 +01:00
if shared . args . extensions is not None :
2023-02-23 18:55:21 +01:00
extensions_module . create_extensions_block ( )
2023-01-29 13:48:18 +01:00
2023-03-15 15:04:30 +01:00
shared . input_params = [ shared . gradio [ k ] for k in [ ' textbox ' , ' max_new_tokens ' , ' do_sample ' , ' temperature ' , ' top_p ' , ' typical_p ' , ' repetition_penalty ' , ' encoder_repetition_penalty ' , ' top_k ' , ' min_length ' , ' no_repeat_ngram_size ' , ' num_beams ' , ' penalty_alpha ' , ' length_penalty ' , ' early_stopping ' ] ]
2023-02-24 21:31:23 +01:00
output_params = [ shared . gradio [ k ] for k in [ ' textbox ' , ' markdown ' , ' html ' ] ]
2023-02-25 04:23:51 +01:00
gen_events . append ( shared . gradio [ ' Generate ' ] . click ( generate_reply , shared . input_params , output_params , show_progress = shared . args . no_stream , api_name = ' textgen ' ) )
gen_events . append ( shared . gradio [ ' textbox ' ] . submit ( generate_reply , shared . input_params , output_params , show_progress = shared . args . no_stream ) )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Stop ' ] . click ( None , None , None , cancels = gen_events )
2023-01-19 02:44:47 +01:00
else :
2023-03-09 13:44:08 +01:00
with gr . Blocks ( css = ui . css , analytics_enabled = False , title = title ) as shared . gradio [ ' interface ' ] :
2023-01-09 00:10:31 +01:00
gr . Markdown ( description )
2023-01-07 02:05:37 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' textbox ' ] = gr . Textbox ( value = default_text , lines = 15 , label = ' Input ' )
shared . gradio [ ' max_new_tokens ' ] = gr . Slider ( minimum = shared . settings [ ' max_new_tokens_min ' ] , maximum = shared . settings [ ' max_new_tokens_max ' ] , step = 1 , label = ' max_new_tokens ' , value = shared . settings [ ' max_new_tokens ' ] )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Generate ' ] = gr . Button ( ' Generate ' )
2023-01-19 02:44:47 +01:00
with gr . Row ( ) :
with gr . Column ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Continue ' ] = gr . Button ( ' Continue ' )
2023-01-19 02:44:47 +01:00
with gr . Column ( ) :
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Stop ' ] = gr . Button ( ' Stop ' )
2023-02-08 02:08:21 +01:00
2023-03-02 15:25:04 +01:00
create_settings_menus ( default_preset )
2023-02-23 16:05:25 +01:00
if shared . args . extensions is not None :
2023-02-23 18:55:21 +01:00
extensions_module . create_extensions_block ( )
2023-01-29 13:48:18 +01:00
2023-01-07 02:05:37 +01:00
with gr . Column ( ) :
with gr . Tab ( ' Raw ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' output_textbox ' ] = gr . Textbox ( lines = 15 , label = ' Output ' )
2023-01-07 02:05:37 +01:00
with gr . Tab ( ' Markdown ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' markdown ' ] = gr . Markdown ( )
2023-01-07 03:14:08 +01:00
with gr . Tab ( ' HTML ' ) :
2023-02-24 20:46:50 +01:00
shared . gradio [ ' html ' ] = gr . HTML ( )
2023-01-07 02:05:37 +01:00
2023-03-15 15:04:30 +01:00
shared . input_params = [ shared . gradio [ k ] for k in [ ' textbox ' , ' max_new_tokens ' , ' do_sample ' , ' temperature ' , ' top_p ' , ' typical_p ' , ' repetition_penalty ' , ' encoder_repetition_penalty ' , ' top_k ' , ' min_length ' , ' no_repeat_ngram_size ' , ' num_beams ' , ' penalty_alpha ' , ' length_penalty ' , ' early_stopping ' ] ]
2023-02-24 20:46:50 +01:00
output_params = [ shared . gradio [ k ] for k in [ ' output_textbox ' , ' markdown ' , ' html ' ] ]
2023-02-25 04:23:51 +01:00
gen_events . append ( shared . gradio [ ' Generate ' ] . click ( generate_reply , shared . input_params , output_params , show_progress = shared . args . no_stream , api_name = ' textgen ' ) )
gen_events . append ( shared . gradio [ ' textbox ' ] . submit ( generate_reply , shared . input_params , output_params , show_progress = shared . args . no_stream ) )
gen_events . append ( shared . gradio [ ' Continue ' ] . click ( generate_reply , [ shared . gradio [ ' output_textbox ' ] ] + shared . input_params [ 1 : ] , output_params , show_progress = shared . args . no_stream ) )
2023-02-24 21:31:23 +01:00
shared . gradio [ ' Stop ' ] . click ( None , None , None , cancels = gen_events )
2022-12-21 17:27:31 +01:00
2023-02-24 20:46:50 +01:00
shared . gradio [ ' interface ' ] . queue ( )
2023-02-23 16:05:25 +01:00
if shared . args . listen :
2023-03-11 12:50:59 +01:00
shared . gradio [ ' interface ' ] . launch ( prevent_thread_lock = True , share = shared . args . share , server_name = ' 0.0.0.0 ' , server_port = shared . args . listen_port , inbrowser = shared . args . auto_launch )
2023-01-21 03:45:16 +01:00
else :
2023-03-11 12:50:59 +01:00
shared . gradio [ ' interface ' ] . launch ( prevent_thread_lock = True , share = shared . args . share , server_port = shared . args . listen_port , inbrowser = shared . args . auto_launch )
2023-02-08 02:08:21 +01:00
# I think that I will need this later
while True :
time . sleep ( 0.5 )