diff --git a/matrix_pygmalion_bot/ai/koboldcpp.py b/matrix_pygmalion_bot/ai/koboldcpp.py index 5214589..1172819 100644 --- a/matrix_pygmalion_bot/ai/koboldcpp.py +++ b/matrix_pygmalion_bot/ai/koboldcpp.py @@ -13,6 +13,7 @@ import base64 from PIL import Image, PngImagePlugin from .pygmalion_helpers import get_full_prompt, num_tokens +#from .llama_helpers import get_full_prompt, num_tokens logger = logging.getLogger(__name__) @@ -74,8 +75,10 @@ async def generate_sync( complete_reply += partial_reply tokens += input_data["max_length"] await typing_fn() - if not partial_reply or partial_reply.find('<|endoftext|>') != -1 or partial_reply.find("\nYou:") != -1 or tokens >= max_new_tokens: + if not partial_reply or partial_reply.find('<|endoftext|>') != -1 or partial_reply.find("\nYou:") != -1 or partial_reply.find("\n### Human:") != -1 or tokens >= max_new_tokens: idx = complete_reply.find(f"\nYou:") + if idx == -1: + idx = complete_reply.find(f"\n### Human:") if idx != -1: complete_reply = complete_reply[:idx].strip() else: diff --git a/matrix_pygmalion_bot/ai/llama_helpers.py b/matrix_pygmalion_bot/ai/llama_helpers.py new file mode 100644 index 0000000..93294cb --- /dev/null +++ b/matrix_pygmalion_bot/ai/llama_helpers.py @@ -0,0 +1,45 @@ +import asyncio +import os, tempfile +import logging + +import json +import requests + +from transformers import AutoTokenizer, AutoConfig +from huggingface_hub import hf_hub_download + +import io +import base64 +from PIL import Image, PngImagePlugin + +logger = logging.getLogger(__name__) + + +async def get_full_prompt(simple_prompt: str, bot, chat_history): + + # 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 = 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" + prompt += f"{bot.name} is always coherent and makes sense, but if she isn't sure if what she is saying is correct she will ask the user for help.\n" + prompt += f"{bot.name} is a very helpful AI and will help the user with anything they need, she is also very friendly and will try to make the user feel better if they are sad.\n" + prompt += f"{bot.name} is also very curious and will ask the user a lot of questions about themselves and their life, she will also try to make the user like her.\n" + prompt += f"\n" + prompt += f"{bot.user_name}: " + simple_prompt + "\n" + prompt += f"{bot.name}:" + + + return prompt + + +async def num_tokens(input_text: str): + return estimate_num_tokens(input_text) + + +async def estimate_num_tokens(input_text: str): + return len(input_text)//4+1 diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py index 377e2b4..1621cba 100644 --- a/matrix_pygmalion_bot/core.py +++ b/matrix_pygmalion_bot/core.py @@ -81,6 +81,8 @@ class Callbacks(object): if not (event.sender == self.bot.owner or chat_message.is_own_message): return + self.bot.user_name = room.user_name(event.sender) + if event.body.startswith('!replybot'): print(event) await self.bot.send_message(self.client, room.room_id, "Hello World!") @@ -281,6 +283,7 @@ class ChatBot(object): self.owner = None self.translate = None + self.user_name = "You" self.name = None self.persona = None