1#!/usr/bin/env python2"""Prepare the general-language corpus for the 103M pilot.34Source: HuggingFaceFW/fineweb-edu, subset "sample-10BT" (English web text, ODC-By 1.0), streamed in its5canonical shard order (deterministic). Documents are exact-hash deduplicated, assigned to train / val / test6by a hash of their content (never by position), and tokenized with a byte-level BPE (vocab 32768, digits split7individually) trained on the TRAIN split only. The test split is locked: it is written but never read by any8training or development script (see protocol/preregistration.md).910Outputs (data/processed/<name>/): tokenizer.json, {train,val,test}.bin (uint16, documents separated by <|endoftext|>),11manifest.json (counts, sha256, parameters).12"""13from __future__ import annotations1415import argparse, hashlib, json, os, sys, time16from pathlib import Path1718import numpy as np192021def doc_hash(text: str) -> str:22 return hashlib.sha256(text.encode("utf-8", "ignore")).hexdigest()232425def split_of(h: str, val_pct: float, test_pct: float) -> str:26 u = int(h[:12], 16) / 16 ** 1227 if u < test_pct:28 return "test"29 if u < test_pct + val_pct:30 return "val"31 return "train"323334def main() -> int:35 ap = argparse.ArgumentParser(description=__doc__)36 ap.add_argument("--out", default="data/processed/fineweb_edu_300m")37 ap.add_argument("--target-train-tokens", type=int, default=300_000_000)38 ap.add_argument("--vocab-size", type=int, default=32768)39 ap.add_argument("--tokenizer-docs", type=int, default=60_000)40 ap.add_argument("--val-pct", type=float, default=0.01)41 ap.add_argument("--test-pct", type=float, default=0.01)42 ap.add_argument("--max-doc-chars", type=int, default=40_000)43 args = ap.parse_args()44 out = Path(args.out); out.mkdir(parents=True, exist_ok=True)45 raw_dir = Path("data/raw/fineweb_edu"); raw_dir.mkdir(parents=True, exist_ok=True)4647 from datasets import load_dataset48 from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers4950 t0 = time.time()51 ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)52 # ---- pass 1: stream documents, dedup, assign splits, spill text to jsonl per split --------------------53 seen = set()54 counts = {"train": 0, "val": 0, "test": 0}55 chars = {"train": 0, "val": 0, "test": 0}56 files = {s: open(raw_dir / f"{s}.jsonl", "w") for s in counts}57 tok_sample = []58 # ~4.3 chars per BPE token on English web text: stream until the train split has enough characters.59 target_chars = int(args.target_train_tokens * 4.3)60 n_seen = 061 for ex in ds:62 n_seen += 163 text = ex["text"]64 if len(text) > args.max_doc_chars:65 text = text[: args.max_doc_chars]66 h = doc_hash(text)67 if h in seen:68 continue69 seen.add(h)70 s = split_of(h, args.val_pct, args.test_pct)71 files[s].write(json.dumps({"id": ex.get("id"), "h": h, "text": text}) + "\n")72 counts[s] += 173 chars[s] += len(text)74 if s == "train" and len(tok_sample) < args.tokenizer_docs:75 tok_sample.append(text)76 if n_seen % 20000 == 0:77 print(f"[pass1] seen={n_seen} train_chars={chars['train']/1e6:.1f}M ({time.time()-t0:.0f}s)", flush=True)78 if chars["train"] >= target_chars:79 break80 for f in files.values():81 f.close()82 print("pass1 done", counts, chars, f"{time.time()-t0:.0f}s", flush=True)8384 # ---- tokenizer (train split only) ------------------------------------------------------------------85 tok = Tokenizer(models.BPE(unk_token=None))86 tok.pre_tokenizer = pre_tokenizers.Sequence([pre_tokenizers.Digits(individual_digits=True),87 pre_tokenizers.ByteLevel(add_prefix_space=False)])88 tok.decoder = decoders.ByteLevel()89 trainer = trainers.BpeTrainer(vocab_size=args.vocab_size, special_tokens=["<|endoftext|>", "<|pad|>"],90 initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)91 tok.train_from_iterator(tok_sample, trainer=trainer)92 tok.save(str(out / "tokenizer.json"))93 eot = tok.token_to_id("<|endoftext|>")94 assert tok.get_vocab_size() == args.vocab_size, tok.get_vocab_size()95 print("tokenizer trained", tok.get_vocab_size(), f"{time.time()-t0:.0f}s", flush=True)9697 # ---- pass 2: tokenize each split to uint16 ---------------------------------------------------------98 manifest = {"source": "HuggingFaceFW/fineweb-edu sample-10BT (streaming, canonical shard order)",99 "license": "ODC-By 1.0 (https://opendatacommons.org/licenses/by/1-0/)",100 "prepared_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "args": vars(args), "docs_streamed": n_seen,101 "docs": counts, "chars": chars, "splits": {}, "tokenizer_sha256": hashlib.sha256((out / "tokenizer.json").read_bytes()).hexdigest(),102 "eot_id": eot, "dtype": "uint16"}103 for s in counts:104 ids_all = []105 with open(raw_dir / f"{s}.jsonl") as f:106 batch = []107 for line in f:108 batch.append(json.loads(line)["text"])109 if len(batch) == 2000:110 for enc in tok.encode_batch(batch):111 ids_all.append(np.asarray(enc.ids + [eot], dtype=np.uint16))112 batch = []113 if batch:114 for enc in tok.encode_batch(batch):115 ids_all.append(np.asarray(enc.ids + [eot], dtype=np.uint16))116 arr = np.concatenate(ids_all) if ids_all else np.zeros(0, dtype=np.uint16)117 arr.tofile(out / f"{s}.bin")118 manifest["splits"][s] = {"tokens": int(arr.size), "docs": len(ids_all),119 "sha256": hashlib.sha256(arr.tobytes()).hexdigest(), "locked": s == "test"}120 print(f"[pass2] {s}: {arr.size:,} tokens ({time.time()-t0:.0f}s)", flush=True)121 (out / "manifest.json").write_text(json.dumps(manifest, indent=2))122 print("done", json.dumps(manifest["splits"]))123 return 0124125126if __name__ == "__main__":127 raise SystemExit(main())128