DeepInfer/serve/prompt_tokenizer.py

"""Prompt tokenization for the serving frontends: exact token ids, much faster than `tokenizer.encode` on long prompts.

Agentic traffic re-sends the whole conversation every turn, so the frontend re-tokenizes 100K+ token prompts on
nearly every request. HF `tokenizer.encode` takes ~1.3 ms per 1K tokens (GB300 Grace core), which becomes the largest
part of TTFT once the prefix cache hits. Backends, all required to produce the ids of
`tokenizer.encode(text, add_special_tokens=False)`:

- `fastokens` (optional `pip install fastokens`): a faster Rust implementation of the HF pipeline, built from the
  tokenizer exactly as transformers configured it (`backend_tokenizer.to_str()`, not the file on disk).
- `segment`: no extra dependency. HF tokenizers first split the input at non-normalized added tokens and encode the
  text between them independently. Splitting there ourselves and caching each piece's ids by its text is exact, and
  a new turn only pays for the pieces it adds (the new tool result, the new user message). It is enabled only when
  the tokenizer's config shows the split is exact (see `_segment_split_tokens`).
- `hf`: `tokenizer.encode`.

`auto` takes the first of fastokens, segment, hf that is available and matches `tokenizer.encode` on a set of probes
built from the tokenizer's own added tokens. At runtime a sample of prompts is re-encoded with `tokenizer.encode`
on a background thread; the first mismatch switches the frontend to `hf`.
"""

from __future__ import annotations

import array
import collections
import contextlib
import json
import queue
import random
import re
import threading
import time
from typing import Any

BACKENDS = ("auto", "fastokens", "segment", "hf")


