| """Train the ZENITH BPE tokenizer on a text sample and export the fast (HF) copy. |
| |
| Usage: python -m zenith.tokenizer.train_tokenizer --sample data/interim/tokenizer_sample.txt \ |
| --out zenith/tokenizer/zenith-16k.json |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import random |
| from pathlib import Path |
|
|
| from zenith.tokenizer.bpe import train_bpe |
| from zenith.tokenizer.export import encode_fast, export_hf, load_fast |
|
|
|
|
| def main(argv=None) -> None: |
| ap = argparse.ArgumentParser(prog="python -m zenith.tokenizer.train_tokenizer") |
| ap.add_argument("--sample", type=Path, required=True) |
| ap.add_argument("--out", type=Path, required=True, help="Python tokenizer path; the HF copy goes next to it as .hf.json") |
| ap.add_argument("--vocab-size", type=int, default=16384) |
| ap.add_argument("--parity-lines", type=int, default=10_000) |
| args = ap.parse_args(argv) |
|
|
| text = args.sample.read_text() |
| docs = [d for d in text.split("\n\n") if d.strip()] |
| tok = train_bpe(docs, vocab_size=args.vocab_size) |
| tok.save(args.out) |
| hf_path = args.out.with_name(args.out.stem + ".hf.json") |
| export_hf(tok, hf_path) |
|
|
| lines = [l for l in text.splitlines() if l.strip()] |
| random.Random(0).shuffle(lines) |
| lines = lines[: args.parity_lines] |
| fast_ids = encode_fast(load_fast(hf_path), lines) |
| for line, ids in zip(lines, fast_ids): |
| if ids != tok.encode(line): |
| raise SystemExit(f"parity failure on line: {line[:80]!r}") |
| n_tokens = sum(len(ids) for ids in fast_ids) |
| n_bytes = sum(len(l.encode()) for l in lines) |
| print(f"parity ok: {len(lines)} lines") |
| print(f"bytes/token: {n_bytes / max(n_tokens, 1):.2f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
