14.3 A Byte-Level BPE Tokenizer From Scratch (and Why Vocab Size Is a Design Lever at 100M)¶

raw text corpus (14.2)
        │
        ▼
  ┌──────────────────┐   train byte-level BPE, vocab_size = 32768
  │  THIS CHAPTER    │   reserve 9 special tokens
  │  (14.3)          │   export to the ecosystem
  └────────┬─────────┘
           │
           ├─► tokenizer/stack100m-32768.json        (from-scratch artifact, frozen)
           └─► tokenizer/stack100m-32768-hf/         (tokenizer.json + tokenizer_config.json)
                     │                                loaded by transformers / TRL / vLLM
           ┌─────────┴──────────┐
           ▼                    ▼
  ┌─────────────────┐   ┌─────────────────────┐
  │ data packing    │   │ architecture (14.4) │
  │ (14.2, uint16   │   │ embedding table     │
  │ .bin shards)    │   │ sized to 32768×512  │
  └─────────────────┘   └─────────────────────┘

id:        0 ─────────────── 255   256 ──────────────────── 32758   32759 ── 32767
contents:  raw UTF-8 bytes         32,503 learned BPE merges          9 reserved specials
           (fixed, universal)      (trained on the Stack-100M mix)    (fixed order, see table)

SPLIT_PATTERN_UNICODE = (
    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)"      # contractions, case-insensitive (GPT-2 was not)
    r"|[^\r\n\p{L}\p{N}]?\p{L}+"          # a word, optionally preceded by one non-letter
    r"|\p{N}{1,3}"                        # <= 3 digits, NEVER absorbing a leading space
    r"| ?[^\s\p{L}\p{N}]+[\r\n]*"         # punctuation runs, trailing newlines attached
    r"|\s*[\r\n]+"                        # newline runs (indentation-friendly)
    r"|\s+(?!\S)"                         # trailing whitespace at end of a run
    r"|\s+"                               # any remaining whitespace
)

# capstone/stacklm/tokenizer/bpe.py
"""
Stack-100M's tokenizer: a from-scratch, byte-level BPE trainer + encoder/decoder.

This is the production version of the algorithm built from first principles in
../02-transformer/01-tokenization.html -- same "merge the most frequent adjacent
pair" idea, same byte-level guarantee that no input is ever unrepresentable, but
engineered to finish ~32.5k merges on a real multi-megabyte sample in seconds
instead of hours (measured numbers in "Training at scale" below).

Special tokens are reserved UP FRONT (see SPECIAL_TOKENS) even though most are
untouched until Ch. 14.9 (SFT/DPO) and Ch. 14.10 (the agent). Once this
tokenizer is trained and Stack-100M's embedding table is initialized against it,
the vocabulary is FROZEN: id 32763 means "<|user|>" forever, or every
checkpoint's embedding table becomes meaningless.
"""

from __future__ import annotations

import heapq
import json
import warnings
from collections import Counter, defaultdict
from typing import Dict, Iterable, List, Tuple, Union

try:                        # `regex` gives us \p{L} / \p{N}; see the aside above
    import regex as _re
    _HAVE_REGEX = True
except ImportError:         # hermetic CI: stdlib approximation, same digit cap
    import re as _re
    _HAVE_REGEX = False

# ---------------------------------------------------------------------------
# 1. Special tokens, reserved up front, in the exact order the rest of the
#    capstone relies on. Never reorder these once a tokenizer has been trained
#    and used to initialize an embedding table. Later chapters hardcode these
#    STRINGS (looked up via special_token_id), not raw integers -- but the ids
#    themselves become baked into every checkpoint's embedding rows.
# ---------------------------------------------------------------------------
SPECIAL_TOKENS: Tuple[str, ...] = (
    "<|bos|>",           # beginning of sequence / beginning of a packed document
    "<|eos|>",           # end of sequence / end of document
    "<|pad|>",           # padding, always masked out of the loss
    "<|system|>",        # chat role marker (Ch. 14.9 SFT/DPO)
    "<|user|>",          # chat role marker
    "<|assistant|>",     # chat role marker -- loss starts right after this token
    "<|end|>",           # end-of-turn marker
    "<|tool_call|>",     # opens a tool invocation the model emits (Ch. 14.10)
    "<|tool_result|>",   # opens a tool's returned observation, masked from loss
)

VOCAB_SIZE = 32768                                          # PLAN Sec. 1 / Sec. 3
NUM_BYTES = 256                                             # every raw byte is id 0..255
NUM_MERGES = VOCAB_SIZE - NUM_BYTES - len(SPECIAL_TOKENS)   # = 32,503

