"""Shared web-corpus streaming + quality filtering for Phase-1 language data. Used by BOTH the offline sharder (scripts/prepare_pretrain.py) and the online streaming loader (cortex/stream_data.py), so the mixture/filters live in one place. Default open mixture (token-weighted by the consumers): FineWeb-Edu 40% (edu int_score >= 4, English, length/symbol heuristics) DCLM 60% (fastText quality-score knob; already top-decile, English) Nemotron-CC is gated: accept its HF license, set HF_TOKEN, and add a Source here. """ from __future__ import annotations import hashlib from dataclasses import dataclass from pathlib import Path @dataclass class Source: name: str hf_id: str weight: float config: str | None = None split: str = "train" text_field: str = "text" score_field: str | None = None # numeric quality score (higher = better) min_score: float = 0.0 en_field: str | None = None # field holding an English-language probability en_min: float = 0.0 text_fn: object = None # optional ex->str builder (structured sources, e.g. Wikipedia) def default_sources() -> list[Source]: return [ Source("fineweb-edu", "HuggingFaceFW/fineweb-edu", 0.40, config="sample-100BT", score_field="int_score", min_score=4, en_field="language_score", en_min=0.85), Source("dclm", "mlfoundations/dclm-baseline-1.0", 0.60, score_field="fasttext_openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train_prob", min_score=0.0), ] # ---- Phase-2 knowledge/math/reasoning mixture -------------------------------------- # FineMath-4+ and UltraData-Math (synthetic QA + textbook exercises) for math; Structured # Wikipedia for dense facts; Cosmopedia-v2 (synthetic textbooks) for reasoning -- promoted # from val into TRAIN here; FineWeb-Edu kept at 10% as a REPLAY buffer against forgetting. # (No code yet -- Phase 3, needs long context.) Token-weighted by the loader; the loader # also caps a source's weight once it loops (see StreamingMixLoader) so the small math sets # don't over-repeat. _WIKI_SKIP = {"url", "id", "type", "language", "wikidata_id", "version", "infoboxes", "links"} def _collect_strings(o) -> str: """Recursively collect all human-text strings from a nested structure, skipping metadata keys -- robust to the exact structured-wikipedia section schema.""" if isinstance(o, str): return o if len(o) > 1 else "" if isinstance(o, dict): return "\n".join(s for k, v in o.items() if k not in _WIKI_SKIP for s in [_collect_strings(v)] if s) if isinstance(o, list): return "\n".join(s for x in o for s in [_collect_strings(x)] if s) return "" def _wiki_text(ex) -> str: """structured-wikipedia (enwiki_namespace_0): title + abstract + flattened sections.""" title = ex.get("name") or "" body = _collect_strings({"abstract": ex.get("abstract"), "sections": ex.get("sections")}) return (title + "\n\n" + body).strip() def phase2_sources() -> list[Source]: return [ Source("finemath-4plus", "HuggingFaceTB/finemath", 0.20, config="finemath-4plus"), Source("ultradata-l3-qa", "openbmb/UltraData-Math", 0.12, config="UltraData-Math-L3-QA-Synthetic", text_field="content"), Source("ultradata-l3-textbook", "openbmb/UltraData-Math", 0.12, config="UltraData-Math-L3-Textbook-Exercise-Synthetic", text_field="content"), Source("structured-wiki", "wikimedia/structured-wikipedia", 0.23, config="enwiki_namespace_0", text_fn=_wiki_text), Source("cosmopedia-v2", "HuggingFaceTB/smollm-corpus", 0.23, config="cosmopedia-v2"), Source("fineweb-edu", "HuggingFaceFW/fineweb-edu", 0.10, config="sample-100BT", score_field="int_score", min_score=4, en_field="language_score", en_min=0.85), ] def cosmopedia_source() -> Source: """Cosmopedia v2 (synthetic textbooks/stories) -- used as an OOD VALIDATION set only; it is never part of the training mixture.""" return Source("cosmopedia-v2", "HuggingFaceTB/smollm-corpus", 1.0, config="cosmopedia-v2") def _gsm8k_text(ex) -> str: """GSM8K example -> question + chain-of-thought answer (keeps the '#### ' line).""" return (ex.get("question", "") + "\n" + ex.get("answer", "")).strip() def gsm8k_source() -> Source: """GSM8K grade-school math word problems (test split) -- a held-out math-REASONING benchmark, never in the training mixture, so it tracks Phase-2's math objective directly.""" return Source("gsm8k", "openai/gsm8k", 1.0, config="main", split="test", text_fn=_gsm8k_text) # ---- deterministic train/val holdout ------------------------------------------------- # A content hash splits each train source into disjoint train/val partitions, stable # across shuffles, resumes and sessions: the validation slice is the docs whose hash # falls in bucket 0; training skips exactly those. No ordering/skip alignment needed. def holdout_is_val(text: str, mod: int = 100) -> bool: h = int.from_bytes(hashlib.md5(text[:512].encode("utf-8", "ignore")).digest()[:8], "big") return h % mod == 0 MIN_CHARS, MAX_CHARS = 200, 200_000 def _en_ok(ex, src: Source) -> bool: if not src.en_field: return True v = ex.get(src.en_field) return v is not None and float(v) >= src.en_min def _text_ok(text: str) -> bool: n = len(text) if n < MIN_CHARS or n > MAX_CHARS: return False head = text[:2000] alpha = sum(c.isalpha() or c.isspace() for c in head) return alpha / max(len(head), 1) >= 0.60 # drop mostly-symbol/markup dumps def make_decontam(path: str | None): """Return a `text -> bool` membership test against benchmark test strings, or None.""" if not path: return None needles = [ln.strip().lower() for ln in Path(path).read_text(encoding="utf-8").splitlines() if len(ln.strip()) >= 32] print(f"[decontam] {len(needles)} benchmark needles loaded from {path}") def hit(text: str) -> bool: t = text.lower() return any(nd in t for nd in needles) return hit def stream_texts(src: Source, *, shuffle_buf: int = 0, seed: int = 0, decontam=None, loop: bool = False, stats: dict | None = None, holdout: str | None = None, holdout_mod: int = 100, resume_state: dict | None = None, on_ds=None): """Yield filtered text strings from one HF source. If loop, restart forever (reshuffling each pass) so a training stream never starves. If stats is given, increment stats['seen']/stats['kept'] for filter-rate visibility. holdout='train' skips the validation partition; holdout='val' yields ONLY it (see holdout_is_val) -- disjoint by construction. resume_state (an HF IterableDataset.state_dict) fast-forwards the stream to where a prior session stopped; on_ds(ds) is called with each freshly built dataset so the caller can snapshot ds.state_dict() for the next resume.""" from datasets import load_dataset p = 0 while True: kw = dict(streaming=True, split=src.split) if src.config: kw["name"] = src.config ds = load_dataset(src.hf_id, **kw) if shuffle_buf: ds = ds.shuffle(seed=seed + p, buffer_size=shuffle_buf) if resume_state is not None: ds.load_state_dict(resume_state) # HF-native fast resume (exact position) resume_state = None # only the first pass resumes if on_ds is not None: on_ds(ds) for ex in ds: if stats is not None: stats["seen"] = stats.get("seen", 0) + 1 text = (src.text_fn(ex) if src.text_fn else ex.get(src.text_field)) or "" if src.score_field is not None and float(ex.get(src.score_field, 0) or 0) < src.min_score: continue if not _en_ok(ex, src) or not _text_ok(text): continue if holdout == "train" and holdout_is_val(text, holdout_mod): continue # reserve the val partition if holdout == "val" and not holdout_is_val(text, holdout_mod): continue # only the val partition if decontam is not None and decontam(text): continue if stats is not None: stats["kept"] = stats.get("kept", 0) + 1 yield text.strip() if not loop: return p += 1