Partial sync of codebase

Author: hauntsaninja

Cargo.toml

@@ -14,13 +14,13 @@ python = [
 ]
 
 [dependencies]
-pyo3 = { version = "0.26", default-features = false, features = [
+pyo3 = { version = "0.26.0", default-features = false, features = [
     "extension-module",
     "macros",
 ], optional = true }
 
 # tiktoken dependencies
-fancy-regex = "0.16"
+fancy-regex = "0.13.0"
 regex = "1.10.3"
 rustc-hash = "2"
 bstr = "1.5.0"

pyproject.toml

@@ -31,7 +31,6 @@ skip = [
   "*-manylinux_i686",
   "*-musllinux_i686",
   "*-win32",
-  "*-musllinux_aarch64",
 ]
 macos.archs = ["x86_64", "arm64"]
 # When cross-compiling on Intel, it is not possible to test arm64 wheels.

src/py.rs

@@ -28,7 +28,7 @@ impl CoreBPE {
 
     #[pyo3(name = "encode_ordinary")]
     fn py_encode_ordinary(&self, py: Python, text: &str) -> Vec<Rank> {
-        py.allow_threads(|| self.encode_ordinary(text))
+        py.detach(|| self.encode_ordinary(text))
     }
 
     #[pyo3(name = "encode")]
@@ -38,7 +38,7 @@ impl CoreBPE {
         text: &str,
         allowed_special: HashSet<PyBackedStr>,
     ) -> PyResult<Vec<Rank>> {
-        py.allow_threads(|| {
+        py.detach(|| {
             let allowed_special: HashSet<&str> =
                 allowed_special.iter().map(|s| s.as_ref()).collect();
             match self.encode(text, &allowed_special) {
@@ -54,7 +54,7 @@ impl CoreBPE {
         text: &str,
         allowed_special: HashSet<PyBackedStr>,
     ) -> PyResult<Py<PyAny>> {
-        let tokens_res = py.allow_threads(|| {
+        let tokens_res = py.detach(|| {
             let allowed_special: HashSet<&str> =
                 allowed_special.iter().map(|s| s.as_ref()).collect();
             self.encode(text, &allowed_special)
@@ -70,7 +70,7 @@ impl CoreBPE {
     }
 
     fn _encode_bytes(&self, py: Python, bytes: &[u8]) -> Vec<Rank> {
-        py.allow_threads(|| {
+        py.detach(|| {
             match std::str::from_utf8(bytes) {
                 // Straightforward case
                 Ok(text) => self.encode_ordinary(text),
@@ -121,7 +121,7 @@ impl CoreBPE {
         text: &str,
         allowed_special: HashSet<PyBackedStr>,
     ) -> PyResult<(Vec<Rank>, Py<PyList>)> {
-        let (tokens, completions): (Vec<Rank>, HashSet<Vec<Rank>>) = py.allow_threads(|| {
+        let (tokens, completions): (Vec<Rank>, HashSet<Vec<Rank>>) = py.detach(|| {
             let allowed_special: HashSet<&str> =
                 allowed_special.iter().map(|s| s.as_ref()).collect();
             self._encode_unstable_native(text, &allowed_special)
@@ -155,7 +155,7 @@ impl CoreBPE {
 
     #[pyo3(name = "decode_bytes")]
     fn py_decode_bytes(&self, py: Python, tokens: Vec<Rank>) -> Result<Py<PyBytes>, PyErr> {
-        match py.allow_threads(|| self.decode_bytes(&tokens)) {
+        match py.detach(|| self.decode_bytes(&tokens)) {
             Ok(bytes) => Ok(PyBytes::new(py, &bytes).into()),
             Err(e) => Err(pyo3::exceptions::PyKeyError::new_err(format!("{}", e))),
         }

tiktoken/core.py

@@ -4,11 +4,11 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, AbstractSet, Collection, Literal, NoReturn, Sequence
 
-import regex
-
 from tiktoken import _tiktoken
 
 if TYPE_CHECKING:
+    import re
+
     import numpy as np
     import numpy.typing as npt
 
@@ -391,6 +391,9 @@ def _encode_single_piece(self, text_or_bytes: str | bytes) -> list[int]:
 
     def _encode_only_native_bpe(self, text: str) -> list[int]:
         """Encodes a string into tokens, but do regex splitting in Python."""
+        # We need specifically `regex` in order to compile pat_str due to e.g. \p
+        import regex
+
         _unused_pat = regex.compile(self._pat_str)
         ret = []
         for piece in regex.findall(_unused_pat, text):