make ai_name exchangable

matrix_pygmalion_bot/ai/llama_helpers.py

@@ -16,13 +16,15 @@ logger = logging.getLogger(__name__)
 
 
 async def get_full_prompt(simple_prompt: str, bot, chat_history):
+    ai_name = "### Assistant" # bot.name
+    user_name = "### Human" # bot.user_name
 
     # https://github.com/ggerganov/llama.cpp/tree/master/examples
 ##    prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n"
 #    prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n"
 #    #"BEGINNING OF CONVERSATION:"
-#    prompt += "### Human: " + simple_prompt + "\n"
-#    prompt += "### Assistant:"
+#    prompt += user_name + ": " + simple_prompt + "\n"
+#    prompt += ai_name + ":"
 
     #prompt = f"This is a transcript of a 1000 page, never ending conversation between {bot.user_name} and the cute and helpful AI assistant {bot.name}. {bot.name} is a girl who is an AI running on the users computer.\n"
     #prompt += f"{bot.name} can think for herself without the user seeing her thoughts by adding a /think prefix to her output. She uses this to reason about the world and to think about what she should say next.\n"
@@ -39,18 +41,18 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history):
     prompt += "### Response:\n"
 
     for dialogue_item in bot.example_dialogue:
-        dialogue_item = dialogue_item.replace('{{user}}', f"### Human")
-        dialogue_item = dialogue_item.replace('{{char}}', f"### Assistant")
+        dialogue_item = dialogue_item.replace('{{user}}', user_name)
+        dialogue_item = dialogue_item.replace('{{char}}', ai_name)
         prompt += dialogue_item + "\n\n"
     prompt += "<START>" + "\n"
-    #prompt += bot.name + ": " + bot.greeting + "\n"
-    #prompt += f"### Human: " + simple_prompt + "\n"
-    #prompt += f"### Assistant:"
+    #prompt += f"{ai_name}: {bot.greeting}\n"
+    #prompt += f"{user_name}: {simple_prompt}\n"
+    #prompt += f"{ai_name}:"
 
     MAX_TOKENS = 2048
     max_new_tokens = 200
     total_num_tokens = await num_tokens(prompt)
-    total_num_tokens += await num_tokens(f"### Human: " + simple_prompt + f"\n### Assistant:")
+    total_num_tokens += await num_tokens(f"{user_name}: {simple_prompt}\n{ai_name}:")
     visible_history = []
     current_message = True
     for key, chat_item in reversed(chat_history.chat_history.items()):
@@ -66,7 +68,7 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history):
         #if chat_item.message["en"] == bot.greeting:
         #    continue
         if chat_item.num_tokens == None:
-            chat_item.num_tokens = await num_tokens("{}: {}".format("### Human", chat_item.message["en"]))
+            chat_item.num_tokens = await num_tokens("{}: {}".format(user_name, chat_item.message["en"]))
         # TODO: is it MAX_TOKENS or MAX_TOKENS - max_new_tokens??
         logger.debug(f"History: " + str(chat_item) + " [" + str(chat_item.num_tokens) + "]")
         if total_num_tokens + chat_item.num_tokens < MAX_TOKENS - max_new_tokens:
@@ -79,15 +81,15 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history):
     if not hasattr(bot, "greeting_num_tokens"):
         bot.greeting_num_tokens = await num_tokens(bot.greeting)
     if total_num_tokens + bot.greeting_num_tokens < MAX_TOKENS - max_new_tokens:
-        prompt += "### Assistant: " + bot.greeting + "\n"
+        prompt += f"{ai_name}: {bot.greeting}\n"
 
     for chat_item in visible_history:
         if chat_item.is_own_message:
-            prompt += "### Assistant: " + chat_item.message["en"] + "\n"
+            prompt += f"{ai_name}: {chat_item.message['en']}\n"
         else:
-            prompt += f"### Human: " + chat_item.message["en"] + "\n"
-    prompt += f"### Human: " + simple_prompt + "\n"
-    prompt += f"### Assistant:"
+            prompt += f"{user_name}: {chat_item.message['en']}\n"
+    prompt += f"{user_name}: {simple_prompt}\n"
+    prompt += f"{ai_name}:"
 
     return prompt