Merge pull request #2 from pamelafox/tokencounting

Count tokens better

Author: pamelafox

.pre-commit-config.yaml

@@ -1,15 +1,15 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.6.0
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.280
+    rev: v0.4.1
     hooks:
     -   id: ruff
 -   repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 24.4.0
     hooks:
     -   id: black

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "llm-messages-token-helper"
 description = "A helper library for estimating tokens used by messages."
-version = "0.0.2"
+version = "0.0.3"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"
@@ -33,6 +33,8 @@ dev = [
     "ruff",
     "black",
     "flit",
+    "azure-identity",
+    "python-dotenv"
 ]
 
 [build-system]
@@ -42,9 +44,11 @@ build-backend = "flit_core.buildapi"
 [tool.ruff]
 line-length = 120
 target-version = "py39"
+output-format = "full"
+
+[tool.ruff.lint]
 select = ["E", "F", "I", "UP"]
 ignore = ["D203", "E501"]
-show-source = true
 
 [tool.black]
 line-length = 120

src/llm_messages_token_helper/model_helper.py

@@ -35,7 +35,9 @@ def get_token_limit(model: str) -> int:
 
 def count_tokens_for_message(model: str, message: Mapping[str, object]) -> int:
     """
-    Calculate the number of tokens required to encode a message.
+    Calculate the number of tokens required to encode a message. Based off cookbook:
+    https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+
     Args:
         model (str): The name of the model to use for encoding.
         message (Mapping): The message to encode, in a dictionary-like object.
@@ -45,16 +47,19 @@ def count_tokens_for_message(model: str, message: Mapping[str, object]) -> int:
     >> model = 'gpt-3.5-turbo'
     >> message = {'role': 'user', 'content': 'Hello, how are you?'}
     >> count_tokens_for_message(model, message)
-    11
+    13
     """
 
     encoding = tiktoken.encoding_for_model(get_oai_chatmodel_tiktok(model))
-    num_tokens = 2  # For "role" and "content" keys
-    for value in message.values():
+    # Assumes we're using a recent model
+    tokens_per_message = 3
+
+    num_tokens = tokens_per_message
+    for key, value in message.items():
         if isinstance(value, list):
             # For GPT-4-vision support, based on https://github.com/openai/openai-cookbook/pull/881/files
             for item in value:
-                num_tokens += len(encoding.encode(item["type"]))
+                # Note: item[type] does not seem to be counted in the token count
                 if item["type"] == "text":
                     num_tokens += len(encoding.encode(item["text"]))
                 elif item["type"] == "image_url":
@@ -63,6 +68,9 @@ def count_tokens_for_message(model: str, message: Mapping[str, object]) -> int:
             num_tokens += len(encoding.encode(value))
         else:
             raise ValueError(f"Could not encode unsupported message value type: {type(value)}")
+        if key == "name":
+            num_tokens += 1
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
     return num_tokens
 
 

tests/__init__.py

@@ -0,0 +1,65 @@
+user_message = {
+    "message": {
+        "role": "user",
+        "content": "Hello, how are you?",
+    },
+    "count": 13,
+}
+
+user_message_unicode = {
+    "message": {
+        "role": "user",
+        "content": "á",
+    },
+    "count": 8,
+}
+
+system_message_short = {
+    "message": {
+        "role": "system",
+        "content": "You are a bot.",
+    },
+    "count": 12,
+}
+
+system_message = {
+    "message": {
+        "role": "system",
+        "content": "You are a helpful, pattern-following assistant that translates corporate jargon into plain English.",
+    },
+    "count": 25,
+}
+
+system_message_unicode = {
+    "message": {
+        "role": "system",
+        "content": "á",
+    },
+    "count": 8,
+}
+
+system_message_with_name = {
+    "message": {
+        "role": "system",
+        "name": "example_user",
+        "content": "New synergies will help drive top-line growth.",
+    },
+    "count": 20,  # Less tokens in older vision preview models
+}
+
+text_and_image_message = {
+    "message": {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Describe this picture:"},
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg==",
+                    "detail": "auto",
+                },
+            },
+        ],
+    },
+    "count": 266,
+}

tests/messages.py

@@ -1,43 +1,35 @@
 from llm_messages_token_helper import build_messages, count_tokens_for_message
 
+from .messages import system_message_short, system_message_unicode, user_message, user_message_unicode
+
 
 def test_messagebuilder():
-    messages = build_messages("gpt-35-turbo", "You are a bot.")
-    assert messages == [
-        # 1 token, 1 token, 1 token, 5 tokens
-        {"role": "system", "content": "You are a bot."}
-    ]
-    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == 8
+    messages = build_messages("gpt-35-turbo", system_message_short["message"]["content"])
+    assert messages == [system_message_short["message"]]
+    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_short["count"]
 
 
 def test_messagebuilder_append():
-    messages = build_messages("gpt-35-turbo", "You are a bot.", new_user_message="Hello, how are you?")
-    assert messages == [
-        # 1 token, 1 token, 1 token, 5 tokens
-        {"role": "system", "content": "You are a bot."},
-        # 1 token, 1 token, 1 token, 6 tokens
-        {"role": "user", "content": "Hello, how are you?"},
-    ]
-    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == 8
-    assert count_tokens_for_message("gpt-35-turbo", messages[1]) == 9
+    messages = build_messages(
+        "gpt-35-turbo", system_message_short["message"]["content"], new_user_message=user_message["message"]["content"]
+    )
+    assert messages == [system_message_short["message"], user_message["message"]]
+    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_short["count"]
+    assert count_tokens_for_message("gpt-35-turbo", messages[1]) == user_message["count"]
 
 
 def test_messagebuilder_unicode():
