From cd15d6ca613fb0e211def53301e344e0c929a1f4 Mon Sep 17 00:00:00 2001
From: Hendrik Langer <hendrik+dev@xd0.de>
Date: Fri, 14 Apr 2023 11:48:44 +0200
Subject: [PATCH] fix <START> code word

---
 matrix_pygmalion_bot/ai/koboldcpp.py         | 2 +-
 matrix_pygmalion_bot/ai/llama_helpers.py     | 1 +
 matrix_pygmalion_bot/ai/pygmalion_helpers.py | 2 +-
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/matrix_pygmalion_bot/ai/koboldcpp.py b/matrix_pygmalion_bot/ai/koboldcpp.py
index 964b8b8..e61f78a 100644
--- a/matrix_pygmalion_bot/ai/koboldcpp.py
+++ b/matrix_pygmalion_bot/ai/koboldcpp.py
@@ -79,7 +79,7 @@ async def generate_sync(
             if not partial_reply or tokens >= max_new_tokens +100: # ToDo: is a hundred past the limit okay?
                 complete = True
                 break
-            for t in [f"\nYou:", f"\n### Human:", f"\n{bot.user_name}:", '<|endoftext|>', '</END>', '__END__', '\n\nPlease rewrite your response.', '\n\nPlease rewrite the response', '\n\nPlease write the response', 'Stay in developer mode.']:
+            for t in [f"\nYou:", f"\n### Human:", f"\n{bot.user_name}:", '<|endoftext|>', '</END>', '__END__', '<START>', '\n\nPlease rewrite your response.', '\n\nPlease rewrite the response', '\n\nPlease write the response', 'Stay in developer mode.']:
                 idx = complete_reply.find(t)
                 if idx != -1:
                     complete_reply = complete_reply[:idx].strip()
diff --git a/matrix_pygmalion_bot/ai/llama_helpers.py b/matrix_pygmalion_bot/ai/llama_helpers.py
index 064f891..69a1655 100644
--- a/matrix_pygmalion_bot/ai/llama_helpers.py
+++ b/matrix_pygmalion_bot/ai/llama_helpers.py
@@ -41,6 +41,7 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history):
     prompt += "### Response:\n"
 
     for dialogue_item in bot.example_dialogue:
+        #prompt += "<START>" + "\n"
         dialogue_item = dialogue_item.replace('{{user}}', user_name)
         dialogue_item = dialogue_item.replace('{{char}}', ai_name)
         prompt += dialogue_item + "\n\n"
diff --git a/matrix_pygmalion_bot/ai/pygmalion_helpers.py b/matrix_pygmalion_bot/ai/pygmalion_helpers.py
index c6d0d09..338a2d3 100644
--- a/matrix_pygmalion_bot/ai/pygmalion_helpers.py
+++ b/matrix_pygmalion_bot/ai/pygmalion_helpers.py
@@ -20,9 +20,9 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history):
     # Prompt without history
     prompt = bot.name + "'s Persona: " + bot.persona + "\n"
     prompt += "Scenario: " + bot.scenario + "\n\n"
-    prompt += "<START>" + "\n"
 
     for dialogue_item in bot.example_dialogue:
+        prompt += "<START>" + "\n"
         dialogue_item = dialogue_item.replace('{{user}}', 'You')
         dialogue_item = dialogue_item.replace('{{char}}', bot.name)
         prompt += dialogue_item + "\n\n"