class PromptTokenizer:
    def __init__(
        self,
        tokenizer: Any,
        backend: str = "auto",
        verify_rate: float = 0.005,
        cache_bytes: int = 512 << 20,
        extra_probes: list[str] | None = None,
    ) -> None:
        assert backend in BACKENDS, backend
        self.tokenizer = tokenizer
        self.verify_rate = verify_rate
        self.stats = dict(calls=0, errors=0, verified=0, verify_mismatches=0)
        self.notes: list[str] = []
        self._times: collections.deque[float] = collections.deque(maxlen=4096)
        self._probes = _probes(tokenizer) + list(extra_probes or [])

        candidates = ("fastokens", "segment") if backend == "auto" else (() if backend == "hf" else (backend,))
        self.name, self._encode = "hf", self._hf_encode
        for name in candidates:
            try:
                encode = _FastokensBackend(tokenizer).encode if name == "fastokens" else _SegmentBackend(tokenizer, cache_bytes).encode
                if (bad := self._first_mismatch(encode)) is not None:
                    raise ValueError(f"ids differ from tokenizer.encode on probe {bad!r}")
            except Exception as e:  # noqa: BLE001 - a backend that cannot be built or fails parity is skipped
                self.notes.append(f"{name}: {type(e).__name__}: {e}"[:300])
                if backend != "auto":
                    raise RuntimeError(f"tokenizer backend {name} unavailable: {self.notes[-1]}") from e
                continue
            self.name, self._encode = name, encode
            break

        self._verify_q: queue.Queue[tuple[str, list[int]]] = queue.Queue(maxsize=8)
        if self.name != "hf" and verify_rate > 0:
            threading.Thread(target=self._verify_forever, name="prompt-tokenizer-verify", daemon=True).start()

    def encode(self, text: str) -> list[int]:
        t0 = time.perf_counter()
        name = self.name
        try:
            ids = self._encode(text)
        except Exception:  # noqa: BLE001 - e.g. a regex engine limit on pathological input
            self.stats["errors"] += 1
            name, ids = "hf", self._hf_encode(text)
        self._times.append(time.perf_counter() - t0)
        self.stats["calls"] += 1
        if name != "hf" and self.verify_rate > 0 and random.random() < self.verify_rate:
            with contextlib.suppress(queue.Full):  # the checker is behind: skip this sample
                self._verify_q.put_nowait((text, ids))
        return ids

    def metrics(self) -> dict[str, float]:
        """Counters plus encode latency percentiles (ms) over the last 4096 calls."""
        out: dict[str, float] = dict(self.stats)
        times = sorted(self._times)
        for q in (50, 90, 99):
            out[f"encode_ms_p{q}"] = round(times[min(len(times) - 1, len(times) * q // 100)] * 1e3, 3) if times else 0.0
        return out

    def describe(self) -> str:
        skipped = "; ".join(self.notes)
        return f"prompt tokenizer backend: {self.name}" + (f" (skipped {skipped})" if skipped else "")

    def _hf_encode(self, text: str) -> list[int]:
        return self.tokenizer.encode(text, add_special_tokens=False)

    def _first_mismatch(self, encode: Any) -> str | None:
        for p in self._probes:
            if list(encode(p)) != self._hf_encode(p):
                return p[:80]
        return None

    def _verify_forever(self) -> None:
        while True:
            text, ids = self._verify_q.get()
            ref = self._hf_encode(text)
            self.stats["verified"] += 1
            if ids != ref:
                self.stats["verify_mismatches"] += 1
                at = next((i for i, (a, b) in enumerate(zip(ids, ref, strict=False)) if a != b), min(len(ids), len(ref)))
                print(
                    f"[server] prompt tokenizer {self.name} differs from tokenizer.encode at token {at} of {len(ref)} "
                    f"(text {text[:60]!r}...); switching to hf",
                    flush=True,
                )
                self.name, self._encode = "hf", self._hf_encode


class _FastokensBackend:
    def __init__(self, tokenizer: Any) -> None:
        from fastokens._native import Tokenizer  # optional dependency

        if not getattr(tokenizer, "is_fast", False):
            raise ValueError("needs a fast (tokenizer.json) tokenizer")
        self._tok = Tokenizer.from_json_str(tokenizer.backend_tokenizer.to_str())

    def encode(self, text: str) -> list[int]:
        return self._tok.encode(text, add_special_tokens=False).ids


class _SegmentBackend:
    """Split at added tokens, encode only the pieces not seen before (one batched call), LRU-cache pieces by text."""

    def __init__(self, tokenizer: Any, cache_bytes: int) -> None:
        if not getattr(tokenizer, "is_fast", False):
            raise ValueError("needs a fast (tokenizer.json) tokenizer")
        self._backend = tokenizer.backend_tokenizer
        split = _segment_split_tokens(self._backend)
        # Longest first: Python's alternation is leftmost-first, HF's matcher is leftmost-longest. lstrip/rstrip
        # tokens swallow the whitespace next to them, as in HF (Unicode White_Space, Rust's char::is_whitespace).
        alts = [
            f"{_WS}*" * t.lstrip + re.escape(t.content) + f"{_WS}*" * t.rstrip
            for t in sorted(split.values(), key=lambda t: -len(t.content))
        ]
        self._split = re.compile("(" + "|".join(alts) + ")")
        self._ids = {c: t.id for c, t in split.items()}
        self._stripped = any(t.lstrip or t.rstrip for t in split.values())
        # ids as uint32 arrays: 4 B/token instead of ~40 B for a list of Python ints
        self._cache: collections.OrderedDict[str, array.array] = collections.OrderedDict()
        self._cache_bytes = 0
        self._max_bytes = cache_bytes
        self._lock = threading.Lock()

    def encode(self, text: str) -> list[int]:
        parts = self._split.split(text)  # odd indices are the split tokens
        pieces: list[Any] = [None] * len(parts)
        missing: list[int] = []
        with self._lock:
            for i, p in enumerate(parts):
                if i % 2:
                    pieces[i] = [self._ids[_strip_ws(p) if self._stripped else p]]
                elif p:
                    ids = self._cache.get(p)
                    if ids is None:
                        missing.append(i)
                    else:
                        self._cache.move_to_end(p)
                        pieces[i] = ids
        if missing:
            encs = self._backend.encode_batch([parts[i] for i in missing], add_special_tokens=False)
            with self._lock:
                for i, e in zip(missing, encs, strict=True):
                    pieces[i] = e.ids
                    if parts[i] not in self._cache:
                        self._cache[parts[i]] = array.array("I", e.ids)
                        self._cache_bytes += _cost(parts[i], e.ids)
                while self._cache_bytes > self._max_bytes and self._cache:
                    k, v = self._cache.popitem(last=False)
                    self._cache_bytes -= _cost(k, v)
        out: list[int] = []
        for ids in pieces:
            if ids:
                out.extend(ids)
        return out


_WS = "[\\t\\n\\x0b\\x0c\\r \\x85\\xa0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000]"
_WS_CHARS = "\t\n\x0b\x0c\r \x85\xa0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000"


def _cost(text: str, ids: Any) -> int:
    return len(text) + 4 * len(ids)


def _strip_ws(s: str) -> str:
    return s.strip(_WS_CHARS)


_SplitToken = collections.namedtuple("_SplitToken", "content id lstrip rstrip")


def _segment_split_tokens(backend: Any) -> dict[str, _SplitToken]:
    """Added tokens at which splitting the input and encoding the pieces separately gives tokenizer.encode's ids.

    HF's AddedVocabulary first matches the non-normalized added tokens on the raw input (leftmost-longest, an
    lstrip/rstrip token also taking the whitespace next to it), and the rest of the pipeline (normalizer,
    pre-tokenizer, model) then runs on each piece in between on its own. Raises when the config can make a piece
    depend on its neighbours.
    """
    cfg = json.loads(backend.to_str())
    added = cfg.get("added_tokens") or []
    if any(t["single_word"] for t in added):
        raise ValueError("single_word added tokens depend on the characters around them")
    if _has_first_only_prefix(cfg.get("pre_tokenizer")) or _has_first_only_prefix(cfg.get("normalizer")):
        raise ValueError("prepend_scheme='first' prefixes only the first piece of the whole input")
    split = {t["content"]: _SplitToken(t["content"], t["id"], t["lstrip"], t["rstrip"]) for t in added if not t["normalized"]}
    if not split:
        raise ValueError("no non-normalized added tokens to split at")
    if any(_strip_ws(c) != c for c in split):
        raise ValueError("an added token starts or ends with whitespace")
    return split


def _has_first_only_prefix(node: Any) -> bool:
    if isinstance(node, dict):
        if node.get("prepend_scheme") == "first":
            return True
        return any(_has_first_only_prefix(v) for v in node.values())
    if isinstance(node, list):
        return any(_has_first_only_prefix(v) for v in node)
    return False


def _probes(tokenizer: Any) -> list[str]:
    """Inputs that exercise added-token boundaries, scripts, digits and whitespace for this tokenizer."""
    added = [t.content for t in tokenizer.added_tokens_decoder.values()]
    rng = random.Random(0)
    sample = added if len(added) <= 64 else added[:16] + rng.sample(added, 48)
    texts = [
        "",
        "hello world",
        " leading space",
        "trailing space ",
        "tabs\tand\r\nwindows\nlines\n\n\n",
        "    indented code:\n        return x  # comment\n",
        "数字 12345678 和中文、日本語のテキスト、カタカナ",
        "emoji 😀👍🏽 and combining é é ñ",
        '{"name": "read_file", "arguments": {"path": "/tmp/a b.py", "n": 3.14}}',
        "a" * 3000,
        " " * 500 + "x",
        "-" * 2000,
    ]
    probes = list(texts)
    for i, tok in enumerate(sample):
        a, b = texts[(i + 1) % len(texts)][:120], texts[(i + 5) % len(texts)][:120]
        probes += [tok, f"{a}{tok}{b}", f"{a} {tok} {b}", f"{tok}{tok}", f"\n{tok}\n"]
    probes.append("".join(f"{t}{texts[i % len(texts)][:40]}" for i, t in enumerate(sample)))
    return probes