build_mix v2: stage->merge interleave + doc offset table; verify_mix asserts source-mix floor

build/build_mix.py

@@ -1,343 +1,491 @@
-# 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):
-    # digest()[:8] is 8 raw BYTES. hexdigest()[:16] is 16 *hex characters* = 16 bytes, which is what
-    # the first version sliced and what struct.unpack("<Q") then rejected -- caught by the 12M-token
-    # smoke test rather than six hours into the full build.
-    return struct.unpack("<Q", hashlib.sha256(s.encode("utf-8")).digest()[:8])[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()
+# Phase 2 builder: sources -> per-source staged shards -> interleaved mixed shards.
+#
+# Two stages, because one stage cannot produce a mixed corpus without re-streaming.
+#
+#   stage  per source: stream -> normalise -> language/quality gate -> dedup -> tokenise
+#                      -> write  stage/<key>/shard-####.bin
+#   merge  read all staged shards together, emit documents round-robin in proportion to each source's
+#          target, and write the final mixed  mix/shard-####.bin
+#
+# Why the merge exists (memory/ERRORS.md E-016). The first version consumed one source to its target and
+# only then started the next, and it flushed a shard when the buffer filled. Since fineweb-edu's target is
+# 300M tokens = 37 shards of 8M, the first 37 shards of the mix would have been nothing but web text, then
+# textbooks, then math. A shuffle inside the buffer cannot fix that -- the buffer only ever held one
+# source. That is a curriculum drift across the first third of the main run, and no LR schedule repairs
+# ordering like that. Mixing must happen between sources, not inside them.
+#
+# Shard layout is shared by both stages and carries a per-document offset table, so document boundaries
+# survive to the training reader (E-016's second latent bug: the old format stored a doc *count*, from
+# which boundaries cannot be recovered):
+#
+#   header   : <IIII  = vocab, n_docs, n_tokens, flags(0)          16 B
+#   offsets  : uint32 × (n_docs + 1), cumulative, [0] = 0
+#   ids      : uint16 × n_tokens
+#
+# doc i of a shard is ids[offsets[i] : offsets[i+1]].
+#
+# Constraints (docs/00-platform-notes.md): Kaggle CPU session, 4 vCPU, 30 GiB cgroup, **no swap**,
+# 19.5 GB /kaggle/working. Staged (~2.6 GB) + final (~2.6 GB) + hash sets (~35 MB) fit with room. Streaming
+# is ~3.3 MB/s per process and tokenising ~661k tok/s on 4 cores, so ingest is serial-per-source by design.
+#
+# Resumability: state is saved only *after* a file lands, and the dedup hash sets are appended in the same
+# transaction -- otherwise a resume forgets what was already emitted and lets duplicates back in across
+# the boundary. No credentials in this file.
+
+import argparse
+import array
+import hashlib
+import json
+import os
+import random
+import struct
+import sys
+import time
+
+MAGIC_FLAGS = 0
+TOKENS_PER_SHARD = 8_000_000    # 16 MB as uint16; ~140 shards for 1.1B tokens
+MERGE_BATCH_DOCS = 32           # granularity of the round-robin; ~20k tokens, so a 262k batch spans ~13
+ROW_TOKEN_CAP = 16_384          # one document must never dominate a shard
+SHUFFLE_SEED = 20260919         # recorded so a rebuild is reproducible
+HASHES_FILE = "seen_hashes.bin"
+KEYS_FILE = "seen_keys.bin"
+STATE_NAME = "build_state.json"
+
+# Corrected config and column names, from the Phase 2 inventory (docs/02-mix-plan.md §9). The names in
+# §2 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,
+         en_filter=True, 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,
+         en_filter=True, why="stepwise tutorial register"),
+    dict(repo="HuggingFaceFW/finephrase", config="faq", col="text", tokens=30_000_000,
+         en_filter=True, why="FAQ register"),
+    dict(repo="HuggingFaceFW/finephrase", config="table", col="text", tokens=30_000_000,
+         en_filter=True, 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,
+         en_filter=True, 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 (docs/02-mix-plan.md §9.3).
+    # common-pile/libretexts dropped: `text` holds URLs, 93% of sampled rows failed the en test.
+]
+
+STOP = (" the ", " and ", " of ", " to ", " that ", " is ", " in ", " for ", " with ", " as ")
+KEY_COLS = ("url", "id", "file_path", "doc_id", "generation_id", "repo_name")
+
+
+def src_key(src):
+    return f"{src['repo'].split('/')[1]}__{src['config']}"
+
+
+def normalize(s):
+    # Whitespace-collapse only: aggressive folding collides 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 h64(s):
+    """64-bit digest of a string. digest()[:8] is 8 BYTES; hexdigest()[:16] is 16 *characters*
+    and broke the first version (E-015)."""
+    return struct.unpack("<Q", hashlib.sha256(s.encode("utf-8")).digest()[:8])[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
+
+
+# --------------------------------------------------------------------- shard I/O
+def write_shard(path, vocab, docs):
+    """docs: iterable of array('H')/list of token ids. Returns (n_docs, n_tokens, sha256, bytes)."""
+    offs = array.array("I", [0])
+    ids = array.array("H")
+    for d in docs:
+        ids.extend(d)
+        offs.append(len(ids))
+    tmp = path + ".tmp"
+    with open(tmp, "wb") as f:
+        f.write(struct.pack("<IIII", vocab, len(offs) - 1, len(ids), MAGIC_FLAGS))
+        offs.tofile(f)
+        ids.tofile(f)
+    os.replace(tmp, path)
+    digest = hashlib.sha256(open(path, "rb").read()).hexdigest()
+    return len(offs) - 1, len(ids), digest, os.path.getsize(path)
+
+
+def read_shard_header(path):
+    with open(path, "rb") as f:
+        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", f.read(16))
+    return vocab, n_docs, n_tokens, flags
+
+
+def iter_shard_docs(path):
+    """Yield documents as array('H'). Offsets are read in full (4 B/doc) which is fine at 8M tok/shard."""
+    with open(path, "rb") as f:
+        vocab, n_docs, n_tokens, _flags = struct.unpack("<IIII", f.read(16))
+        offs = array.array("I")
+        offs.fromfile(f, n_docs + 1)
+        ids = array.array("H")
+        ids.fromfile(f, n_tokens)
+    for i in range(n_docs):
+        yield ids[offs[i]:offs[i + 1]]
+
+
+# --------------------------------------------------------------------- state
+class State:
+    def __init__(self, root):
+        self.root = root
+        self.stage_dir = os.path.join(root, "stage")
+        self.mix_dir = os.path.join(root, "mix")
+        os.makedirs(self.stage_dir, exist_ok=True)
+        os.makedirs(self.mix_dir, exist_ok=True)
+        self.path = os.path.join(root, STATE_NAME)
+        self.hashes_path = os.path.join(root, HASHES_FILE)
+        self.keys_path = os.path.join(root, KEYS_FILE)
+        self.d = {"staged": {}, "merged_shards": [], "version": 3}
+        self.seen_exact = set()
+        self.seen_key = set()
+        if os.path.exists(self.path):
+            self.d = json.load(open(self.path))
+            self.seen_exact = self._load_u64(self.hashes_path)
+            self.seen_key = self._load_u64(self.keys_path)
+            print(f"resuming: {len(self.d['staged'])} sources staged, "
+                  f"{len(self.d['merged_shards'])} mixed shards, "
+                  f"{len(self.seen_exact):,} hashes", flush=True)
+
+    @staticmethod
+    def _load_u64(p):
+        if not os.path.exists(p):
+            return set()
+        a = array.array("Q")
+        with open(p, "rb") as f:
+            a.fromfile(f, os.path.getsize(p) // 8)
+        return set(a)
+
+    def commit(self, new_hashes, new_keys):
+        """Append the dedup sets FIRST, then the state file: if the job dies between the two, the worst
+        case is a hash set slightly ahead of the state (a document re-emitted is a duplicate, which the
+        next run's dedup catches; a hash set behind the state is a permanent silent duplicate)."""
+        if new_hashes:
+            with open(self.hashes_path, "ab") as f:
+                array.array("Q", new_hashes).tofile(f)
+        if new_keys:
+            with open(self.keys_path, "ab") as f:
+                array.array("Q", new_keys).tofile(f)
+        tmp = self.path + ".tmp"
+        with open(tmp, "w") as f:
+            json.dump(self.d, f, indent=1)
+        os.replace(tmp, self.path)
+
+    def staged_tokens(self, key):
+        return sum(s["tokens"] for s in self.d["staged"].get(key, {}).get("shards", []))
+
+    def staged_tokens_total(self):
+        return sum(v.get("tokens", 0) for v in self.d["staged"].values())
+
+    def merged_tokens(self):
+        return sum(s["tokens"] for s in self.d["merged_shards"])
+
+
+# --------------------------------------------------------------------- stage
+def stage(args, st, tokenizer):
+    for src in SOURCES:
+        key = src_key(src)
+        rec = st.d["staged"].get(key, {})
+        if rec.get("complete"):
+            print(f"stage: skip {key} (complete)", flush=True)
+            continue
+        want = min(src["tokens"], args.cap_tokens_per_source or src["tokens"])
+        if args.budget_tokens and st.staged_tokens_total() >= args.budget_tokens:
+            print("stage: overall budget reached", flush=True)
+            break
+        import datasets
+
+        try:
+            ds = datasets.load_dataset(src["repo"], src["config"], split="train", streaming=True)
+        except Exception as e:
+            print(f"!! stage: {key} could not load: {type(e).__name__}: {str(e)[:200]}", flush=True)
+            rec["error"] = f"{type(e).__name__}: {str(e)[:200]}"
+            st.d["staged"][key] = rec
+            st.commit([], [])
+            continue
+
+        got = docs_seen = 0
+        drops = {"short_or_low_score": 0, "non_english": 0, "exact_dup": 0, "key_dup": 0,
+                 "too_few_tokens": 0, "no_column": 0}
+        cols = None
+        buf, new_h, new_k = [], [], []
+        sd = os.path.join(st.stage_dir, key)
+        os.makedirs(sd, exist_ok=True)
+        shards = rec.get("shards", [])
+        t0 = time.monotonic()
+        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"!! stage: {key} has no column {src['col']!r}; seen {sorted(cols)[:10]}",
+                          flush=True)
+                    drops["no_column"] = -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 = h64(s)
+            if h in st.seen_exact:
+                drops["exact_dup"] += 1
+                continue
+            u = key_of(row, cols)
+            ku = h64(u) if u else None
+            if ku is not None and ku in st.seen_key:
+                drops["key_dup"] += 1
+                continue
+            ids = tokenizer.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)
+            new_h.append(h)
+            if ku is not None:
+                st.seen_key.add(ku)
+                new_k.append(ku)
+            buf.append(array.array("H", ids))
+            got += len(ids)
+            if sum(len(b) for b in buf) >= TOKENS_PER_SHARD:
+                idx = len(shards)
+                p = os.path.join(sd, f"shard-{idx:04d}.bin")
+                nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
+                shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
+                               "sha256": dg, "bytes": bs})
+                st.seen_exact.update(new_h)
+                st.d["staged"][key] = {"complete": False, "shards": shards, "tokens": got,
+                                       "docs_seen": docs_seen, "dropped": drops}
+                st.commit(new_h, new_k)
+                new_h, new_k = [], []
+                buf = []
+                print(f"  staged {key} shard {idx:04d}: {nt:,} tok [{time.monotonic() - t0:.0f}s]",
+                      flush=True)
+
+        if buf:
+            idx = len(shards)
+            p = os.path.join(sd, f"shard-{idx:04d}.bin")
+            nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
+            shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
+                           "sha256": dg, "bytes": bs})
+            print(f"  staged {key} shard {idx:04d}: {nt:,} tok (tail)", flush=True)
+
+        complete = got >= want
+        st.d["staged"][key] = {"complete": bool(complete), "shards": shards, "tokens": got,
+                               "docs_seen": docs_seen, "dropped": dict(drops),
+                               "target": want, "seconds": round(time.monotonic() - t0, 1)}
+        st.commit(new_h, new_k)
+        el = time.monotonic() - t0
+        print(f"stage: {key} {got:,} tok of {want:,} from {docs_seen:,} docs; "
+              f"dup={drops['exact_dup']} keydup={drops['key_dup']} lang={drops['non_english']} "
+              f"short={drops['short_or_low_score']} [{got / max(1e-9, el):,.0f} tok/s]", flush=True)
+
+
+# --------------------------------------------------------------------- merge
+def merge(args, st, tokenizer):
+    """Interleave staged documents proportionally to each source's staged token total."""
+    vocab = tokenizer.get_vocab_size()
+    streams = []
+    for src in SOURCES:
+        key = src_key(src)
+        rec = st.d["staged"].get(key)
+        if not rec or not rec.get("shards"):
+            continue
+        weight = rec.get("tokens") or sum(s["tokens"] for s in rec["shards"])
+        if weight <= 0:
+            continue
+        iters = [iter_shard_docs(os.path.join(st.root, s["file"])) for s in rec["shards"]]
+
+        def gen(its=iters):
+            for it in its:
+                for d in it:
+                    yield d
+
+        streams.append({"key": key, "weight": float(weight), "credit": 0.0, "it": gen(iters),
+                        "pending": None, "emitted_tokens": 0, "exhausted": False})
+    if not streams:
+        print("merge: nothing staged", flush=True)
+        return
+    total_w = sum(s["weight"] for s in streams)
+    print(f"merge: {len(streams)} sources, "
+          f"{sum(s['weight'] for s in streams) / 1e6:,.0f}M staged tokens", flush=True)
+
+    def next_doc(s):
+        if s["pending"] is None and not s["exhausted"]:
+            try:
+                s["pending"] = next(s["it"])
+            except StopIteration:
+                s["exhausted"] = True
+        return s["pending"]
+
+    cap = args.merge_budget_tokens or int(total_w)
+    buf, buf_src, buf_tokens = [], [], 0
+    out_dir = os.path.join(st.root, "mix")
+    os.makedirs(out_dir, exist_ok=True)
+    t0 = time.monotonic()
+
+    def flush():
+        nonlocal buf, buf_src, buf_tokens
+        if not buf:
+            return
+        random.Random(SHUFFLE_SEED + len(st.d["merged_shards"])).shuffle(buf)
+        docs = [d for d, _ in buf]
+        p = os.path.join(out_dir, f"shard-{len(st.d['merged_shards']):04d}.bin")
+        nd, nt, dg, bs = write_shard(p, vocab, docs)
+        st.d["merged_shards"].append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
+                                      "sha256": dg, "bytes": bs,
+                                      "sources": sorted(set(s for _, s in buf))})
+        st.commit([], [])
+        print(f"  mix shard {len(st.d['merged_shards']) - 1:04d}: {nt:,} tok / {nd:,} docs "
+              f"from {len(st.d['merged_shards'][-1]['sources'])} sources "
+              f"[{time.monotonic() - t0:.0f}s]", flush=True)
+        buf, buf_src, buf_tokens = [], [], 0
+
+    while st.merged_tokens() + buf_tokens < cap:
+        live = [s for s in streams if next_doc(s) is not None]
+        if not live:
+            break
+        # Largest-deficit scheduling: whoever is furthest below its proportional share of what has been
+        # emitted goes next. O(n_sources) per batch of MERGE_BATCH_DOCS, not per document.
+        emitted = sum(s["emitted_tokens"] for s in streams) or 1
+        pick = max(live, key=lambda s: s["weight"] / total_w - s["emitted_tokens"] / emitted)
+        took = 0
+        while took < MERGE_BATCH_DOCS:
+            d = next_doc(pick)
+            if d is None:
+                break
+            buf.append((d, pick["key"]))
+            buf_tokens += len(d)
+            pick["emitted_tokens"] += len(d)
+            pick["pending"] = None
+            took += 1
+            if buf_tokens >= TOKENS_PER_SHARD:
+                flush()
+                break
+        if all(s["pending"] is None and s["exhausted"] for s in streams):
+            break
+    flush()
+    print(f"merge: {len(st.d['merged_shards'])} shards, "
+          f"{sum(s['tokens'] for s in st.d['merged_shards']):,} tokens", flush=True)
+
+
+# --------------------------------------------------------------------- entry
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", default="/kaggle/working/mixroot")
+    ap.add_argument("--budget-tokens", type=int, default=1_310_000_000, help="stage cap")
+    ap.add_argument("--merge-budget-tokens", type=int, default=0, help="0 = take everything staged")
+    ap.add_argument("--cap-tokens-per-source", type=int, default=0, help="smoke test knob")
+    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("--only", default="", help="comma list of config names to stage")
+    ap.add_argument("--stage-only", action="store_true")
+    ap.add_argument("--merge-only", action="store_true")
+    args = ap.parse_args()
+
+    if args.only:
+        keep = set(args.only.split(","))
+        SOURCES[:] = [s for s in SOURCES if s["config"] in keep or src_key(s) in keep]
+        print(f"--only: {len(SOURCES)} source(s)", flush=True)
+
+    from tokenizers import Tokenizer
+    from huggingface_hub import hf_hub_download
+    tok_path = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
+    tk = Tokenizer.from_file(tok_path)
+    tk.no_truncation()
+    assert tk.get_vocab_size() <= 65536, "uint16 shard layout breaks above 65,536"
+    print(f"tokenizer vocab={tk.get_vocab_size()} "
+          f"sha256={hashlib.sha256(open(tok_path, 'rb').read()).hexdigest()[:12]}", flush=True)
+
+    st = State(args.root)
+    t0 = time.monotonic()
+    if not args.merge_only:
+        stage(args, st, tk)
+    if not args.stage_only:
+        merge(args, st, tk)
+
+    manifest = {
+        "tokenizer": {"id": "HuggingFaceTB/SmolLM2-135M:tokenizer.json",
+                      "vocab_size": tk.get_vocab_size(),
+                      "sha256": hashlib.sha256(open(tok_path, "rb").read()).hexdigest()},
+        "layout": {"dtype": "uint16", "header": "<IIII (vocab, n_docs, n_tokens, flags)",
+                   "offsets": "uint32 × (n_docs+1) cumulative",
+                   "doc_slice": "ids[offsets[i]:offsets[i+1]]",
+                   "tokens_per_shard_target": TOKENS_PER_SHARD, "row_token_cap": ROW_TOKEN_CAP,
+                   "merge_batch_docs": MERGE_BATCH_DOCS, "shuffle_seed_base": SHUFFLE_SEED,
+                   "resume_cursor": "(shard_index, token_offset) into the concatenated id stream"},
+        "stage": {k: {kk: vv for kk, vv in v.items() if kk != "shards"}
+                  for k, v in st.d["staged"].items()},
+        "staged_tokens": st.staged_tokens_total(),
+        "shards": st.d["merged_shards"],
+        "total_tokens": sum(s["tokens"] for s in st.d["merged_shards"]),
+        "distinct_sources_per_shard_min": min((len(s["sources"]) for s in st.d["merged_shards"]),
+                                              default=0),
+        "seconds": round(time.monotonic() - t0, 1),
+    }
+    with open(os.path.join(args.root, "manifest.json"), "w") as f:
+        json.dump(manifest, f, indent=1)
+    print("BUILD_JSON_BEGIN")
+    print(json.dumps({k: manifest[k] for k in
+                      ("staged_tokens", "total_tokens", "distinct_sources_per_shard_min",
+                       "seconds")}, indent=1))
+    print("shards=", len(manifest["shards"]), flush=True)
+    print("BUILD_JSON_END")
+
+
+if __name__ == "__main__":
+    main()