-    messages = build_messages("gpt-35-turbo", "a\u0301")
-    assert messages == [
-        # 1 token, 1 token, 1 token, 1 token
-        {"role": "system", "content": "á"}
-    ]
-    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == 4
+    messages = build_messages("gpt-35-turbo", system_message_unicode["message"]["content"])
+    assert messages == [system_message_unicode["message"]]
+    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_unicode["count"]
 
 
 def test_messagebuilder_unicode_append():
-    messages = build_messages("gpt-35-turbo", "a\u0301", new_user_message="a\u0301")
-    assert messages == [
-        # 1 token, 1 token, 1 token, 1 token
-        {"role": "system", "content": "á"},
-        # 1 token, 1 token, 1 token, 1 token
-        {"role": "user", "content": "á"},
-    ]
-    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == 4
-    assert count_tokens_for_message("gpt-35-turbo", messages[1]) == 4
+    messages = build_messages(
+        "gpt-35-turbo",
+        system_message_unicode["message"]["content"],
+        new_user_message=user_message_unicode["message"]["content"],
+    )
+    assert messages == [system_message_unicode["message"], user_message_unicode["message"]]
+    assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_unicode["count"]
+    assert count_tokens_for_message("gpt-35-turbo", messages[1]) == user_message_unicode["count"]

tests/test_messagebuilder.py

@@ -1,6 +1,8 @@
 import pytest
 from llm_messages_token_helper import count_tokens_for_message, get_token_limit
 
+from .messages import system_message, system_message_with_name, text_and_image_message, user_message
+
 
 def test_get_token_limit():
     assert get_token_limit("gpt-35-turbo") == 4000
@@ -29,52 +31,26 @@ def test_get_token_limit_error():
         "gpt-4v",
     ],
 )
-def test_count_tokens_for_message(model: str):
-    message = {
-        # 1 token : 1 token
-        "role": "user",
-        # 1 token : 5 tokens
-        "content": "Hello, how are you?",
-    }
-    assert count_tokens_for_message(model, message) == 9
-
-
-def test_count_tokens_for_message_gpt4():
-    message = {
-        # 1 token : 1 token
-        "role": "user",
-        # 1 token : 5 tokens
-        "content": "Hello, how are you?",
-    }
-    model = "gpt-4"
-    assert count_tokens_for_message(model, message) == 9
+@pytest.mark.parametrize(
+    "message",
+    [
+        user_message,
+        system_message,
+        system_message_with_name,
+    ],
+)
+def test_count_tokens_for_message(model: str, message: dict):
+    assert count_tokens_for_message(model, message["message"]) == message["count"]
 
 
 def test_count_tokens_for_message_list():
-    message = {
-        # 1 token : 1 token
-        "role": "user",
-        # 1 token : 262 tokens
-        "content": [
-            {"type": "text", "text": "Describe this picture:"},  # 1 token  # 4 tokens
-            {
-                "type": "image_url",  # 2 tokens
-                "image_url": {
-                    "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg==",  # 255 tokens
-                    "detail": "auto",
-                },
-            },
-        ],
-    }
     model = "gpt-4"
-    assert count_tokens_for_message(model, message) == 265
+    assert count_tokens_for_message(model, text_and_image_message["message"]) == text_and_image_message["count"]
 
 
 def test_count_tokens_for_message_error():
     message = {
-        # 1 token : 1 token
         "role": "user",
-        # 1 token : 5 tokens
         "content": {"key": "value"},
     }
     model = "gpt-35-turbo"
@@ -85,7 +61,7 @@ def test_count_tokens_for_message_error():
 def test_get_oai_chatmodel_tiktok_error():
     message = {
         "role": "user",
-        "content": {"key": "value"},
+        "content": "hello",
     }
     with pytest.raises(ValueError, match="Expected valid OpenAI GPT model name"):
         count_tokens_for_message("", message)

tests/test_modelhelper.py

@@ -0,0 +1,55 @@
+import os
+
+import azure.identity
+import openai
+from dotenv import load_dotenv
+from messages import (
+    system_message,
+    system_message_short,
+    system_message_unicode,
+    system_message_with_name,
+    text_and_image_message,
+    user_message,
+    user_message_unicode,
+)
+
+# Setup the OpenAI client to use either Azure OpenAI or OpenAI API
+load_dotenv()
+API_HOST = os.getenv("API_HOST")
+
+if API_HOST == "azure":
+    token_provider = azure.identity.get_bearer_token_provider(
+        azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+    )
+    client = openai.AzureOpenAI(
+        api_version=os.getenv("AZURE_OPENAI_VERSION"),
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        azure_ad_token_provider=token_provider,
+    )
+    MODEL_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+else:
+    client = openai.OpenAI(api_key=os.getenv("OPENAI_KEY"))
+    MODEL_NAME = os.getenv("OPENAI_MODEL")
+
+# Test the token count for each message
+for message_count_pair in [
+    user_message,
+    user_message_unicode,
+    system_message,
+    system_message_short,
+    system_message_unicode,
+    system_message_with_name,
+    text_and_image_message,
+]:
+    response = client.chat.completions.create(
+        model=MODEL_NAME,
+        temperature=0.7,
+        n=1,
+        messages=[message_count_pair["message"]],
+    )
+
+    print(message_count_pair["message"])
+    expected_tokens = message_count_pair["count"]
+    assert (
+        response.usage.prompt_tokens == expected_tokens
+    ), f"Expected {expected_tokens} tokens, got {response.usage.prompt_tokens}"