File size: 886 Bytes
ae9e4fe | 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 | #!/usr/bin/env python3
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from mla.tokenizer import Tokenizer
TRAIN = Path("data/splits/train.txt")
OUT = Path("data/tokenizer/tokenizer.json")
VOCAB_SIZE = 4096
def main():
text = TRAIN.read_text(encoding="utf-8")
tok = Tokenizer()
print(f"training byte-BPE: target vocab={VOCAB_SIZE}")
tok.train(text, VOCAB_SIZE)
OUT.parent.mkdir(parents=True, exist_ok=True)
tok.save(OUT)
print(f"vocab_size={tok.vocab_size} merges={len(tok.merges)}")
print(f"wrote {OUT}")
sample = "<bos><|user|>hey how are you?<eos><|assistant|>i'm good, you?<eos>"
ids = tok.encode(sample)
back = tok.decode(ids)
print(f"sample_ids_len={len(ids)}")
print(f"roundtrip_ok={back == sample}")
print(f"decoded={back!r}")
if __name__ == "__main__":
main()
|