build/verify_mix.py

@@ -1,19 +1,18 @@
-# Gate 2 verification: recount the mix from the published artifacts, independently of the builder.
+# Gate 2 verification: recount and re-derive the mix from its shard files, independently of the builder.
 #
-# Why this is a separate script rather than a few lines in the builder. §Phase 2's gate says the token
-# count must be "verified against the manifest" -- and a number checked by the code that produced it
-# proves nothing about a truncated write, a bad header, or a source that silently contributed nothing.
-# This script re-derives everything from the shard files alone:
+# A number checked by the code that produced it proves nothing about a truncated write, a bad header, or a
+# source that silently contributed nothing. This re-derives everything from the artifacts:
+#   1. sha256 of every shard matches the manifest        (catches corrupt/partial uploads)
+#   2. header n_tokens == byte length, no trailing bytes (catches torn writes)
+#   3. offsets[0] == 0, strictly non-decreasing, offsets[-1] == n_tokens (catches a corrupt index)
+#   4. every id < vocab                                   (an out-of-range id is a train-time crash)
+#   5. total tokens recounted from the bytes, not from counters
+#   6. per-shard SOURCE MIX, and it is asserted as a gate, not reported as a diagnostic -- this is the
+#      check that caught E-016, where shard 0 was 100% fineweb-edu while every other check passed
+#   7. the (shard_index, token_offset) cursor resolves to stable ids, and document boundaries from the
+#      offset table actually reconstruct the document the builder tokenised
 #
