build_v2_blend.py

#!/usr/bin/env python3
"""v2 corpus builder: multi-source FineWeb-Edu-dominant blend -> decontam -> tokenize -> .bin.

S8/Latarnik blend (synthetic <= 25%, anti-collapse):
  FineWeb-Edu 58 | DCLM-baseline 15 | FineMath-4plus 5 | Cosmopedia-v2 18 | Nemotron-HQ-DQA 4
  (synthetic = Cosmopedia + Nemotron = 22%). Nemotron optional -> fold into Cosmopedia if unavailable.

Decontam:
  - 13-gram shingles (ALL sources) vs BLiMP/ARC/WikiText (decontam_index.json).
  - SYNTHETIC sources additionally use STRICTER 8-gram shingles (paraphrase-proxy: synthetic
    (Cosmopedia/Nemotron) paraphrases test-sets; short-shingle catches near-dupes the 13-gram misses).
    NOTE: this is a lightweight proxy for full perplexity-variance-decon (reference-model gate =
    follow-up before v2-train if Arek attaches a scoring model).

Per-source token budget = frac * --target-tokens. Streams each source until its budget, decontam+tokenize,
concatenates (eot-separated) + shuffles at doc-granularity boundaries. Writes train.bin/val.bin uint16 + meta.

Vocab 12288 (v1-consistent). --target-tokens ~12.5B fits pod-2 30GB disk (~26GB). Morning-upgrade to
20-25B needs a bigger volume; 32k re-tokenize is a separate Arek-approved step.
"""
import argparse
import hashlib
import json
import random
import re
from pathlib import Path

import numpy as np

# v2 blend recipe (REVISED per JugnuLM R3 honest-negative: diversity/DCLM DILUTES ARC; educational-
# distribution IS the ARC signal). DROP DCLM+FineMath (proven ARC-diluters). FWE-heavy + Cosmopedia
# (synthetic educational textbooks = ARC-aligned bet, untested by JugnuLM). Nemotron-sample too small
# (exhausts ~1.6M/shard) -> dropped. synthetic=22% (Cosmopedia), under 25% anti-collapse limit.
RECIPE = [
    {"name": "fineweb-edu", "dataset": "HuggingFaceFW/fineweb-edu", "config": "default",
     "split": "train", "content_field": "text", "frac": 0.78, "synthetic": False},
    {"name": "cosmopedia-v2", "dataset": "HuggingFaceTB/smollm-corpus", "config": "cosmopedia-v2",
     "split": "train", "content_field": "text", "frac": 0.22, "synthetic": True},
]

N_HARD = 13   # standard shingle (all sources)
N_SYNTH = 8   # stricter shingle (synthetic paraphrase-proxy)


def normalize(text):
    text = text.lower()
    text = re.sub(r"[^\w\s]", " ", text)      # match build_decontam_index.py (shingle-hash parity)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def shingles(text, n):
    words = normalize(text).split()
    return {hashlib.blake2b(" ".join(words[i:i + n]).encode("utf-8"), digest_size=8).hexdigest()
            for i in range(len(words) - n + 1)}


