Phase 2: build_mix -- resumable, cross-source shards, seeded intra-shuffle, mandatory en filter

build/build_mix.py

@@ -0,0 +1,340 @@
+# Phase 2 builder: turn the selected sources into tokenised, deduplicated, resumable shards.
+#
+# Constraints this is written against (all measured, docs/00-platform-notes.md):
+#   * Kaggle **CPU** session: 4 vCPU, 30 GiB cgroup, **no swap**, 19.5 GB /kaggle/working. Dedup
+#     structures are sized by *document count* (~2.1M docs), never by tokens.
+#   * Hub download is 35-89 MB/s and tokenisation ~661k tok/s on 4 cores. Neither is the bottleneck;
+#     `datasets` streaming at ~3.3 MB/s per process is. So sources go one at a time, streamed, never
+#     materialised.
+#   * The session can be killed at any time (§3.1). Every completed shard is written and recorded before
+#     the next begins, and the exact-duplicate hash set is persisted alongside it -- otherwise "resume"
+#     silently means "start forgetting what you already emitted", which would let duplicates back in
+#     across the resume boundary. State is saved only *after* a shard lands on disk.
+#
+# Dedup policy, with its limits stated:
+#   * exact, global: 64-bit hash of the normalised document.
+#   * cross-source, on the upstream url/id/file_path field -- this is what actually catches FinePhrase's
+#     configs being rewrites of the same FineWeb-Edu documents (§02-mix-plan §2).
+#   * near-dup MinHash is NOT run here. It needs every signature resident at once (~1.1 GB) and the
+#     within-source benefit is already largely captured by the exact hash. It is a v2 candidate, decided
+#     by the measured exact-dedup loss this script reports, not by convention.
+#
+# Language: §3.4 is English-only and the inventory measured 8.5% of finepdfs `eng_Latn` and 7.5% of
+# open-web-math failing a crude English test, so the filter below is mandatory, not decorative.
+#
+# No credentials in this file. Publish is a separate script.
+
+import argparse
+import array
+import hashlib
+import json
+import os
+import random
+import struct
+import time
+
+TOKENS_PER_SHARD = 8_000_000    # 16 MB as uint16 -> ~140 shards for 1.1B tokens, ~2.3 GB total
+ROW_TOKEN_CAP = 16_384          # one document must never dominate a shard
+SHUFFLE_SEED = 20260919         # recorded in the manifest so a rebuild is reproducible
+HASHES_FILE = "seen_hashes.bin"
+STATE_NAME = "build_state.json"
+
+# Corrected config and column names, from the Phase 2 inventory (docs/02-mix-plan.md §9). The §2 names
+# came from dataset cards and five of them did not exist.
+SOURCES = [
+    dict(repo="HuggingFaceFW/fineweb-edu", config="sample-10BT", col="text", tokens=300_000_000,
+         min_score=4, why="quality-scored general English"),
+    dict(repo="HuggingFaceFW/finepdfs-edu", config="eng_Latn", col="text", tokens=150_000_000,
+         en_filter=True, why="long-form educational PDFs; needs own en filter"),
+    dict(repo="HuggingFaceTB/cosmopedia", config="stanford", col="text", tokens=60_000_000,
+         why="textbook expository"),
+    dict(repo="HuggingFaceTB/cosmopedia", config="openstax", col="text", tokens=30_000_000,
+         why="textbook science"),
+    dict(repo="HuggingFaceTB/cosmopedia", config="khanacademy", col="text", tokens=10_000_000,
+         why="worked steps"),
+    dict(repo="HuggingFaceTB/cosmopedia", config="auto_math_text", col="text", tokens=30_000_000,
+         why="synthetic math exposition"),
+    dict(repo="HuggingFaceTB/cosmopedia", config="wikihow", col="text", tokens=30_000_000,
+         why="procedural / physical-world steps"),
+    dict(repo="HuggingFaceFW/finewiki", config="en", col="rewritten", tokens=80_000_000,
+         why="rewritten encyclopedic prose"),
+    dict(repo="omarkamali/wikipedia-monthly", config="20260101.en", col="text", tokens=60_000_000,
+         why="current encyclopedic breadth"),
+    dict(repo="HuggingFaceTB/finemath", config="finemath-4plus", col="text", tokens=130_000_000,
+         why="math; upstream-decontaminated against GSM8K/MMLU/ARC"),
+    dict(repo="open-web-math/open-web-math", config="default", col="text", tokens=70_000_000,
+         en_filter=True, why="forum/webbook math prose"),
+    dict(repo="HuggingFaceFW/finephrase", config="tutorial", col="text", tokens=60_000_000,
+         why="stepwise tutorial register"),
+    dict(repo="HuggingFaceFW/finephrase", config="faq", col="text", tokens=30_000_000,
+         why="FAQ register"),
+    dict(repo="HuggingFaceFW/finephrase", config="table", col="text", tokens=30_000_000,
+         why="tabular->prose"),
+    dict(repo="SimpleStories/SimpleStories", config="default", col="story", tokens=40_000_000,
+         why="long-range simple narrative coherence"),
+    dict(repo="common-pile/arxiv_abstracts", config="default", col="text", tokens=20_000_000,
+         why="CC0 scientific abstracts"),
+    # No code source yet: stack-v3-train has only a `default` config, and its rows are whole repos with a
+    # nested `files` struct, so it needs its own probe first (docs/02-mix-plan.md §9.3).
+    # libretexts dropped outright: its `text` column holds URLs and 93% of sampled rows failed en.
+]
+
+STOP = (" the ", " and ", " of ", " to ", " that ", " is ", " in ", " for ", " with ", " as ")
+KEY_COLS = ("url", "id", "file_path", "doc_id", "generation_id", "repo_name")
+
+
+def normalize(s):
+    # Whitespace-collapse only. Aggressive folding would collide genuinely different documents (tables,
+    # code) and a false merge is worse than a missed near-duplicate.
+    return " ".join(s.split())
+
+
+def looks_english(s):
+    low = s[:600].lower()
+    return sum(w in low for w in STOP) >= 3
+
+
+def doc_hash(s):
+    return struct.unpack("<Q", hashlib.sha256(s.encode("utf-8")).hexdigest()[:16].encode())[0]
+
+
+def key_of(row, cols):
+    for c in KEY_COLS:
+        if c in cols:
+            v = row.get(c)
+            if isinstance(v, str) and v:
+                return v
+    return None
+
+
+class State:
+    def __init__(self, out_dir):
+        self.dir = out_dir
+        self.path = os.path.join(out_dir, STATE_NAME)
+        self.hashes_path = os.path.join(out_dir, HASHES_FILE)
+        self.d = {"shards": [], "per_source": {}, "started": time.time(), "version": 2}
+        self.seen_exact = set()
+        self.seen_url = set()
+        if os.path.exists(self.path):
+            self.d = json.load(open(self.path))
+            if os.path.exists(self.hashes_path):
+                a = array.array("Q")
+                with open(self.hashes_path, "rb") as f:
+                    a.fromfile(f, os.path.getsize(self.hashes_path) // 8)
+                self.seen_exact = set(a)
+            print(f"resuming: {len(self.d['shards'])} shards, "
+                  f"{len(self.seen_exact):,} exact hashes, {self.tokens_done():,} tokens", flush=True)
+            for s in self.d["per_source"].values():
+                self.seen_url.update(s.get("keys", []))
+
+    def tokens_done(self):
+        return sum(s["tokens"] for s in self.d["shards"])
+
+    def save(self, new_hashes, new_keys):
+        with open(self.hashes_path, "ab") as f:
+            array.array("Q", new_hashes).tofile(f)
+        d = dict(self.d)
+        d["per_source"] = {k: {kk: vv for kk, vv in v.items() if kk != "keys"}
+                           for k, v in self.d["per_source"].items()}
+        tmp = self.path + ".tmp"
+        with open(tmp, "w") as f:
+            json.dump(d, f, indent=1)
+        os.replace(tmp, self.path)
+
+    def record_source(self, key, done, tokens, docs_seen, drops, seconds, keys):
+        self.d["per_source"][key] = {"complete": bool(done), "tokens": tokens,
+                                     "docs_seen": docs_seen, "dropped": drops,
+                                     "seconds": round(seconds, 1), "keys_sample": keys[:2000]}
+        self.seen_url.update(keys[:2000])
+
+
+def build(args):
+    t_start = time.monotonic()
+    from tokenizers import Tokenizer
+    from huggingface_hub import hf_hub_download
+
+    os.makedirs(args.out, exist_ok=True)
+    tok_path = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
+    tk = Tokenizer.from_file(tok_path)
+    tk.no_truncation()
+    vocab = tk.get_vocab_size()
+    assert vocab <= 65536, f"uint16 shard layout breaks above 65,536 (vocab={vocab})"
+    tok_sha = hashlib.sha256(open(tok_path, "rb").read()).hexdigest()
+    print(f"tokenizer vocab={vocab} sha256={tok_sha[:12]} budget={args.budget_tokens:,}", flush=True)
+
+    st = State(args.out)
+    buf, buf_src = [], []
+    buf_tokens = 0
+    pending_hashes, pending_keys = [], []
+
+    def flush_shard():
+        nonlocal buf, buf_src, buf_tokens, pending_hashes, pending_keys
+        if not buf:
+            pending_hashes, pending_keys = [], []
+            return
+        ids = array.array("H")
+        # The buffer deliberately spans source boundaries (nothing flushes at end-of-source), so a shard
+        # already mixes domains; shuffling the documents first removes the residual "all fineweb-edu, then
+        # all cosmopedia" ordering inside a pass. Seeded and recorded so a rebuild is reproducible.
+        random.Random(SHUFFLE_SEED + len(st.d["shards"])).shuffle(buf)
+        docs = 0
+        for doc in buf:
+            ids.extend(doc)
+            docs += 1
+        idx = len(st.d["shards"])
+        path = os.path.join(args.out, f"shard-{idx:04d}.bin")
+        tmp = path + ".tmp"
+        with open(tmp, "wb") as f:
+            f.write(struct.pack("<II", vocab, len(ids)))
+            ids.tofile(f)
+        os.replace(tmp, path)
+        digest = hashlib.sha256(open(path, "rb").read()).hexdigest()
+        st.d["shards"].append({"file": os.path.basename(path), "tokens": len(ids), "docs": docs,
+                               "sha256": digest, "bytes": os.path.getsize(path),
+                               "sources": sorted(set(buf_src))})
+        st.save(pending_hashes, pending_keys)
+        print(f"  shard {idx:04d}: {len(ids):,} tok / {docs:,} docs sha256={digest[:12]} "
+              f"[{time.monotonic() - t_start:.0f}s]", flush=True)
+        buf, buf_src, buf_tokens, pending_hashes, pending_keys = [], [], 0, [], []
+
+    if args.max_shards and len(st.d["shards"]) >= args.max_shards:
+        print(f"max_shards={args.max_shards} reached; exiting without work", flush=True)
+        return {"total_tokens": st.tokens_done(), "shards": len(st.d["shards"])}
+
+    stats = []
+    for src in SOURCES:
+        key = f"{src['repo']}::{src['config']}"
+        if st.d["per_source"].get(key, {}).get("complete"):
+            print(f"skip {key} (already complete)", flush=True)
+            continue
+        if st.tokens_done() >= args.budget_tokens:
+            print("budget reached", flush=True)
+            break
+        if args.max_shards and len(st.d["shards"]) >= args.max_shards:
+            print(f"max_shards={args.max_shards} reached", flush=True)
+            break
+        import datasets
+
+        try:
+            ds = datasets.load_dataset(src["repo"], src["config"], split="train", streaming=True)
+        except Exception as e:
+            # A wrong config name must be loud, must be recorded, and must not abort later sources.
+            print(f"!! {key} could not load: {type(e).__name__}: {str(e)[:200]}", flush=True)
+            stats.append({"source": key, "error": f"{type(e).__name__}: {str(e)[:200]}"})
+            continue
+
+        want = min(src["tokens"], args.budget_tokens - st.tokens_done())
+        got = docs_seen = 0
+        drops = {"short_or_low_score": 0, "non_english": 0, "exact_dup": 0, "url_dup": 0, "too_few_tokens": 0}
+        cols = None
+        t0 = time.monotonic()
+        src_keys = []
+        for row in ds:
+            if got >= want:
+                break
+            if args.sec_per_source and time.monotonic() - t0 > args.sec_per_source:
+                break
+            if cols is None:
+                cols = set(row.keys())
+                if src["col"] not in cols:
+                    print(f"!! {key} has no column {src['col']!r}; seen {sorted(cols)[:10]}", flush=True)
+                    drops["short_or_low_score"] = -1
+                    break
+            raw = row.get(src["col"])
+            if not isinstance(raw, str):
+                continue
+            docs_seen += 1
+            s = normalize(raw)
+            if len(s) < args.min_chars:
+                drops["short_or_low_score"] += 1
+                continue
+            if src.get("min_score") is not None:
+                sc = row.get("int_score", row.get("score"))
+                try:
+                    if sc is not None and float(sc) < src["min_score"]:
+                        drops["short_or_low_score"] += 1
+                        continue
+                except (TypeError, ValueError):
+                    pass
+            if src.get("en_filter") and not looks_english(s):
+                drops["non_english"] += 1
+                continue
+            h = doc_hash(s)
+            if h in st.seen_exact:
+                drops["exact_dup"] += 1
+                continue
+            u = key_of(row, cols)
+            if u and u in st.seen_url:
+                drops["url_dup"] += 1
+                continue
+            ids = tk.encode(s[:args.max_chars], add_special_tokens=False).ids[:ROW_TOKEN_CAP]
+            if len(ids) < 16:
+                drops["too_few_tokens"] += 1
+                continue
+            st.seen_exact.add(h)
+            pending_hashes.append(h)
+            if u:
+                st.seen_url.add(u)
+                src_keys.append(u)
+                pending_keys.append(u)
+            buf.append(ids)
+            buf_src.append(key)
+            buf_tokens += len(ids)
+            got += len(ids)
+            if buf_tokens >= TOKENS_PER_SHARD:
+                flush_shard()
+
+        elapsed = time.monotonic() - t0
+        complete = got >= want
+        st.record_source(key, complete, got, docs_seen, drops, elapsed, src_keys)
+        stats.append({"source": key, "tokens": got, "docs_seen": docs_seen, "dropped": dict(drops),
+                      "seconds": round(elapsed, 1), "tok_per_s": round(got / max(1e-9, elapsed)),
+                      "complete": bool(complete)})
+        print(f"  {key}: {got:,} tok of {want:,} target from {docs_seen:,} docs; dup={drops['exact_dup']}"
+              f"/{drops['url_dup']} lang={drops['non_english']} short={drops['short_or_low_score']} "
+              f"[{stats[-1]['tok_per_s']:,} tok/s, {elapsed:.0f}s]", flush=True)
+
+    flush_shard()
+    total_docs = sum(s.get("docs_seen", 0) for s in stats) or 1
+    total_dupes = sum(d["exact_dup"] + d["url_dup"]
+                      for s in stats for d in [s.get("dropped") or {}] if d)
+    manifest = {
+        "tokenizer": {"id": "HuggingFaceTB/SmolLM2-135M:tokenizer.json", "vocab_size": vocab,
+                      "sha256": tok_sha},
+        "layout": {"dtype": "uint16", "header": "<II (vocab, n_tokens)",
+                   "resume_cursor": "(shard_index, token_offset) into the concatenated id stream",
+                   "tokens_per_shard_target": TOKENS_PER_SHARD, "row_token_cap": ROW_TOKEN_CAP,
+                   "shuffle_seed_base": SHUFFLE_SEED,
+                   "shuffle_note": "documents within shard idx are shuffled with seed base+idx; "
+                                   "shard order is the training order"},
+        "total_tokens": st.tokens_done(), "shards": st.d["shards"], "sources": stats,
+        "budget_tokens": args.budget_tokens,
+        "exact_dedup_drop_rate": round(total_dupes / total_docs, 4),
+        "seconds": round(time.monotonic() - t_start, 1),
+    }
+    with open(os.path.join(args.out, "manifest.json"), "w") as f:
+        json.dump(manifest, f, indent=1)
+    return manifest
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="/kaggle/working/mix")
+    ap.add_argument("--budget-tokens", type=int, default=1_310_000_000)
+    ap.add_argument("--min-chars", type=int, default=300)
+    ap.add_argument("--max-chars", type=int, default=100_000)
+    ap.add_argument("--sec-per-source", type=int, default=0)
+    ap.add_argument("--max-shards", type=int, default=0, help="0 = unlimited")
+    a = ap.parse_args()
+    if a.max_shards == 0:
+        a.max_shards = None
+    m = build(a)
+    print("BUILD_JSON_BEGIN")
+    print(json.dumps({"total_tokens": m["total_tokens"], "shards": len(m["shards"]),
+                      "exact_dedup_drop_rate": m.get("exact_dedup_drop_rate"),
+                      "seconds": m["seconds"]}))
+    print("BUILD_JSON_END")
+
+
+if __name__ == "__main__":
+    main()