#!/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()