CoolFace
Apppublic

aphilippov/python-server-api

sourceHugging Faceapache-2.0updated 3y agoView on Hugging Face
0likes
token_count_utils.py189 linesDownload Raw Back to autogen
1from typing import List, Union, Dict2import logging3import json4import tiktoken5import re6 7 8logger = logging.getLogger(__name__)9 10 11def get_max_token_limit(model="gpt-3.5-turbo-0613"):12    # Handle common azure model names/aliases13    model = re.sub(r"^gpt\-?35", "gpt-3.5", model)14    model = re.sub(r"^gpt4", "gpt-4", model)15 16    max_token_limit = {17        "gpt-3.5-turbo": 4096,18        "gpt-3.5-turbo-0301": 4096,19        "gpt-3.5-turbo-0613": 4096,20        "gpt-3.5-turbo-instruct": 4096,21        "gpt-3.5-turbo-16k": 16385,22        "gpt-3.5-turbo-16k-0613": 16385,23        "gpt-3.5-turbo-1106": 16385,24        "gpt-4": 8192,25        "gpt-4-32k": 32768,26        "gpt-4-32k-0314": 32768,  # deprecate in Sep27        "gpt-4-0314": 8192,  # deprecate in Sep28        "gpt-4-0613": 8192,29        "gpt-4-32k-0613": 32768,30        "gpt-4-1106-preview": 128000,31        "gpt-4-vision-preview": 128000,32    }33    return max_token_limit[model]34 35 36def percentile_used(input, model="gpt-3.5-turbo-0613"):37    return count_token(input) / get_max_token_limit(model)38 39 40def token_left(input: Union[str, List, Dict], model="gpt-3.5-turbo-0613") -> int:41    """Count number of tokens left for an OpenAI model.42 43    Args:44        input: (str, list, dict): Input to the model.45        model: (str): Model name.46 47    Returns:48        int: Number of tokens left that the model can use for completion.49    """50    return get_max_token_limit(model) - count_token(input, model=model)51 52 53def count_token(input: Union[str, List, Dict], model: str = "gpt-3.5-turbo-0613") -> int:54    """Count number of tokens used by an OpenAI model.55    Args:56        input: (str, list, dict): Input to the model.57        model: (str): Model name.58 59    Returns:60        int: Number of tokens from the input.61    """62    if isinstance(input, str):63        return _num_token_from_text(input, model=model)64    elif isinstance(input, list) or isinstance(input, dict):65        return _num_token_from_messages(input, model=model)66    else:67        raise ValueError("input must be str, list or dict")68 69 70def _num_token_from_text(text: str, model: str = "gpt-3.5-turbo-0613"):71    """Return the number of tokens used by a string."""72    try:73        encoding = tiktoken.encoding_for_model(model)74    except KeyError:75        logger.warning(f"Model {model} not found. Using cl100k_base encoding.")76        encoding = tiktoken.get_encoding("cl100k_base")77    return len(encoding.encode(text))78 79 80def _num_token_from_messages(messages: Union[List, Dict], model="gpt-3.5-turbo-0613"):81    """Return the number of tokens used by a list of messages.82 83    retrieved from https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb/84    """85    if isinstance(messages, dict):86        messages = [messages]87 88    try:89        encoding = tiktoken.encoding_for_model(model)90    except KeyError:91        print("Warning: model not found. Using cl100k_base encoding.")92        encoding = tiktoken.get_encoding("cl100k_base")93    if model in {94        "gpt-3.5-turbo-0613",95        "gpt-3.5-turbo-16k-0613",96        "gpt-4-0314",97        "gpt-4-32k-0314",98        "gpt-4-0613",99        "gpt-4-32k-0613",100    }:101        tokens_per_message = 3102        tokens_per_name = 1103    elif model == "gpt-3.5-turbo-0301":104        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n105        tokens_per_name = -1  # if there's a name, the role is omitted106    elif "gpt-3.5-turbo" in model:107        logger.info("gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")108        return _num_token_from_messages(messages, model="gpt-3.5-turbo-0613")109    elif "gpt-4" in model:110        logger.info("gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")111        return _num_token_from_messages(messages, model="gpt-4-0613")112    else:113        raise NotImplementedError(114            f"""_num_token_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens."""115        )116    num_tokens = 0117    for message in messages:118        num_tokens += tokens_per_message119        for key, value in message.items():120            if value is None:121                continue122 123            # function calls124            if not isinstance(value, str):125                try:126                    value = json.dumps(value)127                except TypeError:128                    logger.warning(129                        f"Value {value} is not a string and cannot be converted to json. It is a type: {type(value)} Skipping."130                    )131                    continue132 133            num_tokens += len(encoding.encode(value))134            if key == "name":135                num_tokens += tokens_per_name136    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>137    return num_tokens138 139 140def num_tokens_from_functions(functions, model="gpt-3.5-turbo-0613") -> int:141    """Return the number of tokens used by a list of functions.142 143    Args:144        functions: (list): List of function descriptions that will be passed in model.145        model: (str): Model name.146 147    Returns:148        int: Number of tokens from the function descriptions.149    """150    try:151        encoding = tiktoken.encoding_for_model(model)152    except KeyError:153        print("Warning: model not found. Using cl100k_base encoding.")154        encoding = tiktoken.get_encoding("cl100k_base")155 156    num_tokens = 0157    for function in functions:158        function_tokens = len(encoding.encode(function["name"]))159        function_tokens += len(encoding.encode(function["description"]))160        function_tokens -= 2161        if "parameters" in function:162            parameters = function["parameters"]163            if "properties" in parameters:164                for propertiesKey in parameters["properties"]:165                    function_tokens += len(encoding.encode(propertiesKey))166                    v = parameters["properties"][propertiesKey]167                    for field in v:168                        if field == "type":169                            function_tokens += 2170                            function_tokens += len(encoding.encode(v["type"]))171                        elif field == "description":172                            function_tokens += 2173                            function_tokens += len(encoding.encode(v["description"]))174                        elif field == "enum":175                            function_tokens -= 3176                            for o in v["enum"]:177                                function_tokens += 3178                                function_tokens += len(encoding.encode(o))179                        else:180                            print(f"Warning: not supported field {field}")181                function_tokens += 11182                if len(parameters["properties"]) == 0:183                    function_tokens -= 2184 185        num_tokens += function_tokens186 187    num_tokens += 12188    return num_tokens189