Add tokenizer utils

.gitignore

@@ -2,3 +2,6 @@
 __pycache__
 .DS_Store
 **/.DS_Store
+
+anthropic.egg-info/
+build/

README.md

@@ -11,4 +11,5 @@ pip install .
 export ANTHROPIC_API_KEY=<insert token here>
 python examples/basic_sync.py
 python examples/basic_stream.py
+python examples/count_tokens.py
 ```

anthropic/__init__.py

@@ -1,2 +1,3 @@
 from .api import Client
 from .constants import HUMAN_PROMPT, AI_PROMPT, ANTHROPIC_CLIENT_VERSION
+from .tokenizer import count_tokens, get_tokenizer

anthropic/tokenizer.py

@@ -0,0 +1,36 @@
+import os
+import tempfile
+
+import httpx
+from tokenizers import Tokenizer
+
+CLAUDE_TOKENIZER_REMOTE_FILE = "https://public-json-tokenization-0d8763e8-0d7e-441b-a1e2-1c73b8e79dc3.storage.googleapis.com/claude-v1-tokenization.json"
+
+claude_tokenizer = None
+
+def _get_cached_tokenizer_file_as_str() -> str:
+    cache_dir = os.path.join(tempfile.gettempdir(), "anthropic")
+
+    tokenizer_file = os.path.join(cache_dir, 'claude_tokenizer_file.json')
+    if not os.path.exists(tokenizer_file):
+        os.makedirs(cache_dir, exist_ok=True)
+        response = httpx.get(CLAUDE_TOKENIZER_REMOTE_FILE)
+        with open(tokenizer_file, 'w') as f:
+            f.write(response.text)
+
+    with open(tokenizer_file, 'r') as f:
+        return f.read()
+
+def get_tokenizer() -> Tokenizer:
+    global claude_tokenizer
+
+    if not claude_tokenizer:
+        tokenizer_data = _get_cached_tokenizer_file_as_str()
+        claude_tokenizer = Tokenizer.from_str(tokenizer_data)
+
+    return claude_tokenizer
+
+def count_tokens(text: str) -> int:
+    tokenizer = get_tokenizer()
+    encoded_text = tokenizer.encode(text)
+    return len(encoded_text.ids)

examples/count_tokens.py

@@ -0,0 +1,9 @@
+import anthropic
+
+def main(sample_str: str = "Hello world!"):
+    num_tokens = anthropic.count_tokens(sample_str)
+    print(f"Number of tokens: {num_tokens}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -21,8 +21,10 @@ classifiers = [
   "Operating System :: OS Independent"
 ]
 dependencies = [
-    "requests",
+    "httpx",
     'importlib-metadata; python_version<"3.8"',
+    "requests",
+    "tokenizers",
 ]
 
 [project.urls]
@@ -35,4 +37,4 @@ repository = "https://github.com/anthropics/anthropic-sdk-python.git"
 dev = [
     "black >= 22.3.0",
     "pytest",
-]
+]