mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Refactor chat functions (#2003)

Browse source
This commit is contained in:
oobabooga 2023-05-11 15:37:04 -03:00 • committed by GitHub
parent 4e9da22c58
commit 638c6a65a2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 138 additions and 157 deletions
Download patch file Download diff file Expand all files Collapse all files

74
modules/chat.py
View file

@ -188,7 +188,7 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False):
# Generate
for i in range(state['chat_generation_attempts']):
reply = None
for j, reply in enumerate(generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", state, eos_token=eos_token, stopping_strings=stopping_strings)):
for j, reply in enumerate(generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", state, eos_token=eos_token, stopping_strings=stopping_strings, is_chat=True)):
reply = cumulative_reply + reply
# Extracting the reply
@ -242,7 +242,7 @@ def impersonate_wrapper(text, state):
cumulative_reply = text
for i in range(state['chat_generation_attempts']):
reply = None
for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", state, eos_token=eos_token, stopping_strings=stopping_strings):
for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", state, eos_token=eos_token, stopping_strings=stopping_strings, is_chat=True):
reply = cumulative_reply + reply
reply, next_character_found = extract_message_from_reply(reply, state)
yield reply
@ -255,35 +255,31 @@ def impersonate_wrapper(text, state):
yield reply
def cai_chatbot_wrapper(text, state):
for history in chatbot_wrapper(text, state):
def generate_chat_reply(text, state, regenerate=False, _continue=False):
if regenerate or _continue:
text = ''
if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
yield shared.history['visible']
return
for history in chatbot_wrapper(text, state, regenerate=regenerate, _continue=_continue):
yield history
# Same as above but returns HTML
def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False):
for history in generate_chat_reply(text, state, regenerate, _continue):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def regenerate_wrapper(text, state):
if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'], state['chat_style'])
else:
for history in chatbot_wrapper('', state, regenerate=True):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def continue_wrapper(text, state):
if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
yield chat_html_wrapper(shared.history['visible'], state['name1'], state['name2'], state['mode'], state['chat_style'])
else:
for history in chatbot_wrapper('', state, _continue=True):
yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'])
def remove_last_message(name1, name2, mode, style):
def remove_last_message():
if len(shared.history['visible']) > 0 and shared.history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>':
last = shared.history['visible'].pop()
shared.history['internal'].pop()
else:
last = ['', '']
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style), last[0]
return last[0]
def send_last_reply_to_input():
@ -293,35 +289,27 @@ def send_last_reply_to_input():
return ''
def replace_last_reply(text, name1, name2, mode, style):
def replace_last_reply(text):
if len(shared.history['visible']) > 0:
shared.history['visible'][-1][1] = text
shared.history['internal'][-1][1] = apply_extensions("input", text)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def send_dummy_message(text, name1, name2, mode, style):
def send_dummy_message(text):
shared.history['visible'].append([text, ''])
shared.history['internal'].append([apply_extensions("input", text), ''])
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def send_dummy_reply(text, name1, name2, mode, style):
def send_dummy_reply(text):
if len(shared.history['visible']) > 0 and not shared.history['visible'][-1][1] == '':
shared.history['visible'].append(['', ''])
shared.history['internal'].append(['', ''])
shared.history['visible'][-1][1] = text
shared.history['internal'][-1][1] = apply_extensions("input", text)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def clear_html():
return chat_html_wrapper([], "", "")
def clear_chat_log(name1, name2, greeting, mode, style):
def clear_chat_log(greeting, mode):
shared.history['visible'] = []
shared.history['internal'] = []
@ -332,14 +320,12 @@ def clear_chat_log(name1, name2, greeting, mode, style):
save_history(mode)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def redraw_html(name1, name2, mode, style, reset_cache=False):
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style, reset_cache=reset_cache)
def redraw_html(name1, name2, mode, style):
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
def tokenize_dialogue(dialogue, name1, name2, mode, style):
def tokenize_dialogue(dialogue, name1, name2):
history = []
messages = []
dialogue = re.sub('<START>', '', dialogue)
@ -447,7 +433,7 @@ def generate_pfp_cache(character):
return None
def load_character(character, name1, name2, mode, style):
def load_character(character, name1, name2, mode):
shared.character = character
context = greeting = turn_template = ""
greeting_field = 'greeting'
@ -521,7 +507,7 @@ def load_character(character, name1, name2, mode, style):
# Create .json log files since they don't already exist
save_history(mode)
return name1, name2, picture, greeting, context, repr(turn_template)[1:-1], chat_html_wrapper(shared.history['visible'], name1, name2, mode, style)
return name1, name2, picture, greeting, context, repr(turn_template)[1:-1]
def upload_character(json_file, img, tavern=False):
@ -556,7 +542,7 @@ def upload_tavern_character(img, name1, name2):
return upload_character(json.dumps(_json), img, tavern=True)
def upload_your_profile_picture(img, name1, name2, mode, style):
def upload_your_profile_picture(img):
cache_folder = Path("cache")
if not cache_folder.exists():
cache_folder.mkdir()
@ -568,5 +554,3 @@ def upload_your_profile_picture(img, name1, name2, mode, style):
img = make_thumbnail(img)
img.save(Path('cache/pfp_me.png'))
logging.info('Profile picture saved to "cache/pfp_me.png"')
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, style, reset_cache=True)