Upload src/xscript/data/fineweb.py with huggingface_hub

src/xscript/data/fineweb.py

@@ -0,0 +1,326 @@
+"""Model-training text pools from FineWeb-HQ (EN) / FineWeb2-HQ (DE/FR/AR/ZH).
+
+Both datasets come from the same lab and the same XLM-R-embedding quality
+classifier (top-10%% selection), keeping filtering methodology constant across
+all five languages -- with two documented exceptions: FineWeb2-HQ's arb_Arab
+split is only ~29GB total (85 parquet files, confirmed exhausted 2026-07-13),
+and fra_Latn -HQ fully exhausted at 124.1GB (2026-07-14), both far short of
+what our token budgets need. `FALLBACK_SOURCES` lets a language keep pulling
+from a second, less-strict source (FineWeb2's own dedup+quality-filtered
+"train" split, one filter step short of the top-decile HQ cut -- still same
+lab/pipeline) once its primary source is exhausted, rather than epoch heavily
+over a small pool. This is a real, deliberate deviation from strict
+filter-parity for these two languages; note it in the thesis.
+
+We stream parquet with column pruning (only `text`), so the large `embeddings`
+column in FineWeb2-HQ is never downloaded. Output pools are zstd jsonl shards
+of ~1GB uncompressed text. The first parquet file of the primary source is
+reserved exclusively for the in-domain eval holdout (never enters the pool).
+"""
+import json
+from pathlib import Path
+
+from ..langs import LANGS
+from ..paths import MANIFEST_CACHE, POOLS, HOLDOUT, pool_dir, ensure
+
+POOL_SHARD_BYTES = 1 << 30      # uncompressed text per pool shard
+HOLDOUT_BYTES = 30 * (1 << 20)  # 30MB per language
+
+FALLBACK_SOURCES: dict[str, tuple[str, str]] = {
+    "ar": ("HuggingFaceFW/fineweb-2", "data/arb_Arab/train"),
+    # fra_Latn -HQ (epfml/FineWeb2-HQ) fully exhausted 2026-07-14 at 124.1GB,
+    # short of the 152.2GB budget -- same fallback pattern as ar.
+    "fr": ("HuggingFaceFW/fineweb-2", "data/fra_Latn/train"),
+}
+
+
+def _sources_for(lang: str) -> list[tuple[str, str]]:
+    """[(repo, subdir), ...] in priority order: primary (quality-filtered -HQ)
+    first, then any FALLBACK_SOURCES entry once the primary is exhausted."""
+    L = LANGS[lang]
+    srcs = [(L.fineweb_repo, L.fineweb_subdir)]
+    if lang in FALLBACK_SOURCES:
+        srcs.append(FALLBACK_SOURCES[lang])
+    return srcs
+
+
+def _list_parquets(repo: str, subdir: str) -> list[str]:
+    """Manifest of parquet files under repo/subdir, round-robin across CC dumps
+    for raw FineWeb-HQ (EN)."""
+    from huggingface_hub import HfApi
+    cache = ensure(MANIFEST_CACHE / "fineweb_hq_pools") / \
+        f"{repo.replace('/', '__')}__{subdir.replace('/', '_')}.json"
+    if cache.exists():
+        return json.loads(cache.read_text())
+    api = HfApi()
+    files = [e.path for e in api.list_repo_tree(repo, subdir,
+                                                repo_type="dataset", recursive=True)
+             if e.__class__.__name__ == "RepoFile" and e.path.endswith(".parquet")]
+    if repo == "epfml/FineWeb-HQ":
+        # data/CC-MAIN-YYYY-WW/000_xxxxx.parquet: interleave dumps so the pool
+        # spans the full 2013-2024 crawl range like FineWeb2-HQ does.
+        by_dump: dict[str, list[str]] = {}
+        for f in sorted(files):
+            by_dump.setdefault(f.split("/")[1], []).append(f)
+        out = []
+        for i in range(max(len(v) for v in by_dump.values())):
+            for d in sorted(by_dump):
+                if i < len(by_dump[d]):
+                    out.append(by_dump[d][i])
+        files = out
+    else:
+        files = sorted(files)
+    cache.write_text(json.dumps(files))
+    return files
+
+
+def _iter_texts(repo: str, path_in_repo: str):
+    import pyarrow.parquet as pq
+    from huggingface_hub import HfFileSystem
+    fs = HfFileSystem()
+    with fs.open(f"datasets/{repo}/{path_in_repo}", "rb") as f:
+        pf = pq.ParquetFile(f)
+        for rg in range(pf.num_row_groups):
+            tbl = pf.read_row_group(rg, columns=["text"])
+            for t in tbl.column("text").to_pylist():
+                if t:
+                    yield t
+
+
+class _PoolWriter:
+    def __init__(self, out_dir: Path, prefix: str = "pool", start_idx: int = -1,
+                 total_bytes: int = 0, total_docs: int = 0):
+        import zstandard
+        self.dir = ensure(out_dir)
+        self.prefix = prefix
+        self.zstd = zstandard
+        self.idx = start_idx
+        self.cur = None
+        self.cur_bytes = 0
+        self.total_bytes = total_bytes
+        self.total_docs = total_docs
+        self._roll()
+
+    def _roll(self):
+        if self.cur:
+            self.cur.close()
+        self.idx += 1
+        self.cur_bytes = 0
+        raw = open(self.dir / f"{self.prefix}_{self.idx:05d}.jsonl.zst", "wb")
+        self.cur = self.zstd.ZstdCompressor(level=3).stream_writer(raw, closefd=True)
+
+    def write(self, text: str):
+        line = json.dumps({"text": text}, ensure_ascii=False) + "\n"
+        b = line.encode("utf-8")
+        self.cur.write(b)
+        n = len(text.encode("utf-8"))
+        self.cur_bytes += n
+        self.total_bytes += n
+        self.total_docs += 1
+        if self.cur_bytes >= POOL_SHARD_BYTES:
+            self._roll()
+
+    def close(self):
+        if self.cur:
+            self.cur.close()
+            self.cur = None
+
+
+CHECKPOINT_EVERY_N_FILES = 3   # bound on re-downloaded work if the process dies
+
+
+def _next_shard_idx(out_dir: Path, prefix: str = "pool") -> int:
+    existing = sorted(out_dir.glob(f"{prefix}_*.jsonl.zst"))
+    if not existing:
+        return -1
+    stem = existing[-1].name[len(prefix) + 1:]        # "NNNNN.jsonl.zst"
+    return int(stem.split(".", 1)[0])
+
+
+def build_pool(lang: str, budget_bytes: float, holdout_bytes: int = HOLDOUT_BYTES) -> dict:
+    """Build (or resume) a language's text pool.
+
+    Crash-resumable: every CHECKPOINT_EVERY_N_FILES source files, the current
+    shard is closed (a truncated zstd frame from a mid-write kill can't be
+    decoded, so a checkpoint never leaves a shard half-written) and
+    `stats.json` records which manifest files are already consumed plus the
+    index of that last cleanly-closed shard (`shard_idx`). A crash *between*
+    checkpoints (e.g. mid-way through one large parquet file) can still leave
+    a higher-numbered shard on disk with an unterminated zstd frame -- since
+    none of its bytes were counted into `text_bytes`/`docs` at the last
+    checkpoint, resuming deletes any shard past `shard_idx` before writing
+    starts again, so a stray corrupt shard never lingers for `pack()` to trip
+    on. Re-running `xscript pool` after an interruption (SSH drop, node
+    hiccup, ^C, session teardown) picks up from there instead of
+    re-downloading from scratch. `files_consumed` entries are tagged
+    "repo::path" so a language with a FALLBACK_SOURCES entry can track
+    consumption across sources without collision; old untagged checkpoints
+    (single-source) are migrated on load by assuming the primary repo.
+    """
+    sources = _sources_for(lang)
+    primary_repo, primary_subdir = sources[0]
+    first_files = _list_parquets(primary_repo, primary_subdir)
+    if not first_files:
+        raise RuntimeError(f"no parquet files found for {lang}")
+    out = pool_dir(lang)
+    stats_path = out / "stats.json"
+    resume = None
+    if stats_path.exists():
+        st = json.loads(stats_path.read_text())
+        if st["text_bytes"] >= budget_bytes * 0.99:
+            print(f"[pool] {lang}: cached ({st['text_bytes']/1e9:.1f}GB)")
+            return st
+        if st.get("files_consumed"):
+            resume = st
+            print(f"[pool] {lang}: resuming from checkpoint "
+                  f"({st['text_bytes']/1e9:.1f}/{budget_bytes/1e9:.1f}GB, "
+                  f"{len(st['files_consumed'])} files already consumed)")
+
+    if resume is None:
+        # holdout from the primary source's first file only; pool starts at the second
+        hw = _PoolWriter(HOLDOUT, prefix=lang)
+        got = 0
+        for t in _iter_texts(primary_repo, first_files[0]):
+            hw.write(t)
+            got += len(t.encode("utf-8"))
+            if got >= holdout_bytes:
+                break
+        hw.close()
+        used: list[str] = []
+        pw = _PoolWriter(out)
+    else:
+        got = resume["holdout_bytes"]
+        used = [u if "::" in u else f"{primary_repo}::{u}" for u in resume["files_consumed"]]
+        # shard_idx is missing on checkpoints written before this field existed;
+        # best-effort fall back to whatever's on disk (pre-existing behaviour).
+        last_good_idx = resume.get("shard_idx")
+        if last_good_idx is None:
+            last_good_idx = _next_shard_idx(out)
+        else:
+            for stray in out.glob("pool_*.jsonl.zst"):
+                idx = int(stray.name[len("pool_"):].split(".", 1)[0])
+                if idx > last_good_idx:
+                    stray.unlink()   # never checkpointed -- may be a truncated zstd frame
+            # The file *at* shard_idx is the one _checkpoint()'s _roll() had just
+            # opened (empty) when that checkpoint was written -- everything written
+            # into it since then, up to a crash, was never counted in text_bytes/
+            # docs. It's fine (already fully closed) if the run reached this point
+            # via the final `pw.close()` on graceful completion, but indistinguishable
+            # from a crash-mid-write from stats.json alone -- so verify by decoding.
+            at_idx = out / f"pool_{last_good_idx:05d}.jsonl.zst"
+            if at_idx.exists():
+                import io
+                import zstandard
+                try:
+                    with open(at_idx, "rb") as raw:
+                        reader = zstandard.ZstdDecompressor().stream_reader(raw)
+                        for jline in io.TextIOWrapper(reader, encoding="utf-8"):
+                            if jline.strip():
+                                json.loads(jline)
+                except Exception as exc:
+                    print(f"[pool] {lang}: {at_idx.name} failed validation ({exc}) "
+                          f"-- discarding (never counted in checkpointed totals)")
+                    at_idx.unlink()
+        pw = _PoolWriter(out, start_idx=last_good_idx,
+                         total_bytes=resume["text_bytes"], total_docs=resume["docs"])
+
+    def _checkpoint():
+        pw._roll()   # close the current shard so it's a complete, valid zstd frame
+        st = {"lang": lang, "budget_bytes": budget_bytes, "text_bytes": pw.total_bytes,
+              "docs": pw.total_docs, "holdout_bytes": got, "holdout_file": first_files[0],
+              "files_consumed": used, "shard_idx": pw.idx, "exhausted": False}
+        stats_path.write_text(json.dumps(st, indent=2))
+
+    done = False
+    for i, (repo, subdir) in enumerate(sources):
+        files = first_files if i == 0 else _list_parquets(repo, subdir)
+        pool_files = files[1:] if i == 0 else files   # only the primary source reserves a holdout file
+        for f in pool_files:
+            tag = f"{repo}::{f}"
+            if tag in used:
+                continue
+            used.append(tag)
+            try:
+                for t in _iter_texts(repo, f):
+                    pw.write(t)
+                    if pw.total_bytes >= budget_bytes:
+                        break
+            except Exception as exc:
+                print(f"[pool] WARN {tag}: {exc}")
+                used.pop()   # not actually consumed -- retry it on the next run
+            if pw.total_bytes >= budget_bytes:
+                done = True
+                _checkpoint()
+                break
+            if len(used) % CHECKPOINT_EVERY_N_FILES == 0:
+                _checkpoint()
+            if len(used) % 20 == 0:
+                print(f"[pool] {lang}: {pw.total_bytes/1e9:.1f}/{budget_bytes/1e9:.1f}GB "
+                      f"({len(used)} files, source {i+1}/{len(sources)}: {repo})")
+        if done:
+            break
+        if i + 1 < len(sources):
+            print(f"[pool] {lang}: source {i+1}/{len(sources)} ({repo}) exhausted at "
+                  f"{pw.total_bytes/1e9:.1f}GB -> falling back to {sources[i+1][0]}")
+    pw.close()
+    st = {"lang": lang, "budget_bytes": budget_bytes, "text_bytes": pw.total_bytes,
+          "docs": pw.total_docs, "holdout_bytes": got, "holdout_file": first_files[0],
+          "files_consumed": used, "shard_idx": pw.idx,
+          "exhausted": pw.total_bytes < budget_bytes * 0.99}
+    stats_path.write_text(json.dumps(st, indent=2))
+    if st["exhausted"]:
+        print(f"[pool] WARNING {lang}: corpus exhausted at {pw.total_bytes/1e9:.1f}GB "
+              f"< budget {budget_bytes/1e9:.1f}GB -> training will epoch over this pool")
+    print(f"[pool] {lang}: {pw.total_bytes/1e9:.2f}GB text, {pw.total_docs} docs")
+    return st
+
+
+def _measured_bytes_per_token(flavor: str = "unigram", condition: str = "destarved") -> dict[str, float]:
+    """Real bytes/token per study language, measured on FLORES+ dev with the
+    tokenizer that will actually pack the pool text. Destarved is the more
+    byte-hungry of the two conditions (better fertility -> more input bytes
+    needed per token), so it's the binding case for sizing. Scripts vary a
+    lot (Arabic ~6.8 bytes/token vs Chinese ~4.0) -- returns {} (caller falls
+    back to a flat estimate) if the tokenizer or FLORES+ aren't ready yet.
+    """
+    from .. import flores
+    from ..langs import tok_name
+    from ..paths import tokenizer_dir
+    from ..tok.wrapper import Tok
+
+    tdir = tokenizer_dir(tok_name(flavor, condition))
+    if not (tdir / "meta.json").exists():
+        return {}
+    try:
+        tok = Tok(tdir)
+        par = flores.load_parallel(list(LANGS), "dev")
+    except Exception as exc:
+        print(f"[pool] WARN: couldn't measure real bytes/token ({exc}); "
+              f"falling back to the flat estimate")
+        return {}
+    out = {}
+    for l, sents in par.items():
+        b = sum(len(s.encode("utf-8")) for s in sents)
+        t = sum(len(tok.encode(s)) for s in sents)
+        out[l] = b / t
+    return out
+
+
+def plan_budgets(tokens_per_run: float = 30e9, est_bytes_per_token: float = 4.5,
+                 safety: float = 1.15) -> dict[str, float]:
+    """Per-language pool byte budgets.
+
+    Monolingual runs need the full token budget in one language; bilingual
+    runs need half. The pool must cover the *max* need across planned runs
+    under the worst-case (most byte-hungry, i.e. destarved) tokenizer.
+
+    Sized from each language's REAL measured bytes/token (destarved
+    tokenizer on FLORES+), not a flat guess -- a single constant badly
+    under/over-shoots per language given how much bytes/token varies by
+    script. `est_bytes_per_token` is only a fallback for languages the
+    measurement can't cover yet (e.g. before the tokenizer gate).
+    """
+    need_tokens = {l: tokens_per_run for l in LANGS}  # monolingual dominates
+    measured = _measured_bytes_per_token()
+    return {l: need_tokens[l] * measured.get(l, est_bytes_per_token) * safety
+            for l in LANGS}