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

allow custom stopping strings in all modes (#903)

Browse source
This commit is contained in:
catalpaaa 2023-04-11 08:30:06 -07:00 • committed by GitHub
parent 0f212093a3
commit 78bbc66fc4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 37 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

26
modules/chat.py
View file

@ -74,8 +74,18 @@ def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat
return prompt
def get_stopping_strings(state):
if state['mode'] == 'instruct':
stopping_strings = [f"\n{state['name1']}", f"\n{state['name2']}"]
else:
stopping_strings = [f"\n{state['name1']}:", f"\n{state['name2']}:"]
stopping_strings += state['custom_stopping_strings']
return stopping_strings
def extract_message_from_reply(reply, state):
next_character_found = False
stopping_strings = get_stopping_strings(state)
if state['stop_at_newline']:
lines = reply.split('\n')
@ -83,7 +93,7 @@ def extract_message_from_reply(reply, state):
if len(lines) > 1:
next_character_found = True
else:
for string in [f"\n{state['name1']}:", f"\n{state['name2']}:"]:
for string in stopping_strings:
idx = reply.find(string)
if idx != -1:
reply = reply[:idx]
@ -92,7 +102,7 @@ def extract_message_from_reply(reply, state):
# If something like "\nYo" is generated just before "\nYou:"
# is completed, trim it
if not next_character_found:
for string in [f"\n{state['name1']}:", f"\n{state['name2']}:"]:
for string in stopping_strings:
for j in range(len(string) - 1, 0, -1):
if reply[-j:] == string[:j]:
reply = reply[:-j]
@ -106,10 +116,6 @@ def extract_message_from_reply(reply, state):
def chatbot_wrapper(text, state, regenerate=False, _continue=False):
if state['mode'] == 'instruct':
stopping_strings = [f"\n{state['name1']}", f"\n{state['name2']}"]
else:
stopping_strings = [f"\n{state['name1']}:", f"\n{state['name2']}:"]
# Defining some variables
cumulative_reply = ''
@ -117,6 +123,7 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False):
just_started = True
visible_text = custom_generate_chat_prompt = None
eos_token = '\n' if state['stop_at_newline'] else None
stopping_strings = get_stopping_strings(state)
# Check if any extension wants to hijack this function call
for extension, _ in extensions_module.iterator():
@ -186,15 +193,12 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False):
def impersonate_wrapper(text, state):
if state['mode'] == 'instruct':
stopping_strings = [f"\n{state['name1']}", f"\n{state['name2']}"]
else:
stopping_strings = [f"\n{state['name1']}:", f"\n{state['name2']}:"]
# Defining some variables
cumulative_reply = ''
eos_token = '\n' if state['stop_at_newline'] else None
prompt = generate_chat_prompt(text, state['max_new_tokens'], state['name1'], state['name2'], state['context'], state['chat_prompt_size'], end_of_turn=state['end_of_turn'], impersonate=True)
stopping_strings = get_stopping_strings(state)
# Yield *Is typing...*
yield shared.processing_message
@ -498,4 +502,4 @@ def upload_your_profile_picture(img, name1, name2, mode):
img.save(Path('cache/pfp_me.png'))
print('Profile picture saved to "cache/pfp_me.png"')
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, reset_cache=True)
return chat_html_wrapper(shared.history['visible'], name1, name2, mode, reset_cache=True)