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
|
"""FineWeb-Edu -> GPT-2 BPE uint16 shards (train.bin / val.bin).
python scripts/prepare_data.py --tokens 3e9 --out data/fineweb [--dataset HuggingFaceFW/fineweb-edu --name sample-10BT]"""
import os, sys, argparse
import numpy as np
p = argparse.ArgumentParser()
p.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
p.add_argument("--name", default="sample-10BT")
p.add_argument("--tokens", type=float, default=3e9)
p.add_argument("--val_tokens", type=float, default=5e6)
p.add_argument("--out", default="data/fineweb")
a = p.parse_args()
os.makedirs(a.out, exist_ok=True)
import tiktoken
from datasets import load_dataset
enc = tiktoken.get_encoding("gpt2")
ds = load_dataset(a.dataset, name=a.name, split="train", streaming=True)
train_path, val_path = os.path.join(a.out, "train.bin"), os.path.join(a.out, "val.bin")
# write to .partial and rename on completion, so an interrupted prep never leaves a
# half-written train.bin that run_ladder.sh's existence check would mistake for done
ftr, fva = open(train_path + ".partial", "wb"), open(val_path + ".partial", "wb")
n_tr = n_va = 0
target_tr, target_va = int(a.tokens), int(a.val_tokens)
buf = []
for i, ex in enumerate(ds):
ids = enc.encode_ordinary(ex["text"]) + [enc.eot_token]
arr = np.array(ids, dtype=np.uint16)
if n_va < target_va and i % 100 == 0: # every 100th doc to val until filled
fva.write(arr.tobytes()); n_va += len(arr)
else:
ftr.write(arr.tobytes()); n_tr += len(arr)
if n_tr % 50_000_000 < len(arr):
print(f"train {n_tr/1e6:.0f}M val {n_va/1e6:.1f}M tokens", flush=True)
if n_tr >= target_tr and n_va >= target_va:
break
ftr.close(); fva.close()
os.rename(train_path + ".partial", train_path); os.rename(val_path + ".partial", val_path)
print(f"DONE train {n_tr/1e6:.1f}M val {n_va/1e6:.1f}M -> {a.out}", flush=True)
sys.stdout.flush(); sys.stderr.flush()
os._exit(0) # hard-exit: HF streaming leaves non-daemon prefetch threads after `break` that hang interpreter shutdown
|
