Partial sync of codebase

Author: hauntsaninja

CHANGELOG.md

@@ -2,6 +2,17 @@
 
 This is the changelog for the open source version of tiktoken.
 
+## [v0.11.0]
+- Support for `GPT-5`
+- Update version of `pyo3`
+- Use new Rust edition
+- Fix special token handling in `encode_to_numpy`
+- Improvements to private APIs
+
+## [v0.10.0]
+- Support for newer models
+- Improvements to private APIs
+
 ## [v0.9.0]
 - Support for `o1` and `o3` models
 - Better error messages when loading invalid vocabulary files

Cargo.toml

@@ -1,8 +1,7 @@
 [package]
 name = "tiktoken"
-version = "0.10.0"
-edition = "2021"
-rust-version = "1.57.0"
+version = "0.11.0"
+edition = "2024"
 
 [lib]
 name = "tiktoken"
@@ -15,7 +14,7 @@ python = [
 ]
 
 [dependencies]
-pyo3 = { version = "0.22.2", default-features = false, features = [
+pyo3 = { version = "0.24.1", default-features = false, features = [
     "extension-module",
     "macros",
 ], optional = true }

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tiktoken"
-version = "0.10.0"
+version = "0.11.0"
 description = "tiktoken is a fast BPE tokeniser for use with OpenAI's models"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -22,7 +22,7 @@ requires = ["setuptools>=62.4", "wheel", "setuptools-rust>=1.5.2"]
 build-frontend = "build"
 build-verbosity = 1
 
-linux.before-all = "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal"
+linux.before-all = "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y"
 linux.environment = { PATH = "$PATH:$HOME/.cargo/bin" }
 macos.before-all = "rustup target add aarch64-apple-darwin x86_64-apple-darwin"
 macos.environment = { MACOSX_DEPLOYMENT_TARGET = "10.12" }

src/lib.rs

@@ -481,7 +481,9 @@ impl CoreBPE {
 
         assert!(
             encoder.len() == decoder.len(),
-            "Encoder and decoder must be of equal length; maybe you had duplicate token indices in your encoder?"
+            "Encoder and decoder must be of equal length. Encoder length: {}, decoder length: {}.\nMaybe you had duplicate token indices in your encoder?",
+            encoder.len(),
+            decoder.len()
         );
 
         let special_tokens_decoder: HashMap<Rank, Vec<u8>> = special_tokens_encoder
@@ -524,7 +526,7 @@ mod tests {
     use fancy_regex::Regex;
     use rustc_hash::FxHashMap as HashMap;
 
-    use crate::{byte_pair_split, Rank};
+    use crate::{Rank, byte_pair_split};
 
     fn setup_ranks() -> HashMap<Vec<u8>, Rank> {
         HashMap::from_iter([(b"ab".to_vec(), 0), (b"cd".to_vec(), 1)])

src/py.rs

@@ -1,15 +1,14 @@
 use std::collections::HashSet;
 
 use pyo3::{
-    exceptions,
+    PyResult, exceptions,
     prelude::*,
     pybacked::PyBackedStr,
     types::{PyBytes, PyList, PyTuple},
-    PyResult,
 };
 use rustc_hash::FxHashMap as HashMap;
 
-use crate::{byte_pair_encode, CoreBPE, Rank};
+use crate::{CoreBPE, Rank, byte_pair_encode};
 
 #[pymethods]
 impl CoreBPE {
@@ -19,12 +18,8 @@ impl CoreBPE {
         special_tokens_encoder: HashMap<String, Rank>,
         pattern: &str,
     ) -> PyResult<Self> {
-        Self::new_internal(
-            encoder,
-            special_tokens_encoder,
-            pattern,
-        )
-        .map_err(|e| PyErr::new::<exceptions::PyValueError, _>(e.to_string()))
+        Self::new_internal(encoder, special_tokens_encoder, pattern)
+            .map_err(|e| PyErr::new::<exceptions::PyValueError, _>(e.to_string()))
     }
 
     // ====================
@@ -178,7 +173,7 @@ impl CoreBPE {
     fn token_byte_values(&self, py: Python) -> Vec<Py<PyBytes>> {
         self.sorted_token_bytes
             .iter()
-            .map(|x| PyBytes::new_bound(py, x).into())
+            .map(|x| PyBytes::new(py, x).into())
             .collect()
     }
 }
@@ -204,39 +199,47 @@ impl TiktokenBuffer {
                 "Object is not writable",
             ));
         }
-
-        (*view).obj = slf.clone().into_any().into_ptr();
-
-        let data = &slf.borrow().tokens;
-        (*view).buf = data.as_ptr() as *mut std::os::raw::c_void;
-        (*view).len = (data.len() * std::mem::size_of::<Rank>()) as isize;
-        (*view).readonly = 1;
-        (*view).itemsize = std::mem::size_of::<Rank>() as isize;
-        (*view).format = if (flags & pyo3::ffi::PyBUF_FORMAT) == pyo3::ffi::PyBUF_FORMAT {
-            let msg = std::ffi::CString::new("I").unwrap();
-            msg.into_raw()
-        } else {
-            std::ptr::null_mut()
-        };
-        (*view).ndim = 1;
-        (*view).shape = if (flags & pyo3::ffi::PyBUF_ND) == pyo3::ffi::PyBUF_ND {
-            &mut (*view).len
-        } else {
-            std::ptr::null_mut()
-        };
-        (*view).strides = if (flags & pyo3::ffi::PyBUF_STRIDES) == pyo3::ffi::PyBUF_STRIDES {
-            &mut (*view).itemsize
-        } else {
-            std::ptr::null_mut()
-        };
-        (*view).suboffsets = std::ptr::null_mut();
-        (*view).internal = std::ptr::null_mut();
+        unsafe {
+            let view_ref = &mut *view;
+            view_ref.obj = slf.clone().into_any().into_ptr();
+
+            let data = &slf.borrow().tokens;
+            view_ref.buf = data.as_ptr() as *mut std::os::raw::c_void;
+            view_ref.len = (data.len() * std::mem::size_of::<Rank>()) as isize;
+            view_ref.readonly = 1;
+            view_ref.itemsize = std::mem::size_of::<Rank>() as isize;
+            view_ref.format = if (flags & pyo3::ffi::PyBUF_FORMAT) == pyo3::ffi::PyBUF_FORMAT {
+                let msg = std::ffi::CString::new("I").unwrap();
+                msg.into_raw()
+            } else {
+                std::ptr::null_mut()
+            };
+            view_ref.ndim = 1;
+            view_ref.shape = if (flags & pyo3::ffi::PyBUF_ND) == pyo3::ffi::PyBUF_ND {
+                &mut view_ref.len
+            } else {
+                std::ptr::null_mut()
+            };
+            view_ref.strides = if (flags & pyo3::ffi::PyBUF_STRIDES) == pyo3::ffi::PyBUF_STRIDES {
+                &mut view_ref.itemsize
+            } else {
+                std::ptr::null_mut()
+            };
+            view_ref.suboffsets = std::ptr::null_mut();
+            view_ref.internal = std::ptr::null_mut();
+        }
 
         Ok(())
     }
 
     unsafe fn __releasebuffer__(&self, view: *mut pyo3::ffi::Py_buffer) {
-        std::mem::drop(std::ffi::CString::from_raw((*view).format));
+        // Note that Py_buffer doesn't have a Drop impl
+        unsafe {
+            let view_ref = &mut *view;
+            if !view_ref.format.is_null() {
+                std::mem::drop(std::ffi::CString::from_raw(view_ref.format));
+            }
+        }
     }
 }
 

tests/test_encoding.py

@@ -49,6 +49,13 @@ def test_simple_repeated():
     assert enc.encode("00000000000000000") == [8269, 10535, 830]
 
 
+def test_large_repeated():
+    enc = tiktoken.get_encoding("o200k_base")
+
+    with pytest.raises(ValueError):
+        enc.encode("x" * 1_000_000)
+
+
 def test_simple_regex():
     enc = tiktoken.get_encoding("cl100k_base")
     assert enc.encode("rer") == [38149]
@@ -85,7 +92,7 @@ def test_encode_bytes():
 
 @pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
 @hypothesis.given(bytestring=st.binary())
-@hypothesis.settings(deadline=None)
+@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
 def test_hyp_encode_bytes(make_enc: Callable[[], tiktoken.Encoding], bytestring: bytes):
     enc = make_enc()
     assert enc.decode_bytes(enc._encode_bytes(bytestring)) == bytestring
@@ -140,7 +147,7 @@ def test_basic_roundtrip(make_enc):
 
 @pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
 @hypothesis.given(text=st.text())
-@hypothesis.settings(deadline=None)
+@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
 def test_hyp_roundtrip(make_enc: Callable[[], tiktoken.Encoding], text):
     enc = make_enc()
 
@@ -246,11 +253,11 @@ def test_batch_encode(make_enc: Callable[[], tiktoken.Encoding]):
 
 @pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
 @hypothesis.given(batch=st.lists(st.text()))
-@hypothesis.settings(deadline=None)
+@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
 def test_hyp_batch_roundtrip(make_enc: Callable[[], tiktoken.Encoding], batch):
     enc = make_enc()
 
-    encoded = enc.encode_batch(batch)
-    assert encoded == [enc.encode(t) for t in batch]
+    encoded = enc.encode_batch(batch, allowed_special="all")
+    assert encoded == [enc.encode(t, allowed_special="all") for t in batch]
     decoded = enc.decode_batch(encoded)
     assert decoded == batch

tests/test_misc.py

@@ -17,6 +17,8 @@ def test_encoding_for_model():
     assert enc.name == "cl100k_base"
     enc = tiktoken.encoding_for_model("gpt-4o")
     assert enc.name == "o200k_base"
+    enc = tiktoken.encoding_for_model("gpt-oss-120b")
+    assert enc.name == "o200k_harmony"
 
 
 def test_optional_blobfile_dependency():

tiktoken/__init__.py

@@ -5,4 +5,4 @@
 from .registry import get_encoding as get_encoding
 from .registry import list_encoding_names as list_encoding_names
 
-__version__ = "0.10.0"
+__version__ = "0.11.0"

tiktoken/core.py

@@ -155,7 +155,7 @@ def encode_to_numpy(
 
         import numpy as np
 
-        buffer = self._core_bpe.encode_to_tiktoken_buffer(text, self.special_tokens_set)
+        buffer = self._core_bpe.encode_to_tiktoken_buffer(text, allowed_special)
         return np.frombuffer(buffer, dtype=np.uint32)
 
     def encode_ordinary_batch(self, text: list[str], *, num_threads: int = 8) -> list[list[int]]:
@@ -394,7 +394,7 @@ def _encode_only_native_bpe(self, text: str) -> list[int]:
         _unused_pat = regex.compile(self._pat_str)
         ret = []
         for piece in regex.findall(_unused_pat, text):
-            ret.extend(self._core_bpe.encode_single_piece(piece))
+            ret.extend(self._core_bpe.encode_single_piece(piece.encode("utf-8")))
         return ret
 
     def _encode_bytes(self, text: bytes) -> list[int]:

tiktoken/load.py

@@ -16,7 +16,7 @@ def read_file(blobpath: str) -> bytes:
         with blobfile.BlobFile(blobpath, "rb") as f:
             return f.read()
 
-    # avoiding blobfile for public files helps avoid auth issues, like MFA prompts
+    # avoiding blobfile for public files helps avoid auth issues, like MFA prompts.
     import requests
 
     resp = requests.get(blobpath)
@@ -88,6 +88,7 @@ def data_gym_to_mergeable_bpe_ranks(
     encoder_json_file: str,
     vocab_bpe_hash: str | None = None,
     encoder_json_hash: str | None = None,
+    clobber_one_byte_tokens: bool = False,
 ) -> dict[bytes, int]:
     # NB: do not add caching to this function
     rank_to_intbyte = [b for b in range(2**8) if chr(b).isprintable() and chr(b) != " "]
@@ -109,7 +110,10 @@ def decode_data_gym(value: str) -> bytes:
         return bytes(data_gym_byte_to_byte[b] for b in value)
 
     # add the single byte tokens
+    # if clobber_one_byte_tokens is True, we'll replace these with ones from the encoder json
     bpe_ranks = {bytes([b]): i for i, b in enumerate(rank_to_intbyte)}
+    del rank_to_intbyte
+
     # add the merged tokens
     n = len(bpe_ranks)
     for first, second in bpe_merges:
@@ -126,6 +130,12 @@ def decode_data_gym(value: str) -> bytes:
     # drop these two special tokens if present, since they're not mergeable bpe tokens
     encoder_json_loaded.pop(b"<|endoftext|>", None)
     encoder_json_loaded.pop(b"<|startoftext|>", None)
+
+    if clobber_one_byte_tokens:
+        for k in encoder_json_loaded:
+            if len(k) == 1:
+                bpe_ranks[k] = encoder_json_loaded[k]
+
     assert bpe_ranks == encoder_json_loaded
 
     return bpe_ranks