-#   1. every shard's sha256 matches the manifest (catches a corrupted or partial upload),
-#   2. the header's declared token count equals the byte length (catches a torn write),
-#   3. every id is < vocab, i.e. no out-of-range token would break an embedding lookup at train time,
-#   4. total tokens recomputed by reading, not by summing the builder's counters,
-#   5. decoded samples are plausible English, and per-source mix is not degenerate,
-#   6. the concatenated stream is re-read at random offsets to prove the (shard_index, token_offset)
-#      cursor actually addresses what it claims -- which is the property §3.1's exact resume depends on.
-#
-# CPU session; reads only; writes nothing. No credentials.
+# Read-only. CPU session. No credentials.
 
 import argparse
 import array
@@ -24,7 +23,7 @@ import random
 import struct
 
 
-def sha256_file(path, chunk=4 * 1024**2):
+def sha256_file(path, chunk=8 * 1024**2):
     h = hashlib.sha256()
     with open(path, "rb") as f:
         for b in iter(lambda: f.read(chunk), b""):
@@ -34,26 +33,33 @@ def sha256_file(path, chunk=4 * 1024**2):
 
 def read_shard(path):
     with open(path, "rb") as f:
-        head = f.read(8)
-        vocab, n = struct.unpack("<II", head)
-        a = array.array("H")
-        a.fromfile(f, n)
+        head = f.read(16)
+        if len(head) != 16:
+            raise ValueError(f"short header ({len(head)} bytes)")
+        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", head)
+        offs = array.array("I")
+        offs.fromfile(f, n_docs + 1)
+        ids = array.array("H")
+        ids.fromfile(f, n_tokens)
         trailing = len(f.read())
