conjurer/ai_functions.py

import tiktoken


def num_tokens_from_string(message, model):
    """
    The function takes a string message and a model as input and returns the number of tokens in the
    message according to the given model.

    :param message: A string containing the message or text from which you want to count the number of
    tokens
    :param model: The model parameter refers to a language model or tokenizer that can be used to
    tokenize the input string. It could be a pre-trained model or a custom tokenizer
    """
    tokens_per_message = 3
    tokens_per_name = 1
    chat_gpt_encoding = tiktoken.encoding_for_model(model)

    num_tokens = 0
    num_tokens += tokens_per_message
    for keys, values in message.items():
        num_tokens += len(chat_gpt_encoding.encode(values))
        if keys == "role":
            num_tokens += tokens_per_name
    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
    return num_tokens