allow custom stopping strings in all modes (#903)

modules/chat.py

@@ -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)

modules/shared.py

@@ -34,6 +34,7 @@ settings = {
     'context': 'This is a conversation with your Assistant. The Assistant is very helpful and is eager to chat with you and answer your questions.',
     'greeting': 'Hello there!',
     'end_of_turn': '',
+    'custom_stopping_strings': '',
     'stop_at_newline': False,
     'add_bos_token': True,
     'chat_prompt_size': 2048,

modules/text_generation.py

@@ -174,10 +174,14 @@ def generate_reply(question, state, eos_token=None, stopping_strings=[]):
     eos_token_ids = [shared.tokenizer.eos_token_id] if shared.tokenizer.eos_token_id is not None else []
     if eos_token is not None:
         eos_token_ids.append(int(encode(eos_token)[0][-1]))
+
+    # Handling the stopping strings
     stopping_criteria_list = transformers.StoppingCriteriaList()
-    if type(stopping_strings) is list and len(stopping_strings) > 0:
-        t = [encode(string, 0, add_special_tokens=False) for string in stopping_strings]
-        stopping_criteria_list.append(_SentinelTokenStoppingCriteria(sentinel_token_ids=t, starting_idx=len(input_ids[0])))
+    for st in [stopping_strings, state['custom_stopping_strings']]:
+        if type(st) is list and len(st) > 0:
+            sentinel_token_ids = [encode(string, 0, add_special_tokens=False) for string in st]
+            stopping_criteria_list.append(_SentinelTokenStoppingCriteria(sentinel_token_ids=sentinel_token_ids, starting_idx=len(input_ids[0])))
+            break
 
     if not shared.args.flexgen:
         for k in ['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']: