CELL 1 rewritten as an INCREMENTAL materializer: the fused set is ~160GB and the campaign needs ~9k rows — shards visited in sha256(name) order one at a time, gated+deduped+slimmed into the pool, source deleted per shard, hard stop at TRAIN_TARGET+EVAL_N (8192+512). Download bounded to ~10GB instead of the corpus. Split logic executed green against a real shard.

colab/anima_closeout.ipynb

@@ -171,126 +171,138 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "# \u2550\u2550\u2550 CELL 1 \u2014 materialize the gated held-out split + write TOMLs \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
-    "# Row pipeline, exactly the campaign record:\n",
-    "#   AUDIT GATE (pod2/dexp011:98-105): a row passes when its age_audit/audit\n",
-    "#     value contains \"approved\" or \"pass\" (case-insensitive); gated rows are\n",
-    "#     counted and reported loudly.\n",
-    "#   DEDUP by id (first occurrence wins) \u2014 guards the stray rank0.part00\n",
-    "#     shard alongside the shard_r0_* series.\n",
-    "#   CAPTION = caption_joycaption or prompt_fused fallback (dexp011:154),\n",
-    "#     materialized as a `caption` column so the fork reads one key.\n",
-    "#   SPLIT: sha256(id) ascending; first EVAL_N surviving rows are eval.\n",
-    "import hashlib, json\n",
+    "# \u2550\u2550\u2550 CELL 1 \u2014 materialize the gated split INCREMENTALLY \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# The fused dataset is ~160GB across 200+ shards; the campaign needs ~9k\n",
+    "# rows. So: NEVER download the corpus. Shards are visited in sha256(name)\n",
+    "# order (deterministic, free of curation-order bias), one at a time; each\n",
+    "# is gated + deduped + slimmed into the pool, then its source file is\n",
+    "# DELETED; we stop the moment the row target is met. Disk stays bounded at\n",
+    "# roughly one source shard + the materialized output.\n",
+    "#\n",
+    "# Row rules, verbatim from the campaign record:\n",
+    "#   AUDIT GATE (pod2/dexp011:98-105): age_audit/audit must contain\n",
+    "#     \"approved\" or \"pass\" (case-insensitive); gated rows counted loudly.\n",
+    "#   CAPTION = caption_joycaption or prompt_fused fallback (dexp011:154;\n",
+    "#     cond = \"caption_joycaption nl77\"), materialized as `caption`.\n",
+    "#   DEDUP by id (first occurrence wins).\n",
+    "#   SPLIT: within the pool, sha256(id) ascending; first EVAL_N are eval.\n",
+    "import hashlib, json, os, shutil\n",
     "import pyarrow as pa\n",
     "import pyarrow.parquet as pq\n",
     "from huggingface_hub import hf_hub_download\n",
     "\n",
-    "EVAL_N = 512\n",
+    "TRAIN_TARGET = 8192      # ~2.7x exp004's criticized 2978; 10k steps @ mbs8\n",
+    "EVAL_N = 512             # = 80k samples ~= 9 epochs over this pool\n",
+    "KEEP_COLS = [\"id\", \"image\", \"image_width\", \"image_height\",\n",
+    "             \"caption\", \"fused_json\", \"age_audit\"]   # slim rows only\n",
+    "\n",
     "TRAIN_DIR = DATA / \"train\"; TRAIN_DIR.mkdir(exist_ok=True)\n",
     "EVAL_PARQUET = DATA / \"eval.parquet\"\n",
     "MANIFEST = DATA / \"split_manifest.json\"\n",
+    "SRC = DATA / \"src\"; SRC.mkdir(exist_ok=True)\n",
     "\n",
-    "def _passes_gate(row_get):\n",
-    "    vals = [str(row_get(c)) for c in (\"age_audit\", \"audit\")\n",
-    "            if row_get(c) is not None]\n",
+    "def _passes_gate(val_of):\n",
+    "    vals = [str(val_of(c)) for c in (\"age_audit\", \"audit\")\n",
+    "            if val_of(c) is not None]\n",
     "    return (not vals) or any((\"approved\" in v.lower() or \"pass\" in v.lower())\n",
     "                             for v in vals)\n",
     "\n",
     "if not MANIFEST.exists():\n",
-    "    _paths = [hf_hub_download(ANIMA[\"DATASET\"], sh, repo_type=\"dataset\")\n",
-    "              for sh in ANIMA[\"SHARDS\"]]\n",
-    "\n",
-    "    # pass 1: gate + dedup + rank survivors by sha256(id)\n",
-    "    _key_cols = [c for c in (\"id\", \"caption_joycaption\", \"prompt_fused\",\n",
-    "                             \"age_audit\", \"audit\")\n",
-    "                 if True]\n",
-    "    seen_ids, ranked = set(), []\n",
+    "    # deterministic pseudo-random shard order \u2014 never curation order\n",
+    "    shard_order = sorted(\n",
+    "        ANIMA[\"SHARDS\"],\n",
+    "        key=lambda n: hashlib.sha256(n.encode()).hexdigest())\n",
+    "\n",
+    "    pool = []                 # slim row dicts\n",
+    "    seen_ids = set()\n",
     "    n_seen = n_gated = n_dup = 0\n",
-    "    for si, pth in enumerate(_paths):\n",
-    "        pf = pq.ParquetFile(pth)\n",
-    "        avail = [c for c in _key_cols if c in pf.schema_arrow.names]\n",
-    "        ri = 0\n",
-    "        for batch in pf.iter_batches(batch_size=256, columns=avail):\n",
-    "            cols = {c: batch.column(c) for c in avail}\n",
+    "    shards_used = []\n",
+    "    target = TRAIN_TARGET + EVAL_N\n",
+    "\n",
+    "    for sh in shard_order:\n",
+    "        if len(pool) >= target:\n",
+    "            break\n",
+    "        local = hf_hub_download(ANIMA[\"DATASET\"], sh, repo_type=\"dataset\",\n",
+    "                                local_dir=str(SRC))\n",
+    "        pf = pq.ParquetFile(local)\n",
+    "        names = pf.schema_arrow.names\n",
+    "        cols = [c for c in (\"id\", \"image\", \"image_width\", \"image_height\",\n",
+    "                            \"caption_joycaption\", \"prompt_fused\",\n",
+    "                            \"fused_json\", \"age_audit\", \"audit\")\n",
+    "                if c in names]\n",
+    "        for batch in pf.iter_batches(batch_size=32, columns=cols):\n",
+    "            if len(pool) >= target:\n",
+    "                break\n",
+    "            data = {c: batch.column(c) for c in\n",
+    "                    (set(cols) & set(batch.schema.names))}\n",
     "            for j in range(batch.num_rows):\n",
+    "                if len(pool) >= target:\n",
+    "                    break\n",
     "                n_seen += 1\n",
-    "                get = lambda c: (cols[c][j].as_py() if c in cols else None)\n",
-    "                rid = str(get(\"id\"))\n",
+    "                get = lambda c: (data[c][j].as_py() if c in data else None)\n",
     "                if not _passes_gate(get):\n",
     "                    n_gated += 1\n",
-    "                elif rid in seen_ids:\n",
+    "                    continue\n",
+    "                rid = str(get(\"id\"))\n",
+    "                if rid in seen_ids:\n",
     "                    n_dup += 1\n",
-    "                else:\n",
-    "                    seen_ids.add(rid)\n",
-    "                    ranked.append((hashlib.sha256(rid.encode()).hexdigest(),\n",
-    "                                   si, ri, rid))\n",
-    "                ri += 1\n",
-    "    ranked.sort()\n",
-    "    assert len(ranked) > 4 * EVAL_N, (\n",
-    "        f\"only {len(ranked)} rows survive the audit gate \"\n",
-    "        f\"(saw {n_seen}, gated {n_gated}, dup {n_dup}) \u2014 inspect before \"\n",
-    "        \"training on so little\")\n",
-    "    eval_set = {(si, ri) for _h, si, ri, _k in ranked[:EVAL_N]}\n",
-    "    keep_set = {(si, ri) for _h, si, ri, _k in ranked}\n",
-    "\n",
-    "    # pass 2: route rows, appending the derived caption column\n",
-    "    ew = tw = None\n",
-    "    for si, pth in enumerate(_paths):\n",
-    "        pf = pq.ParquetFile(pth)\n",
-    "        names = pf.schema_arrow.names\n",
-    "        ri = 0\n",
-    "        for batch in pf.iter_batches(batch_size=64):\n",
-    "            tbl = pa.Table.from_batches([batch])\n",
-    "            jc = tbl.column(\"caption_joycaption\")\n",
-    "            pfu = (tbl.column(\"prompt_fused\")\n",
-    "                   if \"prompt_fused\" in names else None)\n",
-    "            caps, kmask, emask = [], [], []\n",
-    "            for j in range(tbl.num_rows):\n",
-    "                c = jc[j].as_py() or (pfu[j].as_py() if pfu is not None\n",
-    "                                      else None) or \"\"\n",
-    "                caps.append(str(c))\n",
-    "                kmask.append((si, ri + j) in keep_set)\n",
-    "                emask.append((si, ri + j) in eval_set)\n",
-    "            ri += tbl.num_rows\n",
-    "            tbl = tbl.append_column(\"caption\", pa.array(caps))\n",
-    "            esub = tbl.filter(pa.array(emask))\n",
-    "            tsub = tbl.filter(pa.array([k and not e for k, e\n",
-    "                                        in zip(kmask, emask)]))\n",
-    "            if esub.num_rows:\n",
-    "                if ew is None:\n",
-    "                    ew = pq.ParquetWriter(EVAL_PARQUET, esub.schema)\n",
-    "                ew.write_table(esub)\n",
-    "            if tsub.num_rows:\n",
-    "                if tw is None:\n",
-    "                    tw = pq.ParquetWriter(TRAIN_DIR / \"train.parquet\",\n",
-    "                                          tsub.schema)\n",
-    "                tw.write_table(tsub)\n",
-    "    for w in (ew, tw):\n",
-    "        if w:\n",
-    "            w.close()\n",
+    "                    continue\n",
+    "                cap = get(\"caption_joycaption\") or get(\"prompt_fused\") or \"\"\n",
+    "                if not str(cap).strip():\n",
+    "                    n_gated += 1          # empty caption: unusable row\n",
+    "                    continue\n",
+    "                seen_ids.add(rid)\n",
+    "                pool.append({\n",
+    "                    \"id\": rid, \"image\": get(\"image\"),\n",
+    "                    \"image_width\": get(\"image_width\"),\n",
+    "                    \"image_height\": get(\"image_height\"),\n",
+    "                    \"caption\": str(cap),\n",
+    "                    \"fused_json\": str(get(\"fused_json\") or \"\"),\n",
+    "                    \"age_audit\": str(get(\"age_audit\") or \"\")})\n",
+    "        shards_used.append(sh)\n",
+    "        shutil.rmtree(SRC, ignore_errors=True)   # bound the disk\n",
+    "        SRC.mkdir(exist_ok=True)\n",
+    "        print(f\"[split] {sh}: pool {len(pool)}/{target} \"\n",
+    "              f\"(seen {n_seen}, gated {n_gated}, dup {n_dup})\", flush=True)\n",
+    "\n",
+    "    assert len(pool) >= target, (\n",
+    "        f\"exhausted ALL shards with only {len(pool)} usable rows \"\n",
+    "        f\"(target {target}) \u2014 inspect the gate before training\")\n",
+    "\n",
+    "    # split within the pool: sha256(id) ascending, first EVAL_N are eval\n",
+    "    pool.sort(key=lambda r: hashlib.sha256(r[\"id\"].encode()).hexdigest())\n",
+    "    eval_rows, train_rows = pool[:EVAL_N], pool[EVAL_N:target]\n",
+    "\n",
+    "    def _write(rows, path):\n",
+    "        tbl = pa.table({c: [r[c] for r in rows] for c in KEEP_COLS})\n",
+    "        pq.write_table(tbl, path)\n",
+    "        return tbl.num_rows\n",
+    "\n",
+    "    ne = _write(eval_rows, EVAL_PARQUET)\n",
+    "    nt = _write(train_rows, TRAIN_DIR / \"train.parquet\")\n",
     "    MANIFEST.write_text(json.dumps(\n",
