File size: 3,439 Bytes
4397e12 | 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 86 87 | """Train a 32k byte-level BPE on a sample of the actual mix; also write the Engram compressed-id map.
Usage: source env.sh && $TA_PY scripts/train_tokenizer.py
Outputs: $TA_DATA/tokenizer.json, $TA_DATA/cid_map.npy
"""
import argparse
import itertools
import unicodedata
import numpy as np
from tokenizers import Regex, Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
from tiny_agent.text import DATA, SPECIAL_TOKENS, source_files
# GPT-4 style split (contractions, letters, 1 digit at a time, punctuation runs, newline runs).
SPLIT = (r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""")
# MB of text sampled per source for training the tokenizer
SAMPLE_MB = {"english": 120, "math": 90, "code_python": 70, "code_shell": 25, "code_markdown": 20,
"code_javascript": 15, "code_typescript": 15, "code_go": 10, "code_rust": 10, "code_c": 10}
def sample_texts():
srcs = source_files()
for name, mb in SAMPLE_MB.items():
budget, got = mb * 2**20, 0
files = srcs.get(name, [])
if not files:
print("missing source", name, flush=True)
continue
f, it = files[0]
for t in it(f):
t = t[:20000]
got += len(t)
yield t
if got >= budget:
break
print(name, got // 2**20, "MB", flush=True)
def build_tokenizer() -> Tokenizer:
tok = Tokenizer(models.BPE(byte_fallback=False))
tok.normalizer = normalizers.NFC()
tok.pre_tokenizer = pre_tokenizers.Sequence([
pre_tokenizers.Split(Regex(SPLIT), behavior="isolated"),
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
])
tok.decoder = decoders.ByteLevel()
return tok
def compressed_map(tok: Tokenizer) -> np.ndarray:
"""Engram tokenizer compression: tokens that normalize to the same surface (NFKC, accents
stripped, lowercased, whitespace stripped) share one compressed id."""
V = tok.get_vocab_size()
keys = {}
out = np.zeros(V, dtype=np.int32)
for i in range(V):
s = tok.decode([i], skip_special_tokens=False)
if i >= len(SPECIAL_TOKENS):
s = unicodedata.normalize("NFKC", s)
s = "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn")
s = s.lower().strip() or "<ws>"
out[i] = keys.setdefault(s, len(keys))
print(f"compressed vocab {len(keys)} / {V} ({1 - len(keys) / V:.1%} smaller)")
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--vocab", type=int, default=32768)
a = ap.parse_args()
tok = build_tokenizer()
trainer = trainers.BpeTrainer(vocab_size=a.vocab, special_tokens=SPECIAL_TOKENS, min_frequency=2,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)
tok.train_from_iterator(sample_texts(), trainer=trainer)
assert tok.token_to_id("<|endoftext|>") == 0 and tok.token_to_id("<tool_call>") == 6
tok.save(f"{DATA}/tokenizer.json")
np.save(f"{DATA}/cid_map.npy", compressed_map(tok))
for s in ["def add(a, b):\n return a + b\n", "The answer is 12345.", "grep -rn 'foo' src/ | head"]:
ids = tok.encode(s).ids
print(len(s), "chars ->", len(ids), "tokens:", [tok.decode([i]) for i in ids][:16])
if __name__ == "__main__":
main()
|