File size: 2,892 Bytes
f31f59b | 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 72 | """Training data for the paper: FineWeb-Edu (sample-10BT), Gemma tokenizer (262,144 tokens), uint32 shards.
2.05B training tokens + 10M validation tokens, documents separated by <eos>. Also writes the token counts used
for the unigram prior (counted on the training tokens only).
python prep_data.py --out data [--train 2.05e9] [--val 1e7]
"""
import argparse, pathlib, time
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer
SHARD = 100_000_000
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="data")
ap.add_argument("--train", type=float, default=2.05e9)
ap.add_argument("--val", type=float, default=1e7)
a = ap.parse_args()
out = pathlib.Path(a.out)
out.mkdir(parents=True, exist_ok=True)
tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
eos = tok.eos_token_id
need_val, need_train = int(a.val), int(a.train)
ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True)
counts = np.zeros(262144, np.int64)
buf = np.empty(SHARD + 1_000_000, np.uint32)
n_buf, shard, done_val, total, t0, batch = 0, 0, False, 0, time.time(), []
def write(arr, name):
np.save(out / name, arr)
print(time.strftime("%H:%M:%S"), f"wrote {name} ({len(arr) / 1e6:.0f}M tokens)", flush=True)
for ex in ds:
batch.append(ex["text"])
if len(batch) < 512:
continue
for ids in tok(batch, add_special_tokens=False)["input_ids"]:
ids.append(eos)
buf[n_buf:n_buf + len(ids)] = ids
n_buf += len(ids)
batch = []
if not done_val and n_buf >= need_val:
write(buf[:need_val].copy(), "val.npy")
rest = buf[need_val:n_buf].copy()
buf[:len(rest)] = rest
n_buf, done_val = len(rest), True
elif done_val and n_buf >= SHARD:
take = min(SHARD, need_train - total)
arr = buf[:take].copy()
counts += np.bincount(arr, minlength=262144)
write(arr, f"train_{shard:03d}.npy")
shard, total = shard + 1, total + take
rest = buf[SHARD:n_buf].copy()
buf[:len(rest)] = rest
n_buf = len(rest)
if total >= need_train:
break
print(f" {total / 1e9:.2f}B / {need_train / 1e9:.2f}B {total / (time.time() - t0) / 1e6:.2f}M tok/s", flush=True)
if total < need_train and n_buf:
arr = buf[:min(n_buf, need_train - total)].copy()
counts += np.bincount(arr, minlength=262144)
write(arr, f"train_{shard:03d}.npy")
total += len(arr)
np.save(out / "unigram_counts.npy", counts)
print(time.strftime("%H:%M:%S"), f"done: {total / 1e9:.3f}B train tokens in {time.time() - t0:.0f} s", flush=True)
if __name__ == "__main__":
main()
|