#!/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 -- -- 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())