spb/forge Public MIT

1# Author: Simon-Pierre Boucher — contact@spboucher.ai2"""Stream any mix of Hugging Face pretraining datasets into Forge's3train.bin / val.bin format (uint16 tokens, llm.c-style header — same as4prepare_data.py, whose tokenizer + parallel encoder this reuses).56Built for the sub-1B recipes in SMALL_MODELS_RESEARCH.md: classifier-filtered7web (FineWeb-Edu, DCLM), synthetic textbooks (Cosmopedia), math8(FineMath, OpenWebMath), reference text (Wikipedia), instructions (SmolTalk),9and weighted mixtures of them (e.g. SmolLM2's 60/40 FineWeb-Edu/DCLM).10Everything is STREAMED — no full-corpus downloads; you pay only for the11megabytes you keep.1213Usage:14  # single source15  python3 tools/prepare_hf_data.py --source fineweb-edu --out data/fwe \16      --max-train-mb 500 --vocab-size 163841718  # weighted mixture (name:weight, comma-separated)19  python3 tools/prepare_hf_data.py --mix fineweb-edu:0.6,dclm:0.4 \20      --out data/web-mix --max-train-mb 20002122  # research-backed presets23  python3 tools/prepare_hf_data.py --preset smollm-web --out data/smollm-web24  python3 tools/prepare_hf_data.py --list          # show sources & presets2526Requires: pip install datasets   (numpy already required by prepare_data.py)27"""28import argparse29import os30import random31import sys3233sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))34import prepare_data as pd  # tokenizer training + parallel BPE encoder + .bin writer3536SEP = "<|endoftext|>"  # document separator; same convention as TinyStories373839def messages_to_text(row):40    """Flatten a chat-format row (SmolTalk-style `messages`) to plain text."""41    parts = []42    for m in row.get("messages", []):43        role = m.get("role", "user")44        parts.append(f"<|{role}|>\n{m.get('content', '')}")45    return "\n".join(parts)464748# name -> (repo, config, split, extractor). All public, all stream-capable.49SOURCES = {50    # -- narrow-domain / tiny-model corpora ------------------------------------51    "tinystories":      ("roneneldan/TinyStories", None, "train",52                         lambda r: r["text"]),53    # -- classifier-filtered web (the highest-leverage curation known) ---------54    "fineweb-edu":      ("HuggingFaceFW/fineweb-edu", "sample-10BT", "train",55                         lambda r: r["text"]),56    "fineweb":          ("HuggingFaceFW/fineweb", "sample-10BT", "train",57                         lambda r: r["text"]),58    "fineweb-edu-dedup":("HuggingFaceTB/smollm-corpus", "fineweb-edu-dedup", "train",59                         lambda r: r["text"]),60    "dclm":             ("mlfoundations/dclm-baseline-1.0", None, "train",61                         lambda r: r["text"]),62    # -- synthetic textbooks (phi/Cosmopedia recipe) ----------------------------63    "cosmopedia":       ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", "train",64                         lambda r: r["text"]),65    # -- math (upsample during LR decay per SmolLM2/3) --------------------------66    "finemath":         ("HuggingFaceTB/finemath", "finemath-4plus", "train",67                         lambda r: r["text"]),68    "openwebmath":      ("open-web-math/open-web-math", None, "train",69                         lambda r: r["text"]),70    # -- reference / encyclopedic ----------------------------------------------71    "wikipedia-en":     ("wikimedia/wikipedia", "20231101.en", "train",72                         lambda r: r["text"]),73    "wikipedia-fr":     ("wikimedia/wikipedia", "20231101.fr", "train",74                         lambda r: r["text"]),75    # -- generic web baseline ---------------------------------------------------76    "c4":               ("allenai/c4", "en", "train",77                         lambda r: r["text"]),78    # -- instructions (SFT-style text; useful for a decay-phase blend) ----------79    "smoltalk":         ("HuggingFaceTB/smoltalk", "all", "train",80                         messages_to_text),81}8283# Research-backed mixtures (SMALL_MODELS_RESEARCH.md §1). Weights are document84# sampling probabilities over the interleaved stream.85PRESETS = {86    # SmolLM2's best web mix: 60% FineWeb-Edu + 40% DCLM87    "smollm-web":   {"fineweb-edu": 0.6, "dclm": 0.4},88    # phi-style: filtered web + synthetic textbooks89    "textbooks":    {"fineweb-edu": 0.6, "cosmopedia": 0.4},90    # single-stage high-quality blend for <500M models (uniform, no staging)91    "smol-full":    {"fineweb-edu": 0.55, "cosmopedia": 0.25,92                     "finemath": 0.10, "openwebmath": 0.10},93    # LR-decay-phase blend: premium math/instructions upsampled94    "decay-anneal": {"fineweb-edu": 0.4, "finemath": 0.3, "smoltalk": 0.3},95    # bilingual reference blend96    "wiki-bilingue":{"wikipedia-en": 0.5, "wikipedia-fr": 0.5},97}9899100def open_stream(name, seed, shuffle_buffer):101    try:102        from datasets import load_dataset103    except ImportError:104        sys.exit("prepare_hf_data: pip install datasets")105    repo, config, split, extract = SOURCES[name]106    ds = load_dataset(repo, config, split=split, streaming=True)107    if shuffle_buffer > 0:108        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)109    it = iter(ds)110111    def docs():112        for row in it:113            text = extract(row)114            if text and len(text) >= 64:  # drop empty/near-empty docs115                yield text116117    return docs()118119120def stream_mixture(weights, seed, shuffle_buffer):121    """Yield documents, choosing the source of each by weight. A source that122    runs dry is dropped and the rest renormalize (streamed sets are huge —123    this mostly matters for small ones like tinystories)."""124    rng = random.Random(seed)125    streams = {n: open_stream(n, seed + i, shuffle_buffer)126               for i, n in enumerate(weights)}127    names = list(weights)128    while names:129        r = rng.random() * sum(weights[n] for n in names)130        acc = 0.0131        pick = names[-1]132        for n in names:133            acc += weights[n]134            if r <= acc:135                pick = n136                break137        try:138            yield next(streams[pick])139        except StopIteration:140            print(f"  source exhausted: {pick}")141            names.remove(pick)142143144def write_corpus(docs, train_path, val_path, train_mb, val_mb, val_every):145    """Split the stream into train/val text files (every Nth doc to val until146    its budget fills), separated by <|endoftext|>."""147    train_budget = int(train_mb * 1024 * 1024)148    val_budget = int(val_mb * 1024 * 1024)149    train_n = val_n = 0150    report_every = max(1, train_budget // (20 * 1024 * 1024) or 1) * 10_000_000151    next_report = report_every152    with open(train_path, "w") as ftr, open(val_path, "w") as fva:153        for i, doc in enumerate(docs):154            piece = doc.rstrip() + "\n" + SEP + "\n"155            if val_n < val_budget and i % val_every == val_every - 1:156                fva.write(piece)157                val_n += len(piece)158            else:159                ftr.write(piece)160                train_n += len(piece)161            if train_n >= next_report:162                print(f"  {train_n / 1e6:.0f} MB train / {val_n / 1e6:.1f} MB val")163                next_report += report_every164            if train_n >= train_budget and val_n >= val_budget:165                break166    print(f"corpus: {train_n / 1e6:.1f} MB train, {val_n / 1e6:.1f} MB val")167168169def main():170    ap = argparse.ArgumentParser(171        description="Stream HF datasets (single, mix, or preset) into Forge .bin files")172    g = ap.add_mutually_exclusive_group()173    g.add_argument("--source", choices=sorted(SOURCES), help="single dataset")174    g.add_argument("--mix", help="name:weight,name:weight,…")175    g.add_argument("--preset", choices=sorted(PRESETS), help="research-backed mixture")176    g.add_argument("--list", action="store_true", help="list sources & presets")177    ap.add_argument("--out", help="output dir (train.bin/val.bin/corpus)")178    ap.add_argument("--vocab-size", type=int, default=16384)179    ap.add_argument("--tokenizer", default=None,180                    help="existing forgebpe .model to reuse (e.g. data/tinystories/"181                         "tok4096.model); otherwise trained on this corpus")182    ap.add_argument("--tokenizer-mb", type=float, default=20.0)183    ap.add_argument("--max-train-mb", type=float, default=500.0)184    ap.add_argument("--max-val-mb", type=float, default=5.0)185    ap.add_argument("--val-every", type=int, default=200,186                    help="route every Nth document to val until its budget fills")187    ap.add_argument("--seed", type=int, default=1337)188    ap.add_argument("--shuffle-buffer", type=int, default=10000,189                    help="streaming shuffle buffer per source (0 disables)")190    ap.add_argument("--keep-text", action="store_true",191                    help="keep the intermediate train.txt/val.txt corpus files")192    args = ap.parse_args()193194    if args.list or not (args.source or args.mix or args.preset):195        print("sources:")196        for n, (repo, cfg, _, _) in sorted(SOURCES.items()):197            print(f"  {n:<18} {repo}" + (f" [{cfg}]" if cfg else ""))198        print("presets:")199        for n, w in PRESETS.items():200            print(f"  {n:<18} " + ", ".join(f"{k}:{v}" for k, v in w.items()))201        return202    if not args.out:203        ap.error("--out is required")204205    if args.source:206        weights = {args.source: 1.0}207    elif args.preset:208        weights = PRESETS[args.preset]209    else:210        weights = {}211        for part in args.mix.split(","):212            name, _, w = part.partition(":")213            if name not in SOURCES:214                ap.error(f"unknown source '{name}' (see --list)")215            weights[name] = float(w) if w else 1.0216217    os.makedirs(args.out, exist_ok=True)218    print("mixture: " + ", ".join(f"{n}:{w}" for n, w in weights.items()))219220    train_txt = os.path.join(args.out, "train.txt")221    val_txt = os.path.join(args.out, "val.txt")222    docs = stream_mixture(weights, args.seed, args.shuffle_buffer)223    write_corpus(docs, train_txt, val_txt, args.max_train_mb, args.max_val_mb,224                 args.val_every)225226    model = args.tokenizer or os.path.join(args.out, f"tok{args.vocab_size}.model")227    if not os.path.exists(model):228        import train_tokenizer as tt229        with open(train_txt, "rb") as f:230            sample = f.read(int(args.tokenizer_mb * 1024 * 1024))231        print(f"training {args.vocab_size}-vocab BPE on {len(sample) / 1e6:.1f} MB")232        tt.write_model(model, tt.train(sample, args.vocab_size))233        print(f"wrote {model}")234235    pd.tokenize_file(train_txt, model, os.path.join(args.out, "train.bin"),236                     args.max_train_mb + 1)237    pd.tokenize_file(val_txt, model, os.path.join(args.out, "val.bin"),238                     args.max_val_mb + 1)239240    if not args.keep_text:241        os.remove(train_txt)242        os.remove(val_txt)243    print("done — point `forge train --data` at", args.out,244          f"(set model vocab_size to match the tokenizer: {model})")245246247if __name__ == "__main__":248    main()249