File size: 2,381 Bytes
9c6799a | 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 | """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()
|