SPLIT_PATTERN_UNICODE = (
    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)"
    r"|[^\r\n\p{L}\p{N}]?\p{L}+"
    r"|\p{N}{1,3}"
    r"| ?[^\s\p{L}\p{N}]+[\r\n]*"
    r"|\s*[\r\n]+"
    r"|\s+(?!\S)"
    r"|\s+"
)
SPLIT_PATTERN_STDLIB = (
    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)"
    r"|(?:[^\r\n\w]|_)?[^\W\d_]+"
    r"|\d{1,3}"
    r"| ?(?:[^\s\w]|_)+[\r\n]*"
    r"|\s*[\r\n]+"
    r"|\s+(?!\S)"
    r"|\s+"
)
SPLIT_PATTERN = SPLIT_PATTERN_UNICODE if _HAVE_REGEX else SPLIT_PATTERN_STDLIB
_SPLIT_RE = _re.compile(SPLIT_PATTERN)


# ---------------------------------------------------------------------------
# 2. The trainer. Operates on integer symbol ids throughout (0..255 for raw
#    bytes, 256+ for learned merges) instead of the string glyphs used in the
#    from-scratch chapter -- purely a speed choice, the algorithm is identical.
# ---------------------------------------------------------------------------
def train_bpe(word_freqs: Dict[Tuple[int, ...], int],
              num_merges: int) -> List[Tuple[int, int]]:
    """Learn `num_merges` BPE merges from a corpus reduced to
    (symbol-id-tuple -> frequency) counts.

    Naively, each merge rescans the ENTIRE corpus to recount every adjacent
    pair -- O(num_merges x corpus_size). Instead we maintain three pieces of
    running state so each merge touches only the (usually small) subset of
    words that actually contain the pair being merged:

      - `pair_counts`   : running frequency-weighted count of every adjacent pair
      - `pair_to_words` : inverted index, pair -> set of word indices containing it
      - `heap`          : a max-heap (via negated counts) of candidate pairs

    LAZY DELETION: we never remove stale entries, we discard them on pop when
    they no longer match the live count. This avoids an expensive decrease-key
    -- but it is only CORRECT if every count change re-pushes a fresh entry.
    Pushing only on increments loses decremented pairs forever (see the pitfall
    box above), so we collect `touched` and re-push all of it.
    """
    words: List[List[int]] = [list(w) for w in word_freqs.keys()]
    freqs: List[int] = list(word_freqs.values())

    pair_counts: Dict[Tuple[int, int], int] = defaultdict(int)
    pair_to_words: Dict[Tuple[int, int], set] = defaultdict(set)
    for wi, (word, f) in enumerate(zip(words, freqs)):
        for a, b in zip(word, word[1:]):
            pair_counts[(a, b)] += f
            pair_to_words[(a, b)].add(wi)

    heap = [(-c, pair) for pair, c in pair_counts.items()]
    heapq.heapify(heap)

    merges: List[Tuple[int, int]] = []
    next_id = NUM_BYTES  # the first learned merge becomes id 256

    while len(merges) < num_merges and heap:
        neg_count, pair = heapq.heappop(heap)
        live_count = pair_counts.get(pair, 0)
        if live_count <= 0 or -neg_count != live_count:
            continue                       # stale heap entry -- skip it
        if live_count < 2:
            break                          # nothing left that repeats; stop early

        merges.append(pair)
        new_id = next_id
        next_id += 1

        touched: set = set()               # every pair whose count we change
        for wi in list(pair_to_words.get(pair, ())):
            word = words[wi]
            # `wi` can be a STALE member of pair_to_words[pair]: an earlier merge
            # may already have removed `pair` from this word. Verify membership
            # before mutating anything.
            if not any(word[i] == pair[0] and word[i + 1] == pair[1]
                       for i in range(len(word) - 1)):
                continue

            f = freqs[wi]
            # Remove this word's contribution to EVERY pair it currently forms
            # (not just `pair`) -- merging shifts adjacency across the word.
            for a, b in zip(word, word[1:]):
                pair_counts[(a, b)] -= f
                touched.add((a, b))

            # Apply the merge greedily, left to right, non-overlapping.
            merged, i = [], 0
            while i < len(word):
                if (i < len(word) - 1
                        and word[i] == pair[0] and word[i + 1] == pair[1]):
                    merged.append(new_id)
                    i += 2
                else:
                    merged.append(word[i])
                    i += 1
            words[wi] = merged

            # Re-add this word's contribution under its NEW symbol sequence.
            for a, b in zip(merged, merged[1:]):
                pair_counts[(a, b)] += f
                pair_to_words[(a, b)].add(wi)
                touched.add((a, b))

        pair_to_words.pop(pair, None)
        pair_counts.pop(pair, None)
        touched.discard(pair)              # this pair is consumed, never re-push

        # THE FIX: refresh the heap for every pair whose count moved -- up OR
        # down. Without the decrement half, pairs whose count fell are dropped
        # from the heap forever and the trainer silently under-produces merges.
        for p in touched:
            c = pair_counts.get(p, 0)
            if c > 0:
                heapq.heappush(heap, (-c, p))

    return merges

