# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Stream any mix of Hugging Face pretraining datasets into Forge's
train.bin / val.bin format (uint16 tokens, llm.c-style header — same as
prepare_data.py, whose tokenizer + parallel encoder this reuses).

Built for the sub-1B recipes in SMALL_MODELS_RESEARCH.md: classifier-filtered
web (FineWeb-Edu, DCLM), synthetic textbooks (Cosmopedia), math
(FineMath, OpenWebMath), reference text (Wikipedia), instructions (SmolTalk),
and weighted mixtures of them (e.g. SmolLM2's 60/40 FineWeb-Edu/DCLM).
Everything is STREAMED — no full-corpus downloads; you pay only for the
megabytes you keep.

Usage:
  # single source
  python3 tools/prepare_hf_data.py --source fineweb-edu --out data/fwe \
      --max-train-mb 500 --vocab-size 16384

  # weighted mixture (name:weight, comma-separated)
  python3 tools/prepare_hf_data.py --mix fineweb-edu:0.6,dclm:0.4 \
      --out data/web-mix --max-train-mb 2000

  # research-backed presets
  python3 tools/prepare_hf_data.py --preset smollm-web --out data/smollm-web
  python3 tools/prepare_hf_data.py --list          # show sources & presets

Requires: pip install datasets   (numpy already required by prepare_data.py)
"""
import argparse
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import prepare_data as pd  # tokenizer training + parallel BPE encoder + .bin writer

SEP = "<|endoftext|>"  # document separator; same convention as TinyStories


def messages_to_text(row):
    """Flatten a chat-format row (SmolTalk-style `messages`) to plain text."""
    parts = []
    for m in row.get("messages", []):
        role = m.get("role", "user")
        parts.append(f"<|{role}|>\n{m.get('content', '')}")
    return "\n".join(parts)


# name -> (repo, config, split, extractor). All public, all stream-capable.
SOURCES = {
    # -- narrow-domain / tiny-model corpora ------------------------------------
    "tinystories":      ("roneneldan/TinyStories", None, "train",
                         lambda r: r["text"]),
    # -- classifier-filtered web (the highest-leverage curation known) ---------
    "fineweb-edu":      ("HuggingFaceFW/fineweb-edu", "sample-10BT", "train",
                         lambda r: r["text"]),
    "fineweb":          ("HuggingFaceFW/fineweb", "sample-10BT", "train",
                         lambda r: r["text"]),
    "fineweb-edu-dedup":("HuggingFaceTB/smollm-corpus", "fineweb-edu-dedup", "train",
                         lambda r: r["text"]),
    "dclm":             ("mlfoundations/dclm-baseline-1.0", None, "train",
                         lambda r: r["text"]),
    # -- synthetic textbooks (phi/Cosmopedia recipe) ----------------------------
    "cosmopedia":       ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", "train",
                         lambda r: r["text"]),
    # -- math (upsample during LR decay per SmolLM2/3) --------------------------
    "finemath":         ("HuggingFaceTB/finemath", "finemath-4plus", "train",
                         lambda r: r["text"]),
    "openwebmath":      ("open-web-math/open-web-math", None, "train",
                         lambda r: r["text"]),
    # -- reference / encyclopedic ----------------------------------------------
    "wikipedia-en":     ("wikimedia/wikipedia", "20231101.en", "train",
                         lambda r: r["text"]),
    "wikipedia-fr":     ("wikimedia/wikipedia", "20231101.fr", "train",
                         lambda r: r["text"]),
    # -- generic web baseline ---------------------------------------------------
    "c4":               ("allenai/c4", "en", "train",
                         lambda r: r["text"]),
    # -- instructions (SFT-style text; useful for a decay-phase blend) ----------
    "smoltalk":         ("HuggingFaceTB/smoltalk", "all", "train",
                         messages_to_text),
}

# Research-backed mixtures (SMALL_MODELS_RESEARCH.md §1). Weights are document
# sampling probabilities over the interleaved stream.
PRESETS = {
    # SmolLM2's best web mix: 60% FineWeb-Edu + 40% DCLM
    "smollm-web":   {"fineweb-edu": 0.6, "dclm": 0.4},
    # phi-style: filtered web + synthetic textbooks
    "textbooks":    {"fineweb-edu": 0.6, "cosmopedia": 0.4},
    # single-stage high-quality blend for <500M models (uniform, no staging)
    "smol-full":    {"fineweb-edu": 0.55, "cosmopedia": 0.25,
                     "finemath": 0.10, "openwebmath": 0.10},
    # LR-decay-phase blend: premium math/instructions upsampled
    "decay-anneal": {"fineweb-edu": 0.4, "finemath": 0.3, "smoltalk": 0.3},
    # bilingual reference blend
    "wiki-bilingue":{"wikipedia-en": 0.5, "wikipedia-fr": 0.5},
}


def open_stream(name, seed, shuffle_buffer):
    try:
        from datasets import load_dataset
    except ImportError:
        sys.exit("prepare_hf_data: pip install datasets")
    repo, config, split, extract = SOURCES[name]
    ds = load_dataset(repo, config, split=split, streaming=True)
    if shuffle_buffer > 0:
        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
    it = iter(ds)

    def docs():
        for row in it:
            text = extract(row)
            if text and len(text) >= 64:  # drop empty/near-empty docs
                yield text

    return docs()


def stream_mixture(weights, seed, shuffle_buffer):
    """Yield documents, choosing the source of each by weight. A source that
    runs dry is dropped and the rest renormalize (streamed sets are huge —
    this mostly matters for small ones like tinystories)."""
    rng = random.Random(seed)
    streams = {n: open_stream(n, seed + i, shuffle_buffer)
               for i, n in enumerate(weights)}
    names = list(weights)
    while names:
        r = rng.random() * sum(weights[n] for n in names)
        acc = 0.0
        pick = names[-1]
        for n in names:
            acc += weights[n]
            if r <= acc:
                pick = n
                break
        try:
            yield next(streams[pick])
        except StopIteration:
            print(f"  source exhausted: {pick}")
            names.remove(pick)


def write_corpus(docs, train_path, val_path, train_mb, val_mb, val_every):
    """Split the stream into train/val text files (every Nth doc to val until
    its budget fills), separated by <|endoftext|>."""
    train_budget = int(train_mb * 1024 * 1024)
    val_budget = int(val_mb * 1024 * 1024)
    train_n = val_n = 0
    report_every = max(1, train_budget // (20 * 1024 * 1024) or 1) * 10_000_000
    next_report = report_every
    with open(train_path, "w") as ftr, open(val_path, "w") as fva:
        for i, doc in enumerate(docs):
            piece = doc.rstrip() + "\n" + SEP + "\n"
            if val_n < val_budget and i % val_every == val_every - 1:
                fva.write(piece)
                val_n += len(piece)
            else:
                ftr.write(piece)
                train_n += len(piece)
            if train_n >= next_report:
                print(f"  {train_n / 1e6:.0f} MB train / {val_n / 1e6:.1f} MB val")
                next_report += report_every
            if train_n >= train_budget and val_n >= val_budget:
                break
    print(f"corpus: {train_n / 1e6:.1f} MB train, {val_n / 1e6:.1f} MB val")


def main():
    ap = argparse.ArgumentParser(
        description="Stream HF datasets (single, mix, or preset) into Forge .bin files")
    g = ap.add_mutually_exclusive_group()
    g.add_argument("--source", choices=sorted(SOURCES), help="single dataset")
    g.add_argument("--mix", help="name:weight,name:weight,…")
    g.add_argument("--preset", choices=sorted(PRESETS), help="research-backed mixture")
    g.add_argument("--list", action="store_true", help="list sources & presets")
    ap.add_argument("--out", help="output dir (train.bin/val.bin/corpus)")
    ap.add_argument("--vocab-size", type=int, default=16384)
    ap.add_argument("--tokenizer", default=None,
                    help="existing forgebpe .model to reuse (e.g. data/tinystories/"
                         "tok4096.model); otherwise trained on this corpus")
    ap.add_argument("--tokenizer-mb", type=float, default=20.0)
    ap.add_argument("--max-train-mb", type=float, default=500.0)
    ap.add_argument("--max-val-mb", type=float, default=5.0)
    ap.add_argument("--val-every", type=int, default=200,
                    help="route every Nth document to val until its budget fills")
    ap.add_argument("--seed", type=int, default=1337)
    ap.add_argument("--shuffle-buffer", type=int, default=10000,
                    help="streaming shuffle buffer per source (0 disables)")
    ap.add_argument("--keep-text", action="store_true",
                    help="keep the intermediate train.txt/val.txt corpus files")
    args = ap.parse_args()

    if args.list or not (args.source or args.mix or args.preset):
        print("sources:")
        for n, (repo, cfg, _, _) in sorted(SOURCES.items()):
            print(f"  {n:<18} {repo}" + (f" [{cfg}]" if cfg else ""))
        print("presets:")
        for n, w in PRESETS.items():
            print(f"  {n:<18} " + ", ".join(f"{k}:{v}" for k, v in w.items()))
        return
    if not args.out:
        ap.error("--out is required")

    if args.source:
        weights = {args.source: 1.0}
    elif args.preset:
        weights = PRESETS[args.preset]
    else:
        weights = {}
        for part in args.mix.split(","):
            name, _, w = part.partition(":")
            if name not in SOURCES:
                ap.error(f"unknown source '{name}' (see --list)")
            weights[name] = float(w) if w else 1.0

    os.makedirs(args.out, exist_ok=True)
    print("mixture: " + ", ".join(f"{n}:{w}" for n, w in weights.items()))

    train_txt = os.path.join(args.out, "train.txt")
    val_txt = os.path.join(args.out, "val.txt")
    docs = stream_mixture(weights, args.seed, args.shuffle_buffer)
    write_corpus(docs, train_txt, val_txt, args.max_train_mb, args.max_val_mb,
                 args.val_every)

    model = args.tokenizer or os.path.join(args.out, f"tok{args.vocab_size}.model")
    if not os.path.exists(model):
        import train_tokenizer as tt
        with open(train_txt, "rb") as f:
            sample = f.read(int(args.tokenizer_mb * 1024 * 1024))
        print(f"training {args.vocab_size}-vocab BPE on {len(sample) / 1e6:.1f} MB")
        tt.write_model(model, tt.train(sample, args.vocab_size))
        print(f"wrote {model}")

    pd.tokenize_file(train_txt, model, os.path.join(args.out, "train.bin"),
                     args.max_train_mb + 1)
    pd.tokenize_file(val_txt, model, os.path.join(args.out, "val.bin"),
                     args.max_val_mb + 1)

    if not args.keep_text:
        os.remove(train_txt)
        os.remove(val_txt)
    print("done — point `forge train --data` at", args.out,
          f"(set model vocab_size to match the tokenizer: {model})")


if __name__ == "__main__":
    main()