diff --git a/README.md b/README.md index e9fe141..90caa97 100644 --- a/README.md +++ b/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). diff --git a/python/fastokens/__init__.py b/python/fastokens/__init__.py index f03d25e..223b057 100644 --- a/python/fastokens/__init__.py +++ b/python/fastokens/__init__.py @@ -1,3 +1,6 @@ +from __future__ import annotations + +from fastokens._compat import _TokenizerShim from fastokens._native import Tokenizer __all__ = ["Tokenizer", "patch_transformers", "unpatch_transformers"] @@ -5,6 +8,17 @@ _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 diff --git a/python/fastokens/_compat.py b/python/fastokens/_compat.py index 6d4dbd6..7c70fb4 100644 --- a/python/fastokens/_compat.py +++ b/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,6 +87,7 @@ def __getstate__(self): self._fast.truncation, self._fast.padding, self._encode_special_tokens, + self._pcre2_options, ) def __setstate__(self, state) -> None: @@ -66,8 +95,12 @@ def __setstate__(self, state) -> None: # 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,12 +128,12 @@ 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( @@ -107,17 +141,18 @@ def from_pretrained( 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 -------------------------------------------------- diff --git a/python/fastokens/_native.pyi b/python/fastokens/_native.pyi index 0292a16..4be5bfd 100644 --- a/python/fastokens/_native.pyi +++ b/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: ... diff --git a/python/src/lib.rs b/python/src/lib.rs index 45ff6ab..1d6f1e6 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -382,6 +382,22 @@ struct TokenizerState { post_processor_json: Option, } +fn tokenizer_options( + pcre2_match_limit: Option, + pcre2_depth_limit: Option, + pcre2_heap_limit: Option, + pcre2_max_jit_stack_size: Option, +) -> fastokens::TokenizerOptions { + fastokens::TokenizerOptions { + pcre2_limits: fastokens::Pcre2Limits { + match_limit: pcre2_match_limit, + depth_limit: pcre2_depth_limit, + heap_limit: pcre2_heap_limit, + max_jit_stack_size: pcre2_max_jit_stack_size, + }, + } +} + impl TokenizerState { fn do_truncate(&self, ids: &mut Vec) { let Some(ref t) = self.trunc else { return }; @@ -438,7 +454,11 @@ impl PyTokenizer { /// Build from a raw JSON string, extracting the post-processor field so /// the getter can return it without needing to re-serialize. - fn build_from_str(json: &str, py: Python<'_>) -> PyResult { + fn build_from_str( + json: &str, + py: Python<'_>, + options: fastokens::TokenizerOptions, + ) -> PyResult { let value: Value = serde_json::from_str(json).map_err(|e| PyValueError::new_err(e.to_string()))?; let post_processor_json = value @@ -446,7 +466,10 @@ impl PyTokenizer { .filter(|v| !v.is_null()) .map(|v| v.to_string()); let inner = py - .allow_threads(|| fastokens::Tokenizer::from_json(value).map_err(|e| e.to_string())) + .allow_threads(|| { + fastokens::Tokenizer::from_json_with_options(value, options) + .map_err(|e| e.to_string()) + }) .map_err(PyValueError::new_err)?; Ok(Self { state: RwLock::new(TokenizerState { @@ -466,34 +489,100 @@ impl PyTokenizer { /// /// (This is an alias for Tokenizer.from_model) #[new] - fn new(model: &str, py: Python<'_>) -> PyResult { - Self::from_model(model, py) + #[pyo3(signature = (model, pcre2_match_limit = None, pcre2_depth_limit = None, pcre2_heap_limit = None, pcre2_max_jit_stack_size = None))] + fn new( + model: &str, + pcre2_match_limit: Option, + pcre2_depth_limit: Option, + pcre2_heap_limit: Option, + pcre2_max_jit_stack_size: Option, + py: Python<'_>, + ) -> PyResult { + Self::from_model( + model, + pcre2_match_limit, + pcre2_depth_limit, + pcre2_heap_limit, + pcre2_max_jit_stack_size, + py, + ) } /// Create a tokenizer from a `tokenizer.json` file. #[staticmethod] - fn from_file(path: &str, py: Python<'_>) -> PyResult { + #[pyo3(signature = (path, pcre2_match_limit = None, pcre2_depth_limit = None, pcre2_heap_limit = None, pcre2_max_jit_stack_size = None))] + fn from_file( + path: &str, + pcre2_match_limit: Option, + pcre2_depth_limit: Option, + pcre2_heap_limit: Option, + pcre2_max_jit_stack_size: Option, + py: Python<'_>, + ) -> PyResult { let json = std::fs::read_to_string(path) .map_err(|e| PyValueError::new_err(format!("cannot read {path}: {e}")))?; - Self::build_from_str(&json, py) + Self::build_from_str( + &json, + py, + tokenizer_options( + pcre2_match_limit, + pcre2_depth_limit, + pcre2_heap_limit, + pcre2_max_jit_stack_size, + ), + ) } /// Create a tokenizer from a raw JSON string for `tokenizer.json`. #[staticmethod] - fn from_json_str(json: &str, py: Python<'_>) -> PyResult { - Self::build_from_str(json, py) + #[pyo3(signature = (json, pcre2_match_limit = None, pcre2_depth_limit = None, pcre2_heap_limit = None, pcre2_max_jit_stack_size = None))] + fn from_json_str( + json: &str, + pcre2_match_limit: Option, + pcre2_depth_limit: Option, + pcre2_heap_limit: Option, + pcre2_max_jit_stack_size: Option, + py: Python<'_>, + ) -> PyResult { + Self::build_from_str( + json, + py, + tokenizer_options( + pcre2_match_limit, + pcre2_depth_limit, + pcre2_heap_limit, + pcre2_max_jit_stack_size, + ), + ) } /// Download `tokenizer.json` from HuggingFace Hub for the given model /// (e.g. `"meta-llama/Llama-3.1-8B"`) and create a tokenizer with it. #[staticmethod] - fn from_model(model: &str, py: Python<'_>) -> PyResult { + #[pyo3(signature = (model, pcre2_match_limit = None, pcre2_depth_limit = None, pcre2_heap_limit = None, pcre2_max_jit_stack_size = None))] + fn from_model( + model: &str, + pcre2_match_limit: Option, + pcre2_depth_limit: Option, + pcre2_heap_limit: Option, + pcre2_max_jit_stack_size: Option, + py: Python<'_>, + ) -> PyResult { let json = py .allow_threads(|| { fastokens::Tokenizer::download_tokenizer_json(model).map_err(|e| e.to_string()) }) .map_err(PyValueError::new_err)?; - Self::build_from_str(&json, py) + Self::build_from_str( + &json, + py, + tokenizer_options( + pcre2_match_limit, + pcre2_depth_limit, + pcre2_heap_limit, + pcre2_max_jit_stack_size, + ), + ) } // ── Post-processor ──────────────────────────────────────────────── diff --git a/python/tests/test_pcre2_limits.py b/python/tests/test_pcre2_limits.py new file mode 100644 index 0000000..0b0706d --- /dev/null +++ b/python/tests/test_pcre2_limits.py @@ -0,0 +1,49 @@ +import copy +import json +import pickle + +import pytest + +from fastokens._compat import _TokenizerShim +from fastokens._native import Tokenizer + + +def _tokenizer_json() -> str: + return json.dumps( + { + "model": { + "type": "BPE", + "vocab": {"a": 0, "!": 1}, + "merges": [], + }, + "pre_tokenizer": { + "type": "Split", + "pattern": {"Regex": "^(a+)+$"}, + "behavior": "Isolated", + "invert": False, + }, + } + ) + + +def test_native_from_json_str_applies_pcre2_match_limit(): + tok = Tokenizer.from_json_str(_tokenizer_json(), pcre2_match_limit=1) + + with pytest.raises(ValueError, match="match limit"): + tok.encode("aaaaaaaaaaaaaaaa!") + + +def test_shim_deepcopy_preserves_pcre2_limits(): + shim = _TokenizerShim(_tokenizer_json(), pcre2_match_limit=1) + cloned = copy.deepcopy(shim) + + with pytest.raises(ValueError, match="match limit"): + cloned.encode("aaaaaaaaaaaaaaaa!") + + +def test_shim_pickle_preserves_pcre2_limits(): + shim = _TokenizerShim(_tokenizer_json(), pcre2_match_limit=1) + restored = pickle.loads(pickle.dumps(shim)) + + with pytest.raises(ValueError, match="match limit"): + restored.encode("aaaaaaaaaaaaaaaa!") diff --git a/src/json_structs.rs b/src/json_structs.rs index a9aaa5e..7795e19 100644 --- a/src/json_structs.rs +++ b/src/json_structs.rs @@ -116,7 +116,7 @@ pub enum PreTokenizerConfig { ByteLevel(pre_tokenizers::ByteLevel), Whitespace, WhitespaceSplit, - Split(pre_tokenizers::Split), + Split(pre_tokenizers::SplitConfig), Punctuation { #[serde(default)] behavior: String, diff --git a/src/lib.rs b/src/lib.rs index b931528..1ba511a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,7 +22,7 @@ pub use self::{ models::Model, normalizers::{Nfc, Normalizer}, post_processors::PostProcessor, - pre_tokenizers::{ByteLevel, PreTokenizer, Split, SplitBehavior}, + pre_tokenizers::{ByteLevel, Pcre2Limits, PreTokenizer, Split, SplitBehavior}, }; use self::{ @@ -35,7 +35,7 @@ use self::{ mod hf_hub_support { pub use hf_hub::api::sync::ApiError; - use super::{Error, Tokenizer, TokenizerJson}; + use super::{Error, Tokenizer, TokenizerJson, TokenizerOptions}; use hf_hub::api::sync::{Api, ApiBuilder}; use std::fs; @@ -62,13 +62,21 @@ mod hf_hub_support { /// Used by `Tokenizer::from_model` and `Tokenizer::from_model_with_token` to fetch /// `tokenizer.json` from the HuggingFace Hub and build a `Tokenizer`. pub fn from_model_with_token(model: &str, token: Option<&str>) -> Result { + from_model_with_token_and_options(model, token, TokenizerOptions::default()) + } + + pub fn from_model_with_token_and_options( + model: &str, + token: Option<&str>, + options: TokenizerOptions, + ) -> Result { validate_model_id(model)?; let api = make_api(token)?; let repo = api.model(model.to_string()); let json_path = repo.get("tokenizer.json")?; let raw = fs::read_to_string(json_path)?; let json: TokenizerJson = serde_json::from_str(&raw)?; - Tokenizer::build(json) + Tokenizer::build_with_options(json, options) } /// Used by the Python layer to fetch `tokenizer.json` from the HuggingFace Hub and @@ -114,6 +122,12 @@ pub enum Error { InvalidIdentifier(String), } +/// Options applied while constructing a [`Tokenizer`]. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct TokenizerOptions { + pub pcre2_limits: Pcre2Limits, +} + /// An LLM tokenizer backed by `tokenizer.json`. pub struct Tokenizer { added_tokens: Option, @@ -130,11 +144,15 @@ pub struct Tokenizer { impl Tokenizer { /// Build the pipeline steps from a parsed JSON config. fn build(json: TokenizerJson) -> Result { + Self::build_with_options(json, TokenizerOptions::default()) + } + + fn build_with_options(json: TokenizerJson, options: TokenizerOptions) -> Result { let added_tokens = AddedTokens::from_configs(&json.added_tokens).map_err(Error::Model)?; let normalizer = json.normalizer.map(Normalizer::from_config).transpose()?; let pre_tokenizer = json .pre_tokenizer - .map(PreTokenizer::from_config) + .map(|config| PreTokenizer::from_config_with_limits(config, options.pcre2_limits)) .transpose()?; let model = Model::from_config(json.model).map_err(Error::Model)?; let post_processor = json @@ -181,12 +199,24 @@ impl Tokenizer { Self::build(json) } + /// Create a tokenizer from a raw JSON value for `tokenizer.json` with construction options. + pub fn from_json_with_options(json: Value, options: TokenizerOptions) -> Result { + let json: TokenizerJson = serde_json::from_value(json)?; + Self::build_with_options(json, options) + } + /// Create a tokenizer from a `tokenizer.json` file. pub fn from_file(path: &Path) -> Result { let json: TokenizerJson = serde_json::from_str(&fs::read_to_string(path)?)?; Self::build(json) } + /// Create a tokenizer from a `tokenizer.json` file with construction options. + pub fn from_file_with_options(path: &Path, options: TokenizerOptions) -> Result { + let json: TokenizerJson = serde_json::from_str(&fs::read_to_string(path)?)?; + Self::build_with_options(json, options) + } + /// Download `tokenizer.json` from HuggingFace Hub for the given model (e.g. /// `"meta-llama/Llama-3.1-8B"`) and create a tokenizer with it. /// @@ -198,6 +228,12 @@ impl Tokenizer { Self::from_model_with_token(model, None) } + /// Like [`Self::from_model`] but accepts construction options. + #[cfg(feature = "hf-hub")] + pub fn from_model_with_options(model: &str, options: TokenizerOptions) -> Result { + Self::from_model_with_token_and_options(model, None, options) + } + /// Like [`Self::from_model`] but accepts an explicit HuggingFace token, /// overriding the credential cache. Pass `None` to use the credential /// cache (`~/.cache/huggingface/token`, set via `huggingface-cli login`). @@ -206,6 +242,16 @@ impl Tokenizer { hf_hub_support::from_model_with_token(model, token) } + /// Like [`Self::from_model_with_token`] but accepts construction options. + #[cfg(feature = "hf-hub")] + pub fn from_model_with_token_and_options( + model: &str, + token: Option<&str>, + options: TokenizerOptions, + ) -> Result { + hf_hub_support::from_model_with_token_and_options(model, token, options) + } + /// Download `tokenizer.json` and return its raw content without building /// the tokenizer. Used by the Python layer to extract fields (such as /// `post_processor`) before handing the JSON off to [`Self::from_json`]. @@ -593,6 +639,45 @@ pub fn decode_stream_step( } } +#[cfg(test)] +mod local_tests { + use serde_json::json; + + use super::*; + + #[test] + fn from_json_with_options_propagates_pcre2_limits() { + let tokenizer = Tokenizer::from_json_with_options( + json!({ + "model": { + "type": "BPE", + "vocab": {"a": 0, "!": 1}, + "merges": [] + }, + "pre_tokenizer": { + "type": "Split", + "pattern": {"Regex": "^(a+)+$"}, + "behavior": "Isolated", + "invert": false + } + }), + TokenizerOptions { + pcre2_limits: Pcre2Limits { + match_limit: Some(1), + ..Default::default() + }, + }, + ) + .unwrap(); + + let err = tokenizer.encode("aaaaaaaaaaaaaaaa!").unwrap_err(); + assert!( + err.to_string().contains("match limit"), + "expected match limit error, got {err}" + ); + } +} + #[cfg(all(test, feature = "hf-hub"))] mod tests { use crate::hf_hub_support::make_api; diff --git a/src/pre_tokenizers.rs b/src/pre_tokenizers.rs index 30cc65e..1e1ce27 100644 --- a/src/pre_tokenizers.rs +++ b/src/pre_tokenizers.rs @@ -8,7 +8,7 @@ use crate::{ pub use self::{ byte_level::ByteLevel, - split::{Split, SplitBehavior}, + split::{Pcre2Limits, Split, SplitBehavior, SplitConfig}, }; pub(crate) use self::byte_level::BYTE_TO_CHAR; @@ -42,13 +42,22 @@ pub enum PreTokenizer { impl PreTokenizer { /// Build a pre-tokenizer from its JSON configuration. pub fn from_config(config: PreTokenizerConfig) -> Result { + Self::from_config_with_limits(config, Pcre2Limits::default()) + } + + pub fn from_config_with_limits( + config: PreTokenizerConfig, + limits: Pcre2Limits, + ) -> Result { match config { PreTokenizerConfig::ByteLevel(bl) => Ok(Self::ByteLevel(bl)), - PreTokenizerConfig::Split(s) => Ok(Self::Split(s)), + PreTokenizerConfig::Split(s) => Ok(Self::Split(Split::from_split_config_with_limits( + s, limits, + )?)), PreTokenizerConfig::Sequence { pretokenizers } => { let steps = pretokenizers .into_iter() - .map(Self::from_config) + .map(|config| Self::from_config_with_limits(config, limits)) .collect::, _>>()?; Ok(Self::Sequence(steps)) } diff --git a/src/pre_tokenizers/split.rs b/src/pre_tokenizers/split.rs index 2008e47..5145bdd 100644 --- a/src/pre_tokenizers/split.rs +++ b/src/pre_tokenizers/split.rs @@ -2,6 +2,7 @@ use std::cell::RefCell; use std::sync::atomic::{AtomicUsize, Ordering}; use fancy_regex::Regex; +use memchr::{memchr, memchr2}; use rayon::prelude::*; use serde::Deserialize; use serde_json::Value; @@ -27,8 +28,50 @@ struct SplitCache { /// Minimum shared prefix length (bytes) before incremental re-use kicks in. const INCREMENTAL_MIN_PREFIX: usize = 4096; +/// PCRE2 resource limits applied while compiling Split regexes. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Pcre2Limits { + pub match_limit: Option, + pub depth_limit: Option, + /// Heap limit in KiB, matching PCRE2's `(*LIMIT_HEAP=...)` units. + pub heap_limit: Option, + pub max_jit_stack_size: Option, +} + +impl Pcre2Limits { + fn has_any(self) -> bool { + self.match_limit.is_some() + || self.depth_limit.is_some() + || self.heap_limit.is_some() + || self.max_jit_stack_size.is_some() + } + + fn apply_to_source(self, source: &str) -> std::string::String { + if self.match_limit.is_none() && self.depth_limit.is_none() && self.heap_limit.is_none() { + return source.to_string(); + } + + let mut prefixed = String::new(); + if let Some(limit) = self.heap_limit { + prefixed.push_str(&format!("(*LIMIT_HEAP={limit})")); + } + if let Some(limit) = self.match_limit { + prefixed.push_str(&format!("(*LIMIT_MATCH={limit})")); + } + if let Some(limit) = self.depth_limit { + prefixed.push_str(&format!("(*LIMIT_DEPTH={limit})")); + } + prefixed.push_str(source); + prefixed + } +} + /// Wrapper around a JIT-compiled PCRE2 regex for the Llama-3 pattern. -struct Pcre2Regex(pcre2::bytes::Regex); +struct Pcre2Regex { + regex: pcre2::bytes::Regex, + source: std::string::String, + max_jit_stack_size: Option, +} // Safety: PCRE2 JIT-compiled regexes are thread-safe for matching. // Each thread uses independent match data internally via pcre2 crate. @@ -43,15 +86,17 @@ impl std::fmt::Debug for Pcre2Regex { impl Clone for Pcre2Regex { fn clone(&self) -> Self { - // Re-compile for independent match state. - Self( - pcre2::bytes::RegexBuilder::new() - .utf(true) - .ucp(true) - .jit_if_available(true) - .build(self.0.as_str()) - .expect("re-compile PCRE2 regex"), - ) + let mut builder = pcre2::bytes::RegexBuilder::new(); + builder + .utf(true) + .ucp(true) + .jit_if_available(true) + .max_jit_stack_size(self.max_jit_stack_size); + Self { + regex: builder.build(&self.source).expect("re-compile PCRE2 regex"), + source: self.source.clone(), + max_jit_stack_size: self.max_jit_stack_size, + } } } @@ -118,13 +163,13 @@ pub enum SplitBehavior { } /// Raw deserialization helper for [`Split`]. -#[derive(Deserialize)] -struct SplitRaw { - pattern: Pattern, +#[derive(Clone, Debug, Deserialize)] +pub struct SplitConfig { + pub pattern: Pattern, #[serde(default)] - behavior: SplitBehavior, + pub behavior: SplitBehavior, #[serde(default)] - invert: bool, + pub invert: bool, } /// Monotonic counter for unique Split instance IDs. @@ -138,7 +183,7 @@ static SPLIT_ID_COUNTER: AtomicUsize = AtomicUsize::new(1); /// /// [`PreTokenizerConfig::Split`]: crate::PreTokenizerConfig::Split #[derive(Clone, Debug, Deserialize)] -#[serde(try_from = "SplitRaw")] +#[serde(try_from = "SplitConfig")] pub struct Split { /// Unique identity for cache invalidation across different Split instances. #[serde(skip)] @@ -157,21 +202,96 @@ pub struct Split { /// Compile PCRE2 JIT regexes from `source`, returning `None` if PCRE2 cannot /// handle the pattern (e.g. unsupported syntax). -fn try_compile_pcre2_regexes(source: &str, n: usize) -> Option> { - if source.contains("&&") { - return None; +fn try_compile_pcre2_regexes( + source: &str, + n: usize, + limits: Pcre2Limits, +) -> Result>, Error> { + if has_class_intersection(source) { + if limits.has_any() { + return Err(Error::Unsupported( + "PCRE2 limits require a PCRE2-compatible regex pattern".into(), + )); + } + return Ok(None); } + let source = limits.apply_to_source(source); let mut regexes = Vec::with_capacity(n); for _ in 0..n { - let re = pcre2::bytes::RegexBuilder::new() + let mut builder = pcre2::bytes::RegexBuilder::new(); + builder .utf(true) .ucp(true) .jit_if_available(true) - .build(source) - .ok()?; - regexes.push(Pcre2Regex(re)); + .max_jit_stack_size(limits.max_jit_stack_size); + let re = match builder.build(&source) { + Ok(re) => re, + Err(e) if limits.has_any() => { + return Err(Error::Unsupported(format!( + "PCRE2 limits require a PCRE2-compatible regex pattern: {e}" + ))); + } + Err(_) => return Ok(None), + }; + regexes.push(Pcre2Regex { + regex: re, + source: source.clone(), + max_jit_stack_size: limits.max_jit_stack_size, + }); } - Some(regexes) + Ok(Some(regexes)) +} + +fn has_class_intersection(source: &str) -> bool { + // && inside a regex character class is the fancy-regex set-intersection syntax, e.g. [a&&b], + // which PCRE2 does not support. + let bytes = source.as_bytes(); + + if memchr::memmem::find(bytes, b"&&").is_none() { + return false; + } + + let mut search_start = 0; + 'classes: while let Some(open_rel) = memchr(b'[', &bytes[search_start..]) { + let open = search_start + open_rel; + if is_escaped(bytes, open) { + search_start = open + 1; + continue; + } + + let mut pos = open + 1; + loop { + let Some(rel) = memchr2(b'&', b']', &bytes[pos..]) else { + return false; + }; + let idx = pos + rel; + if is_escaped(bytes, idx) { + pos = idx + 1; + continue; + } + + match bytes[idx] { + b'&' if bytes.get(idx + 1) == Some(&b'&') => return true, + b']' if idx > open + 1 => { + search_start = idx + 1; + continue 'classes; + } + _ => pos = idx + 1, + } + } + } + + false +} + +fn is_escaped(bytes: &[u8], pos: usize) -> bool { + let mut slash_count = 0; + let mut i = pos; + while i > 0 && bytes[i - 1] == b'\\' { + slash_count += 1; + i -= 1; + } + slash_count % 2 == 1 } /// Compile `n` independent copies of a regex from `source`. @@ -183,41 +303,62 @@ fn compile_regexes(source: &str, n: usize) -> Result, Error> { Ok(regexes) } -impl TryFrom for Split { +impl TryFrom for Split { type Error = Error; - fn try_from(raw: SplitRaw) -> Result { - let source = raw.pattern.source(); - let pcre2_regexes = try_compile_pcre2_regexes(&source, max_parallel()); + fn try_from(raw: SplitConfig) -> Result { + Self::from_split_config_with_limits(raw, Pcre2Limits::default()) + } +} + +impl Split { + fn from_parts( + source: std::string::String, + behavior: SplitBehavior, + invert: bool, + limits: Pcre2Limits, + ) -> Result { let regexes = compile_regexes(&source, max_parallel())?; + let pcre2_regexes = try_compile_pcre2_regexes(&source, max_parallel(), limits)?; Ok(Self { id: SPLIT_ID_COUNTER.fetch_add(1, Ordering::Relaxed), regexes, - behavior: raw.behavior, - invert: raw.invert, + behavior, + invert, pcre2_regexes, }) } -} -impl Split { + pub fn from_split_config_with_limits( + config: SplitConfig, + limits: Pcre2Limits, + ) -> Result { + Self::from_parts( + config.pattern.source(), + config.behavior, + config.invert, + limits, + ) + } + /// Build a [`Split`] from raw JSON fields. /// /// `pattern` must be a JSON object with either a `"String"` key (literal, /// will be regex-escaped) or a `"Regex"` key (used as-is). pub fn from_config(pattern: &Value, behavior: &str, invert: bool) -> Result { + Self::from_config_with_limits(pattern, behavior, invert, Pcre2Limits::default()) + } + + pub fn from_config_with_limits( + pattern: &Value, + behavior: &str, + invert: bool, + limits: Pcre2Limits, + ) -> Result { let pattern: Pattern = serde_json::from_value(pattern.clone())?; let source = pattern.source(); - let regexes = compile_regexes(&source, max_parallel())?; let behavior: SplitBehavior = serde_json::from_value(Value::String(behavior.to_string()))?; - let pcre2_regexes = try_compile_pcre2_regexes(&source, max_parallel()); - Ok(Self { - id: SPLIT_ID_COUNTER.fetch_add(1, Ordering::Relaxed), - regexes, - behavior, - invert, - pcre2_regexes, - }) + Self::from_parts(source, behavior, invert, limits) } /// Refine the splits of a [`PreTokenizedString`] in place. @@ -435,7 +576,7 @@ impl Split { if p >= bytes.len() { return Ok(None); } - match regex.0.find_at(bytes, p) { + match regex.regex.find_at(bytes, p) { Ok(Some(m)) => { if m.start() == m.end() { p = m.end() + 1; @@ -831,7 +972,7 @@ fn find_matches_pcre2( let bytes = input.as_bytes(); let mut pos = 0; while pos < bytes.len() { - match regex.0.find_at(bytes, pos) { + match regex.regex.find_at(bytes, pos) { Ok(Some(m)) => { if m.start() == m.end() { pos = m.end() + 1; @@ -1015,6 +1156,61 @@ mod tests { assert!(matches!(err, Error::Json(_))); } + #[test] + fn pcre2_match_limit_is_enforced() { + let split = Split::from_config_with_limits( + &json!({"Regex": "^(a+)+$"}), + "Isolated", + false, + Pcre2Limits { + match_limit: Some(1), + ..Default::default() + }, + ) + .unwrap(); + + let err = split.split("aaaaaaaaaaaaaaaa!").unwrap_err(); + assert!( + err.to_string().contains("match limit"), + "expected match limit error, got {err}" + ); + } + + #[test] + fn pcre2_limits_require_pcre2_compatible_pattern() { + let err = Split::from_config_with_limits( + &json!({"Regex": "[a&&b]"}), + "Isolated", + false, + Pcre2Limits { + match_limit: Some(1_000), + ..Default::default() + }, + ) + .unwrap_err(); + + assert!( + err.to_string().contains("PCRE2 limits require"), + "expected PCRE2 compatibility error, got {err}" + ); + } + + #[test] + fn pcre2_limits_allow_literal_ampersands() { + let split = Split::from_config_with_limits( + &json!({"String": "&&"}), + "Isolated", + false, + Pcre2Limits { + match_limit: Some(1_000), + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(split.split("a&&b").unwrap(), vec!["a", "&&", "b"]); + } + // ── Real-world tokenizer patterns ────────────────────── /// The Llama-3 / GPT-4 tokenizer pre-tokenization pattern.