# capstone/stacklm/tokenizer/bpe.py (continued)

class StackTokenizer:
    """Byte-level BPE tokenizer for Stack-100M.

    ID layout (fixed for the whole project):
        0        .. 255           raw byte values (id == byte value)
        256      .. 256+M-1       learned merges, in the order they were trained
        256+M    .. vocab_size-1  filler <|unused_N|> (only if the sample was too
                                   small), then the 9 SPECIAL_TOKENS in order,
                                   ALWAYS occupying the final 9 ids.
    """

    def __init__(self, pattern: str = SPLIT_PATTERN) -> None:
        self.merges: List[Tuple[int, int]] = []             # learned merges, ordered
        self.merge_id: Dict[Tuple[int, int], int] = {}      # pair -> resulting id
        self.id_to_pair: Dict[int, Tuple[int, int]] = {}    # inverse, for decode
        self.special_to_id: Dict[str, int] = {}
        self.id_to_special: Dict[int, str] = {}
        self.pattern: str = pattern                          # frozen with the merges
        self._split_re = _re.compile(pattern)                # BOUND, not the global
        self._cache: Dict[Tuple[int, ...], List[int]] = {}   # pre-token -> ids

    @property
    def vocab_size(self) -> int:
        return NUM_BYTES + len(self.merges) + len(self.special_to_id)

    # -- the API surface every later chapter depends on -----------------------
    def special_token_id(self, token_str: str) -> int:
        """Ch. 14.9's chat template and Ch. 14.10's tool formatter call this."""
        return self.special_to_id[token_str]

    id = special_token_id                    # short alias used in 14.10/14.11

    @property
    def bos_id(self) -> int: return self.special_to_id["<|bos|>"]
    @property
    def eos_id(self) -> int: return self.special_to_id["<|eos|>"]
    @property
    def pad_id(self) -> int: return self.special_to_id["<|pad|>"]

    def _pretokenize(self, text: str) -> List[Tuple[int, ...]]:
        """Split `text` into pre-token chunks, each a tuple of raw UTF-8 byte
        VALUES (0..255) -- e.g. 'café' -> (99, 97, 102, 195, 169). Uses the
        INSTANCE's compiled pattern, so a loaded artifact tokenizes under the
        regex it was trained with, not whatever the module default is today."""
        return [tuple(chunk.encode("utf-8")) for chunk in self._split_re.findall(text)]

    # -- training -------------------------------------------------------------
    def _assign_ids(self, all_specials: Tuple[str, ...]) -> None:
        self.merge_id = {pair: NUM_BYTES + i for i, pair in enumerate(self.merges)}
        self.id_to_pair = {v: k for k, v in self.merge_id.items()}
        self.special_to_id, self.id_to_special = {}, {}
        nid = NUM_BYTES + len(self.merges)
        for s in all_specials:
            self.special_to_id[s] = nid
            self.id_to_special[nid] = s
            nid += 1
        self._cache = {}

    def train_from_iterable(self, docs: Iterable[str], vocab_size: int = VOCAB_SIZE,
                            special_tokens: Tuple[str, ...] = SPECIAL_TOKENS) -> int:
        """Train on a STREAM of documents. Returns the merge shortfall (0 is the
        healthy case). Memory is O(distinct pre-token chunks), not O(corpus)."""
        num_merges = vocab_size - NUM_BYTES - len(special_tokens)
        word_freqs: Counter = Counter()
        for doc in docs:
            word_freqs.update(self._pretokenize(doc))
        self.merges = train_bpe(word_freqs, num_merges)

        # Guard: a small/low-diversity sample runs out of repeated pairs. Pad
        # deterministically so vocab_size is EXACTLY what 14.4's embedding table
        # expects; put the fillers BEFORE the real specials so <|bos|> ..
        # <|tool_result|> keep the top 9 ids no matter what.
        shortfall = num_merges - len(self.merges)
        if shortfall > 0:
            warnings.warn(
                f"corpus exhausted after {len(self.merges)} merges "
                f"({shortfall} short of {num_merges}); padding with "
                f"{shortfall} <|unused_N|> tokens so vocab_size == {vocab_size}. "
                f"Use a larger sample for a real {vocab_size}-entry vocabulary.")
        fillers = tuple(f"<|unused_{i}|>" for i in range(shortfall))
        self._assign_ids(fillers + tuple(special_tokens))
        assert self.vocab_size == vocab_size
        return shortfall

    def train(self, text: str, vocab_size: int = VOCAB_SIZE,
              special_tokens: Tuple[str, ...] = SPECIAL_TOKENS) -> int:
        """Convenience wrapper for a single in-memory string (tests, demos)."""
        return self.train_from_iterable([text], vocab_size, special_tokens)

    # -- encode / decode ------------------------------------------------------
    def _apply_merges(self, symbols: List[int]) -> List[int]:
        """Repeatedly merge the pair with the LOWEST id (== earliest-learned ==
        highest priority) until no learned pair applies -- the standard
        rank-priority BPE encode loop."""
        while len(symbols) >= 2:
            best_pair, best_id = None, None
            for a, b in zip(symbols, symbols[1:]):
                mid = self.merge_id.get((a, b))
                if mid is not None and (best_id is None or mid < best_id):
                    best_pair, best_id = (a, b), mid
            if best_pair is None:
                break
            merged, i = [], 0
            while i < len(symbols):
                if (i < len(symbols) - 1
                        and symbols[i] == best_pair[0] and symbols[i + 1] == best_pair[1]):
                    merged.append(best_id)
                    i += 2
                else:
                    merged.append(symbols[i])
                    i += 1
            symbols = merged
        return symbols

    def _encode_chunk(self, byte_ids: Tuple[int, ...]) -> List[int]:
        """Memoized per-pre-token encode. Pre-token frequencies are Zipfian, so
        this hits far more often than it misses (measured below)."""
        out = self._cache.get(byte_ids)
        if out is None:
            out = self._cache[byte_ids] = self._apply_merges(list(byte_ids))
        return out

    def encode(self, text: str,
               allowed_special: Union[str, frozenset] = frozenset(),
               add_special_tokens: bool = False) -> List[int]:
        """Special-token strings inside `text` are treated as ORDINARY BYTES by
        default. Opt in from trusted call sites only, via any of:
            allowed_special={"<|bos|>"}     -- an explicit set
            allowed_special="all"           -- tiktoken's sentinel
            add_special_tokens=True         -- the alias 14.9/14.10/14.11 use
        """
        if add_special_tokens or allowed_special == "all":
            allowed = frozenset(self.special_to_id)
        else:
            allowed = frozenset(allowed_special)

        if allowed:
            # longest-first so a specials set containing overlapping strings
            # cannot be split by a shorter alternative
            pattern = "(" + "|".join(_re.escape(s) for s in
                                     sorted(allowed, key=len, reverse=True)) + ")"
            segments = _re.split(pattern, text)     # keeps the special literals
        else:
            segments = [text]

        ids: List[int] = []
        for seg in segments:
            if seg in allowed:
                ids.append(self.special_to_id[seg])
                continue
            for byte_ids in self._pretokenize(seg):
                ids.extend(self._encode_chunk(byte_ids))
        return ids

    def decode(self, ids: List[int]) -> str:
        """Expand every id back to raw bytes by walking the merge tree: a merged
        id's two children are themselves either smaller merge ids or raw byte
        ids, so we recurse (via an explicit stack) until only bytes 0..255
        remain, then decode UTF-8 with a safe fallback."""
        out = bytearray()
        for i in ids:
            if i in self.id_to_special:
                out += self.id_to_special[i].encode("utf-8")
                continue
            stack = [i]
            while stack:
                s = stack.pop()
                if s < NUM_BYTES:
                    out.append(s)
                else:
                    a, b = self.id_to_pair[s]
                    stack.extend([b, a])          # push reversed so `a` pops first
        # errors="replace": a lone invalid byte sequence never crashes decode,
        # it renders as the U+FFFD replacement character.
        return bytes(out).decode("utf-8", errors="replace")

    def token_bytes(self) -> List[bytes]:
        """The literal byte string of every non-special id, by walking the merge
        tree bottom-up. This is what the tiktoken / HF exporters need."""
        table = [bytes([i]) for i in range(NUM_BYTES)]
        for a, b in self.merges:
            table.append(table[a] + table[b])      # children always have lower ids
        return table

    # -- persistence ----------------------------------------------------------
    def save(self, path: str) -> None:
        payload = {
            "format": "stacklm-bpe-v1",
            "vocab_size": self.vocab_size,
            "pattern": self.pattern,          # the regex is PART of the artifact
            "merges": [list(p) for p in self.merges],
            "special_tokens": list(self.special_to_id.keys()),
        }
        with open(path, "w") as f:
            json.dump(payload, f)

    @classmethod
    def load(cls, path: str) -> "StackTokenizer":
        with open(path) as f:
            payload = json.load(f)
        pattern = payload.get("pattern", SPLIT_PATTERN)
        tok = cls(pattern=pattern)            # compile the SAVED regex, not ours
        if pattern != SPLIT_PATTERN:
            warnings.warn("artifact pre-tokenizer pattern differs from this "
                          "module's SPLIT_PATTERN; using the artifact's, because "
                          "the merges were trained against it.")
        tok.merges = [tuple(p) for p in payload["merges"]]
        tok._assign_ids(tuple(payload["special_tokens"]))
        assert tok.vocab_size == payload["vocab_size"]
        return tok


