cortex/webdata.py

"""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 '#### <final>' 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