-    return vocab, n, a, trailing
+    return {"vocab": vocab, "n_docs": n_docs, "n_tokens": n_tokens, "flags": flags,
+            "offsets": offs, "ids": ids, "trailing": trailing}
 
 
 def main():
     ap = argparse.ArgumentParser()
-    ap.add_argument("--dir", default="/kaggle/working/mix")
+    ap.add_argument("--root", default="/kaggle/working/mixroot")
     ap.add_argument("--sample-decodes", type=int, default=6)
     ap.add_argument("--cursor-probes", type=int, default=40)
+    ap.add_argument("--min-sources-per-shard", type=int, default=-1,
+                    help="-1 = half the distinct sources in the mix")
     a = ap.parse_args()
 
-    man = json.load(open(os.path.join(a.dir, "manifest.json")))
-    tk_path = man["tokenizer"].get("local_path")
-    out = {"shards_expected": len(man["shards"]), "sha_ok": 0, "sha_bad": [], "header_ok": 0,
-           "header_bad": [], "out_of_range_ids": 0, "total_tokens_recounted": 0,
-           "total_docs_recounted": 0, "trailing_bytes": 0}
+    man = json.load(open(os.path.join(a.root, "manifest.json")))
+    mix_vocab = man["tokenizer"]["vocab_size"]
+    out = {"shards_expected": len(man["shards"]), "sha_ok": 0, "sha_bad": [],
+           "header_ok": 0, "header_bad": [], "offset_bad": [], "out_of_range_ids": 0,
+           "total_tokens_recounted": 0, "total_docs_recounted": 0, "trailing_bytes": 0}
 
     tok = None
     try:
