fetch_coco_subset.py

"""Build the PixelModel v3 train/eval data from MS-COCO.

Same source as v1: `sayakpaul/coco-30-val-2014` (MS-COCO val2014 caption/image
pairs). The set is split *deterministically by image hash* into a train subset
and an eval subset so the two are provably disjoint and reproducible - exactly
the regression-safe protocol v1 used.

Outputs (into --out, default ../pm-work):
  coco_train.npz   images (uint8, N x S x S x 3) + token ids (int32, N x T)
  coco_eval.npz    eval images + raw caption strings (for FID/CLIP)
  vocab.json       word-level tokenizer vocabulary (built from TRAIN only)
  split_manifest.json  hashes + counts, so disjointness is auditable

Run on the training machine (needs `datasets` + network):
    python fetch_coco_subset.py --out ../pm-work
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import time

import numpy as np
from PIL import Image

from model import build_vocab, encode_caption, PAD_ID


def center_square_resize(img: Image.Image, size: int) -> np.ndarray:
    """Center-crop to a square, resize to size x size, return uint8 HWC RGB."""
    img = img.convert("RGB")
    w, h = img.size
    s = min(w, h)
    left = (w - s) // 2
    top = (h - s) // 2
    img = img.crop((left, top, left + s, top + s)).resize((size, size), Image.BICUBIC)
    return np.asarray(img, dtype=np.uint8)


def image_hash(img: Image.Image) -> str:
    """Stable content hash of an image, independent of dataset row order."""
    buf = io.BytesIO()
    img.convert("RGB").resize((64, 64), Image.BICUBIC).save(buf, format="PNG")
    return hashlib.md5(buf.getvalue()).hexdigest()


def caption_of(example) -> str:
    """COCO rows expose captions under a few different keys across mirrors."""
    for key in ("caption", "captions", "text", "sentences", "annotations_captions"):
        if key in example and example[key]:
            v = example[key]
            if isinstance(v, (list, tuple)):
                return str(v[0])
            return str(v)
    raise KeyError(f"no caption field found; row keys = {list(example.keys())}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="../pm-work", help="output directory")
    ap.add_argument("--dataset", default="sayakpaul/coco-30-val-2014",
                    help="HF dataset id (MS-COCO caption/image pairs)")
    ap.add_argument("--split", default="train", help="HF split to pull from")
    ap.add_argument("--n-train", type=int, default=20000)
    ap.add_argument("--n-eval", type=int, default=5000)
    ap.add_argument("--store-size", type=int, default=144,
                    help="square size images are stored at (must be >= train crop)")
    ap.add_argument("--vocab-size", type=int, default=8192, help="base model, incl. PAD+UNK")
    ap.add_argument("--vocab-size-xl", type=int, default=16384, help="XL variant vocab cap")
    ap.add_argument("--min-freq", type=int, default=2)
    ap.add_argument("--max-tokens", type=int, default=20)
    ap.add_argument("--eval-hash-mod", type=int, default=6,
                    help="1/eval-hash-mod of images are reserved for eval")
    ap.add_argument("--streaming", action="store_true",
                    help="stream the split instead of downloading it first. "
                         "Faster to start, but a long collection loop can die "
                         "to a dropped connection; off by default.")
    ap.add_argument("--load-retries", type=int, default=6)
    ap.add_argument("--max-reconnects", type=int, default=8,
                    help="resume attempts if iteration dies mid-collection")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    from datasets import load_dataset

    def open_dataset():
        """Load the split, retrying the network part. Default is NON-streaming:
        the shards are downloaded once (HF retries/resumes those internally) and
        iteration is then purely local, so a long collection loop cannot die
        halfway to a dropped HTTP connection."""
        last = None
        for t in range(args.load_retries):
            try:
                return load_dataset(args.dataset, split=args.split,
                                    streaming=args.streaming)
            except Exception as e:
                last = e
                wait = 5 * (t + 1)
                print(f"[data] load attempt {t + 1}/{args.load_retries} failed "
                      f"({type(e).__name__}: {e}); retrying in {wait}s", flush=True)
                time.sleep(wait)
        raise RuntimeError(f"could not load {args.dataset}: {last}")

    train_imgs, train_caps_raw, train_hashes = [], [], []
    eval_imgs, eval_caps, eval_hashes = [], [], []
    seen = set()

    def need_more():
        return len(train_imgs) < args.n_train or len(eval_imgs) < args.n_eval

    last_report = 0
    reconnects = 0
    print(f"[data] loading {args.dataset} split={args.split} "
          f"(streaming={args.streaming})", flush=True)
    while need_more():
        ds = open_dataset()
        try:
            for ex in ds:
                if not need_more():
                    break
                try:
                    img = ex.get("image") or ex.get("img")
                    if img is None:
                        continue
                    cap = caption_of(ex)
                    h = image_hash(img)
                    if h in seen:
                        continue
                    seen.add(h)
                    is_eval = (int(h, 16) % args.eval_hash_mod == 0)
                    if is_eval:
                        if len(eval_imgs) >= args.n_eval:
                            continue
                        eval_imgs.append(center_square_resize(img, 256))
                        eval_caps.append(cap)
                        eval_hashes.append(h)
                    else:
                        if len(train_imgs) >= args.n_train:
                            continue
                        train_imgs.append(center_square_resize(img, args.store_size))
                        train_caps_raw.append(cap)
                        train_hashes.append(h)
                except Exception:
                    continue
                total = len(train_imgs) + len(eval_imgs)
                if total - last_report >= 1000:
                    last_report = total
                    print(f"[data] train={len(train_imgs)} eval={len(eval_imgs)}", flush=True)
            break
        except Exception as e:
            reconnects += 1
            print(f"[data] iteration died ({type(e).__name__}: {e}); "
                  f"reconnect {reconnects}/{args.max_reconnects} with "
                  f"train={len(train_imgs)} eval={len(eval_imgs)}", flush=True)
            if reconnects >= args.max_reconnects:
                print("[data] giving up on reconnects; proceeding with what we have",
                      flush=True)
                break
            time.sleep(5)

    assert set(train_hashes).isdisjoint(set(eval_hashes)), "train/eval overlap!"
    print(f"[data] collected train={len(train_imgs)} eval={len(eval_imgs)} "
          f"(disjoint: {set(train_hashes).isdisjoint(set(eval_hashes))})")

    vocab = build_vocab(train_caps_raw, args.vocab_size, args.min_freq)
    with open(os.path.join(args.out, "vocab.json"), "w") as f:
        json.dump(vocab, f)
    vocab_xl = build_vocab(train_caps_raw, args.vocab_size_xl, args.min_freq)
    with open(os.path.join(args.out, "vocab_xl.json"), "w") as f:
        json.dump(vocab_xl, f)
    print(f"[data] vocab size = {len(vocab)} (cap {args.vocab_size}), "
          f"vocab_xl = {len(vocab_xl)} (cap {args.vocab_size_xl})")

    train_tokens = np.stack(
        [encode_caption(c, vocab, args.max_tokens) for c in train_caps_raw])
    np.savez_compressed(
        os.path.join(args.out, "coco_train.npz"),
        images=np.stack(train_imgs),
        tokens=train_tokens,
        captions=np.array(train_caps_raw, dtype=object),
        store_size=args.store_size,
        max_tokens=args.max_tokens,
    )
    np.savez_compressed(
        os.path.join(args.out, "coco_eval.npz"),
        images=np.stack(eval_imgs),
        captions=np.array(eval_caps, dtype=object),
    )
    with open(os.path.join(args.out, "split_manifest.json"), "w") as f:
        json.dump({
            "dataset": args.dataset,
            "n_train": len(train_imgs),
            "n_eval": len(eval_imgs),
            "eval_hash_mod": args.eval_hash_mod,
            "vocab_size": len(vocab),
            "store_size": args.store_size,
            "max_tokens": args.max_tokens,
            "disjoint": bool(set(train_hashes).isdisjoint(set(eval_hashes))),
            "train_hash_sample": train_hashes[:8],
            "eval_hash_sample": eval_hashes[:8],
        }, f, indent=2)

    mb = os.path.getsize(os.path.join(args.out, "coco_train.npz")) / 1e6
    print(f"[data] wrote coco_train.npz ({mb:.1f} MB), coco_eval.npz, vocab.json")
    print(f"[data] done -> {args.out}")


if __name__ == "__main__":
    main()