| """Byte-level BPE tokenizer, vocab 8192 (D-007). | |
| Byte-level means every string can be encoded (unknown characters fall back to | |
| raw bytes), and decode(encode(text)) == text exactly. | |
| """ | |
| from collections.abc import Iterable, Iterator | |
| from pathlib import Path | |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers | |
| VOCAB_SIZE = 8192 | |
| EOS = "<|endoftext|>" # separates stories in the token stream; id 0 | |
| TOKENIZER_PATH = Path(__file__).resolve().parents[2] / "tokenizer" / "tinystories-bpe8192.json" | |
| def train_tokenizer(texts: Iterable[str] | Iterator[list[str]], vocab_size: int = VOCAB_SIZE) -> Tokenizer: | |
| tok = Tokenizer(models.BPE()) | |
| # GPT-2-style split into words/punctuation/spaces before BPE, no prefix space | |
| # (so the roundtrip is exact). | |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) | |
| tok.decoder = decoders.ByteLevel() | |
| trainer = trainers.BpeTrainer( | |
| vocab_size=vocab_size, | |
| min_frequency=2, | |
| special_tokens=[EOS], | |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), # all 256 bytes always present | |
| show_progress=True, | |
| ) | |
| tok.train_from_iterator(texts, trainer=trainer) | |
| return tok | |
| def load_tokenizer(path: Path = TOKENIZER_PATH) -> Tokenizer: | |
| return Tokenizer.from_file(str(path)) | |