File size: 1,730 Bytes
0cf8154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """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()
|