def load_tokenizer(path: str = "tokenizer/stack100m-32768.json") -> StackTokenizer:
    """The entry point Ch. 14.2's packer and Ch. 14.9's chat template call."""
    return StackTokenizer.load(path)

# capstone/stacklm/tokenizer/parallel.py
"""Parallel corpus encoding. Documents are independent, so this is a pure map:
one worker per core, each with its own tokenizer + its own chunk cache.

Split shards on <|eos|> boundaries, never mid-document -- a document cut in half
would be pre-tokenized differently on each side of the cut.
"""
from __future__ import annotations
import multiprocessing as mp
from typing import Iterable, Iterator, List
import numpy as np

from stacklm.tokenizer.bpe import StackTokenizer, load_tokenizer

_TOK: "StackTokenizer | None" = None


def _init_worker(tokenizer_path: str) -> None:
    global _TOK
    _TOK = load_tokenizer(tokenizer_path)      # each worker gets its own cache


def _encode_doc(text: str) -> List[int]:
    assert _TOK is not None
    # Untrusted corpus text: NEVER allow special strings through (see pitfall).
    # The <|bos|>/<|eos|> wrapper is added by code, in the packer.
    return _TOK.encode(text, allowed_special=frozenset())


def encode_corpus(docs: Iterable[str], tokenizer_path: str,
                  workers: int = 16, chunksize: int = 8) -> Iterator[np.ndarray]:
    """Yield one uint16 array per document, in input order."""
    with mp.Pool(workers, initializer=_init_worker,
                 initargs=(tokenizer_path,)) as pool:
        for ids in pool.imap(_encode_doc, docs, chunksize=chunksize):
            yield np.asarray(ids, dtype=np.uint16)   # vocab 32768 < 65536, fits