-    "        {\"dataset\": ANIMA[\"DATASET\"], \"rows_seen\": n_seen,\n",
-    "         \"gated_out\": n_gated, \"duplicate_ids\": n_dup,\n",
-    "         \"rows_kept\": len(ranked), \"eval_n\": EVAL_N,\n",
-    "         \"eval_ids\": [k for _h, _s, _r, k in ranked[:EVAL_N]],\n",
+    "        {\"dataset\": ANIMA[\"DATASET\"],\n",
+    "         \"shard_order\": \"sha256(shard_name) ascending\",\n",
+    "         \"shards_used\": shards_used, \"rows_seen\": n_seen,\n",
+    "         \"gated_or_empty\": n_gated, \"duplicate_ids\": n_dup,\n",
+    "         \"train_rows\": nt, \"eval_n\": ne,\n",
+    "         \"eval_ids\": [r[\"id\"] for r in eval_rows],\n",
     "         \"caption_rule\": \"caption_joycaption or prompt_fused fallback \"\n",
     "                         \"(dexp011:154; cond = caption_joycaption nl77)\",\n",
     "         \"gate_rule\": \"age_audit/audit must contain approved|pass \"\n",
     "                      \"(dexp011:98-105)\",\n",
-    "         \"split_rule\": \"sha256(id) ascending; first EVAL_N are eval\"},\n",
+    "         \"split_rule\": \"sha256(id) ascending within the pool; first \"\n",
+    "                       \"EVAL_N are eval\"},\n",
     "        indent=1))\n",
-    "    print(f\"split: seen {n_seen} -> gated {n_gated}, dup {n_dup}, kept \"\n",
-    "          f\"{len(ranked)} = {len(ranked) - EVAL_N} train + {EVAL_N} eval\")\n",
+    "    print(f\"split DONE: {nt} train + {ne} eval from \"\n",
+    "          f\"{len(shards_used)}/{len(ANIMA['SHARDS'])} shards \"\n",
+    "          f\"(seen {n_seen}, gated/empty {n_gated}, dup {n_dup})\")\n",
     "else:\n",
     "    _m = json.loads(MANIFEST.read_text())\n",
-    "    print(f\"split already materialized: kept {_m['rows_kept']} \"\n",
-    "          f\"(gated {_m['gated_out']}, dup {_m['duplicate_ids']})\")\n",
+    "    print(f\"split already materialized: {_m['train_rows']} train + \"\n",
+    "          f\"{_m['eval_n']} eval from {len(_m['shards_used'])} shards\")\n",
     "\n",
     "# -- dataset TOMLs (fused recipe: derived caption column) -----------------\n",
-    "_wh = (\"image_width_column = 'image_width'\\n\"\n",
-    "       \"image_height_column = 'image_height'\\n\") if ANIMA[\"HAS_WH\"] else \"\"\n",
-    "\n",
     "def _dataset_toml(files):\n",
     "    return f\"\"\"# generated by anima_closeout.ipynb\n",
     "resolutions = [512]\n",
@@ -305,7 +317,9 @@
     "type = 'parquet'\n",
     "parquet_files = '{files}'\n",
     "image_column = 'image'\n",
-    "{_wh}caption_column = 'caption'\n",
+    "image_width_column = 'image_width'\n",
+    "image_height_column = 'image_height'\n",
+    "caption_column = 'caption'\n",
     "caption_type = 'text'\n",
     "num_repeats = 1\n",
     "skip_empty_caption = true\n",