Bump version, sync codebase

Author: hauntsaninja

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 This is the changelog for the open source version of tiktoken.
 
+## [v0.2.0]
+- Add ``tiktoken.encoding_for_model`` to get the encoding for a specific model
+- Improve portability of caching logic
+
 ## [v0.1.2]
 - Avoid use of `blobfile` for public files
 - Add support for Python 3.8

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tiktoken"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 rust-version = "1.57.0"
 

Makefile

@@ -7,6 +7,9 @@ OpenAI's models.
 import tiktoken
 enc = tiktoken.get_encoding("gpt2")
 assert enc.decode(enc.encode("hello world")) == "hello world"
+
+# To get the tokeniser corresponding to a specific model in the OpenAI API:
+enc = tiktoken.encoding_for_model("text-davinci-003")
 ```
 
 The open source version of `tiktoken` can be installed from PyPI:
@@ -16,7 +19,9 @@ pip install tiktoken
 
 The tokeniser API is documented in `tiktoken/core.py`.
 
-Example code using `tiktoken` can be found in the [OpenAI Cookbook](https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb).
+Example code using `tiktoken` can be found in the
+[OpenAI Cookbook](https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb).
+
 
 ## Performance
 
@@ -28,3 +33,72 @@ Performance measured on 1GB of text using the GPT-2 tokeniser, using `GPT2Tokeni
 `tokenizers==0.13.2` and `transformers==4.24.0`.
 
 
+## Getting help
+
+Please post questions in the [issue tracker](https://github.com/openai/tiktoken/issues).
+
+If you work at OpenAI, make sure to check the internal documentation or feel free to contact
+@shantanu.
+
+
+## Extending tiktoken
+
+You may wish to extend `tiktoken` to support new encodings. There are two ways to do this.
+
+
+**Create your `Encoding` object exactly the way you want and simply pass it around.**
+
+```python
+cl100k_base = tiktoken.get_encoding("cl100k_base")
+
+# In production, load the arguments directly instead of accessing private attributes
+# See openai_public.py for examples of arguments for specific encodings
+enc = tiktoken.Encoding(
+    # If you're changing the set of special tokens, make sure to use a different name
+    # It should be clear from the name what behaviour to expect.
+    name="cl100k_im",
+    pat_str=cl100k_base._pat_str,
+    mergeable_ranks=cl100k_base._mergeable_ranks,
+    special_tokens={
+        **cl100k_base._special_tokens,
+        "<|im_start|>": 100264,
+        "<|im_end|>": 100265,
+    }
+)
+```
+
+**Use the `tiktoken_ext` plugin mechanism to register your `Encoding` objects with `tiktoken`.**
+
+This is only useful if you need `tiktoken.get_encoding` to find your encoding, otherwise prefer
+option 1.
+
+To do this, you'll need to create a namespace package under `tiktoken_ext`.
+
+Layout your project like this, making sure to omit the `tiktoken_ext/__init__.py` file:
+```
+my_tiktoken_extension
+├── tiktoken_ext
+│   └── my_encodings.py
+└── setup.py
+```
+
+`my_encodings.py` should be a module that contains a variable named `ENCODING_CONSTRUCTORS`.
+This is a dictionary from an encoding name to a function that takes no arguments and returns
+arguments that can be passed to `tiktoken.Encoding` to construct that encoding. For an example, see
+`tiktoken_ext/openai_public.py`. For precise details, see `tiktoken/registry.py`.
+
+Your `setup.py` should look something like this:
+```python
+from setuptools import setup, find_namespace_packages
+
+setup(
+    name="my_tiktoken_extension",
+    packages=find_namespace_packages(include=['tiktoken_ext.*'])
+    install_requires=["tiktoken"],
+    ...
+)
+```
+
+Then simply `pip install my_tiktoken_extension` and you should be able to use your custom encodings!
+Make sure **not** to use an editable install.
+

README.md

@@ -1,12 +1,12 @@
 [project]
 name = "tiktoken"
+version = "0.2.0"
 dependencies = ["blobfile>=2", "regex>=2022.1.18", "requests>=2.26.0"]
-dynamic = ["version"]
 requires-python = ">=3.8"
 
 [build-system]
 build-backend = "setuptools.build_meta"
-requires = ["setuptools>=61", "wheel", "setuptools-rust>=1.3"]
+requires = ["setuptools>=62.4", "wheel", "setuptools-rust>=1.5.2"]
 
 [tool.cibuildwheel]
 build-frontend = "build"

pyproject.toml

@@ -1,14 +1,8 @@
 from setuptools import setup
 from setuptools_rust import Binding, RustExtension
 
-public = True
-
-if public:
-    version = "0.1.2"
-
 setup(
     name="tiktoken",
-    version=version,
     rust_extensions=[
         RustExtension(
             "tiktoken._tiktoken",

setup.py

@@ -17,3 +17,10 @@ def test_simple():
         enc = tiktoken.get_encoding(enc_name)
         for token in range(10_000):
             assert enc.encode_single_token(enc.decode_single_token_bytes(token)) == token
+
+
+def test_encoding_for_model():
+    enc = tiktoken.encoding_for_model("gpt2")
+    assert enc.name == "gpt2"
+    enc = tiktoken.encoding_for_model("text-davinci-003")
+    assert enc.name == "p50k_base"

tests/test_simple_public.py

@@ -1,3 +1,4 @@
 from .core import Encoding as Encoding
+from .model import encoding_for_model as encoding_for_model
 from .registry import get_encoding as get_encoding
 from .registry import list_encoding_names as list_encoding_names

tiktoken/__init__.py

@@ -19,6 +19,21 @@ def __init__(
         special_tokens: dict[str, int],
         explicit_n_vocab: Optional[int] = None,
     ):
+        """Creates an Encoding object.
+
+        See openai_public.py for examples of how to construct an Encoding object.
+
+        Args:
+            name: The name of the encoding. It should be clear from the name of the encoding
+                what behaviour to expect, in particular, encodings with different special tokens
+                should have different names.
+            pat_str: A regex pattern string that is used to split the input text.
+            mergeable_ranks: A dictionary mapping mergeable token bytes to their ranks. The ranks
+                must correspond to merge priority.
+            special_tokens: A dictionary mapping special token strings to their token values.
+            explicit_n_vocab: The number of tokens in the vocabulary. If provided, it is checked
+                that the number of mergeable tokens and special tokens is equal to this number.
+        """
         self.name = name
 
         self._pat_str = pat_str

tiktoken/core.py

@@ -4,6 +4,7 @@
 import hashlib
 import json
 import os
+import tempfile
 import uuid
 
 import blobfile
@@ -24,7 +25,7 @@ def read_file_cached(blobpath: str) -> bytes:
     elif "DATA_GYM_CACHE_DIR" in os.environ:
         cache_dir = os.environ["DATA_GYM_CACHE_DIR"]
     else:
-        cache_dir = "/tmp/data-gym-cache"
+        cache_dir = os.path.join(tempfile.gettempdir(), "data-gym-cache")
 
     if cache_dir == "":
         # disable caching