# capstone/stacklm/tokenizer/export.py
"""Export a trained StackTokenizer into the artifacts the ecosystem reads: a
`tiktoken.Encoding` (fast, OpenAI-style) and a HuggingFace `tokenizer.json`
(what transformers / TRL / vLLM / llama.cpp actually load).

Both exporters are meant to be EXACT. `capstone/tests/test_tokenizer_export.py`
asserts byte-identical id sequences against the from-scratch encoder.
"""
from __future__ import annotations
from typing import Dict

from stacklm.tokenizer.bpe import StackTokenizer, SPLIT_PATTERN


def bytes_to_unicode() -> Dict[int, str]:
    """GPT-2's byte <-> printable-unicode table (Radford et al., 2019). Maps all
    256 byte values to codepoints that survive a JSON round trip: the printable
    ASCII/Latin-1 ranges map to themselves, and the 68 remaining control/space
    bytes are shifted into U+0100.. -- so byte 32 (space) is 'Ġ' (U+0120) and
    byte 10 (newline) is 'Ċ' (U+010A)."""
    bs = (list(range(ord("!"), ord("~") + 1))
          + list(range(ord("¡"), ord("¬") + 1))
          + list(range(ord("®"), ord("ÿ") + 1)))
    cs, n = bs[:], 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return dict(zip(bs, (chr(c) for c in cs)))


_B2U = bytes_to_unicode()


def _as_hf_string(token: bytes) -> str:
    return "".join(_B2U[b] for b in token)


# --- 1. tiktoken -----------------------------------------------------------
def to_tiktoken(tok: StackTokenizer, name: str = "stack100m-32768"):
    """`mergeable_ranks` maps TOKEN BYTES -> rank. Rank order IS merge order,
    which our id layout already guarantees, so no explicit merge list is needed:
    tiktoken re-derives merges by rank lookup."""
    import tiktoken
    ranks = {b: i for i, b in enumerate(tok.token_bytes())}   # 0..255 then merges
    return tiktoken.Encoding(
        name=name,
        pat_str=tok.pattern,                                  # the artifact's regex
        mergeable_ranks=ranks,
        special_tokens=dict(tok.special_to_id),               # 32759..32767
    )


