train_bpe.py

import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from nano_transformer.tokenization import gpt2_bytes_to_unicode, train_bpe


DATA_PATH = ROOT / "data"


def serialize_tokenizer_artifacts(vocab, merges, vocab_path: Path, merges_path: Path) -> None:
    gpt2_byte_encoder = gpt2_bytes_to_unicode()

    def bytes_to_token(b: bytes) -> str:
        return "".join(gpt2_byte_encoder[x] for x in b)

    vocab_encoded = {bytes_to_token(v): k for k, v in vocab.items()}

    with open(vocab_path, "w", encoding="utf-8") as f:
        json.dump(vocab_encoded, f, ensure_ascii=False, indent=4)

    with open(merges_path, "w", encoding="utf-8", newline="\n") as f:
        for b1, b2 in merges:
            f.write(f"{bytes_to_token(b1)} {bytes_to_token(b2)}\n")

def train_bpe_tinystories():
    input_path = DATA_PATH / "TinyStoriesV2-GPT4-train.txt"
    vocab, merges = train_bpe(
        input_path=input_path,
        vocab_size=10000,
        special_tokens=["<|endoftext|>"],
    )
    vocab_path = DATA_PATH / "train-bpe-tinystories-vocab.json"
    merges_path = DATA_PATH / "train-bpe-tinystories-merges.txt"
    serialize_tokenizer_artifacts(vocab, merges, vocab_path, merges_path)


def train_bpe_openwebtext():
    input_path = DATA_PATH / "owt_train.txt"
    vocab, merges = train_bpe(
        input_path=input_path,
        vocab_size=32000,
        special_tokens=["<|endoftext|>"],
    )
    vocab_path = DATA_PATH / "train-bpe-openwebtext-vocab.json"
    merges_path = DATA_PATH / "train-bpe-openwebtext-merges.txt"
    serialize_tokenizer_artifacts(vocab, merges, vocab_path, merges_path)


if __name__ == "__main__":
    # train_bpe_tinystories()
    train_bpe_openwebtext()