| """Tokenize FineWeb-Edu (sample-10BT) with the GPT-2 tokenizer into flat uint16 files. |
| |
| Each document is written as [EOT] + tokens. Train comes from shards 000-006 |
| (~5B tokens), val from the first --val_tokens of shard 013. |
| """ |
|
|
| import argparse |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| import tiktoken |
| from huggingface_hub import hf_hub_download |
|
|
| REPO = "HuggingFaceFW/fineweb-edu" |
| ENC = None |
|
|
|
|
| def encode(texts): |
| global ENC |
| if ENC is None: |
| ENC = tiktoken.get_encoding("gpt2") |
| out = [] |
| for ids in ENC.encode_ordinary_batch(texts): |
| out.append(ENC.eot_token) |
| out.extend(ids) |
| return np.asarray(out, dtype=np.uint16) |
|
|
|
|
| def text_batches(path, batch=1024): |
| for rb in pq.ParquetFile(path).iter_batches(batch_size=batch, columns=["text"]): |
| yield rb.column("text").to_pylist() |
|
|
|
|
| def write_split(shards, dst, max_tokens, workers): |
| n = 0 |
| with open(dst, "wb") as f, Pool(workers) as pool: |
| for shard in shards: |
| local = hf_hub_download(REPO, f"sample/10BT/{shard}_00000.parquet", repo_type="dataset") |
| for arr in pool.imap(encode, text_batches(local), chunksize=4): |
| if max_tokens and n + len(arr) > max_tokens: |
| arr = arr[: max_tokens - n] |
| f.write(arr.tobytes()) |
| n += len(arr) |
| if max_tokens and n >= max_tokens: |
| return n |
| print(f"{dst.name}: {shard} done, {n / 1e9:.3f}B tokens", flush=True) |
| return n |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--out_dir", default="/data/fineweb-edu-gpt2") |
| p.add_argument("--train_shards", default="000,001,002,003,004,005,006") |
| p.add_argument("--val_shard", default="013") |
| p.add_argument("--val_tokens", type=int, default=20_000_000) |
| p.add_argument("--workers", type=int, default=32) |
| args = p.parse_args() |
| out = Path(args.out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| nv = write_split([args.val_shard], out / "val.bin", args.val_tokens, args.workers) |
| print(f"val: {nv} tokens", flush=True) |
| nt = write_split(args.train_shards.split(","), out / "train.bin", 0, args.workers) |
| print(f"train: {nt} tokens", flush=True) |
| (out / "DONE").write_text(f"train={nt}\nval={nv}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
