feat: add PCRE2 resource limits API

README.md

@@ -73,6 +73,23 @@ tokenizer = Tokenizer.from_model("deepseek-ai/DeepSeek-V3.2")
 tokens = tokenizer.encode("A very long prompt that is now lightning fast.")
 ```
 
+PCRE2 resource limits can be set when constructing a tokenizer to guard against
+pathological regex/input combinations in `tokenizer.json`:
+
+```python
+from fastokens._native import Tokenizer
+
+tokenizer = Tokenizer.from_model(
+    "deepseek-ai/DeepSeek-V3.2",
+    pcre2_match_limit=1_000_000,
+)
+```
+
+If a limit is reached during pre-tokenization, encoding returns an error instead
+of continuing an expensive regex match. The same keyword arguments are accepted
+by `Tokenizer(...)`, `Tokenizer.from_file(...)`, `Tokenizer.from_json_str(...)`,
+and `fastokens.patch_transformers(...)`.
+
 ### Dynamo usage
 
 `fastokens` is integrated with NVIDIA Dynamo's frontend, and can be used by passing the flag `--tokenizer fastokens` to the latest version (either build from source or wait for the official release, coming in the next few days).

python/fastokens/__init__.py

@@ -1,10 +1,24 @@
+from __future__ import annotations
+
+from fastokens._compat import _TokenizerShim
 from fastokens._native import Tokenizer
 
 __all__ = ["Tokenizer", "patch_transformers", "unpatch_transformers"]
 
 
 _patched = False
 _originals: dict = {}
+_patch_pcre2_options: dict[str, int | None] = {
+    "pcre2_match_limit": None,
+    "pcre2_depth_limit": None,
+    "pcre2_heap_limit": None,
+    "pcre2_max_jit_stack_size": None,
+}
+
+
+class _ConfiguredTokenizerShim(_TokenizerShim):
+    def __init__(self, src, **kwargs):
+        super().__init__(src, **{**_patch_pcre2_options, **kwargs})
 
 
 def _swap_backend(tokenizer, shim_cls):
@@ -15,7 +29,12 @@ def _swap_backend(tokenizer, shim_cls):
     return tokenizer
 
 
-def patch_transformers() -> None:
+def patch_transformers(
+    pcre2_match_limit: int | None = None,
+    pcre2_depth_limit: int | None = None,
+    pcre2_heap_limit: int | None = None,
+    pcre2_max_jit_stack_size: int | None = None,
+) -> None:
     """
     Monkey-patch ``tokenizers.Tokenizer`` so that the
     ``transformers`` library uses fastokens for encoding.
@@ -33,15 +52,26 @@ def patch_transformers() -> None:
 
     Supports both transformers v4 (``tokenization_utils_fast``)
     and v5+ (``tokenization_utils_tokenizers``).
+
+    PCRE2 resource limits can be supplied to guard against tokenizer regexes
+    with pathological backtracking on adversarial input.
     """
     global _patched
     if _patched:
         print("[fastokens] patch_transformers: already patched.")
         return
 
-    from fastokens._compat import _TokenizerShim
     from fastokens._native import DecodeStream
 
+    _patch_pcre2_options.update(
+        {
+            "pcre2_match_limit": pcre2_match_limit,
+            "pcre2_depth_limit": pcre2_depth_limit,
+            "pcre2_heap_limit": pcre2_heap_limit,
+            "pcre2_max_jit_stack_size": pcre2_max_jit_stack_size,
+        }
+    )
+
     import tokenizers.decoders as _td
 
     # ── v5+: wrap from_pretrained on TokenizersBackend ────────────────
@@ -64,7 +94,7 @@ def patch_transformers() -> None:
         @classmethod
         def _patched_from_pretrained(cls, *args, **kwargs):
             tokenizer = _orig_fp.__func__(cls, *args, **kwargs)
-            return _swap_backend(tokenizer, _TokenizerShim)
+            return _swap_backend(tokenizer, _ConfiguredTokenizerShim)
 
         _originals["TokenizersBackend.from_pretrained"] = _orig_fp
         TokenizersBackend.from_pretrained = _patched_from_pretrained
@@ -80,7 +110,7 @@ def _patched_from_pretrained(cls, *args, **kwargs):
             import transformers.tokenization_utils_fast as _tuf
 
             _originals["tokenization_utils_fast"] = (_tuf, _tuf.TokenizerFast)
-            _tuf.TokenizerFast = _TokenizerShim
+            _tuf.TokenizerFast = _ConfiguredTokenizerShim
         except ImportError:
             pass
 
@@ -139,4 +169,12 @@ def unpatch_transformers() -> None:
     _td.DecodeStream = _originals["DecodeStream"]
 
     _originals.clear()
+    _patch_pcre2_options.update(
+        {
+            "pcre2_match_limit": None,
+            "pcre2_depth_limit": None,
+            "pcre2_heap_limit": None,
+            "pcre2_max_jit_stack_size": None,
+        }
+    )
     _patched = False

python/fastokens/_compat.py

@@ -19,6 +19,19 @@
 # _Encoding directly from this module.
 _Encoding = Encoding
 
+def _pcre2_options(
+    pcre2_match_limit: int | None = None,
+    pcre2_depth_limit: int | None = None,
+    pcre2_heap_limit: int | None = None,
+    pcre2_max_jit_stack_size: int | None = None,
+) -> dict[str, int | None]:
+    return {
+        "pcre2_match_limit": pcre2_match_limit,
+        "pcre2_depth_limit": pcre2_depth_limit,
+        "pcre2_heap_limit": pcre2_heap_limit,
+        "pcre2_max_jit_stack_size": pcre2_max_jit_stack_size,
+    }
+
 
 # ---------------------------------------------------------------------------
 # _TokenizerShim
@@ -33,17 +46,32 @@ class _TokenizerShim:
     original ``tokenizers.Tokenizer`` is kept.
     """
 
-    def __init__(self, src) -> None:
+    def __init__(
+        self,
+        src,
+        pcre2_match_limit: int | None = None,
+        pcre2_depth_limit: int | None = None,
+        pcre2_heap_limit: int | None = None,
+        pcre2_max_jit_stack_size: int | None = None,
+    ) -> None:
+        self._pcre2_options = _pcre2_options(
+            pcre2_match_limit,
+            pcre2_depth_limit,
+            pcre2_heap_limit,
+            pcre2_max_jit_stack_size,
+        )
         if isinstance(src, str):
             self._json = src
-            self._fast = Tokenizer.from_json_str(src)
+            self._fast = Tokenizer.from_json_str(src, **self._pcre2_options)
         elif isinstance(src, _TokenizerShim):
             self._json = src._json
-            self._fast = Tokenizer.from_json_str(src._json)
+            if all(value is None for value in self._pcre2_options.values()):
+                self._pcre2_options = dict(src._pcre2_options)
+            self._fast = Tokenizer.from_json_str(src._json, **self._pcre2_options)
         elif hasattr(src, "to_str"):
             # Accept a real tokenizers.Tokenizer (e.g. from convert_slow_tokenizer).
             self._json = src.to_str()
-            self._fast = Tokenizer.from_json_str(self._json)
+            self._fast = Tokenizer.from_json_str(self._json, **self._pcre2_options)
         else:
             raise TypeError(
                 f"expected JSON string, _TokenizerShim, or tokenizers.Tokenizer; "
@@ -59,15 +87,20 @@ def __getstate__(self):
             self._fast.truncation,
             self._fast.padding,
             self._encode_special_tokens,
+            self._pcre2_options,
         )
 
     def __setstate__(self, state) -> None:
         if isinstance(state, str):
             # Backwards compat: old pickles stored just the JSON string.
             self.__init__(state)  # type: ignore[misc]
         else:
-            json_str, trunc, pad, enc_special = state
-            self.__init__(json_str)  # type: ignore[misc]
+            if len(state) == 4:
+                json_str, trunc, pad, enc_special = state
+                pcre2_options = {}
+            else:
+                json_str, trunc, pad, enc_special, pcre2_options = state
+            self.__init__(json_str, **pcre2_options)  # type: ignore[misc]
             if trunc is not None:
                 self._fast.enable_truncation(**trunc)
             if pad is not None:
@@ -78,7 +111,8 @@ def __deepcopy__(self, memo):
         new = object.__new__(_TokenizerShim)
         memo[id(self)] = new
         new._json = self._json
-        new._fast = Tokenizer.from_json_str(self._json)
+        new._pcre2_options = dict(self._pcre2_options)
+        new._fast = Tokenizer.from_json_str(self._json, **new._pcre2_options)
         trunc = self._fast.truncation
         if trunc is not None:
             new._fast.enable_truncation(**trunc)
@@ -94,30 +128,31 @@ def __deepcopy__(self, memo):
     # -- Factory class methods ------------------------------------------
 
     @classmethod
-    def from_str(cls, json_str: str) -> _TokenizerShim:
-        return cls(json_str)
+    def from_str(cls, json_str: str, **kwargs) -> _TokenizerShim:
+        return cls(json_str, **kwargs)
 
     @classmethod
-    def from_file(cls, path: str) -> _TokenizerShim:
-        return cls(Path(path).read_text(encoding="utf-8"))
+    def from_file(cls, path: str, **kwargs) -> _TokenizerShim:
+        return cls(Path(path).read_text(encoding="utf-8"), **kwargs)
 
     @classmethod
     def from_pretrained(
         cls,
         identifier: str,
         revision: str = "main",
         token: str | None = None,
+        **kwargs,
     ) -> _TokenizerShim:
         from huggingface_hub import hf_hub_download
 
         path = hf_hub_download(
             identifier, "tokenizer.json", revision=revision, token=token,
         )
-        return cls.from_file(path)
+        return cls.from_file(path, **kwargs)
 
     @classmethod
-    def from_buffer(cls, buf: bytes) -> _TokenizerShim:
-        return cls(buf.decode("utf-8"))
+    def from_buffer(cls, buf: bytes, **kwargs) -> _TokenizerShim:
+        return cls(buf.decode("utf-8"), **kwargs)
 
     # -- Serialization --------------------------------------------------
 

python/fastokens/_native.pyi

@@ -72,16 +72,41 @@ class Encoding:
 class Tokenizer:
     """An LLM tokenizer backed by ``tokenizer.json``."""
 
-    def __new__(cls, model: str) -> "Tokenizer": ...
+    def __new__(
+        cls,
+        model: str,
+        pcre2_match_limit: Optional[int] = None,
+        pcre2_depth_limit: Optional[int] = None,
+        pcre2_heap_limit: Optional[int] = None,
+        pcre2_max_jit_stack_size: Optional[int] = None,
+    ) -> "Tokenizer": ...
 
     @staticmethod
-    def from_file(path: str) -> "Tokenizer": ...
+    def from_file(
+        path: str,
+        pcre2_match_limit: Optional[int] = None,
+        pcre2_depth_limit: Optional[int] = None,
+        pcre2_heap_limit: Optional[int] = None,
+        pcre2_max_jit_stack_size: Optional[int] = None,
+    ) -> "Tokenizer": ...
 
     @staticmethod
-    def from_json_str(json: str) -> "Tokenizer": ...
+    def from_json_str(
+        json: str,
+        pcre2_match_limit: Optional[int] = None,
+        pcre2_depth_limit: Optional[int] = None,
+        pcre2_heap_limit: Optional[int] = None,
+        pcre2_max_jit_stack_size: Optional[int] = None,
+    ) -> "Tokenizer": ...
 
     @staticmethod
-    def from_model(model: str) -> "Tokenizer": ...
+    def from_model(
+        model: str,
+        pcre2_match_limit: Optional[int] = None,
+        pcre2_depth_limit: Optional[int] = None,
+        pcre2_heap_limit: Optional[int] = None,
+        pcre2_max_jit_stack_size: Optional[int] = None,
+    ) -> "Tokenizer": ...
 
     @property
     def vocab_size(self) -> int: ...