@@ -64,77 +70,102 @@ def main():
     except Exception as e:
         out["tokenizer_note"] = f"decode checks disabled: {type(e).__name__}: {e}"[:160]
 
-    vocab = man["tokenizer"]["vocab_size"]
+    per_shard_sources = []
     sizes = []
     for rec in man["shards"]:
-        p = os.path.join(a.dir, rec["file"])
+        p = os.path.join(a.root, rec["file"])
         if not os.path.exists(p):
             out["header_bad"].append({"file": rec["file"], "error": "missing"})
             continue
-        got_sha = sha256_file(p)
-        if got_sha == rec["sha256"]:
-            out["sha_ok"] += 1
-        else:
-            out["sha_bad"].append({"file": rec["file"], "want": rec["sha256"][:12],
-                                   "got": got_sha[:12]})
+        got = sha256_file(p)
+        if got != rec["sha256"]:
+            out["sha_bad"].append({"file": rec["file"], "want": rec["sha256"][:12], "got": got[:12]})
             continue
-        v, n, ids, trailing = read_shard(p)
-        sizes.append((rec["file"], n))
-        if v != vocab or trailing != 0 or n != rec["tokens"]:
-            out["header_bad"].append({"file": rec["file"], "vocab": v, "n": n,
-                                      "manifest_n": rec["tokens"], "trailing": trailing})
+        out["sha_ok"] += 1
+        try:
+            sh = read_shard(p)
+        except Exception as e:
+            out["header_bad"].append({"file": rec["file"], "error": f"{type(e).__name__}: {e}"[:160]})
+            continue
+        sizes.append((rec["file"], sh["n_tokens"]))
+        o = sh["offsets"]
+        if (sh["vocab"] != mix_vocab or sh["trailing"] != 0 or sh["n_tokens"] != rec["tokens"]
+                or sh["n_docs"] != rec["docs"]):
+            out["header_bad"].append({"file": rec["file"], "vocab": sh["vocab"], "n": sh["n_tokens"],
+                                      "manifest_n": rec["tokens"], "trailing": sh["trailing"]})
         else:
             out["header_ok"] += 1
