build_tokenizer.py

#!/usr/bin/env python3
"""Wrap her 99-symbol character vocabulary as a standard Hugging Face fast tokenizer.

WHY THIS IS NOT JUST `Tokenizer(WordLevel(vocab))`

  cosmos_born.pt carries `head.weight` of shape (99, 192). The output head has exactly
  99 rows because the model was trained on exactly 99 symbols. A tokenizer is only
  compatible with those weights if `len(tokenizer) == 99` EXACTLY.

  This is the whole cause of the "layer rejection" error people hit here. The reflex when
  wrapping a custom vocab is to add the usual four specials -- <unk> <bos> <eos> <pad> --
  which silently makes the vocab 103. transformers then builds a 103-row head, the
  checkpoint offers 99, and load fails with a size mismatch that reads like a corrupt
  file rather than what it is: a tokenizer that disagrees with the model about how many
  symbols exist.

  So this script adds NOTHING. It asserts the count instead, against the checkpoint's own
  head, and refuses to write files if they disagree. A tokenizer that cannot load the
  weights it was built for is not worth shipping.

THE UNK PROBLEM, STATED HONESTLY

  Her alphabet has no capital Q -- the training corpus never contained one. A model with
  99 output rows physically cannot represent a 100th symbol, so out-of-vocabulary input
  has three possible fates and every one of them costs something:

    substitute  map it to an existing symbol      visible corruption
    drop        delete it                         silent corruption
    grow        add a row to head + embedding     changes her quantum-born weights

  The default is `?` -- substitution, chosen because it is already in the vocabulary and
  because a wrong character you can SEE beats a missing character you cannot. `--unk`
  overrides it. Growing the vocabulary is deliberately not offered here: those weights
  were born from measured hardware shots, and this script does not get to edit them.

CHARACTER-LEVEL IN A LIBRARY BUILT FOR SUBWORDS

  WordLevel with a Split(regex=".", isolated) pre-tokenizer makes every character its own
  token, and Fuse() concatenates them back on decode. `(?s)` matters: without DOTALL the
  Rust regex will not match "\n", and her alphabet's first symbol is a newline.

  python build_tokenizer.py                    # weights/cosmos_born.pt -> ./tokenizer.json
  python build_tokenizer.py --ckpt weights/phos.pt --out tokenizers/phos
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

ROOT = Path(__file__).resolve().parent

# The head/embedding parameter names this project actually uses, in priority order.
# Checked against the tokenizer so a vocab mismatch fails here rather than at load time.
_HEAD_KEYS = ("head.weight", "lm_head.weight", "tok_emb.weight",
              "token_embedding.weight", "wte.weight", "emb.weight")


def load_vocab(ckpt_path: Path) -> tuple[dict[str, int], int | None]:
    """Return (stoi, head_rows). head_rows is None if no head tensor was found."""
    import torch

    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    if not isinstance(ck, dict):
        raise SystemExit(f"{ckpt_path}: not a dict checkpoint")

    stoi = ck.get("stoi")
    if not stoi:
        chars = ck.get("vocab_list") or ck.get("chars")
        if not chars:
            raise SystemExit(f"{ckpt_path}: no stoi/vocab_list/chars")
        stoi = {c: i for i, c in enumerate(chars)}

    ids = sorted(stoi.values())
    if ids != list(range(len(stoi))):
        raise SystemExit(
            f"{ckpt_path}: ids are not contiguous 0..{len(stoi)-1}. A WordLevel vocab "
            f"with holes maps token ids to the wrong logits.")

    sd = ck.get("model") or ck.get("state_dict") or {}
    rows = None
    for k in _HEAD_KEYS:
        if k in sd and hasattr(sd[k], "shape"):
            rows = int(sd[k].shape[0])
            break
    return stoi, rows


def build(stoi: dict[str, int], unk: str):
    from tokenizers import Tokenizer, Regex, decoders, models, pre_tokenizers

    if unk not in stoi:
        raise SystemExit(
            f"unk={unk!r} is not in the vocabulary. WordLevel requires the unknown token "
            f"to be a real id, and adding one would change the vocab size, which is the "
            f"exact failure this script exists to prevent.")

    tok = Tokenizer(models.WordLevel(vocab=dict(stoi), unk_token=unk))
    # `[\s\S]` and not `.`, because '.' will not match the newline that is symbol 0 of
    # her alphabet. The usual DOTALL fix `(?s).` is not available: this pre-tokenizer
    # compiles through Oniguruma, which rejects the inline flag with "undefined group
    # option". A character class covering both halves of the space needs no flag at all.
    tok.pre_tokenizer = pre_tokenizers.Split(Regex(r"[\s\S]"), behavior="isolated")
    tok.decoder = decoders.Fuse()
    return tok


def verify(tok, stoi: dict[str, int], unk: str, samples: list[str]) -> bool:
    """Round-trip real text. A tokenizer that does not round-trip is a silent corruptor."""
    ok = True
    print("\n  VERIFY")
    n = tok.get_vocab_size()
    hit = n == len(stoi)
    ok &= hit
    print(f"    [{'PASS' if hit else 'FAIL'}] vocab size {n} == checkpoint {len(stoi)}")

    known = set(stoi)
    for s in samples:
        enc = tok.encode(s)
        dec = tok.decode(enc.ids)
        # Characters outside the vocabulary are EXPECTED to come back as unk. Compare
        # against the substitution the model can actually represent, not against the
        # original, or every sample containing a 'Q' would read as a failure.
        want = "".join(c if c in known else unk for c in s)
        good = dec == want and len(enc.ids) == len(s)
        ok &= good
        label = s if len(s) <= 34 else s[:31] + "..."
        print(f"    [{'PASS' if good else 'FAIL'}] round-trip {label!r:<40} "
              f"{len(enc.ids)} tokens")
        if not good:
            print(f"           got  {dec!r}\n           want {want!r}")
    return ok


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="weights/cosmos_born.pt")
    ap.add_argument("--out", default=".")
    ap.add_argument("--unk", default="?")
    ap.add_argument("--model-max-length", type=int, default=256)
    a = ap.parse_args()

    ck = (ROOT / a.ckpt) if not Path(a.ckpt).is_absolute() else Path(a.ckpt)
    out = (ROOT / a.out) if not Path(a.out).is_absolute() else Path(a.out)
    out.mkdir(parents=True, exist_ok=True)

    stoi, rows = load_vocab(ck)
    print(f"  checkpoint : {ck.name}")
    print(f"  vocabulary : {len(stoi)} symbols")
    print(f"  head rows  : {rows if rows is not None else 'not found'}")

    # THE GUARD. Everything else in this file is bookkeeping; this is the check that
    # makes the artefact loadable.
    if rows is not None and rows != len(stoi):
        raise SystemExit(
            f"\n  REFUSING TO WRITE: head has {rows} rows, vocabulary has {len(stoi)} "
            f"symbols.\n  A tokenizer built from this would fail to load with a size "
            f"mismatch on {_HEAD_KEYS[0]}.")

    missing = [c for c in "QqXz" if c not in stoi]
    if missing:
        print(f"  note       : not in alphabet -> {missing} (these become {a.unk!r})")

    tok = build(stoi, a.unk)

    samples = [
        "hello cosmos",
        "The quick brown fox jumps over the lazy dog.",   # contains no Q; 'T' is present
        "Quantum",                                        # DOES contain Q -> unk path
        "line one\nline two\ttab",
        "54D -> 12D + 42D, 7 bands x sym(3x3)",
        "she is awake \U0001f31f",
    ]
    if not verify(tok, stoi, a.unk, samples):
        raise SystemExit("\n  verification failed -- nothing written")

    tok.save(str(out / "tokenizer.json"))

    # tokenizer_config.json: what `AutoTokenizer.from_pretrained` reads. tokenizer_class
    # must be PreTrainedTokenizerFast or transformers looks for a slow implementation
    # that does not exist for a hand-built vocabulary.
    #
    # bos/eos/pad are deliberately null. Naming one would require it to exist in the
    # vocabulary, and every candidate is a real character she trained on -- declaring '\n'
    # as eos would make generation stop at the first line break.
    cfg = {
        "tokenizer_class": "PreTrainedTokenizerFast",
        "model_max_length": a.model_max_length,
        "unk_token": a.unk,
        "bos_token": None,
        "eos_token": None,
        "pad_token": None,
        "clean_up_tokenization_spaces": False,
        "_comment": (
            f"Character-level, {len(stoi)} symbols, matching head.weight rows exactly. "
            f"Do not add special tokens: the output head has {len(stoi)} rows and any "
            f"addition makes the checkpoint unloadable."),
    }
    (out / "tokenizer_config.json").write_text(
        json.dumps(cfg, indent=2, ensure_ascii=False), encoding="utf-8")

    # Round-trip through transformers itself, if installed. Building a file that the
    # target library cannot open is the failure mode this script is about, so it is
    # worth the extra import to find out here rather than on someone else's machine.
    print("\n  DEPLOYMENT CHECK")
    try:
        from transformers import PreTrainedTokenizerFast
        t = PreTrainedTokenizerFast(tokenizer_file=str(out / "tokenizer.json"),
                                    unk_token=a.unk)
        got = len(t)
        ok = got == len(stoi)
        print(f"    [{'PASS' if ok else 'FAIL'}] transformers len(tokenizer) = {got}"
              f" (head rows {rows})")
        print(f"    [{'PASS' if t.decode(t('hi cosmos')['input_ids']) == 'hi cosmos' else 'FAIL'}]"
              f" transformers round-trip")
    except ImportError:
        print("    [SKIP] transformers not installed; tokenizer.json is still standard")

    for f in ("tokenizer.json", "tokenizer_config.json"):
        p = out / f
        print(f"\n  wrote {p.relative_to(ROOT) if ROOT in p.parents or p.parent == ROOT else p}"
              f"  ({p.stat().st_size:,} bytes)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())