From e94ab5dac1f7839b3f5ca0d1b407cb07adbeecfe Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 6 Apr 2023 01:43:10 -0300
Subject: [PATCH] Minor fixes

---
 modules/chat.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/modules/chat.py b/modules/chat.py
index f4ddf427..749ff8c2 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -24,7 +24,6 @@ def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat
     impersonate = kwargs['impersonate'] if 'impersonate' in kwargs else False
     also_return_rows = kwargs['also_return_rows'] if 'also_return_rows' in kwargs else False
 
-    user_input = fix_newlines(user_input)
     rows = [f"{context.strip()}\n"]
 
     # Finding the maximum prompt size
@@ -51,8 +50,8 @@ def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat
         rows.append(f"{prefix1.strip() if not is_instruct else prefix1}")
         limit = 2
     else:
-
         # Adding the user message
+        user_input = fix_newlines(user_input)
         if len(user_input) > 0:
             rows.append(f"{prefix1}{user_input}{end_of_turn}\n")
 
@@ -92,12 +91,14 @@ def extract_message_from_reply(reply, name1, name2, stop_at_newline):
                     if reply[-j:] == string[:j]:
                         reply = reply[:-j]
                         break
+                else:
+                    continue
+                break
 
     reply = fix_newlines(reply)
     return reply, next_character_found
 
 def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn, regenerate=False):
-    just_started = True
     eos_token = '\n' if generate_state['stop_at_newline'] else None
     name1_original = name1
     if 'pygmalion' in shared.model_name.lower():
@@ -129,6 +130,7 @@ def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_tu
 
     # Generate
     cumulative_reply = ''
+    just_started = True
     for i in range(generate_state['chat_generation_attempts']):
         reply = None
         for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_state, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]):
@@ -162,7 +164,6 @@ def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_tu
 
 def impersonate_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
     eos_token = '\n' if generate_state['stop_at_newline'] else None
-
     if 'pygmalion' in shared.model_name.lower():
         name1 = "You"
 
@@ -187,7 +188,7 @@ def impersonate_wrapper(text, generate_state, name1, name2, context, mode, end_o
     yield reply
 
 def cai_chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
-    for history in chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn, regenerate=False):
+    for history in chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
         yield chat_html_wrapper(history, name1, name2, mode)
 
 def regenerate_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):