vllm_mlx.constrained.cache

Cache of TokenEnforcerTokenizerData objects keyed by tokenizer identity.

View the complete module source at #L1-L186.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.constrained.cache

Cache of TokenEnforcerTokenizerData objects keyed by tokenizer identity.

Building TokenEnforcerTokenizerData requires iterating over the entire vocabulary (up to 200k tokens on MiniMax/GLM) and decoding each token. The cost is ~1-2 seconds per model and the result is independent of the JSON schema, so we cache it for the lifetime of the process.

vllm_mlx.constrained.cache.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.constrained.cache._CACHE module-attribute

_CACHE: dict[int, Any] = {}

vllm_mlx.constrained.cache._CACHE_LOCK module-attribute

_CACHE_LOCK = threading.Lock()

vllm_mlx.constrained.cache._resolve_inner_tokenizer

_resolve_inner_tokenizer(tokenizer: Any) -> Any

VLM processors wrap the actual tokenizer under processor.tokenizer. mlx_lm.tokenizer_utils.TokenizerWrapper exposes it via _tokenizer. Return the most-unwrapped tokenizer that still has the HF all_special_ids / eos_token_id surface.

Note: on HF PreTrainedTokenizerFast, _tokenizer points at the rust-level object which lacks all_special_ids; unwrapping to that level would cause every special token (<eos>, <pad>, \n, <|think|> …) to leak into regular_tokens and end up in TokenizerPrefixTree.root as an always-allowed token. We only unwrap when the inner layer still exposes all_special_ids.

Source code in vllm_mlx/constrained/cache.py

def _resolve_inner_tokenizer(tokenizer: Any) -> Any:
    """
    VLM processors wrap the actual tokenizer under ``processor.tokenizer``.
    ``mlx_lm.tokenizer_utils.TokenizerWrapper`` exposes it via ``_tokenizer``.
    Return the most-unwrapped tokenizer that still has the HF
    ``all_special_ids`` / ``eos_token_id`` surface.

    Note: on HF ``PreTrainedTokenizerFast``, ``_tokenizer`` points at the
    rust-level object which lacks ``all_special_ids``; unwrapping to that
    level would cause every special token (``<eos>``, ``<pad>``, ``\\n``,
    ``<|think|>`` …) to leak into ``regular_tokens`` and end up in
    ``TokenizerPrefixTree.root`` as an always-allowed token.  We only
    unwrap when the inner layer still exposes ``all_special_ids``.
    """
    # VLM processor wrapper exposes the HF tokenizer under ``tokenizer``.
    inner = getattr(tokenizer, "tokenizer", None)
    if (
        inner is not None
        and inner is not tokenizer
        and hasattr(inner, "all_special_ids")
    ):
        tokenizer = inner
    # mlx_lm TokenizerWrapper keeps the raw HF tokenizer under ``_tokenizer``.
    # Only unwrap if the inner object still exposes the HF tokenizer surface.
    inner = getattr(tokenizer, "_tokenizer", None)
    if inner is not None and hasattr(inner, "all_special_ids"):
        tokenizer = inner
    return tokenizer

vllm_mlx.constrained.cache._build_regular_tokens_list

_build_regular_tokens_list(tokenizer: Any, vocab_size: int) -> list[tuple[int, str, bool]]

Enumerate the regular (non-special) tokens in the vocabulary and produce the (token_id, decoded_with_leading_space_marker, is_word_start) tuples required by TokenEnforcerTokenizerData.

Mirrors the reference implementation in lmformatenforcer.integrations. transformers but works with the HF tokenizer surface only (so we do not need a hard transformers dependency at the right version).

Source code in vllm_mlx/constrained/cache.py

def _build_regular_tokens_list(
    tokenizer: Any, vocab_size: int
) -> list[tuple[int, str, bool]]:
    """
    Enumerate the regular (non-special) tokens in the vocabulary and produce
    the ``(token_id, decoded_with_leading_space_marker, is_word_start)`` tuples
    required by ``TokenEnforcerTokenizerData``.

    Mirrors the reference implementation in ``lmformatenforcer.integrations.
    transformers`` but works with the HF tokenizer surface only (so we do not
    need a hard transformers dependency at the right version).
    """
    try:
        special_ids = set(tokenizer.all_special_ids)
    except AttributeError:
        special_ids = set()

    try:
        token_0 = tokenizer.encode("0")[-1]
    except Exception:
        token_0 = None

    regular_tokens: list[tuple[int, str, bool]] = []
    for token_idx in range(vocab_size):
        if token_idx in special_ids:
            continue
        try:
            decoded_regular = tokenizer.decode([token_idx])
        except Exception:
            continue
        if token_0 is not None:
            try:
                decoded_after_0 = tokenizer.decode([token_0, token_idx])[1:]
            except Exception:
                decoded_after_0 = decoded_regular
        else:
            decoded_after_0 = decoded_regular
        is_word_start_token = len(decoded_after_0) > len(decoded_regular)
        regular_tokens.append((token_idx, decoded_after_0, is_word_start_token))
    return regular_tokens