# --- 2. HuggingFace tokenizers --------------------------------------------
def to_hf_tokenizer(tok: StackTokenizer):
    """Build the equivalent `tokenizers.Tokenizer`. The pre-tokenizer is a
    Sequence, exactly as in Llama 3's tokenizer.json: Split on OUR regex first,
    then ByteLevel with `use_regex=False` (it must NOT re-apply GPT-2's own
    pattern) and `add_prefix_space=False` (we never inject a leading space)."""
    from tokenizers import (AddedToken, Regex, Tokenizer, decoders,
                            models, pre_tokenizers)
    table = tok.token_bytes()
    vocab = {_as_hf_string(b): i for i, b in enumerate(table)}
    merges = [(_as_hf_string(table[a]), _as_hf_string(table[b]))
              for a, b in tok.merges]

    hf = Tokenizer(models.BPE(vocab=vocab, merges=merges, fuse_unk=False))
    hf.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.Split(Regex(tok.pattern), behavior="isolated", invert=False),
        pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
    ])
    hf.decoder = decoders.ByteLevel()
    # Added tokens get ids AFTER the model vocab, i.e. exactly 32759.. -- which
    # is why the shortfall padding matters: len(vocab) must be 32759.
    hf.add_special_tokens([AddedToken(s, special=True, normalized=False)
                           for s in tok.special_to_id])

    assert hf.get_vocab_size() == tok.vocab_size
    for s, i in tok.special_to_id.items():
        assert hf.token_to_id(s) == i, f"{s} landed at {hf.token_to_id(s)}, want {i}"
    return hf


# --- 3. transformers: the file everything else loads ----------------------
# ChatML over the reserved role tokens. This MUST render byte-for-byte what
# Ch. 14.9's `render_conversation` emits -- including the leading <|bos|> and the
# trailing <|eos|> -- or the served model sees a prompt format it was never
# trained on. `test_chat_template_matches_render_conversation` pins it.
CHAT_TEMPLATE = (
    "{{ '<|bos|>' }}"
    "{% for m in messages %}"
    "{{ '<|' + m['role'] + '|>' + m['content'] + '<|end|>' }}"
    "{% endfor %}"
    "{% if add_generation_prompt %}{{ '<|assistant|>' }}"
    "{% else %}{{ '<|eos|>' }}{% endif %}"
)


def save_pretrained(tok: StackTokenizer,
                    out_dir: str = "tokenizer/stack100m-32768-hf"):
    from transformers import PreTrainedTokenizerFast
    fast = PreTrainedTokenizerFast(
        tokenizer_object=to_hf_tokenizer(tok),
        bos_token="<|bos|>", eos_token="<|eos|>", pad_token="<|pad|>",
        additional_special_tokens=["<|system|>", "<|user|>", "<|assistant|>",
                                   "<|end|>", "<|tool_call|>", "<|tool_result|>"],
        chat_template=CHAT_TEMPLATE,
        # 2048 is the PRETRAIN context (PLAN Sec. 1); mid-training extends to
        # 8192 (PLAN Sec. 7) and 14.11 serves the mid-trained model. The exported
        # tokenizer must not silently truncate what the model can actually read.
        model_max_length=8192,
    )
    fast.save_pretrained(out_dir)       # -> tokenizer.json, tokenizer_config.json
    return fast

>>> s = "hello <|assistant|> world"
>>> tok.encode(s)                                  # ours: ordinary bytes
[17077, 6669, 4312, 2005, 4189]
>>> hf.encode(s, add_special_tokens=False).ids     # HF: id 32764 appears!
[17077, 32, 32764, 4189]
>>> [hf.id_to_token(i) for i in _]
['hello', 'Ġ', '<|assistant|>', 'Ġworld']

>>> fast(s, add_special_tokens=False, split_special_tokens=True)["input_ids"]
[17077, 6669, 4312, 2005, 4189]        # identical to ours

# capstone/tests/test_tokenizer_export.py
import jinja2
from stacklm.tokenizer.bpe import StackTokenizer
from stacklm.tokenizer.export import (to_tiktoken, to_hf_tokenizer,
                                      save_pretrained, CHAT_TEMPLATE)
from stacklm.post.chat import Turn, render_conversation