def contaminated(text, hard, blimp, synth_hard, synthetic):
    """13-gram drop-on-any (hard) + BLiMP fraction>0.5%. Synthetic: also 8-gram drop-on-any vs synth_hard."""
    sh = shingles(text, N_HARD)
    if sh & hard:
        return True
    if blimp:
        hit = len(sh & blimp)
        if sh and hit / len(sh) > 0.005:
            return True
    if synthetic and synth_hard:
        if shingles(text, N_SYNTH) & synth_hard:
            return True
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--target-tokens", type=int, default=12_500_000_000)
    ap.add_argument("--tokenizer", required=True)
    ap.add_argument("--decontam-index", required=True)
    ap.add_argument("--out-dir", default=".")
    ap.add_argument("--val-frac", type=float, default=0.005)
    ap.add_argument("--eot-id", type=int, default=None)
    ap.add_argument("--seed", type=int, default=1337)
    ap.add_argument("--drop-nemotron", action="store_true",
                    help="skip Nemotron (hard access) and reallocate its 4% to Cosmopedia (-> synth 22)")
    ap.add_argument("--sample-n", type=int, default=0, help="test: N docs/source, no write")
    ap.add_argument("--only", default=None, help="build only this source (parallel worker mode)")
    ap.add_argument("--shard-idx", type=int, default=0, help="worker shard index (0..shard-total-1)")
    ap.add_argument("--shard-total", type=int, default=1, help="total shards for --only source (parallelism)")
    a = ap.parse_args()

    from datasets import load_dataset
    from tokenizers import Tokenizer

    tok = Tokenizer.from_file(a.tokenizer)
    vocab = tok.get_vocab_size()
    eos_id = tok.token_to_id("<|endoftext|>")
    if eos_id is None:
        raise SystemExit("tokenizer nie ma <|endoftext|>; podaj --eot-id jawnie")
    eot = a.eot_id if a.eot_id is not None else eos_id
    if eot != eos_id:
        print(f"WARN: --eot-id={eot} rozni sie od <|endoftext|>={eos_id}", flush=True)
    idx = json.loads(Path(a.decontam_index).read_text())
    missing = [k for k in ("hard_hashes", "blimp_hashes", "synth_hard_hashes") if k not in idx]
    if missing:
        # Bramka Harta/Wartownika 2026-09-25: brak klucza = stop, nigdy cichy fallback na 13-gram.
        raise SystemExit(f"decontam index {a.decontam_index} nie ma kluczy: {missing}")
    hard = set(idx["hard_hashes"]); blimp = set(idx["blimp_hashes"])
    synth_hard = set(idx["synth_hard_hashes"])
    print(f"tokenizer vocab={vocab} eot={eot} | decontam hard={len(hard):,} blimp={len(blimp):,} "
          f"synth_hard={len(synth_hard):,}", flush=True)

    recipe = [dict(s) for s in RECIPE]
    if a.drop_nemotron:
        nemo = next(s for s in recipe if s["name"] == "nemotron-dqa")
        cosmo = next(s for s in recipe if s["name"] == "cosmopedia-v2")
        cosmo["frac"] += nemo["frac"]
        recipe = [s for s in recipe if s["name"] != "nemotron-dqa"]
        print(f"drop-nemotron: cosmopedia-v2 frac -> {cosmo['frac']:.2f}", flush=True)
    if a.only:
        recipe = [s for s in recipe if s["name"] == a.only]
        if not recipe:
            print(f"!! --only {a.only}: no such source", flush=True); return

    outd = Path(a.out_dir); outd.mkdir(parents=True, exist_ok=True)
    n_val_target = 0 if (a.sample_n or a.shard_total > 1) else int(a.target_tokens * a.val_frac)
    train_fh = None if a.sample_n else open(outd / "train.bin", "wb")
    val_fh = None if a.sample_n else open(outd / "val.bin", "wb")
    CHUNK = 20_000_000  # flush ~20M tok (40MB) -> bounded RAM, incremental disk-write (12.5B won't fit RAM)
    state = {"chunk": [], "train": 0, "val": 0}

    def flush():
        if a.sample_n or not state["chunk"]:
            state["chunk"] = []
            return
        arr = np.array(state["chunk"], dtype=np.uint16)
        state["chunk"] = []
        room = n_val_target - state["val"]
        if room > 0:
            take = min(room, len(arr))
            arr[:take].tofile(val_fh); state["val"] += take
            arr = arr[take:]
        if len(arr):
            arr.tofile(train_fh); state["train"] += len(arr)

    report = {}
    for s in recipe:
        budget = int(s["frac"] * a.target_tokens) // a.shard_total
        n_tok = n_bytes = kept = drop_dec = seen = 0
        try:
            ds = (load_dataset(s["dataset"], s["config"], split=s["split"], streaming=True)
                  if s["config"] else load_dataset(s["dataset"], split=s["split"], streaming=True))
            if a.shard_total > 1:
                ds = ds.shard(num_shards=a.shard_total, index=a.shard_idx)
        except Exception as e:
            print(f"!! {s['name']} load failed: {e!r} -- SKIP", flush=True)
            report[s["name"]] = {"status": "load-failed", "error": repr(e)[:200]}
            continue
        for ex in ds:
            seen += 1
            content = ex.get(s["content_field"]) or ""
            if contaminated(content, hard, blimp, synth_hard, s["synthetic"]):
                drop_dec += 1
            else:
                ids = [i for i in tok.encode(content).ids if i < vocab]
                if ids:
                    ids.append(eot)
                    state["chunk"].extend(ids)
                    n_tok += len(ids); n_bytes += len(content.encode("utf-8")); kept += 1
                    if len(state["chunk"]) >= CHUNK:
                        flush()
            if a.sample_n and seen >= a.sample_n:
                break
            if not a.sample_n and n_tok >= budget:
                break
            if seen % 20000 == 0:
                print(f"  [{s['name']}] seen={seen:,} kept={kept:,} tok={n_tok:,}/{budget:,} drop_dec={drop_dec:,}", flush=True)
        ratio = n_tok / n_bytes if n_bytes else 0
        report[s["name"]] = {"seen": seen, "kept": kept, "n_tok": n_tok, "n_bytes": n_bytes,
                             "tok_per_byte": round(ratio, 4), "drop_decontam": drop_dec,
                             "budget": budget, "synthetic": s["synthetic"]}
        print(f"=== {s['name']} DONE kept={kept:,} n_tok={n_tok:,} tok/byte={ratio:.4f} drop_dec={drop_dec:,} ===", flush=True)

    flush()
    synth_tok = sum(report[s["name"]]["n_tok"] for s in recipe
                    if s["synthetic"] and report.get(s["name"], {}).get("n_tok"))
    total_written = state["train"] + state["val"]
    print(f"=== BLEND total_tok={total_written:,} synthetic={synth_tok:,} "
          f"({100*synth_tok/max(total_written,1):.1f}% -- limit 25) ===", flush=True)
    if a.sample_n:
        print("SAMPLE mode: no write."); return
    train_fh.close(); val_fh.close()
    meta = {"blend": report, "total_tok": total_written, "synthetic_tok": synth_tok,
            "synthetic_frac": round(synth_tok / max(total_written, 1), 4),
            "vocab": vocab, "eot": eot, "train_tok": state["train"], "val_tok": state["val"], "seed": a.seed,
            "decontam": {"index": str(a.decontam_index), "hard": len(hard), "blimp": len(blimp),
                         "synth_hard": len(synth_hard)}}
    (outd / "mix_meta.json").write_text(json.dumps(meta, indent=2))
    print(f"WROTE {outd}/train.bin ({state['train']:,} tok) + val.bin ({state['val']:,}) + mix_meta.json", flush=True)


if __name__ == "__main__":
    main()