vllm_mlx.constrained.cache._get_eos_token_id

_get_eos_token_id(tokenizer: Any) -> int | list[int]

Source code in vllm_mlx/constrained/cache.py

def _get_eos_token_id(tokenizer: Any) -> int | list[int]:
    # Some tokenizers expose multiple EOS candidates (e.g. Gemma 4 has
    # [1 <eos>, 106 <end_of_turn>, 50 <|think|>] in generation_config.json).
    # Prefer the list form so all stop tokens are treated as EOS by the
    # enforcer; otherwise the model may emit an out-of-schema stop token
    # that the enforcer did not mask (because it's a special token not in
    # ``regular_tokens``), yet the inference runtime still treats as stop.
    eos_list = getattr(tokenizer, "eos_token_ids", None)
    if isinstance(eos_list, (list, tuple)) and eos_list:
        return list(eos_list)
    eos = getattr(tokenizer, "eos_token_id", None)
    if eos is not None:
        return eos
    return 0

vllm_mlx.constrained.cache._get_vocab_size

_get_vocab_size(tokenizer: Any) -> int

Source code in vllm_mlx/constrained/cache.py

def _get_vocab_size(tokenizer: Any) -> int:
    vs = getattr(tokenizer, "vocab_size", None)
    if isinstance(vs, int) and vs > 0:
        return vs
    try:
        return len(tokenizer)
    except TypeError:
        pass
    get_vocab = getattr(tokenizer, "get_vocab", None)
    if callable(get_vocab):
        return len(get_vocab())
    raise ValueError("Cannot determine tokenizer vocab size")

vllm_mlx.constrained.cache._decode_function

_decode_function(tokenizer: Any, tokens: list[int]) -> str

Source code in vllm_mlx/constrained/cache.py

def _decode_function(tokenizer: Any, tokens: list[int]) -> str:
    try:
        decoded = tokenizer.decode(tokens)
    except Exception:
        return ""
    return decoded.rstrip("\ufffd") if isinstance(decoded, str) else ""

vllm_mlx.constrained.cache.get_tokenizer_data

get_tokenizer_data(tokenizer: Any) -> Any | None

Return a cached TokenEnforcerTokenizerData for tokenizer.

Returns None if lm-format-enforcer is not installed or the tokenizer cannot be adapted.

Source code in vllm_mlx/constrained/cache.py

def get_tokenizer_data(tokenizer: Any) -> Any | None:
    """
    Return a cached ``TokenEnforcerTokenizerData`` for ``tokenizer``.

    Returns ``None`` if ``lm-format-enforcer`` is not installed or the
    tokenizer cannot be adapted.
    """
    try:
        from lmformatenforcer.tokenenforcer import TokenEnforcerTokenizerData
    except ImportError:
        return None

    inner = _resolve_inner_tokenizer(tokenizer)
    key = id(inner)
    with _CACHE_LOCK:
        cached = _CACHE.get(key)
        if cached is not None:
            return cached

    try:
        vocab_size = _get_vocab_size(inner)
    except Exception as exc:
        logger.warning(
            "Could not determine vocab size for constrained decoding: %s", exc
        )
        return None

    try:
        regular_tokens = _build_regular_tokens_list(inner, vocab_size)
        decode_fn = functools.partial(_decode_function, inner)
        eos_token_id = _get_eos_token_id(inner)
        data = TokenEnforcerTokenizerData(
            regular_tokens,
            decode_fn,
            eos_token_id,
            use_bitmask=False,
            vocab_size=vocab_size,
        )
    except Exception as exc:
        logger.warning("Failed to build TokenEnforcerTokenizerData: %s", exc)
        return None

    with _CACHE_LOCK:
        _CACHE[key] = data
    return data

vllm_mlx.constrained.cache.clear_cache

clear_cache() -> None

Drop the cache (mainly for tests).

Source code in vllm_mlx/constrained/cache.py

def clear_cache() -> None:
    """Drop the cache (mainly for tests)."""
    with _CACHE_LOCK:
        _CACHE.clear()

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.constrained.cache._resolve_inner_tokenizer · function

vllm_mlx.constrained.cache._resolve_inner_tokenizer(tokenizer: Any) -> Any

VLM processors wrap the actual tokenizer under processor.tokenizer.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Type: Any
• Direct return expressions: tokenizer