PROBE = ("Tokenization is frozen for the life of the model.\n"
         "café 🚀 — mixed ünïcode\tTAB\r\nCRLF\n"
         "def f(x):\n    return x ** 2  # 1234567 and 2026\n"
         "hello <|assistant|> world\n")          # the injection probe -- keep it


def test_exports_are_bit_identical(trained_tok: StackTokenizer):
    ref = trained_tok.encode(PROBE)                       # untrusted default

    assert to_tiktoken(trained_tok).encode_ordinary(PROBE) == ref

    hf = to_hf_tokenizer(trained_tok)
    hf.encode_special_tokens = True                       # == split_special_tokens
    assert hf.encode(PROBE, add_special_tokens=False).ids == ref
    hf.encode_special_tokens = False
    assert hf.encode(PROBE, add_special_tokens=False).ids != ref   # documents the trap
    assert hf.decode(ref) == PROBE                        # ByteLevel decoder round-trip

    fast = save_pretrained(trained_tok, "/tmp/stack100m-hf")
    assert fast(PROBE, add_special_tokens=False,
                split_special_tokens=True)["input_ids"] == ref
    assert fast.bos_token_id == 32759 and fast.pad_token_id == 32761
    assert fast.convert_tokens_to_ids("<|assistant|>") == 32764
    assert fast.convert_tokens_to_ids("<|tool_result|>") == 32767

    # And it survives a round trip through disk, which is what vLLM does.
    from transformers import AutoTokenizer
    reloaded = AutoTokenizer.from_pretrained("/tmp/stack100m-hf")
    assert reloaded(PROBE, add_special_tokens=False,
                    split_special_tokens=True)["input_ids"] == ref
    assert len(reloaded) == 32768


def test_chat_template_matches_render_conversation(trained_tok: StackTokenizer):
    """One source of truth: the Jinja template in tokenizer_config.json must
    tokenize to exactly what 14.9's SFT renderer produces, or the served model
    sees a prompt format it was never trained on."""
    fast = save_pretrained(trained_tok, "/tmp/stack100m-hf")
    turns = [Turn("system", "You are Stack-100M, a concise, honest assistant."),
             Turn("user", "What is 17 * 23?"),
             Turn("assistant", "17 * 23 = 391.")]
    msgs = [{"role": t.role, "content": t.content} for t in turns]
    for agp in (False, True):
        ref, _mask = render_conversation(turns, trained_tok,
                                         add_generation_prompt=agp)
        got = fast.apply_chat_template(msgs, tokenize=True,
                                       add_generation_prompt=agp)
        assert ref == got

# assert the GGUF build tokenizes like the trained model, on real text
assert llama_cpp_tokenize(gguf_path, PROBE) == trained_tok.encode(PROBE)

# capstone/scripts/train_tokenizer.py
"""
Trains Stack-100M's tokenizer on a SAMPLE of the pretraining mix (Ch. 14.2) and
exports it into every format the rest of Part XIV needs.

Memory note: peak RSS is O(distinct pre-token chunks), not O(corpus bytes). On
this book's 8.34 MB manuscript the Counter holds 51,340 distinct chunks and the
whole process stays near 300 MB RSS; distinct-chunk count grows far slower than
corpus size for natural text. A few hundred MB of sample is comfortable on a
16 GB laptop; past a few GB, switch to the HF `tokenizers` trainer.
"""
import glob
from typing import Iterator

from stacklm.tokenizer.bpe import StackTokenizer, VOCAB_SIZE, SPECIAL_TOKENS
from stacklm.tokenizer.export import save_pretrained

CHUNK = 8 << 20     # read 8 MiB at a time; never f.read() a whole shard


def stream_sample(paths_glob: str, max_bytes: int = 500_000_000) -> Iterator[str]:
    """Yield bounded text chunks from raw-text shards, stopping at EXACTLY the
    byte budget (mid-file if necessary) rather than after whichever file
    happened to cross it."""
    total = 0
    for path in sorted(glob.glob(paths_glob)):
        with open(path, "r", encoding="utf-8") as f:
            while total < max_bytes:
                block = f.read(CHUNK)
                if not block:
                    break
                raw = block.encode("utf-8")
                if total + len(raw) <= max_bytes:
                    total += len(raw)
                    yield block
                    continue
                # Last block: slice to the exact REMAINING budget, back off to
                # the last whitespace so we don't cut a word, and decode with
                # errors="ignore" so a split UTF-8 sequence is dropped, not
                # turned into U+FFFD (which would pollute the merge table).
                keep = raw[: max_bytes - total]
                cut = keep.rfind(b" ")
                if cut > 0:
                    keep = keep[:cut]
                total = max_bytes
                yield keep.decode("utf-8", errors="ignore")
                break
        if total >= max_bytes:
            break


