diff --git a/matrix_pygmalion_bot/ai/llama_helpers.py b/matrix_pygmalion_bot/ai/llama_helpers.py index d799134..07e5d2a 100644 --- a/matrix_pygmalion_bot/ai/llama_helpers.py +++ b/matrix_pygmalion_bot/ai/llama_helpers.py @@ -24,12 +24,23 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history): # 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"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"Write a script for a dialogue between the user {bot.user_name} and the fictional character {bot.name}.\n" + prompt += bot.persona + "\n" + prompt += "Scenario: " + bot.scenario + "\n" + + for dialogue_item in bot.example_dialogue: + dialogue_item = dialogue_item.replace('{{user}}', f"{bot.user_name}") + dialogue_item = dialogue_item.replace('{{char}}', bot.name) + prompt += dialogue_item + "\n\n" + prompt += "" + "\n" + #prompt += bot.name + ": " + bot.greeting + "\n" #prompt += f"{bot.user_name}: " + simple_prompt + "\n" #prompt += f"{bot.name}:" @@ -62,7 +73,11 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history): break visible_history = reversed(visible_history) - #prompt += bot.name + ": " + bot.greeting + "\n" + 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 += bot.name + ": " + bot.greeting + "\n" + for chat_item in visible_history: if chat_item.is_own_message: prompt += bot.name + ": " + chat_item.message["en"] + "\n" diff --git a/matrix_pygmalion_bot/ai/pygmalion_helpers.py b/matrix_pygmalion_bot/ai/pygmalion_helpers.py index 1a46879..f594358 100644 --- a/matrix_pygmalion_bot/ai/pygmalion_helpers.py +++ b/matrix_pygmalion_bot/ai/pygmalion_helpers.py @@ -19,15 +19,22 @@ 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" + prompt += "Scenario: " + bot.scenario + "\n\n" prompt += "" + "\n" + + for dialogue_item in bot.example_dialogue: + dialogue_item = dialogue_item.replace('{{user}}', 'You') + dialogue_item = dialogue_item.replace('{{char}}', bot.name) + prompt += dialogue_item + "\n\n" + prompt += "" + "\n" #prompt += bot.name + ": " + bot.greeting + "\n" - prompt += "You: " + simple_prompt + "\n" - prompt += bot.name + ":" + #prompt += "You: " + simple_prompt + "\n" + #prompt += bot.name + ":" MAX_TOKENS = 2048 max_new_tokens = 200 total_num_tokens = await num_tokens(prompt) + total_num_tokens += await num_tokens(f"You: " + simple_prompt + "\n{bot.name}:") visible_history = [] current_message = True for key, chat_item in reversed(chat_history.chat_history.items()): @@ -53,10 +60,11 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history): break visible_history = reversed(visible_history) - prompt = bot.name + "'s Persona: " + bot.persona + "\n" - prompt += "Scenario: " + bot.scenario + "\n" - prompt += "" + "\n" - #prompt += bot.name + ": " + bot.greeting + "\n" + 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 += bot.name + ": " + bot.greeting + "\n" + for chat_item in visible_history: if chat_item.is_own_message: prompt += bot.name + ": " + chat_item.message["en"] + "\n" @@ -68,11 +76,16 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history): return prompt +tokenizer = None + + async def num_tokens(input_text: str): # os.makedirs("./models/pygmalion-6b", exist_ok=True) # hf_hub_download(repo_id="PygmalionAI/pygmalion-6b", filename="config.json", cache_dir="./models/pygmalion-6b") # config = AutoConfig.from_pretrained("./models/pygmalion-6b/config.json") - tokenizer = AutoTokenizer.from_pretrained("PygmalionAI/pygmalion-6b") + global tokenizer + if not tokenizer: + tokenizer = AutoTokenizer.from_pretrained("PygmalionAI/pygmalion-6b") encoding = tokenizer.encode(input_text, add_special_tokens=False) max_input_size = tokenizer.max_model_input_sizes return len(encoding) diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py index 8c6f436..e8f84ff 100644 --- a/matrix_pygmalion_bot/core.py +++ b/matrix_pygmalion_bot/core.py @@ -96,7 +96,7 @@ class Callbacks(object): if m['cmd']: prompt = m['cmd'].strip() if self.bot.image_prompt: - prompt.replace(self.bot.name, self.bot.image_prompt) + prompt = prompt.replace(self.bot.name, self.bot.image_prompt) else: if self.bot.image_prompt: prompt = self.bot.image_prompt @@ -289,6 +289,7 @@ class ChatBot(object): self.persona = None self.scenario = None self.greeting = None + self.example_dialogue = [] self.temperature = 0.90 self.events = [] self.global_tick = 0 @@ -300,11 +301,12 @@ class ChatBot(object): if STORE_PATH and not os.path.isdir(STORE_PATH): os.mkdir(STORE_PATH) - def character_init(self, name, persona, scenario, greeting): + def character_init(self, name, persona, scenario, greeting, example_dialogue=[]): self.name = name self.persona = persona self.scenario = scenario self.greeting = greeting + self.example_dialogue = example_dialogue async def event_loop(self): try: @@ -441,7 +443,11 @@ async def main() -> None: else: device_name = "matrix-nio" bot = ChatBot(homeserver, user_id, password, device_name) - bot.character_init(botname, config[section]['persona'], config[section]['scenario'], config[section]['greeting']) + if config.has_option(section, 'example_dialogue'): + example_dialogue = json.loads(config[section]['example_dialogue']) + else: + example_dialogue = [] + bot.character_init(botname, config[section]['persona'].replace("\\n", "\n"), config[section]['scenario'].replace("\\n", "\n"), config[section]['greeting'].replace("\\n", "\n"), example_dialogue) if config.has_option(section, 'temperature'): bot.temperature = config[section]['temperature'] if config.has_option(section, 'owner'):