Exceptions and behavior

Function _resolve_inner_tokenizer calls getattr, hasattr; returns tokenizer. No direct raise statement appears in this definition.

View source #L26-L53.

vllm_mlx.constrained.cache._build_regular_tokens_list · function

vllm_mlx.constrained.cache._build_regular_tokens_list(tokenizer: Any, vocab_size: int) -> list[tuple[int, str, bool]]

Enumerate the regular (non-special) tokens in the vocabulary and produce the (token_id, decoded_with_leading_space_marker, is_word_start) tuples required by TokenEnforcerTokenizerData.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.
vocab_size	int	yes	none	Required positional or keyword input.

Returns

• Type: list[tuple[int, str, bool]]
• Direct return expressions: regular_tokens

Exceptions and behavior

Function _build_regular_tokens_list calls set, tokenizer.encode, range, tokenizer.decode; returns regular_tokens. No direct raise statement appears in this definition.

View source #L56-L95.

vllm_mlx.constrained.cache._get_eos_token_id · function

vllm_mlx.constrained.cache._get_eos_token_id(tokenizer: Any) -> int | list[int]

Function _get_eos_token_id calls getattr, isinstance, list; has 3 explicit return paths.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Type: int | list[int]
• Direct return expressions: list(eos_list); eos; 0

Exceptions and behavior

Function _get_eos_token_id calls getattr, isinstance, list; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L98-L111.

vllm_mlx.constrained.cache._get_vocab_size · function

vllm_mlx.constrained.cache._get_vocab_size(tokenizer: Any) -> int

Function _get_vocab_size calls getattr, isinstance, len, callable; can raise ValueError; has 3 explicit return paths.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Type: int
• Direct return expressions: vs; len(tokenizer); len(get_vocab())

Exceptions and behavior

Function _get_vocab_size calls getattr, isinstance, len, callable; can raise ValueError; has 3 explicit return paths. Directly raised exceptions: ValueError.

View source #L114-L125.

vllm_mlx.constrained.cache._decode_function · function

vllm_mlx.constrained.cache._decode_function(tokenizer: Any, tokens: list[int]) -> str

Function _decode_function calls tokenizer.decode, isinstance, decoded.rstrip; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.
tokens	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: ''; decoded.rstrip('�') if isinstance(decoded, str) else ''

Exceptions and behavior

Function _decode_function calls tokenizer.decode, isinstance, decoded.rstrip; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L128-L133.

vllm_mlx.constrained.cache.get_tokenizer_data · function

vllm_mlx.constrained.cache.get_tokenizer_data(tokenizer: Any) -> Any | None

Return a cached TokenEnforcerTokenizerData for tokenizer.

Parameters

Name	Type	Required	Default	Description
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Type: Any | None
• Direct return expressions: None; cached; data

Exceptions and behavior

Function get_tokenizer_data calls _resolve_inner_tokenizer, id, _CACHE.get, _get_vocab_size; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L136-L180.

vllm_mlx.constrained.cache.clear_cache · function

vllm_mlx.constrained.cache.clear_cache() -> None

Drop the cache (mainly for tests).

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Function clear_cache calls _CACHE.clear. No direct raise statement appears in this definition.

View source #L183-L186.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
_resolve_inner_tokenizer	function	_resolve_inner_tokenizer(tokenizer: Any) -> Any	VLM processors wrap the actual tokenizer under processor.tokenizer.	#L26-L53
_build_regular_tokens_list	function	_build_regular_tokens_list(tokenizer: Any, vocab_size: int) -> list[tuple[int, str, bool]]	Enumerate the regular (non-special) tokens in the vocabulary and produce the (token_id, decoded_with_leading_space_marker, is_word_start) tuples required by TokenEnforcerTokenizerData.	#L56-L95
_get_eos_token_id	function	_get_eos_token_id(tokenizer: Any) -> int \| list[int]	Function _get_eos_token_id calls getattr, isinstance, list; has 3 explicit return paths.	#L98-L111
_get_vocab_size	function	_get_vocab_size(tokenizer: Any) -> int	Function _get_vocab_size calls getattr, isinstance, len, callable; can raise ValueError; has 3 explicit return paths.	#L114-L125
_decode_function	function	_decode_function(tokenizer: Any, tokens: list[int]) -> str	Function _decode_function calls tokenizer.decode, isinstance, decoded.rstrip; has 2 explicit return paths.	#L128-L133
get_tokenizer_data	function	get_tokenizer_data(tokenizer: Any) -> Any \| None	Return a cached TokenEnforcerTokenizerData for tokenizer.	#L136-L180
clear_cache	function	clear_cache() -> None	Drop the cache (mainly for tests).	#L183-L186