|
|
| """Stream FineWeb-Edu and write GPT-2 token .bin files for nanoGPT training."""
|
|
|
| import argparse
|
| import json
|
| import os
|
| import pickle
|
| import shutil
|
| import time
|
| from typing import Optional
|
|
|
| import numpy as np
|
| import tiktoken
|
| from datasets import load_dataset
|
| from tqdm import tqdm
|
|
|
|
|
| GPT2_VOCAB_SIZE = 50257
|
|
|
|
|
| def parse_token_count(value: Optional[str]) -> Optional[int]:
|
| if value is None:
|
| return None
|
| raw = str(value).strip().replace("_", "").lower()
|
| if raw in {"none", "all", "full", "0"}:
|
| return None
|
|
|
| multipliers = {
|
| "k": 1_000,
|
| "m": 1_000_000,
|
| "b": 1_000_000_000,
|
| }
|
| suffix = raw[-1]
|
| if suffix in multipliers:
|
| return int(float(raw[:-1]) * multipliers[suffix])
|
| return int(raw)
|
|
|
|
|
| def safe_output_paths(out_dir: str, overwrite: bool) -> tuple[str, str, str]:
|
| os.makedirs(out_dir, exist_ok=True)
|
| train_path = os.path.join(out_dir, "train.bin")
|
| val_path = os.path.join(out_dir, "val.bin")
|
| meta_path = os.path.join(out_dir, "meta.pkl")
|
|
|
| existing = [p for p in (train_path, val_path, meta_path) if os.path.exists(p)]
|
| if existing and not overwrite:
|
| joined = "\n ".join(existing)
|
| raise FileExistsError(
|
| "Refusing to overwrite existing output files. Remove them or pass "
|
| f"--overwrite:\n {joined}"
|
| )
|
| return train_path, val_path, meta_path
|
|
|
|
|
| def write_tokens(handle, tokens: list[int]) -> int:
|
| if not tokens:
|
| return 0
|
| arr = np.asarray(tokens, dtype=np.uint16)
|
| arr.tofile(handle)
|
| return int(arr.size)
|
|
|
|
|
| def main() -> None:
|
| parser = argparse.ArgumentParser(
|
| description=(
|
| "Stream HuggingFaceFW/fineweb-edu sample-10BT, tokenize with the "
|
| "GPT-2 BPE tokenizer, and write data/fineweb/train.bin and val.bin."
|
| )
|
| )
|
| parser.add_argument("--out_dir", default="data/fineweb")
|
| parser.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
|
| parser.add_argument("--name", default="sample-10BT")
|
| parser.add_argument("--split", default="train")
|
| parser.add_argument(
|
| "--max_tokens",
|
| default=None,
|
| help=(
|
| "Maximum total tokens to write, e.g. 100M, 1B, 10B. "
|
| "Use all/full/0 or omit for the full streamed split."
|
| ),
|
| )
|
| parser.add_argument(
|
| "--val_tokens",
|
| default="5M",
|
| help="Tokens reserved from the beginning of the stream for training-time val.bin.",
|
| )
|
| parser.add_argument(
|
| "--min_score",
|
| type=float,
|
| default=None,
|
| help="Optional FineWeb-Edu quality score filter, e.g. 3.0.",
|
| )
|
| parser.add_argument("--overwrite", action="store_true")
|
| args = parser.parse_args()
|
|
|
| max_tokens = parse_token_count(args.max_tokens)
|
| val_tokens = parse_token_count(args.val_tokens)
|
| if val_tokens is None or val_tokens <= 0:
|
| raise ValueError("--val_tokens must be a positive token count")
|
| if max_tokens is not None and max_tokens <= val_tokens:
|
| raise ValueError("--max_tokens must be larger than --val_tokens")
|
|
|
| train_path, val_path, meta_path = safe_output_paths(args.out_dir, args.overwrite)
|
| train_part = train_path + ".part"
|
| val_part = val_path + ".part"
|
|
|
| for path in (train_part, val_part):
|
| if os.path.exists(path):
|
| if args.overwrite:
|
| os.remove(path)
|
| else:
|
| raise FileExistsError(
|
| f"Partial output exists: {path}. Remove it or pass --overwrite."
|
| )
|
|
|
| enc = tiktoken.get_encoding("gpt2")
|
| eot = enc.eot_token
|
|
|
| ds = load_dataset(
|
| args.dataset,
|
| name=args.name,
|
| split=args.split,
|
| streaming=True,
|
| )
|
|
|
| total_written = 0
|
| val_written = 0
|
| train_written = 0
|
| docs_seen = 0
|
| docs_used = 0
|
| start = time.time()
|
|
|
| progress_total = max_tokens if max_tokens is not None else None
|
| pbar = tqdm(total=progress_total, unit="tok", smoothing=0.05)
|
|
|
| with open(val_part, "wb") as val_f, open(train_part, "wb") as train_f:
|
| for doc in ds:
|
| docs_seen += 1
|
| if args.min_score is not None and doc.get("score") is not None:
|
| if float(doc["score"]) < args.min_score:
|
| continue
|
|
|
| text = doc.get("text")
|
| if not text:
|
| continue
|
|
|
| tokens = [eot] + enc.encode_ordinary(text)
|
| if max_tokens is not None:
|
| remaining = max_tokens - total_written
|
| if remaining <= 0:
|
| break
|
| tokens = tokens[:remaining]
|
|
|
| cursor = 0
|
| if val_written < val_tokens:
|
| take = min(val_tokens - val_written, len(tokens))
|
| val_written += write_tokens(val_f, tokens[:take])
|
| cursor = take
|
|
|
| if cursor < len(tokens):
|
| train_written += write_tokens(train_f, tokens[cursor:])
|
|
|
| docs_used += 1
|
| total_written = val_written + train_written
|
| pbar.update(len(tokens))
|
|
|
| if max_tokens is not None and total_written >= max_tokens:
|
| break
|
|
|
| pbar.close()
|
|
|
| os.replace(train_part, train_path)
|
| os.replace(val_part, val_path)
|
|
|
| metadata = {
|
| "vocab_size": GPT2_VOCAB_SIZE,
|
| "tokenizer": "tiktoken:gpt2",
|
| "dataset": args.dataset,
|
| "name": args.name,
|
| "split": args.split,
|
| "max_tokens": max_tokens,
|
| "val_tokens_requested": val_tokens,
|
| "val_tokens_written": val_written,
|
| "train_tokens_written": train_written,
|
| "docs_seen": docs_seen,
|
| "docs_used": docs_used,
|
| "min_score": args.min_score,
|
| "elapsed_sec": round(time.time() - start, 2),
|
| }
|
| with open(meta_path, "wb") as f:
|
| pickle.dump({"vocab_size": GPT2_VOCAB_SIZE}, f)
|
| with open(os.path.join(args.out_dir, "data_info.json"), "w") as f:
|
| json.dump(metadata, f, indent=2)
|
|
|
| total_size = os.path.getsize(train_path) + os.path.getsize(val_path)
|
| print("\nFinished FineWeb-Edu data preparation")
|
| print(f" train tokens: {train_written:,} -> {train_path}")
|
| print(f" val tokens: {val_written:,} -> {val_path}")
|
| print(f" disk size: {total_size / 1024**3:.2f} GiB")
|
| print(f" meta: {meta_path}")
|
| print(f" info: {os.path.join(args.out_dir, 'data_info.json')}")
|
|
|
| if shutil.disk_usage(args.out_dir).free < 5 * 1024**3:
|
| print("WARNING: less than 5 GiB free space remains in the output directory.")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
