Fixes chat

gradio_app.py

@@ -170,10 +170,10 @@ class TextGeneration:
                     input_text,
                     **generation_kwargs,
                 )[0]["generated_text"]
-                if generation_kwargs["do_clean"]:
-                    generated_text = cleaner.clean_txt(generated_text)
                 if generated_text.strip().startswith(input_text):
                     generated_text = generated_text.replace(input_text, "", 1).strip()
+                if generation_kwargs["do_clean"]:
+                    generated_text = cleaner.clean_txt(generated_text)
                 if generated_text:
                     if previous_text and previous_text != text:
                         diff = [
@@ -258,9 +258,10 @@ def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k
             break
     context += history_context
     for _ in range(5):
-        response = generator.generate(f"{context}\n\n{user}: {message}.\n", generation_kwargs)[0]
+        prompt = f"{context}\n\n{user}: {message}.\n"
+        response = generator.generate(prompt, generation_kwargs)[0]
         if DEBUG:
-            print("\n-----" + response + "-----\n")
+            print("\n-----\n" + response + "\n-----\n")
         # response = response.split("\n")[-1]
         # if agent in response and response.split(agent)[-1]:
         #     response = response.split(agent)[-1]
@@ -268,7 +269,7 @@ def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k
         #     response = response.split(user)[-1]
         # Take the first response
         response = [
-            r for r in response.split(f"{AGENT}:") if r.strip()
+            r for r in response.replace(prompt, "").split(f"{AGENT}:") if r.strip()
         ][0].split(USER)[0].replace(f"{AGENT}:", "\n").strip()
         if response[0] in string.punctuation:
             response = response[1:].strip()