matrix-pygmalion-bot/matrix_pygmalion_bot/ai/llama_helpers.py

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):
    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 += 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"
    #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"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
    prompt += f"### Instruction:\nGiven the following character description and scenario, write a script for a dialogue between the human user {bot.user_name} and the fictional AI assistant {bot.name}. Play the role of the character {bot.name}\n\n"
    prompt += "### Input:\n"
    prompt += bot.name + "'s Persona: " + bot.persona + "\n"
    prompt += "Scenario: " + bot.scenario + "\n\n"
    prompt += "### Response:\n"

    for dialogue_item in bot.example_dialogue:
        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 += 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"{user_name}: {simple_prompt}\n{ai_name}:")
    visible_history = []
    current_message = True
    for key, chat_item in reversed(chat_history.chat_history.items()):
        if current_message:
            current_message = False
            continue
        if chat_item.message["en"].startswith('!begin'):
            break
        if chat_item.message["en"].startswith('!'):
            continue
        if chat_item.message["en"].startswith('<ERROR>'):
            continue
        #if chat_item.message["en"] == bot.greeting:
        #    continue
        if chat_item.num_tokens == None:
            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:
            visible_history.append(chat_item)
            total_num_tokens += chat_item.num_tokens
        else:
            break
    visible_history = reversed(visible_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 += f"{ai_name}: {bot.greeting}\n"

    for chat_item in visible_history:
        if chat_item.is_own_message:
            prompt += f"{ai_name}: {chat_item.message['en']}\n"
        else:
            prompt += f"{user_name}: {chat_item.message['en']}\n"
    prompt += f"{user_name}: {simple_prompt}\n"
    prompt += f"{ai_name}:"

    return prompt


async def num_tokens(input_text: str):
    return await estimate_num_tokens(input_text)


async def estimate_num_tokens(input_text: str):
    return len(input_text)//4+1