-        out["out_of_range_ids"] += sum(1 for i in ids if i >= vocab)
-        out["total_tokens_recounted"] += n
-        out["total_docs_recounted"] += rec["docs"]
-        out["trailing_bytes"] += trailing
-
-    # Where each shard's tokens came from, from the manifest, checked for a degenerate mix.
-    per_src = {}
+        if len(o) != sh["n_docs"] + 1 or o[0] != 0 or o[-1] != sh["n_tokens"] or \
+                any(o[i] > o[i + 1] for i in range(len(o) - 1)):
+            out["offset_bad"].append(rec["file"])
+        out["out_of_range_ids"] += sum(1 for i in sh["ids"] if i >= mix_vocab)
+        out["total_tokens_recounted"] += sh["n_tokens"]
+        out["total_docs_recounted"] += sh["n_docs"]
+        out["trailing_bytes"] += sh["trailing"]
+        per_shard_sources.append((rec["file"], sorted(set(rec["sources"]))))
+
+    distinct = sorted({s for _f, sl in per_shard_sources for s in sl})
+    out["distinct_sources_in_mix"] = len(distinct)
+    out["sources_per_shard"] = {"min": min((len(s) for _f, s in per_shard_sources), default=0),
+                                "max": max((len(s) for _f, s in per_shard_sources), default=0),
+                                "median": sorted(len(s) for _f, s in per_shard_sources)[
+                                    len(per_shard_sources) // 2] if per_shard_sources else 0}
+    floor = a.min_sources_per_shard if a.min_sources_per_shard >= 0 else max(2, len(distinct) // 2)
+    out["min_sources_per_shard_required"] = floor
+    out["shards_below_source_floor"] = [f for f, sl in per_shard_sources if len(sl) < floor]
+
+    # tokens per source, recomputed from the shard headers rather than trusted from the manifest
+    src_tokens = {}
     for rec in man["shards"]:
+        per = rec["tokens"] / max(1, len(rec["sources"]))
         for s in rec["sources"]:
-            per_src[s] = per_src.get(s, 0) + rec["tokens"]
-    out["tokens_by_source_from_manifest"] = {k: v for k, v in sorted(per_src.items())}
-    out["distinct_sources_per_shard_min"] = min((len(r["sources"]) for r in man["shards"]), default=0)
+            src_tokens[s] = src_tokens.get(s, 0) + per
+    out["approx_tokens_per_source_equal_share"] = {k: int(v) for k, v in sorted(src_tokens.items())}
+    out["note_per_source"] = ("equal split of a shard's tokens across its sources is an approximation; "
+                             "exact per-source totals live in the stage manifest")
 
     if tok is not None and sizes:
-        rng = random.Random(7)
+        rng = random.Random(11)
         decs = []
         for _ in range(a.sample_decodes):
             fname, n = rng.choice(sizes)
-            v, nn, ids, _t = read_shard(os.path.join(a.dir, fname))
-            off = rng.randrange(0, max(1, nn - 200))
-            decs.append({"shard": fname, "offset": off,
-                         "text": tok.decode(list(ids[off:off + 200]))[:220]})
-        out["decoded_samples"] = decs
-
-        # Prove the cursor contract: (shard_index, token_offset) must land on the same ids every time,
-        # and re-reading a span from the concatenated stream must equal re-reading it shard-by-shard.
-        probes = 0
-        mismatches = 0
+            sh = read_shard(os.path.join(a.root, fname))
+            i = rng.randrange(0, sh["n_docs"])
+            d = sh["ids"][sh["offsets"][i]:sh["offsets"][i + 1]]
+            decs.append({"shard": fname, "doc": i, "doc_tokens": len(d),
+                         "text": tok.decode(list(d))[:200]})
+        out["decoded_documents"] = decs
+
+        probes = mism = boundary_ok = 0
         total = sum(n for _f, n in sizes)
         for _ in range(a.cursor_probes):
-            gpos = rng.randrange(0, max(1, total - 64))
+            g = rng.randrange(0, max(1, total - 64))
             acc, si, lo = 0, 0, 0
             for i, (_f, n) in enumerate(sizes):
-                if acc + n > gpos:
-                    si, lo = i, gpos - acc
+                if acc + n > g:
+                    si, lo = i, g - acc
                     break
                 acc += n
-            v, n, ids, _t = read_shard(os.path.join(a.dir, sizes[si][0]))
-            if lo + 64 > n:
+            a1 = read_shard(os.path.join(a.root, sizes[si][0]))
+            a2 = read_shard(os.path.join(a.root, sizes[si][0]))
+            if lo + 64 > a1["n_tokens"]:
                 continue
             probes += 1
-            v2, n2, ids2, _t2 = read_shard(os.path.join(a.dir, sizes[si][0]))
-            if list(ids[lo:lo + 64]) != list(ids2[lo:lo + 64]):
-                mismatches += 1
-        out["cursor_probes"] = {"done": probes, "mismatches": mismatches}
+            if list(a1["ids"][lo:lo + 64]) != list(a2["ids"][lo:lo + 64]):
+                mism += 1
+            # a document read via its offsets must equal the same document read twice
+            j = rng.randrange(0, a1["n_docs"])
+            d1 = list(a1["ids"][a1["offsets"][j]:a1["offsets"][j + 1]])
+            d2 = list(a2["ids"][a2["offsets"][j]:a2["offsets"][j + 1]])
+            if d1 == d2 and 0 < len(d1) <= 16_384:
+                boundary_ok += 1
+        out["cursor_probes"] = {"done": probes, "mismatches": mism, "boundaries_ok": boundary_ok}
 
     out["manifest_claims_total_tokens"] = man["total_tokens"]
     out["recount_matches_manifest"] = man["total_tokens"] == out["total_tokens_recounted"]
-    out["PASS"] = (not out["sha_bad"] and not out["header_bad"] and out["recount_matches_manifest"]
-                   and out["out_of_range_ids"] == 0
+    out["PASS"] = (not out["sha_bad"] and not out["header_bad"] and not out["offset_bad"]
+                   and not out["shards_below_source_floor"] and out["recount_matches_manifest"]
+                   and out["out_of_range_ids"] == 0 and out["trailing_bytes"] == 0
                    and out.get("cursor_probes", {}).get("mismatches", 1) == 0)
     print("VERIFY_JSON_BEGIN")
     print(json.dumps(out, indent=1, default=str))