spb/forge Public MIT

1# Author: Simon-Pierre Boucher — contact@spboucher.ai2"""Download TinyStories, train (or reuse) a forgebpe tokenizer, tokenize,3and write train.bin / val.bin as uint16 token streams with the llm.c-style4header {magic 20240520, version 1, num_tokens} (RESEARCH.md §7).56Encoder equivalence: src/tokenizer/bpe.cpp repeatedly applies the *lowest-id*7merge present. Because a merge (a,b)->c always has c > a and c > b, applying8merge c can only create pairs whose merge id exceeds c — so no lower-id merge9can appear later. Hence one ascending pass over the merge list, applying all10non-overlapping occurrences of each, is equivalent and vectorizes cleanly.11tests/test_tokenizer.cpp checks the two agree on real text.1213Usage:14  python3 tools/prepare_data.py --out data/tinystories [--vocab-size 4096]15      [--max-train-mb 100] [--max-val-mb 5]16"""17import argparse18import multiprocessing as mp19import os20import struct21import urllib.request2223import numpy as np2425BASE = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main"26TRAIN_FILE = "TinyStoriesV2-GPT4-train.txt"27VAL_FILE = "TinyStoriesV2-GPT4-valid.txt"28MAGIC, VERSION = 20240520, 12930_MERGES = None  # worker global: list of (idx, left, right) in ascending id order313233def download(url, dest):34    if os.path.exists(dest):35        print(f"already have {dest}")36        return37    print(f"downloading {url}")38    tmp = dest + ".tmp"39    urllib.request.urlretrieve(url, tmp)40    os.rename(tmp, dest)414243def load_merges(model_path):44    merges = []45    with open(model_path) as f:46        assert f.readline().strip() == "forgebpe v1"47        f.readline()  # vocab size48        for line in f:49            idx, left, right = map(int, line.split())50            merges.append((idx, left, right))51    merges.sort()52    return merges535455def apply_merge(ids, idx, left, right):56    is_pair = (ids[:-1] == left) & (ids[1:] == right)57    pos = np.where(is_pair)[0]58    if pos.size == 0:59        return ids60    if left == right:  # "aaa": keep the first of each run, non-overlapping61        keep, last = [], -262        for p in pos:63            if p != last + 1:64                keep.append(p)65                last = p66        pos = np.asarray(keep, dtype=np.int64)67    ids[pos] = idx68    return np.delete(ids, pos + 1)697071def encode_block(block_bytes):72    ids = np.frombuffer(block_bytes, dtype=np.uint8).astype(np.int32)73    for idx, left, right in _MERGES:74        if ids.size < 2:75            break76        ids = apply_merge(ids, idx, left, right)77    return ids.astype(np.uint16)787980def _init_worker(model_path):81    global _MERGES82    _MERGES = load_merges(model_path)838485def tokenize_file(text_path, model_path, out_path, max_mb, block_mb=2.0):86    with open(text_path, "rb") as f:87        data = f.read(int(max_mb * 1024 * 1024))8889    # split into blocks on story boundaries so merges never straddle documents90    sep = b"<|endoftext|>"91    blocks, cur = [], bytearray()92    limit = int(block_mb * 1024 * 1024)93    for story in data.split(sep):94        if not story.strip():95            continue96        cur += story97        if len(cur) >= limit:98            blocks.append(bytes(cur))99            cur = bytearray()100    if cur:101        blocks.append(bytes(cur))102    print(f"tokenizing {len(data) / 1e6:.1f} MB in {len(blocks)} blocks")103104    with mp.Pool(initializer=_init_worker, initargs=(model_path,)) as pool:105        parts = []106        for i, ids in enumerate(pool.imap(encode_block, blocks)):107            parts.append(ids)108            if (i + 1) % 10 == 0:109                print(f"  block {i + 1}/{len(blocks)}")110    tokens = np.concatenate(parts) if parts else np.zeros(0, dtype=np.uint16)111    print(f"  -> {len(tokens) / 1e6:.2f}M tokens "112          f"({len(data) / max(1, len(tokens)):.2f} bytes/token)")113114    with open(out_path, "wb") as f:115        header = np.zeros(256, dtype=np.int32)116        header[0], header[1], header[2] = MAGIC, VERSION, len(tokens)117        f.write(header.tobytes())118        f.write(tokens.tobytes())119    print(f"wrote {out_path} ({os.path.getsize(out_path) / 1e6:.1f} MB)")120121122def main():123    ap = argparse.ArgumentParser()124    ap.add_argument("--out", required=True)125    ap.add_argument("--vocab-size", type=int, default=4096)126    ap.add_argument("--max-train-mb", type=float, default=100.0)127    ap.add_argument("--max-val-mb", type=float, default=5.0)128    ap.add_argument("--tokenizer-mb", type=float, default=10.0,129                    help="corpus sample size for BPE training")130    ap.add_argument("--tokenizer", default=None,131                    help="existing .model; otherwise trained on the corpus")132    args = ap.parse_args()133134    os.makedirs(args.out, exist_ok=True)135    train_txt = os.path.join(args.out, TRAIN_FILE)136    val_txt = os.path.join(args.out, VAL_FILE)137    download(f"{BASE}/{TRAIN_FILE}?download=true", train_txt)138    download(f"{BASE}/{VAL_FILE}?download=true", val_txt)139140    model = args.tokenizer or os.path.join(args.out, f"tok{args.vocab_size}.model")141    if not os.path.exists(model):142        import sys143        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))144        import train_tokenizer as tt145        with open(train_txt, "rb") as f:146            sample = f.read(int(args.tokenizer_mb * 1024 * 1024))147        print(f"training {args.vocab_size}-vocab BPE on {len(sample) / 1e6:.1f} MB")148        tt.write_model(model, tt.train(sample, args.vocab_size))149        print(f"wrote {model}")150151    tokenize_file(train_txt, model, os.path.join(args.out, "train.bin"),152                  args.max_train_mb)153    tokenize_file(val_txt, model, os.path.join(args.out, "val.bin"), args.max_val_mb)154155156if __name__ == "__main__":157    main()158