File size: 2,792 Bytes
521b329 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Tokenizer: BPE for English (trained here, from scratch), bytes for regex.
English is ~75% of each sequence; BPE shrinks it ~4x. Regex stays byte-level
so every character of the output is an explicit, exact decision.
ids: 0 PAD, 1 SEP, 2 EOS | 3..258 regex bytes | 259.. english BPE
Sequence: <english BPE> SEP <regex bytes> EOS
python src/tokenizer.py --vocab 4096 # train on data/train.jsonl
"""
import argparse
import json
from functools import lru_cache
from pathlib import Path
ROOT = Path(__file__).parents[1]
HERE = Path(__file__).parent
PATH = HERE / "bpe.json" if (HERE / "bpe.json").exists() else ROOT / "data/bpe.json" # flat release or repo
PAD, SEP, EOS = 0, 1, 2
OFFSET = 3
EN = OFFSET + 256
@lru_cache(1)
def _bpe():
from tokenizers import Tokenizer
return Tokenizer.from_file(str(PATH))
def vocab() -> int:
return EN + _bpe().get_vocab_size()
def encode(s: str) -> list[int]:
"""Regex side: raw bytes."""
return [b + OFFSET for b in s.encode("utf-8")]
def decode(ids) -> str:
return bytes(i - OFFSET for i in ids if OFFSET <= i < EN).decode("utf-8", errors="replace")
def english(s: str) -> list[int]:
return [i + EN for i in _bpe().encode(s).ids]
def english_batch(xs: list[str]) -> list[list[int]]:
return [[i + EN for i in e.ids] for e in _bpe().encode_batch(xs)]
def prompt(en: str) -> list[int]:
return english(en) + [SEP]
def pair(en: str, regex: str) -> tuple[list[int], int]:
"""-> (tokens, index of SEP). Loss applies to targets after SEP."""
p = prompt(en)
return p + encode(regex) + [EOS], len(p) - 1
def train(vocab_size: int):
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
tk = Tokenizer(models.BPE())
tk.pre_tokenizer = pre_tokenizers.Sequence([ # one token per digit: counts stay legible
pre_tokenizers.Digits(individual_digits=True), pre_tokenizers.ByteLevel(add_prefix_space=False)])
tk.decoder = decoders.ByteLevel()
tr = trainers.BpeTrainer(vocab_size=vocab_size, min_frequency=2,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
text = (json.loads(l)["en"] for l in open(ROOT / "data/train.jsonl"))
tk.train_from_iterator(text, tr)
tk.save(str(PATH))
return tk
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--vocab", type=int, default=4096)
a = ap.parse_args()
tk = train(a.vocab)
rows = [json.loads(l)["en"] for l in open(ROOT / "data/val.jsonl")]
n_b = sum(len(r.encode()) for r in rows)
n_t = sum(len(e.ids) for e in tk.encode_batch(rows))
print(f"vocab {tk.get_vocab_size()}: {n_b / n_t:.2f} bytes/token on val")
print(tk.encode("validate a ZIP code like 90210-1234, case insensitive").tokens)
|