fetch_data.py

"""
Fetch a slice of curated, license-clean OPEN datasets (real web prose + GitHub code) as extra
distillation input text. Streamed (we don't download whole datasets), tagged by domain, saved
as jsonl shards in data/ alongside the Haiku data.

These datasets are made to be trained on — clean alternative to scraping raw GitHub/web.
Tries several candidates per domain and uses whichever loads (dataset availability shifts).
"""
import json
import os
import time

from datasets import load_dataset

os.makedirs("data", exist_ok=True)
N = 5000  # examples per domain (scaling-sweep dataset; gentler on unauthenticated HF rate limits)

# domain -> list of (dataset, config, split, text_field) candidates, tried in order.
# All license-clean, made-to-train-on open datasets (wiki / GitHub / OpenAssistant / Dolly).
CANDIDATES = {
    "web": [
        ("Salesforce/wikitext", "wikitext-103-raw-v1", "train", "text"),
        ("wikitext", "wikitext-103-raw-v1", "train", "text"),
        ("allenai/c4", "en", "train", "text"),
    ],
    "code": [
        ("codeparrot/codeparrot-clean", None, "train", "content"),       # full, large, parquet
        ("codeparrot/codeparrot-clean-valid", None, "train", "content"),
        ("angie-chen55/python-github-code", None, "train", "text"),
    ],
    # real human conversation (Apache-2.0); fall back to daily_dialog / ultrachat
    "conversation": [
        ("OpenAssistant/oasst1", None, "train", "text"),
        ("daily_dialog", None, "train", "dialog"),
        ("HuggingFaceH4/ultrachat_200k", None, "train_sft", "messages"),
    ],
    # real human-written explanatory answers (CC-BY-SA); fall back to sciq
    "explanation": [
        ("databricks/databricks-dolly-15k", None, "train", "response"),
        ("sciq", None, "train", "support"),
        ("Salesforce/wikitext", "wikitext-103-raw-v1", "train", "text"),
    ],
}


def extract(val):
    """Field values vary: a plain string, a list of utterances, or a list of chat dicts.
    Normalize any of those to a single text string."""
    if isinstance(val, str):
        return val
    if isinstance(val, list):
        parts = []
        for x in val:
            if isinstance(x, str):
                parts.append(x)
            elif isinstance(x, dict):
                parts.append(x.get("content") or x.get("text") or "")
        return " ".join(p for p in parts if p)
    return ""


def clean(s):
    s = extract(s).strip()
    return s if 60 <= len(s) <= 2000 else None


def existing(out):
    return sum(1 for _ in open(out)) if os.path.exists(out) else 0


def pull(name, cfg, split, field, out, domain, retries=3):
    """Stream into a TEMP file; only replace `out` on success. Retry on transient HF errors
    (unauthenticated streaming gets rate-limited). Returns n written, or 0."""
    tmp = out + ".tmp"
    for attempt in range(retries):
        try:
            ds = load_dataset(name, cfg, split=split, streaming=True)
            n = 0
            with open(tmp, "w") as fh:
                for ex in ds:
                    t = clean(ex.get(field))
                    if not t:
                        continue
                    fh.write(json.dumps({"domain": domain, "text": t}, ensure_ascii=False) + "\n")
                    n += 1
                    if n >= N:
                        break
            if n > 0:
                os.replace(tmp, out)            # atomic; never clobbers `out` on failure
                return n
            os.path.exists(tmp) and os.remove(tmp)
            return 0
        except Exception as e:
            os.path.exists(tmp) and os.remove(tmp)
            wait = 5 * (attempt + 1)
            print(f"{domain}: {name} attempt {attempt+1} FAILED ({type(e).__name__}: {str(e)[:90]}) "
                  f"{'retrying in %ds' % wait if attempt < retries-1 else 'giving up'}", flush=True)
            if attempt < retries - 1:
                time.sleep(wait)
    return 0


for domain, cands in CANDIDATES.items():
    out = f"data/{domain}-real.jsonl"
    have = existing(out)
    if have >= N:
        print(f"{domain}: already have {have} >= {N}, skipping (resume)", flush=True)
        continue
    done = False
    for name, cfg, split, field in cands:
        n = pull(name, cfg, split, field, out, domain)
        if n > 0:
            print(f"{domain}: {name} -> {n} examples -> {out}", flush=True)
            done = True
            break
        print(f"{domain}: {name} yielded 0 usable, trying next candidate...", flush=True)
    if not done:
        print(f"{domain}: ALL candidates failed — keeping existing {have} examples", flush=True)