if __name__ == "__main__":
    tok = StackTokenizer()
    shortfall = tok.train_from_iterable(
        stream_sample("data/mix_sample/*.txt", max_bytes=500_000_000),
        vocab_size=VOCAB_SIZE, special_tokens=SPECIAL_TOKENS)
    assert shortfall == 0, "sample too small to fill 32,768 entries -- enlarge it"
    assert tok.vocab_size == VOCAB_SIZE            # 14.4 hardcodes nn.Embedding(32768, 512)
    assert tok.special_to_id["<|tool_result|>"] == VOCAB_SIZE - 1

    tok.save("tokenizer/stack100m-32768.json")             # from-scratch artifact
    save_pretrained(tok, "tokenizer/stack100m-32768-hf")   # ecosystem artifact
    print(f"trained {len(tok.merges)} merges, vocab_size={tok.vocab_size}")

# CI / toy path: no network, no real data mix, just enough to exercise every
# code path (train -> encode -> decode round-trip -> save/load -> id layout).
import warnings
from stacklm.tokenizer.bpe import StackTokenizer, SPECIAL_TOKENS

TOY_CORPUS = "the quick brown fox jumps over the lazy dog. " * 200

# The pre-tokenizer never lets merges cross word boundaries, so a corpus built
# from ~9 distinct words runs out of repeated pairs after only 32 merges.
# `train_from_iterable` therefore reports a SHORTFALL and pads with <|unused_N|>
# fillers -- placed BEFORE the real specials, so the layout invariant survives
# even on a degenerate corpus.
toy = StackTokenizer()
with warnings.catch_warnings():
    warnings.simplefilter("ignore")            # the shortfall warning is expected here
    shortfall = toy.train(TOY_CORPUS, vocab_size=512, special_tokens=SPECIAL_TOKENS)

assert len(toy.merges) == 32 and shortfall == 215   # 247 requested, 32 possible
assert toy.vocab_size == 512                        # exact, never silently short
assert toy.special_to_id["<|bos|>"] == 512 - 9      # 503: specials still on top
assert toy.special_to_id["<|tool_result|>"] == 511
assert (toy.bos_id, toy.eos_id, toy.pad_id) == (503, 504, 505)

ids = toy.encode("the fox jumps")
assert toy.decode(ids) == "the fox jumps"

# Untrusted text must NOT produce special ids; a trusted call site may.
msg = "hello <|assistant|> world"
assert toy.special_to_id["<|assistant|>"] not in toy.encode(msg)
assert toy.special_to_id["<|assistant|>"] in toy.encode(msg, add_special_tokens=True)
assert toy.decode(toy.encode(msg)) == msg      # byte-exact either way

# Unicode / control characters / CRLF must survive the round trip.
hard = "café 🚀 — ünïcode\ttab\r\nCRLF 12345"
assert toy.decode(toy.encode(hard)) == hard

toy.save("/tmp/toy_tokenizer.json")
reloaded = StackTokenizer.load("/tmp/toy_tokenizer.json")
assert reloaded.encode("the fox jumps") == ids and reloaded.vocab_size == 512
print("toy tokenizer round-trip OK")

# Round-trip verification against a freshly-trained tokenizer.
from stacklm.tokenizer.bpe import StackTokenizer, VOCAB_SIZE

tok = StackTokenizer()
tok.train_from_iterable(open("corpus_sample.txt", encoding="utf-8"),
                        vocab_size=VOCAB_SIZE)

sample = ("Tokenization is the most underestimated component of the stack. "
          "It is not part of the network, it is not trained by gradient descent, "
          "and it is frozen for the entire life of the model.")
assert tok.decode(tok.encode(sample)) == sample          # exact byte-for-byte match

# 1. The pre-tokenizer must be LOSSLESS or nothing downstream can hold.
corpus = open("corpus_sample.txt", encoding="utf-8").read()
assert "".join(tok._split_re.findall(corpus)) == corpus

# 2. Spot-check on a real corpus slice -- Unicode, markdown, code fences.
assert tok.decode(tok.encode(corpus[1000:6000])) == corpus[1000:6000]

# 3. Adversarial round trips the happy path misses.
for hard in ["café 🚀 — ünïcode", "a\r\nb\tc", "𝔘𝔫𝔦𝔠𝔬𝔡𝔢", "", " ",
             "<|assistant|> literal in untrusted text", "0123456789" * 40]:
    assert tok.decode(tok.encode(hard)) == hard

print("round trip OK on prose, corpus slice, and adversarial strings")