deploy capit backend

.dockerignore

@@ -0,0 +1,15 @@
+# keep the build context tiny — the Dockerfile only needs pipeline/ + backend/ source.
+# weights come from the Hub at build time, not the context.
+.git
+**/.venv
+**/__pycache__
+**/*.egg-info
+**/.pytest_cache
+**/.ruff_cache
+**/.mypy_cache
+data/
+frontend/
+docs/
+assets/
+**/node_modules
+*.zip

Dockerfile

@@ -0,0 +1,37 @@
+# capit backend — HF Spaces (Docker SDK). Build context = repo root:
+#   docker build -f backend/Dockerfile -t capit-api .
+FROM python:3.12-slim
+
+# uv, pinned to what the lockfile was authored with
+COPY --from=ghcr.io/astral-sh/uv:0.9.9 /uv /uvx /bin/
+
+# HF Spaces require a non-root user with UID 1000; set it up before any COPY
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    HF_HOME=/home/user/.cache/huggingface \
+    UV_PROJECT_ENVIRONMENT=/home/user/venv \
+    UV_LINK_MODE=copy \
+    UV_COMPILE_BYTECODE=1 \
+    CAPIT_ARTIFACT_REPO=Bukunmi2108/capit-sat \
+    PATH=/home/user/venv/bin:$PATH
+
+WORKDIR /home/user/app
+
+# the capit package (model classes) the backend imports, then the backend project
+COPY --chown=user pipeline/ ./pipeline/
+COPY --chown=user backend/ ./backend/
+
+WORKDIR /home/user/app/backend
+RUN uv sync --frozen --no-dev
+
+# bake weights into image layers so the Space has no cold-start downloads
+RUN uv run python -c "from huggingface_hub import hf_hub_download as d; d('Bukunmi2108/capit-sat','capit-sat.pt'); d('Bukunmi2108/capit-sat','vocab.json')" \
+ && uv run python -c "from transformers import BlipForConditionalGeneration as M, BlipProcessor as P; m='Salesforce/blip-image-captioning-base'; P.from_pretrained(m); M.from_pretrained(m)"
+
+# weights are baked above — serve from cache only, no runtime Hub calls (faster cold start)
+ENV HF_HUB_OFFLINE=1 \
+    TRANSFORMERS_OFFLINE=1
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,16 @@
+---
+title: capit
+emoji: 🔎
+colorFrom: red
+colorTo: gray
+sdk: docker
+app_port: 7860
+pinned: false
+short_description: A glass-box image captioner (Show, Attend and Tell) vs BLIP.
+---
+
+# capit — backend
+
+Side-by-side captioning API: a from-scratch *Show, Attend and Tell* model (glass box —
+per-word attention + rejected beams) and BLIP (closed box). Code:
+https://github.com/Bukunmi2108/capit

backend/Dockerfile

@@ -0,0 +1,37 @@
+# capit backend — HF Spaces (Docker SDK). Build context = repo root:
+#   docker build -f backend/Dockerfile -t capit-api .
+FROM python:3.12-slim
+
+# uv, pinned to what the lockfile was authored with
+COPY --from=ghcr.io/astral-sh/uv:0.9.9 /uv /uvx /bin/
+
+# HF Spaces require a non-root user with UID 1000; set it up before any COPY
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    HF_HOME=/home/user/.cache/huggingface \
+    UV_PROJECT_ENVIRONMENT=/home/user/venv \
+    UV_LINK_MODE=copy \
+    UV_COMPILE_BYTECODE=1 \
+    CAPIT_ARTIFACT_REPO=Bukunmi2108/capit-sat \
+    PATH=/home/user/venv/bin:$PATH
+
+WORKDIR /home/user/app
+
+# the capit package (model classes) the backend imports, then the backend project
+COPY --chown=user pipeline/ ./pipeline/
+COPY --chown=user backend/ ./backend/
+
+WORKDIR /home/user/app/backend
+RUN uv sync --frozen --no-dev
+
+# bake weights into image layers so the Space has no cold-start downloads
+RUN uv run python -c "from huggingface_hub import hf_hub_download as d; d('Bukunmi2108/capit-sat','capit-sat.pt'); d('Bukunmi2108/capit-sat','vocab.json')" \
+ && uv run python -c "from transformers import BlipForConditionalGeneration as M, BlipProcessor as P; m='Salesforce/blip-image-captioning-base'; P.from_pretrained(m); M.from_pretrained(m)"
+
+# weights are baked above — serve from cache only, no runtime Hub calls (faster cold start)
+ENV HF_HUB_OFFLINE=1 \
+    TRANSFORMERS_OFFLINE=1
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

backend/app.py

@@ -0,0 +1,127 @@
+"""capit backend: side-by-side captions from the glass-box SAT model and BLIP.
+
+Loads the Stage 4.4 serving artifact (local dir or HF Hub) + BLIP, exposes /health and
+/caption. The SAT path returns per-word attention, the center-crop box, and the rejected beams.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import time
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+import torch
+from fastapi import FastAPI, HTTPException, Query, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image, ImageOps, UnidentifiedImageError
+from transformers import BlipForConditionalGeneration, BlipProcessor
+
+from capit.serving import caption as sat_caption
+from capit.serving import center_crop_box, load_artifact, make_transform
+
+MAX_BYTES = 8 * 1024 * 1024
+BLIP_MODEL = "Salesforce/blip-image-captioning-base"
+ARTIFACT_REPO = os.environ.get("CAPIT_ARTIFACT_REPO")
+ARTIFACT_DIR = Path(os.environ.get("CAPIT_ARTIFACT_DIR", Path(__file__).resolve().parents[1] / "data" / "artifact"))
+ALLOWED_ORIGINS = os.environ.get("CAPIT_CORS", "http://localhost:5173,http://localhost:3000").split(",")
+
+state: dict = {}
+
+
+def _artifact_paths() -> tuple[Path, Path]:
+    """Hub repo (Space; baked into the image cache at build) or a local dir (dev)."""
+    if ARTIFACT_REPO:
+        from huggingface_hub import hf_hub_download
+
+        return (
+            Path(hf_hub_download(ARTIFACT_REPO, "capit-sat.pt")),
+            Path(hf_hub_download(ARTIFACT_REPO, "vocab.json")),
+        )
+    return ARTIFACT_DIR / "capit-sat.pt", ARTIFACT_DIR / "vocab.json"
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    torch.set_num_threads(2)
+    artifact, vocab_path = _artifact_paths()
+    encoder, decoder, vocab, preprocess = load_artifact(artifact, vocab_path)
+    state.update(
+        encoder=encoder,
+        decoder=decoder,
+        vocab=vocab,
+        transform=make_transform(preprocess),
+        preprocess=preprocess,
+        blip_processor=BlipProcessor.from_pretrained(BLIP_MODEL),
+        blip=BlipForConditionalGeneration.from_pretrained(BLIP_MODEL).eval(),
+    )
+    yield
+    state.clear()
+
+
+app = FastAPI(title="Capit", lifespan=lifespan)
+app.add_middleware(
+    CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["GET", "POST"], allow_headers=["*"]
+)
+
+
+async def _ingest(file: UploadFile) -> Image.Image:
+    body = await file.read(MAX_BYTES + 1)
+    if len(body) > MAX_BYTES:
+        raise HTTPException(413, f"image exceeds {MAX_BYTES // (1024 * 1024)} MB")
+    try:
+        image = Image.open(io.BytesIO(body))
+        image.load()
+    except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc:
+        raise HTTPException(422, f"could not decode image: {exc}") from exc
+    return (ImageOps.exif_transpose(image) or image).convert("RGB")
+
+
+@torch.no_grad()
+def _sat(image: Image.Image, beam: int) -> dict:
+    t0 = time.perf_counter()
+    tensor = state["transform"](image)
+    words, alphas, beams = sat_caption(state["encoder"], state["decoder"], state["vocab"], tensor, k=beam)
+    decode = state["vocab"].decode
+    return {
+        "caption": " ".join(words),
+        "words": words,
+        "attention": [[round(a, 5) for a in row] for row in alphas.tolist()],
+        "crop": center_crop_box(image.width, image.height, state["preprocess"]),
+        "beams": [{"caption": " ".join(decode(toks)), "score": round(score, 3)} for toks, score in beams],
+        "decode_ms": round((time.perf_counter() - t0) * 1000),
+    }
+
+
+@torch.no_grad()
+def _blip(image: Image.Image, beam: int) -> dict:
+    t0 = time.perf_counter()
+    inputs = state["blip_processor"](image, return_tensors="pt")
+    out = state["blip"].generate(**inputs, num_beams=beam, max_new_tokens=30)
+    text = state["blip_processor"].decode(out[0], skip_special_tokens=True)
+    return {"caption": text, "decode_ms": round((time.perf_counter() - t0) * 1000)}
+
+
+@app.get("/")
+def root() -> dict:
+    return {"message": "Welcome to the Capit captioning API! Visit /docs for usage details."}
+
+@app.get("/health")
+def health() -> dict:
+    return {"status": "ok"}
+
+
+@app.post("/caption")
+async def caption_endpoint(
+    file: UploadFile,
+    model: str = Query("both", pattern="^(sat|blip|both)$"),
+    beam: int = Query(3, ge=1, le=10),
+) -> dict:
+    image = await _ingest(file)
+    result: dict = {}
+    if model in ("sat", "both"):
+        result["sat"] = _sat(image, beam)
+    if model in ("blip", "both"):
+        result["blip"] = _blip(image, beam)
+    return result

backend/pyproject.toml

@@ -0,0 +1,29 @@
+[project]
+name = "capit-backend"
+version = "0.1.0"
+description = "capit backend — side-by-side SAT + BLIP captioning API"
+requires-python = ">=3.10"
+dependencies = [
+    "capit",
+    "fastapi[standard]",
+    "uvicorn",
+    "python-multipart",
+    "transformers",
+    "huggingface-hub",
+    "pillow",
+    "numpy",
+    "torch",
+    "torchvision",
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "httpx"]
+
+[tool.uv]
+package = false
+
+[tool.uv.sources]
+capit = { path = "../pipeline", editable = true }
+
+[[tool.uv.index]]
+url = "https://download.pytorch.org/whl/cpu"

backend/scripts/bench_blip.py

@@ -0,0 +1,119 @@
+"""Benchmark BLIP caption latency + peak memory on CPU (Phase 5, killer gate #3).
+
+No published CPU latency exists for blip-image-captioning-base; this measures it directly to
+decide whether it ships as-is on the 2-vCPU HF Space, needs int8 quantization, or a fallback.
+Standalone — no capit import; the backend is its own component.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import resource
+import time
+from pathlib import Path
+
+import torch
+from PIL import Image, ImageOps
+from transformers import BlipForConditionalGeneration, BlipProcessor
+
+MODEL = "Salesforce/blip-image-captioning-base"
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _test_images(data_root: Path, n: int) -> list[Path]:
+    records = json.loads((data_root / "dataset_flickr8k.json").read_text())["images"]
+    test = sorted((r for r in records if r["split"] == "test"), key=lambda r: r["filename"])
+    return [data_root / "Images" / r["filename"] for r in test[:n]]
+
+
+def _peak_rss_mb() -> float:
+    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024  # ru_maxrss is KB on Linux
+
+
+def _pct(xs: list[float], q: float) -> float:
+    xs = sorted(xs)
+    k = (len(xs) - 1) * q
+    lo = int(k)
+    hi = min(lo + 1, len(xs) - 1)
+    return xs[lo] + (xs[hi] - xs[lo]) * (k - lo)
+
+
+@torch.no_grad()
+def bench(processor, model, paths: list[Path], num_beams: int) -> dict:
+    latencies: list[float] = []
+    sample = ""
+    for i, p in enumerate(paths):
+        with Image.open(p) as im:
+            image = (ImageOps.exif_transpose(im) or im).convert("RGB")
+        t0 = time.perf_counter()
+        inputs = processor(image, return_tensors="pt")
+        out = model.generate(**inputs, num_beams=num_beams, max_new_tokens=30)
+        dt = time.perf_counter() - t0
+        if i == 0:  # warm-up: first call pays lazy init / caching
+            sample = processor.decode(out[0], skip_special_tokens=True)
+            continue
+        latencies.append(dt)
+    return {
+        "num_beams": num_beams,
+        "n": len(latencies),
+        "mean_s": sum(latencies) / len(latencies),
+        "p95_s": _pct(latencies, 0.95),
+        "min_s": min(latencies),
+        "max_s": max(latencies),
+        "sample_caption": sample,
+    }
+
+
+def _verdict(mean_s: float) -> str:
+    if mean_s <= 5.0:
+        return "SHIP AS-IS (comfortable margin under the 10s Space budget)"
+    if mean_s <= 10.0:
+        return "TIGHT — under budget locally but little Space margin; consider int8 quantization"
+    return "MISS — exceeds budget; int8 quantize (OpenVINO/NNCF) then re-bench, else ViT-GPT2 fallback"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-root", default=str(_REPO_ROOT / "data" / "flickr8k"))
+    parser.add_argument("--n", type=int, default=20)
+    parser.add_argument("--threads", type=int, default=2, help="2 simulates the HF Space's 2 vCPU")
+    parser.add_argument("--beams", type=int, nargs="+", default=[1, 3])
+    parser.add_argument("--out-json", default=str(_REPO_ROOT / "data" / "blip_bench.json"))
+    args = parser.parse_args()
+
+    torch.set_num_threads(args.threads)
+    processor = BlipProcessor.from_pretrained(MODEL)
+    model = BlipForConditionalGeneration.from_pretrained(MODEL).eval()
+    paths = _test_images(Path(args.data_root), args.n + 1)  # +1 for the discarded warm-up
+
+    runs = [bench(processor, model, paths, b) for b in args.beams]
+    peak_rss = _peak_rss_mb()
+    headline = max(runs, key=lambda r: r["num_beams"])  # the slowest config is what the UI serves
+
+    report = {
+        "model": MODEL,
+        "threads": args.threads,
+        "torch_threads_available": torch.get_num_threads(),
+        "peak_rss_mb": round(peak_rss, 1),
+        "runs": runs,
+        "headline_beam": headline["num_beams"],
+        "headline_mean_s": round(headline["mean_s"], 2),
+        "space_budget_s": 10.0,
+        "verdict": _verdict(headline["mean_s"]),
+    }
+    Path(args.out_json).write_text(json.dumps(report, indent=2))
+
+    print(f"model: {MODEL}   threads: {args.threads}   peak RSS: {peak_rss:.0f} MB")
+    for r in runs:
+        print(
+            f"  beams={r['num_beams']}  mean={r['mean_s']:.2f}s  p95={r['p95_s']:.2f}s  "
+            f"min={r['min_s']:.2f}s  max={r['max_s']:.2f}s  (n={r['n']})"
+        )
+    print(f'  sample: "{runs[-1]["sample_caption"]}"')
+    print(f"verdict (beams={headline['num_beams']}, {headline['mean_s']:.2f}s vs 10s budget): {report['verdict']}")
+    print(f"wrote {args.out_json}")
+
+
+if __name__ == "__main__":
+    main()

backend/tests/test_app.py

@@ -0,0 +1,78 @@
+"""Stage 6.1 — backend API. Needs the local artifact + cached BLIP; skips cleanly otherwise."""
+
+import io
+import sys
+from pathlib import Path
+
+import pytest
+from PIL import Image
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from app import ARTIFACT_DIR, app
+
+from fastapi.testclient import TestClient
+
+pytestmark = pytest.mark.skipif(
+    not (ARTIFACT_DIR / "capit-sat.pt").exists(),
+    reason="no local serving artifact; run scripts/export_artifact.py first",
+)
+
+
+@pytest.fixture(scope="module")
+def client():
+    with TestClient(app) as c:
+        yield c
+
+
+def _png(size=(64, 48), mode="RGB", exif=None) -> bytes:
+    img = Image.new(mode, size, "red")
+    buf = io.BytesIO()
+    if exif is not None:
+        img.save(buf, format="JPEG", exif=exif)
+    else:
+        img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_health(client):
+    assert client.get("/health").json() == {"status": "ok"}
+
+
+def test_caption_happy_sat(client):
+    r = client.post("/caption?model=sat&beam=3", files={"file": ("x.png", _png(), "image/png")})
+    assert r.status_code == 200
+    sat = r.json()["sat"]
+    assert sat["caption"] and len(sat["words"]) == len(sat["attention"])
+    assert all(len(row) == 196 for row in sat["attention"])
+    assert set(sat["crop"]) == {"x", "y", "w", "h"}
+    assert sat["beams"][0]["caption"]
+
+
+def test_caption_both_has_blip(client):
+    r = client.post("/caption?model=both", files={"file": ("x.png", _png(), "image/png")})
+    body = r.json()
+    assert r.status_code == 200 and body["blip"]["caption"] and body["sat"]["caption"]
+
+
+def test_junk_bytes_is_422(client):
+    r = client.post("/caption?model=sat", files={"file": ("x.png", b"not an image", "image/png")})
+    assert r.status_code == 422
+
+
+def test_oversized_is_413(client):
+    big = b"\xff" * (8 * 1024 * 1024 + 1)
+    r = client.post("/caption?model=sat", files={"file": ("x.png", big, "image/png")})
+    assert r.status_code == 413
+
+
+def test_rgba_accepted(client):
+    r = client.post("/caption?model=sat", files={"file": ("x.png", _png(mode="RGBA"), "image/png")})
+    assert r.status_code == 200
+
+
+def test_exif_orientation_applied(client):
+    exif = Image.Exif()
+    exif[0x0112] = 6  # rotate 90 CW: a saved 400x200 displays upright as 200x400 (portrait)
+    r = client.post("/caption?model=sat", files={"file": ("x.jpg", _png((400, 200), exif=exif), "image/jpeg")})
+    crop = r.json()["sat"]["crop"]
+    assert crop["w"] > crop["h"]  # portrait-upright crops wider in fractions; ignoring EXIF would flip this

pipeline/.python-version

@@ -0,0 +1 @@
+3.12

pipeline/notebooks/colab_train.ipynb

@@ -0,0 +1,264 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# capit — train on Colab (Stage 3.3)\n",
+    "\n",
+    "A thin launcher: all the logic lives in `train.py` in the repo. This notebook clones it, stages the data, and runs the CLI on the T4.\n",
+    "\n",
+    "**Before you run:**\n",
+    "1. **Runtime → Change runtime type → T4 GPU.**\n",
+    "2. Locally: `python pipeline/scripts/make_train_zip.py` → builds `data/flickr8k_colab.zip` (Images + dataset_flickr8k.json + vocab.json).\n",
+    "3. **Push your latest code** — the notebook runs what is in git, not your local working tree.\n",
+    "\n",
+    "Get the zip onto Drive **once** (either the upload cell below, or drag it into `MyDrive/capit/` via the Drive website — more reliable for ~1 GB). After that, reconnects just re-stage from Drive and `--resume auto` continues the run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
+     ]
+    }
+   ],
+   "source": [
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cloning into '/content/capit'...\n",
+      "remote: Enumerating objects: 149, done.\u001b[K\n",
+      "remote: Counting objects: 100% (149/149), done.\u001b[K\n",
+      "remote: Compressing objects: 100% (89/89), done.\u001b[K\n",
+      "remote: Total 149 (delta 55), reused 132 (delta 38), pack-reused 0 (from 0)\u001b[K\n",
+      "Receiving objects: 100% (149/149), 96.33 KiB | 6.02 MiB/s, done.\n",
+      "Resolving deltas: 100% (55/55), done.\n",
+      "Obtaining file:///content/capit/pipeline\n",
+      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
+      "  Checking if build backend supports build_editable ... \u001b[?25l\u001b[?25hdone\n",
+      "  Getting requirements to build editable ... \u001b[?25l\u001b[?25hdone\n",
+      "  Preparing editable metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
+      "Requirement already satisfied: torch in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (2.11.0+cu128)\n",
+      "Requirement already satisfied: torchvision in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (0.26.0+cu128)\n",
+      "Requirement already satisfied: pillow in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (11.3.0)\n",
+      "Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (2.0.2)\n",
+      "Requirement already satisfied: nltk in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (3.9.1)\n",
+      "Requirement already satisfied: kagglehub in /usr/local/lib/python3.12/dist-packages (from capit==0.1.0) (1.0.1)\n",
+      "Requirement already satisfied: kagglesdk<1.0,>=0.1.22 in /usr/local/lib/python3.12/dist-packages (from kagglehub->capit==0.1.0) (0.1.23)\n",
+      "Requirement already satisfied: packaging in /usr/local/lib/python3.12/dist-packages (from kagglehub->capit==0.1.0) (26.2)\n",
+      "Requirement already satisfied: pyyaml in /usr/local/lib/python3.12/dist-packages (from kagglehub->capit==0.1.0) (6.0.3)\n",
+      "Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from kagglehub->capit==0.1.0) (2.32.4)\n",
+      "Requirement already satisfied: tqdm in /usr/local/lib/python3.12/dist-packages (from kagglehub->capit==0.1.0) (4.67.3)\n",
+      "Requirement already satisfied: click in /usr/local/lib/python3.12/dist-packages (from nltk->capit==0.1.0) (8.4.1)\n",
+      "Requirement already satisfied: joblib in /usr/local/lib/python3.12/dist-packages (from nltk->capit==0.1.0) (1.5.3)\n",
+      "Requirement already satisfied: regex>=2021.8.3 in /usr/local/lib/python3.12/dist-packages (from nltk->capit==0.1.0) (2025.11.3)\n",
+      "Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (3.29.1)\n",
+      "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (4.15.0)\n",
+      "Requirement already satisfied: setuptools<82 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (75.2.0)\n",
+      "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (1.14.0)\n",
+      "Requirement already satisfied: networkx>=2.5.1 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (3.6.1)\n",
+      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (3.1.6)\n",
+      "Requirement already satisfied: fsspec>=0.8.5 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (2025.3.0)\n",
+      "Requirement already satisfied: cuda-toolkit==12.8.1 in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.1)\n",
+      "Requirement already satisfied: cuda-bindings<13,>=12.9.4 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (12.9.7)\n",
+      "Requirement already satisfied: nvidia-cudnn-cu12==9.19.0.56 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (9.19.0.56)\n",
+      "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (0.7.1)\n",
+      "Requirement already satisfied: nvidia-nccl-cu12==2.28.9 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (2.28.9)\n",
+      "Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (3.4.5)\n",
+      "Requirement already satisfied: triton==3.6.0 in /usr/local/lib/python3.12/dist-packages (from torch->capit==0.1.0) (3.6.0)\n",
+      "Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.4.1)\n",
+      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.90)\n",
+      "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (11.3.3.83)\n",
+      "Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (1.13.1.3)\n",
+      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.90)\n",
+      "Requirement already satisfied: nvidia-curand-cu12==10.3.9.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (10.3.9.90)\n",
+      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (11.7.3.90)\n",
+      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.5.8.93)\n",
+      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.93)\n",
+      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.93)\n",
+      "Requirement already satisfied: nvidia-nvtx-cu12==12.8.90.* in /usr/local/lib/python3.12/dist-packages (from cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==12.8.1; platform_system == \"Linux\"->torch->capit==0.1.0) (12.8.90)\n",
+      "Requirement already satisfied: cuda-pathfinder~=1.1 in /usr/local/lib/python3.12/dist-packages (from cuda-bindings<13,>=12.9.4->torch->capit==0.1.0) (1.5.5)\n",
+      "Requirement already satisfied: protobuf in /usr/local/lib/python3.12/dist-packages (from kagglesdk<1.0,>=0.1.22->kagglehub->capit==0.1.0) (5.29.6)\n",
+      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch->capit==0.1.0) (1.3.0)\n",
+      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch->capit==0.1.0) (3.0.3)\n",
+      "Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->kagglehub->capit==0.1.0) (3.4.7)\n",
+      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests->kagglehub->capit==0.1.0) (3.18)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests->kagglehub->capit==0.1.0) (2.5.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests->kagglehub->capit==0.1.0) (2026.5.20)\n",
+      "Building wheels for collected packages: capit\n",
+      "  Building editable for capit (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
+      "  Created wheel for capit: filename=capit-0.1.0-0.editable-py3-none-any.whl size=1307 sha256=fbe47f272d2e771f315dc44349cd36179ae23791926a0b0626981391c9c3976b\n",
+      "  Stored in directory: /tmp/pip-ephem-wheel-cache-630q9_cb/wheels/ff/0a/93/9457afa2ecf0d24429d3eb1754c723aeaf950fdbae03614da9\n",
+      "Successfully built capit\n",
+      "Installing collected packages: capit\n",
+      "  Attempting uninstall: capit\n",
+      "    Found existing installation: capit 0.1.0\n",
+      "    Uninstalling capit-0.1.0:\n",
+      "      Successfully uninstalled capit-0.1.0\n",
+      "Successfully installed capit-0.1.0\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.colab-display-data+json": {
+       "id": "fe04bbce05954a97b04b1aefedeb93e6",
+       "pip_warning": {
+        "packages": [
+         "capit"
+        ]
+       }
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "capit installed\n"
+     ]
+    }
+   ],
+   "source": [
+    "# torch/torchvision are preinstalled on Colab; this pulls the small extras (nltk, ...).\n",
+    "# No -q: a failed install must be visible, not surface 3 hours later as ModuleNotFoundError.\n",
+    "!rm -rf /content/capit && git clone https://github.com/Bukunmi2108/capit.git /content/capit\n",
+    "!pip install -e /content/capit/pipeline\n",
+    "import capit  # fail fast if the install didn't take\n",
+    "print(\"capit installed\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "already on Drive — skipping upload\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os, shutil\n",
+    "dest = '/content/drive/MyDrive/capit/flickr8k_colab.zip'\n",
+    "if os.path.exists(dest):\n",
+    "    print('already on Drive — skipping upload')\n",
+    "else:\n",
+    "    from google.colab import files\n",
+    "    os.makedirs(os.path.dirname(dest), exist_ok=True)\n",
+    "    up = files.upload()  # pick flickr8k_colab.zip\n",
+    "    shutil.move(next(iter(up)), dest)\n",
+    "    print('uploaded and stashed on Drive:', dest)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "8091 images staged\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Copy OFF the Drive mount to local disk, then unzip (never read images over Drive — slow).\n",
+    "# `&&` so unzip only runs if the copy succeeded.\n",
+    "!cp /content/drive/MyDrive/capit/flickr8k_colab.zip /content/flickr8k_colab.zip && unzip -q -o /content/flickr8k_colab.zip -d /content/flickr8k\n",
+    "import os\n",
+    "n = len(os.listdir('/content/flickr8k/Images'))\n",
+    "assert n == 8091, f\"expected 8091 images, got {n} — bad/partial zip\"\n",
+    "print(n, \"images staged\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-11ad3fa6.pth\n",
+      "100% 97.8M/97.8M [00:00<00:00, 147MB/s] \n",
+      "epoch 0  loss 4.8991  val_bleu4 8.18  best 8.18 *  | a boy in a blue shirt and a white shirt is standing on a skateboard\n",
+      "epoch 1  loss 4.0486  val_bleu4 17.58  best 17.58 *  | a boy is sitting on a bench\n",
+      "epoch 2  loss 3.7289  val_bleu4 17.30  best 17.58  | a man in a red shirt and black pants is sitting on a bench\n",
+      "epoch 3  loss 3.5054  val_bleu4 17.85  best 17.85 *  | a man in a red shirt is sitting on a skateboard\n",
+      "epoch 4  loss 3.3269  val_bleu4 19.32  best 19.32 *  | a man in a red shirt is sitting on a skateboard\n",
+      "epoch 5  loss 3.1763  val_bleu4 18.90  best 19.32  | a man in a red shirt is sitting on a skateboard\n",
+      "epoch 6  loss 3.0429  val_bleu4 18.59  best 19.32  | a man in a red shirt is riding on a skateboard\n",
+      "epoch 7  loss 2.9214  val_bleu4 19.62  best 19.62 *  | a man in a red shirt is riding on a swing\n",
+      "epoch 8  loss 2.8085  val_bleu4 18.76  best 19.62  | a man in a red shirt and blue jeans is sitting on a bench in a <unk>\n",
+      "epoch 9  loss 2.7064  val_bleu4 18.01  best 19.62  | a man in a red shirt and blue jeans is sitting on a skateboard\n",
+      "epoch 10  loss 2.6112  val_bleu4 18.01  best 19.62  | a man in a red shirt and blue jeans is sitting on a red bench\n",
+      "epoch 11  loss 2.5185  val_bleu4 18.08  best 19.62  | a man in a red shirt and blue jeans is <unk> down a <unk> <unk> structure\n",
+      "epoch 12  loss 2.4343  val_bleu4 18.57  best 19.62  | a man on a skateboard in a <unk> area\n",
+      "epoch 13  loss 2.3560  val_bleu4 17.20  best 19.62  | a man in a red shirt is sitting on a bench in a <unk> area\n",
+      "epoch 14  loss 2.2798  val_bleu4 16.75  best 19.62  | a man in a red shirt is sitting on a bench in a <unk> area\n",
+      "epoch 15  loss 2.2071  val_bleu4 17.01  best 19.62  | a man on a skateboard on a <unk>\n",
+      "epoch 16  loss 2.1434  val_bleu4 17.04  best 19.62  | a man on a skateboard on a <unk> <unk>\n",
+      "epoch 17  loss 2.0809  val_bleu4 17.65  best 19.62  | a man on a skateboard in a <unk> area\n",
+      "early stop: no val BLEU-4 improvement in 10 epochs (last loss 2.0809)\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python -m capit.train \\\n",
+    "  --data-root /content/flickr8k \\\n",
+    "  --vocab-path /content/flickr8k/vocab.json \\\n",
+    "  --ckpt-dir /content/drive/MyDrive/capit/checkpoints \\\n",
+    "  --resume auto"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## If disconnected\n",
+    "Reconnect → re-run **all** cells. The upload cell skips (zip already on Drive), the data cell re-stages, and `--resume auto` loads `latest.pt` from Drive and continues from the next epoch — checkpoints live on Drive, so a disconnect costs minutes, not the run. (train.py refuses to run if `--ckpt-dir` points at an unmounted Drive.)\n",
+    "\n",
+    "**Exit gate (Stage 3.3):** `best.pt` on Drive with val BLEU-4 (greedy, nltk) ≥ ~14. Download `MyDrive/capit/checkpoints/best.pt` for Phase 4."
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

pipeline/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=61.0.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "capit"
+version = "0.1.0"
+description = "capit — glass-box image captioner (Show, Attend and Tell)"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch",
+    "torchvision",
+    "pillow",
+    "numpy",
+    "nltk",
+    "kagglehub",
+    "huggingface-hub",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "pycocoevalcap",
+    "matplotlib",
+    "scikit-image",
+]
+
+[[tool.uv.index]]
+url = "https://download.pytorch.org/whl/cpu"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

pipeline/scripts/build_vocab.py

@@ -0,0 +1,21 @@
+"""Build data/vocab.json from the full Flickr8k train split."""
+
+from __future__ import annotations
+
+import json
+
+from capit.config import config
+from capit.data.records import train_captions
+from capit.data.vocab import Vocab
+
+
+def build_vocab() -> None:
+    records = json.loads(config.karpathy_json.read_text())["images"]
+    vocab = Vocab.build(train_captions(records), config.min_freq)
+    config.vocab_path.parent.mkdir(parents=True, exist_ok=True)
+    vocab.save(config.vocab_path)
+    print(f"vocab: {len(vocab)} words (min_freq={config.min_freq}) -> {config.vocab_path}")
+
+
+if __name__ == "__main__":
+    build_vocab()

pipeline/scripts/export_artifact.py

@@ -0,0 +1,132 @@
+"""Export the serving artifact: best.pt (training checkpoint) -> capit-sat.pt (inference contract)."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from pathlib import Path
+
+import torch
+
+from capit.checkpoint import load
+from capit.config import config
+from capit.data.vocab import Vocab
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+from capit.serving import build_artifact
+
+
+def _split_counts(karpathy_json: Path) -> dict[str, int]:
+    records = json.loads(karpathy_json.read_text())["images"]
+    counts: dict[str, int] = {}
+    for r in records:
+        counts[r["split"]] = counts.get(r["split"], 0) + 1
+    return counts
+
+
+def _scores_table(metrics: dict[str, dict[str, float]]) -> str:
+    cols = ["BLEU-1", "BLEU-2", "BLEU-3", "BLEU-4", "CIDEr"]
+    lines = ["| beam | " + " | ".join(cols) + " |", "|-----:|" + "|".join(["-------:"] * len(cols)) + "|"]
+    for beam in sorted(metrics, key=int):
+        s = metrics[beam]
+        lines.append(f"| {beam} | " + " | ".join(f"{s[c]:.2f}" for c in cols) + " |")
+    return "\n".join(lines)
+
+
+def _model_card(
+    repo_id: str, splits: dict[str, int], best_bleu4: float, best_epoch: int, metrics: dict[str, dict[str, float]]
+) -> str:
+    return f"""---
+license: mit
+language:
+- en
+library_name: pytorch
+pipeline_tag: image-to-text
+tags:
+- image-captioning
+- show-attend-and-tell
+- visual-attention
+datasets:
+- flickr8k
+metrics:
+- bleu
+- cider
+---
+
+# capit-sat
+
+Show, Attend and Tell image captioner, trained from scratch on Flickr8k (Karpathy split).
+The glass-box half of [capit](https://github.com/Bukunmi2108/capit) — exposes per-word
+attention, beam candidates, and word-by-word playback.
+
+## Test-set scores (pycocoevalcap, Karpathy test = {splits.get('test', '?')} images)
+
+{_scores_table(metrics)}
+
+## Training
+
+- Backbone: frozen ResNet-50 (ImageNet). Decoder trained from scratch.
+- Best val BLEU-4 {best_bleu4:.2f} at epoch {best_epoch} (early-stopped); Colab T4.
+- Splits: train {splits.get('train', '?')}, val {splits.get('val', '?')}, test {splits.get('test', '?')}.
+
+## Known limitation
+
+Attention is effectively 7x7: ResNet-50 at 224px is natively 7x7 and the encoder upsamples
+to 14x14, so heatmaps are coarse (~32px blocks). Captions are grounded; the spots are
+region-level, not pixel-level.
+
+## Use
+
+`huggingface_hub.hf_hub_download("{repo_id}", "capit-sat.pt")` + `vocab.json`, then
+`capit.serving.load_artifact(...)`.
+"""
+
+
+def export(ckpt_path: Path, vocab_path: Path, out_dir: Path, repo_id: str, metrics_json: Path) -> Path:
+    if not metrics_json.exists():
+        raise FileNotFoundError(
+            f"metrics file {metrics_json} not found — generate it first:\n"
+            f"  uv run python -m capit.evaluate --out-json {metrics_json}"
+        )
+    metrics = json.loads(metrics_json.read_text())
+    vocab = Vocab.load(vocab_path)
+    state = load(ckpt_path)
+    if state.vocab_sha256 != vocab.sha256():
+        raise ValueError(f"vocab mismatch: ckpt {state.vocab_sha256[:8]} != vocab {vocab.sha256()[:8]}")
+
+    encoder = Encoder(pretrained=True)
+    decoder = Decoder(vocab_size=len(vocab))
+    decoder.load_state_dict(state.model_state)
+    blob = build_artifact(encoder, decoder, vocab)
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    artifact_path = out_dir / "capit-sat.pt"
+    tmp = artifact_path.with_name(artifact_path.name + ".tmp")
+    torch.save(blob, tmp)
+    os.replace(tmp, artifact_path)
+    shutil.copyfile(vocab_path, out_dir / "vocab.json")
+
+    splits = _split_counts(config.karpathy_json)
+    (out_dir / "README.md").write_text(_model_card(repo_id, splits, state.best_bleu4, state.best_epoch, metrics))
+    return artifact_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", default=str(config.ckpt_dir / "best.pt"))
+    parser.add_argument("--vocab", default=str(config.vocab_path))
+    parser.add_argument("--out-dir", default=str(config.data_root / "artifact"))
+    parser.add_argument("--repo-id", default="Bukunmi2108/capit-sat")
+    parser.add_argument("--metrics-json", default=str(config.data_root / "eval_results.json"))
+    args = parser.parse_args()
+
+    path = export(Path(args.ckpt), Path(args.vocab), Path(args.out_dir), args.repo_id, Path(args.metrics_json))
+    size_mb = path.stat().st_size / 1e6
+    print(f"wrote {path} ({size_mb:.1f} MB), vocab.json, README.md to {args.out_dir}")
+    print(f"push: hf upload {args.repo_id} {args.out_dir} . --repo-type model")
+
+
+if __name__ == "__main__":
+    main()

pipeline/scripts/make_subsample.py

@@ -0,0 +1,39 @@
+"""Build the dev subsample: a small seeded slice of Flickr8k written in the same
+on-disk format (Karpathy JSON + Images/) as the full set."""
+
+from __future__ import annotations
+
+import json
+import shutil
+
+from capit.config import config
+from capit.data.download import _copy_atomic
+from capit.data.records import select_records
+
+
+def make_subsample() -> None:
+    full = json.loads(config.karpathy_json.read_text())
+    selected = select_records(full["images"], config.subsample_counts, config.seed)
+
+    sources = [(config.images_dir / r["filename"], r) for r in selected]
+    missing = [r["filename"] for src, r in sources if not src.is_file()]
+    if missing:
+        raise FileNotFoundError(
+            f"{len(missing)} source images missing under {config.images_dir} "
+            f"(rerun the Stage 0.2 download), e.g. {missing[:3]}"
+        )
+
+    if config.subsample_root.exists():
+        shutil.rmtree(config.subsample_root)
+    config.subsample_images_dir.mkdir(parents=True, exist_ok=True)
+    for src, rec in sources:
+        _copy_atomic(src, config.subsample_images_dir / rec["filename"])
+
+    config.subsample_json.write_text(json.dumps({"images": selected, "dataset": full["dataset"]}))
+
+    tally = {s: sum(r["split"] == s for r in selected) for s in config.subsample_counts}
+    print(f"subsample: {len(selected)} images {tally} -> {config.subsample_root}")
+
+
+if __name__ == "__main__":
+    make_subsample()

pipeline/scripts/make_train_zip.py

@@ -0,0 +1,35 @@
+"""Assemble the full dataset + vocab into one zip for Colab (upload to Drive).
+The zip unpacks to a single CaptionDataset root: Images/ + dataset_flickr8k.json + vocab.json.
+"""
+
+from __future__ import annotations
+
+import json
+import zipfile
+from pathlib import Path
+
+from capit.config import config
+
+
+def make_train_zip(out_path: Path | None = None) -> Path:
+    out = Path(out_path) if out_path else config.data_root / "flickr8k_colab.zip"
+    for required in (config.images_dir, config.karpathy_json, config.vocab_path):
+        if not required.exists():
+            raise FileNotFoundError(f"{required} missing — run the Stage 0.2 download and build_vocab first")
+
+    jpgs = sorted(config.images_dir.glob("*.jpg"))
+    expected = len(json.loads(config.karpathy_json.read_text())["images"])
+    if len(jpgs) < expected:
+        raise ValueError(f"{len(jpgs)} images in {config.images_dir}, expected >= {expected} (partial download?)")
+    with zipfile.ZipFile(out, "w", zipfile.ZIP_STORED) as zf:  # jpgs already compressed
+        for jpg in jpgs:
+            zf.write(jpg, f"Images/{jpg.name}")
+        zf.write(config.karpathy_json, "dataset_flickr8k.json")
+        zf.write(config.vocab_path, "vocab.json")
+
+    print(f"wrote {out} ({len(jpgs)} images + dataset_flickr8k.json + vocab.json)")
+    return out
+
+
+if __name__ == "__main__":
+    make_train_zip()

pipeline/scripts/overfit_one_batch.py

@@ -0,0 +1,20 @@
+"""Overfit one batch (killer gate #1) — run and eyeball the decoded captions."""
+
+from __future__ import annotations
+
+from capit.overfit import run_overfit
+
+
+def main() -> None:
+    result = run_overfit()
+    curve = result["ce_curve"]
+    for s in range(0, len(curve), max(1, len(curve) // 10)):
+        print(f"step {s:4d}  ce {curve[s]:.4f}")
+    print(f"final ce: {result['final_ce']:.4f}\n")
+    for got, want in zip(result["decoded"], result["targets"]):
+        print(f"  target: {' '.join(want)}")
+        print(f"  greedy: {' '.join(got)}\n")
+
+
+if __name__ == "__main__":
+    main()

pipeline/scripts/serve_demo.py

@@ -0,0 +1,34 @@
+"""Exit gate: knowing only a Hub repo id, download the artifact and caption one image.
+
+This is the seed of backend/app.py — it proves capit-sat.pt is a self-sufficient inference
+contract (no checkpoint, no training data, no hardcoded preprocess).
+"""
+
+from __future__ import annotations
+
+import argparse
+from huggingface_hub import hf_hub_download
+from PIL import Image, ImageOps
+
+from capit.serving import caption, load_artifact, make_transform
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("image")
+    parser.add_argument("--repo-id", default="Bukunmi2108/capit-sat")
+    args = parser.parse_args()
+
+    artifact = hf_hub_download(args.repo_id, "capit-sat.pt")
+    vocab_path = hf_hub_download(args.repo_id, "vocab.json")
+    encoder, decoder, vocab, preprocess = load_artifact(artifact, vocab_path)
+    transform = make_transform(preprocess)
+
+    with Image.open(args.image) as img:
+        tensor = transform((ImageOps.exif_transpose(img) or img).convert("RGB"))
+    words, _, _ = caption(encoder, decoder, vocab, tensor)
+    print(" ".join(words))
+
+
+if __name__ == "__main__":
+    main()

pipeline/scripts/visualize_attention.py

@@ -0,0 +1,115 @@
+"""Decode one image at beam 3 and render a per-word attention grid over it."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+from skimage.transform import pyramid_expand
+from torchvision import transforms
+
+from capit.checkpoint import load
+from capit.config import Config, config
+from capit.data.dataset import CaptionDataset, build_transform
+from capit.data.vocab import END, START, Vocab
+from capit.decode import beam_search
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+
+
+def display_transform(cfg: Config = config):
+    return transforms.Compose(
+        [transforms.Resize(cfg.resize), transforms.CenterCrop(cfg.crop), transforms.ToTensor()]
+    )
+
+
+def load_image(path: Path | str, transform) -> torch.Tensor:
+    with Image.open(path) as img:
+        img = ImageOps.exif_transpose(img) or img
+        return transform(img.convert("RGB"))
+
+
+def heatmap(alpha: torch.Tensor, upscale: int = 16, sigma: float = 8.0) -> np.ndarray:
+    grid = config.encoded_size
+    return pyramid_expand(alpha.reshape(grid, grid).numpy(), upscale=upscale, sigma=sigma)
+
+
+def _label(ax, text: str) -> None:
+    ax.text(
+        0.03, 0.95, text, transform=ax.transAxes, fontsize=11, va="top",
+        color="black", bbox=dict(facecolor="white", edgecolor="none", pad=1.5),
+    )
+
+
+def overlay_grid(display_img: torch.Tensor, words: list[str], alphas: torch.Tensor, attn_alpha: float = 0.8):
+    base = display_img.permute(1, 2, 0).numpy()
+    heats = [heatmap(a) for a in alphas]
+    n = len(words) + 1
+    cols = min(n, 5)
+    rows = (n + cols - 1) // cols
+    fig, axes = plt.subplots(rows, cols, figsize=(cols * 2.6, rows * 2.6), squeeze=False, facecolor="white")
+    axes = axes.reshape(-1)
+    axes[0].imshow(base)
+    _label(axes[0], "input")
+    for i, (word, heat) in enumerate(zip(words, heats), start=1):
+        axes[i].imshow(base)
+        axes[i].imshow(heat, alpha=attn_alpha, cmap="Greys_r")
+        _label(axes[i], word)
+    for ax in axes:
+        ax.axis("off")
+    fig.tight_layout()
+    return fig
+
+
+@torch.no_grad()
+def visualize(image_path: Path | str, ckpt_path: Path | str, out_path: Path | str, k: int = config.beam_k) -> list[str]:
+    vocab = Vocab.load(config.vocab_path)
+    encoder = Encoder(pretrained=True).eval()
+    decoder = Decoder(vocab_size=len(vocab))
+    state = load(ckpt_path)
+    if state.vocab_sha256 != vocab.sha256():
+        raise ValueError(f"vocab mismatch: ckpt {state.vocab_sha256[:8]} != vocab {vocab.sha256()[:8]}")
+    decoder.load_state_dict(state.model_state)
+    decoder.eval()
+
+    features = encoder(load_image(image_path, build_transform()).unsqueeze(0))
+    tokens, alphas, _ = beam_search(decoder, features, vocab.word2id[START], vocab.word2id[END], k=k)
+    words = vocab.decode(tokens)
+
+    fig = overlay_grid(load_image(image_path, display_transform()), words, alphas)
+    fig.savefig(out_path, dpi=120, bbox_inches="tight")
+    plt.close(fig)
+    return words
+
+
+def _resolve_image(args) -> Path:
+    if args.image:
+        return Path(args.image)
+    ds = CaptionDataset(args.data_root, "test", Vocab.load(config.vocab_path), build_transform())
+    return Path(args.data_root) / "Images" / ds.records[args.index]["filename"]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", help="explicit image path; overrides --index")
+    parser.add_argument("--index", type=int, default=0, help="index into the test split")
+    parser.add_argument("--data-root", default=str(config.flickr8k_dir))
+    parser.add_argument("--ckpt", default=str(config.ckpt_dir / "best.pt"))
+    parser.add_argument("--out", default=str(config.data_root / "attention.png"))
+    args = parser.parse_args()
+
+    image_path = _resolve_image(args)
+    words = visualize(image_path, args.ckpt, args.out)
+    print(f"{image_path.name}: {' '.join(words)}")
+    print(f"saved {args.out}")
+
+
+if __name__ == "__main__":
+    main()

pipeline/src/capit/checkpoint.py

@@ -0,0 +1,58 @@
+"""Training checkpoints: atomic save/load of model + optimizer state for resumable runs."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from torch.optim import Optimizer
+
+
+@dataclass
+class TrainState:
+    model_state: dict[str, Any]
+    optim_state: dict[str, Any]
+    epoch: int
+    best_bleu4: float
+    best_epoch: int
+    vocab_sha256: str
+
+
+def save(
+    path: Path | str,
+    model: nn.Module,
+    optimizer: Optimizer,
+    epoch: int,
+    best_bleu4: float,
+    best_epoch: int,
+    vocab_sha256: str,
+) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    state = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "epoch": epoch,
+        "best_bleu4": best_bleu4,
+        "best_epoch": best_epoch,
+        "vocab_sha256": vocab_sha256,
+    }
+    tmp = path.with_name(path.name + ".tmp")
+    torch.save(state, tmp)
+    os.replace(tmp, path)
+
+
+def load(path: Path | str) -> TrainState:
+    state = torch.load(Path(path), map_location="cpu", weights_only=False)
+    return TrainState(
+        model_state=state["model"],
+        optim_state=state["optimizer"],
+        epoch=state["epoch"],
+        best_bleu4=state["best_bleu4"],
+        best_epoch=state.get("best_epoch", state["epoch"]),
+        vocab_sha256=state["vocab_sha256"],
+    )

pipeline/src/capit/config.py

@@ -0,0 +1,92 @@
+"""
+Single source of truth for every hyperparameter in the capit project.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+@dataclass(frozen=True)
+class Config:
+    seed: int = 13
+    data_root: Path = _REPO_ROOT / "data"
+    subsample_train: int = 50
+    subsample_val: int = 10
+    subsample_test: int = 10
+    min_freq: int = 5
+    resize: int = 256
+    crop: int = 224
+    imagenet_mean: tuple[float, float, float] = (0.485, 0.456, 0.406)
+    imagenet_std: tuple[float, float, float] = (0.229, 0.224, 0.225)
+    encoded_size: int = 14
+    encoder_dim: int = 2048
+    decoder_dim: int = 512
+    attention_dim: int = 512
+    embed_dim: int = 512
+    dropout: float = 0.5
+    alpha_c: float = 1.0
+    decoder_lr: float = 4e-4
+    batch_size: int = 64
+    max_epochs: int = 50
+    num_workers: int = 2
+    grad_clip: float = 5.0
+    patience: int = 10
+    beam_k: int = 3
+    beam_alpha: float = 0.7
+    max_decode_len: int = 50
+
+
+    @property
+    def flickr8k_dir(self) -> Path:
+        return self.data_root / "flickr8k"
+
+    @property
+    def images_dir(self) -> Path:
+        return self.flickr8k_dir / "Images"
+
+    @property
+    def captions_txt(self) -> Path:
+        return self.flickr8k_dir / "captions.txt"
+
+    @property
+    def karpathy_dir(self) -> Path:
+        return self.data_root / "karpathy"
+
+    @property
+    def karpathy_json(self) -> Path:
+        return self.karpathy_dir / "dataset_flickr8k.json"
+
+    @property
+    def subsample_counts(self) -> dict[str, int]:
+        return {
+            "train": self.subsample_train,
+            "val": self.subsample_val,
+            "test": self.subsample_test,
+        }
+
+    @property
+    def subsample_root(self) -> Path:
+        return self.data_root / "dev_subsample"
+
+    @property
+    def subsample_images_dir(self) -> Path:
+        return self.subsample_root / "Images"
+
+    @property
+    def subsample_json(self) -> Path:
+        return self.subsample_root / "dataset_flickr8k.json"
+
+    @property
+    def vocab_path(self) -> Path:
+        return self.data_root / "vocab.json"
+
+    @property
+    def ckpt_dir(self) -> Path:
+        return self.data_root / "checkpoints"
+
+
+config = Config()

pipeline/src/capit/data/dataset.py

@@ -0,0 +1,78 @@
+"""CaptionDataset: (image, caption) pairs and evaluation views over a Flickr8k root."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import torch
+from PIL import Image, ImageOps
+from torch.utils.data import Dataset
+from torchvision import transforms
+from capit.config import Config, config
+from capit.data.vocab import END, PAD, START, Vocab
+
+
+def build_transform(cfg: Config = config):
+    return transforms.Compose(
+        [
+            transforms.Resize(cfg.resize),
+            transforms.CenterCrop(cfg.crop),
+            transforms.ToTensor(),
+            transforms.Normalize(cfg.imagenet_mean, cfg.imagenet_std),
+        ]
+    )
+
+
+class CaptionDataset(Dataset):
+    def __init__(self, root: Path | str, split: str, vocab: Vocab, transform) -> None:
+        self.root = Path(root)
+        self.split = split
+        self.vocab = vocab
+        self.transform = transform
+
+        records = json.loads((self.root / "dataset_flickr8k.json").read_text())["images"]
+        self.records = [r for r in records if r["split"] == split]
+        if not self.records:
+            available = sorted({r["split"] for r in records})
+            raise ValueError(f"split {split!r} matched 0 records under {self.root}; available: {available}")
+        self.pairs = [(r, s) for r in self.records for s in r["sentences"]]
+        self.max_len = max(len(s["tokens"]) + 2 for r in self.records for s in r["sentences"])
+
+        ordered = sorted(self.records, key=lambda r: r["filename"])
+        self.image_id = {r["filename"]: i for i, r in enumerate(ordered)}
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def __getitem__(self, i: int) -> tuple[torch.Tensor, torch.Tensor, int]:
+        rec, sent = self.pairs[i]
+        image = self._load_image(rec["filename"])
+        ids = [self.vocab.word2id[START], *self.vocab.encode(sent["tokens"]), self.vocab.word2id[END]]
+        caption_len = len(ids)
+        pad = self.max_len - caption_len
+        if pad < 0:
+            raise ValueError(f"caption len {caption_len} exceeds max_len {self.max_len} for {rec['filename']!r}")
+        ids += [self.vocab.word2id[PAD]] * pad
+        return image, torch.tensor(ids), caption_len
+
+    def _load_image(self, filename: str) -> torch.Tensor:
+        with Image.open(self.root / "Images" / filename) as img:
+            img = ImageOps.exif_transpose(img) or img
+            return self.transform(img.convert("RGB"))
+
+    def references(self) -> dict[int, list[list[str]]]:
+        return {
+            self.image_id[r["filename"]]: [s["tokens"] for s in r["sentences"]] for r in self.records
+        }
+
+    def iter_images(self):
+        for r in self.records:
+            yield self.image_id[r["filename"]], self._load_image(r["filename"])
+
+
+def collate_fn(batch):
+    batch = sorted(batch, key=lambda x: x[2], reverse=True)
+    images = torch.stack([img for img, _, _ in batch])
+    captions = torch.stack([ids for _, ids, _ in batch])
+    lengths = torch.tensor([clen for _, _, clen in batch])
+    return images, captions, lengths

pipeline/src/capit/data/download.py

@@ -0,0 +1,105 @@
+"""Dataset acquisition.
+
+Two independent, idempotent fetchers that materialize Flickr8k into the canonical ``data/`` root so the directory is self-contained .
+
+The Karpathy JSON is the source of truth downstream; the ~91 Kaggle images it does not
+reference are never loaded, and captions.txt is downloaded for completeness but unused.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import shutil
+import urllib.request
+import zipfile
+from pathlib import Path
+
+from capit.config import config
+
+KARPATHY_ZIP_URL = (
+    "https://cs.stanford.edu/people/karpathy/deepimagesent/caption_datasets.zip"
+)
+KAGGLE_DATASET = "adityajn105/flickr8k"
+_HTTP_TIMEOUT = 60
+
+
+def _copy_atomic(src: Path, dst: Path) -> None:
+    tmp = dst.with_name(dst.name + ".tmp")
+    shutil.copy2(src, tmp)
+    os.replace(tmp, dst)
+
+
+def download_flickr8k(dest: Path | None = None) -> Path:
+    """Materialize Flickr8k images + captions.txt into ``dest`` (default ``config.flickr8k_dir``)."""
+    dest = Path(dest) if dest is not None else config.flickr8k_dir
+    images_dir = dest / "Images"
+    captions = dest / "captions.txt"
+    sentinel = dest / ".complete"
+
+    if sentinel.exists():
+        return dest
+
+    import kagglehub  # lazy: keeps the module importable without the dep installed
+
+    src = Path(kagglehub.dataset_download(KAGGLE_DATASET))
+    src_images = src / "Images"
+    if not src_images.is_dir():
+        contents = sorted(p.name for p in src.iterdir())
+        raise FileNotFoundError(f"expected Images/ under {src}, found: {contents}")
+    src_captions = src / "captions.txt"
+    if not src_captions.is_file():
+        contents = sorted(p.name for p in src.iterdir())
+        raise FileNotFoundError(f"expected captions.txt under {src}, found: {contents}")
+
+    images_dir.mkdir(parents=True, exist_ok=True)
+    for jpg in src_images.glob("*.jpg"):
+        target = images_dir / jpg.name
+        if not target.exists():
+            _copy_atomic(jpg, target)
+    _copy_atomic(src_captions, captions)
+
+    sentinel.touch()
+    return dest
+
+
+def download_karpathy_splits(dest: Path | None = None) -> Path:
+    """Download Karpathy's caption_datasets.zip and extract dataset_flickr8k.json into ``dest``.
+    """
+    dest = Path(dest) if dest is not None else config.karpathy_dir
+    json_path = dest / "dataset_flickr8k.json"
+    if json_path.is_file():
+        return json_path
+
+    dest.mkdir(parents=True, exist_ok=True)
+    with urllib.request.urlopen(KARPATHY_ZIP_URL, timeout=_HTTP_TIMEOUT) as resp:  # noqa: S310 — trusted https URL
+        archive = resp.read()
+
+    with zipfile.ZipFile(io.BytesIO(archive)) as zf:
+        members = [n for n in zf.namelist() if n.endswith("dataset_flickr8k.json")]
+        if not members:
+            raise FileNotFoundError(
+                f"dataset_flickr8k.json not in archive; contents: {zf.namelist()}"
+            )
+        tmp = json_path.with_name(json_path.name + ".tmp")
+        try:
+            with zf.open(members[0]) as f_in, open(tmp, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+            os.replace(tmp, json_path)
+        except BaseException:
+            tmp.unlink(missing_ok=True)
+            raise
+
+    return json_path
+
+
+def main() -> None:
+    flickr_dir = download_flickr8k()
+    json_path = download_karpathy_splits()
+    n_jpg = len(list((flickr_dir / "Images").glob("*.jpg")))
+    print(f"flickr8k: {n_jpg} jpgs in {flickr_dir / 'Images'}")
+    print(f"karpathy: {json_path}")
+
+
+if __name__ == "__main__":
+    main()

pipeline/src/capit/data/records.py

@@ -0,0 +1,29 @@
+"""Pure operations over Karpathy image records (the `images` list of the split JSON)."""
+
+from __future__ import annotations
+
+import random
+
+
+def train_captions(records: list[dict]) -> list[list[str]]:
+    captions = [s["tokens"] for r in records if r["split"] == "train" for s in r["sentences"]]
+    if not captions:
+        splits = sorted({r["split"] for r in records})
+        raise ValueError(f"no train records found; available splits: {splits}")
+    return captions
+
+
+def select_records(records: list[dict], counts: dict[str, int], seed: int) -> list[dict]:
+    by_split: dict[str, list[dict]] = {split: [] for split in counts}
+    for rec in records:
+        if rec["split"] in by_split:
+            by_split[rec["split"]].append(rec)
+
+    rng = random.Random(seed)
+    selected: list[dict] = []
+    for split, count in counts.items():
+        pool = sorted(by_split[split], key=lambda r: r["filename"])
+        if len(pool) < count:
+            raise ValueError(f"{split}: need {count}, only {len(pool)} available")
+        selected.extend(rng.sample(pool, count))
+    return selected

pipeline/src/capit/data/vocab.py

@@ -0,0 +1,57 @@
+"""Vocabulary: a word<->id mapping with frequency-filtered tokens."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from collections import Counter
+from pathlib import Path
+
+PAD, START, END, UNK = "<pad>", "<start>", "<end>", "<unk>"
+SPECIALS = [PAD, START, END, UNK]
+
+
+class Vocab:
+    def __init__(self, words: list[str]) -> None:
+        if list(words[:4]) != SPECIALS:
+            raise ValueError(f"ids 0-3 must be {SPECIALS}, got {list(words[:4])}")
+        if len(set(words)) != len(words):
+            raise ValueError("vocab contains duplicate words")
+        self.id2word = list(words)
+        self.word2id = {w: i for i, w in enumerate(self.id2word)}
+
+    @classmethod
+    def build(cls, captions: list[list[str]], min_freq: int = 5) -> "Vocab":
+        counts: Counter[str] = Counter()
+        for tokens in captions:
+            counts.update(tokens)
+        kept = sorted((w for w, c in counts.items() if c >= min_freq), key=lambda w: (-counts[w], w))
+        return cls(SPECIALS + kept)
+
+    def encode(self, words: list[str]) -> list[int]:
+        unk = self.word2id[UNK]
+        return [self.word2id.get(w, unk) for w in words]
+
+    def decode(self, ids: list[int]) -> list[str]:
+        n = len(self.id2word)
+        for i in ids:
+            if not 0 <= i < n:
+                raise ValueError(f"decode: id {i} out of range [0, {n})")
+        return [self.id2word[i] for i in ids]
+
+    def __len__(self) -> int:
+        return len(self.id2word)
+
+    def sha256(self) -> str:
+        return hashlib.sha256(json.dumps(self.id2word, ensure_ascii=False).encode()).hexdigest()
+
+    def save(self, path: Path | str) -> None:
+        path = Path(path)
+        tmp = path.with_name(path.name + ".tmp")
+        tmp.write_text(json.dumps(self.id2word, ensure_ascii=False, indent=2))
+        os.replace(tmp, path)
+
+    @classmethod
+    def load(cls, path: Path | str) -> "Vocab":
+        return cls(json.loads(Path(path).read_text()))

pipeline/src/capit/decode.py

@@ -0,0 +1,80 @@
+"""Greedy and beam-search decoding with per-word attention.
+
+Both operate on a single image's encoder features ``[1, L, encoder_dim]`` and return the caption token ids (specials stripped) plus an aligned ``[T, L]`` alpha tensor — the per-word attention that drives the heatmaps. beam_search also returns the full completed set as ``(tokens, normalized_score)`` pairs, winner first ("the road not taken").
+"""
+
+from __future__ import annotations
+
+import torch
+from torch.nn.functional import log_softmax
+from capit.config import config
+from capit.models.decoder import Decoder
+
+
+@torch.no_grad()
+def greedy(
+    decoder: Decoder, features: torch.Tensor, start_id: int, end_id: int, max_len: int = config.max_decode_len
+) -> tuple[list[int], torch.Tensor]:
+    decoder.eval()
+    h, c = decoder.init_hidden(features)
+    token = features.new_full((1,), start_id, dtype=torch.long)
+    ids: list[int] = []
+    alphas: list[torch.Tensor] = []
+    for _ in range(max_len):
+        logits, alpha, h, c = decoder.step(token, h, c, features)
+        token = logits.argmax(dim=1)
+        if int(token) == end_id:
+            break
+        ids.append(int(token))
+        alphas.append(alpha[0])
+    stacked = torch.stack(alphas) if alphas else features.new_zeros(0, features.size(1))
+    return ids, stacked
+
+
+@torch.no_grad()
+def beam_search(
+    decoder: Decoder,
+    features: torch.Tensor,
+    start_id: int,
+    end_id: int,
+    k: int = config.beam_k,
+    max_len: int = config.max_decode_len,
+    alpha_norm: float = config.beam_alpha,
+) -> tuple[list[int], torch.Tensor, list[tuple[list[int], float]]]:
+    decoder.eval()
+    vocab_size, num_pixels = decoder.fc.out_features, features.size(1)
+    feats = features.expand(k, num_pixels, -1)
+    h, c = decoder.init_hidden(feats)
+    seqs = features.new_full((k, 1), start_id, dtype=torch.long)
+    seq_alphas = features.new_zeros(k, 1, num_pixels)
+    scores = features.new_zeros(k, 1)
+    completed: list[tuple[list[int], torch.Tensor, float]] = []
+
+    for step in range(max_len):
+        logits, alpha, h, c = decoder.step(seqs[:, -1], h, c, feats)
+        logp = scores + log_softmax(logits, dim=1)
+        flat = logp[0] if step == 0 else logp.reshape(-1)  # step 0: all beams identical
+        top_scores, top_idx = flat.topk(seqs.size(0))
+        beam_idx, word_idx = top_idx // vocab_size, top_idx % vocab_size
+
+        seqs = torch.cat([seqs[beam_idx], word_idx.unsqueeze(1)], dim=1)
+        seq_alphas = torch.cat([seq_alphas[beam_idx], alpha[beam_idx].unsqueeze(1)], dim=1)
+        scores, h, c, feats = top_scores.unsqueeze(1), h[beam_idx], c[beam_idx], feats[beam_idx]
+
+        ended = word_idx == end_id
+        for i in ended.nonzero().flatten().tolist():
+            completed.append((seqs[i, 1:-1].tolist(), seq_alphas[i, 1:-1], float(scores[i])))
+        keep = (~ended).nonzero().flatten()
+        if len(keep) == 0:
+            break
+        seqs, seq_alphas, scores = seqs[keep], seq_alphas[keep], scores[keep]
+        h, c, feats = h[keep], c[keep], feats[keep]
+
+    if not completed:  # nothing emitted <end> by max_len → fall back to the best live hypothesis
+        i = int(scores.argmax())
+        completed.append((seqs[i, 1:].tolist(), seq_alphas[i, 1:], float(scores[i])))
+
+    ranked = sorted(completed, key=lambda x: x[2] / max(len(x[0]), 1) ** alpha_norm, reverse=True)
+    beams = [(toks, score / max(len(toks), 1) ** alpha_norm) for toks, _, score in ranked]
+    win_tokens, win_alphas, _ = ranked[0]
+    return win_tokens, win_alphas, beams

pipeline/src/capit/evaluate.py

@@ -0,0 +1,89 @@
+"""Official BLEU/CIDEr on the test split.
+
+Generates one caption per test image (beam search) and scores it against the 5 references with the pycocoevalcap Python scorers used directly on pre-tokenized strings. Encoder features are cached once and reused across beam widths.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from pycocoevalcap.bleu.bleu import Bleu
+from pycocoevalcap.cider.cider import Cider
+from capit.checkpoint import load
+from capit.config import config
+from capit.data.dataset import CaptionDataset, build_transform
+from capit.data.vocab import END, START, Vocab
+from capit.decode import beam_search
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+
+METRICS = ["BLEU-1", "BLEU-2", "BLEU-3", "BLEU-4", "CIDEr"]
+
+
+def score(
+    references: dict[int, list[list[str]]], candidates: dict[int, list[str]]
+) -> dict[str, float]:
+    gts = {i: [" ".join(toks) for toks in refs] for i, refs in references.items()}
+    res = {i: [" ".join(candidates[i])] for i in references}
+    bleu = Bleu(4).compute_score(gts, res, verbose=0)[0]
+    cider = Cider().compute_score(gts, res)[0]
+    return {
+        "BLEU-1": float(bleu[0]) * 100,
+        "BLEU-2": float(bleu[1]) * 100,
+        "BLEU-3": float(bleu[2]) * 100,
+        "BLEU-4": float(bleu[3]) * 100,
+        "CIDEr": float(cider) * 100,
+    }
+
+
+@torch.no_grad()
+def evaluate(
+    data_root, ckpt_path, beams: tuple[int, ...] = (1, 3, 5)
+) -> dict[int, dict[str, float]]:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    vocab = Vocab.load(config.vocab_path)
+    encoder = Encoder(pretrained=True).to(device).eval()
+    decoder = Decoder(vocab_size=len(vocab)).to(device)
+    state = load(ckpt_path)
+    if state.vocab_sha256 != vocab.sha256():
+        raise ValueError(f"vocab mismatch: ckpt {state.vocab_sha256[:8]} != vocab {vocab.sha256()[:8]}")
+    decoder.load_state_dict(state.model_state)
+
+    ds = CaptionDataset(data_root, "test", vocab, build_transform())
+    references = ds.references()
+    features = {iid: encoder(img.unsqueeze(0).to(device)) for iid, img in ds.iter_images()}
+
+    sid, eid = vocab.word2id[START], vocab.word2id[END]
+    results: dict[int, dict[str, float]] = {}
+    for k in beams:
+        candidates = {
+            iid: vocab.decode(beam_search(decoder, f, sid, eid, k=k)[0])
+            for iid, f in features.items()
+        }
+        results[k] = score(references, candidates)
+    return results
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-root", default=str(config.flickr8k_dir))
+    parser.add_argument("--ckpt", default=str(config.ckpt_dir / "best.pt"))
+    parser.add_argument("--beams", type=int, nargs="+", default=[1, 3, 5])
+    parser.add_argument("--out-json", help="write the results table to this path for downstream use")
+    args = parser.parse_args()
+
+    results = evaluate(args.data_root, args.ckpt, tuple(args.beams))
+    header = "beam  " + "  ".join(f"{m:>7}" for m in METRICS)
+    print(header)
+    print("-" * len(header))
+    for k, s in results.items():
+        print(f"{k:>4}  " + "  ".join(f"{s[m]:>7.2f}" for m in METRICS))
+    if args.out_json:
+        Path(args.out_json).write_text(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()

pipeline/src/capit/losses.py

@@ -0,0 +1,23 @@
+"""Caption loss: cross-entropy over real positions + doubly-stochastic attention regularizer."""
+
+from __future__ import annotations
+
+import torch
+from torch.nn.functional import cross_entropy
+
+from capit.config import config
+
+
+def word_cross_entropy(logits: torch.Tensor, captions: torch.Tensor) -> torch.Tensor:
+    targets = captions[:, 1 : 1 + logits.size(1)]
+    return cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=0)
+
+
+def caption_loss(
+    logits: torch.Tensor,
+    alphas: torch.Tensor,
+    captions: torch.Tensor,
+    alpha_c: float = config.alpha_c,
+) -> torch.Tensor:
+    reg = alpha_c * ((1 - alphas.sum(dim=1)) ** 2).mean()
+    return word_cross_entropy(logits, captions) + reg

pipeline/src/capit/models/attention.py

@@ -0,0 +1,29 @@
+"""Attention module for processing spatial features."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from capit.config import config
+
+class BahdanauAttention(nn.Module):
+    def __init__(
+        self,
+        encoder_dim: int = config.encoder_dim,
+        decoder_dim: int = config.decoder_dim,
+        attention_dim: int = config.attention_dim,
+    ) -> None:
+        super().__init__()
+        self.encoder_att = nn.Linear(encoder_dim, attention_dim)
+        self.decoder_att = nn.Linear(decoder_dim, attention_dim)
+        self.full_att = nn.Linear(attention_dim, 1)
+        self.relu = nn.ReLU()
+        self.softmax = nn.Softmax(dim=1)
+    
+    def forward(self, features: torch.Tensor, h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        att1 = self.encoder_att(features)
+        att2 = self.decoder_att(h).unsqueeze(1)
+        att = self.full_att(self.relu(att1 + att2)).squeeze(2)
+        alpha = self.softmax(att)
+        context = (features * alpha.unsqueeze(2)).sum(dim=1)
+        return context, alpha

pipeline/src/capit/models/decoder.py

@@ -0,0 +1,93 @@
+"""Decoder: a recurrent neural network for generating captions."""
+
+from __future__ import annotations
+import torch
+from torch import nn
+from capit.config import config
+from capit.models.attention import BahdanauAttention
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_dim: int = config.embed_dim,
+        decoder_dim: int = config.decoder_dim,
+        attention_dim: int = config.attention_dim,
+        dropout: float = config.dropout,
+    ) -> None:
+        super().__init__()
+        self.attention = BahdanauAttention(config.encoder_dim, decoder_dim, attention_dim)
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.decode_step = nn.LSTMCell(embed_dim + config.encoder_dim, decoder_dim)
+        self.init_h = nn.Linear(config.encoder_dim, decoder_dim)
+        self.init_c = nn.Linear(config.encoder_dim, decoder_dim)
+        self.f_beta = nn.Linear(decoder_dim, config.encoder_dim)
+        self.fc = nn.Linear(decoder_dim, vocab_size)
+
+    def init_hidden(self, features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean_features = features.mean(dim=1)
+        h = self.init_h(mean_features)
+        c = self.init_c(mean_features)
+        return h, c
+
+    def step(
+        self, token: torch.Tensor, h: torch.Tensor, c: torch.Tensor, features: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        context, alpha = self.attention(features, h)
+        context = torch.sigmoid(self.f_beta(h)) * context
+        x = torch.cat([self.embedding(token), context], dim=1)
+        h, c = self.decode_step(x, (h, c))
+        return self.fc(self.dropout(h)), alpha, h, c
+
+    def forward(
+        self, features: torch.Tensor, captions: torch.Tensor, lengths: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
+        embeddings = self.embedding(captions)
+        h, c = self.init_hidden(features)
+        decode_lengths = (lengths - 1).tolist()
+        if any(a < b for a, b in zip(decode_lengths, decode_lengths[1:])):
+            raise ValueError(f"forward expects lengths sorted descending; got {lengths.tolist()}")
+        T = max(decode_lengths)
+        batch_size, num_pixels, _ = features.size()
+        vocab_size = self.fc.out_features
+        logits = features.new_zeros(batch_size, T, vocab_size)
+        alphas = features.new_zeros(batch_size, T, num_pixels)
+        for t in range(T):
+            n = sum(l > t for l in decode_lengths)
+            context, alpha = self.attention(features[:n], h[:n])
+            beta = torch.sigmoid(self.f_beta(h[:n]))
+            context = beta * context
+            x = torch.cat([embeddings[:n, t], context], dim=1)
+            h, c = self.decode_step(x, (h[:n], c[:n]))
+            logits[:n, t] = self.fc(self.dropout(h))
+            alphas[:n, t] = alpha
+        return logits, alphas, decode_lengths
+
+    @torch.no_grad()
+    def greedy(
+        self, features: torch.Tensor, start_id: int, end_id: int, max_len: int = 50
+    ) -> list[list[int]]:
+        was_training = self.training
+        self.eval()
+        batch_size = features.size(0)
+        h, c = self.init_hidden(features)
+        prev = features.new_full((batch_size,), start_id, dtype=torch.long)
+        outputs: list[list[int]] = [[] for _ in range(batch_size)]
+        done = [False] * batch_size
+        for _ in range(max_len):
+            logits, _, h, c = self.step(prev, h, c, features)
+            prev = logits.argmax(dim=1)
+            for i in range(batch_size):
+                if not done[i]:
+                    tok = int(prev[i])
+                    if tok == end_id:
+                        done[i] = True
+                    else:
+                        outputs[i].append(tok)
+            if all(done):
+                break
+        if was_training:
+            self.train()
+        return outputs
+

pipeline/src/capit/models/encoder.py

@@ -0,0 +1,34 @@
+"""Encoder: a frozen ResNet-50 backbone producing [B, 196, 2048] spatial features."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torchvision.models import ResNet50_Weights, resnet50
+from capit.config import config
+
+
+class Encoder(nn.Module):
+    def __init__(self, encoded_size: int = config.encoded_size, pretrained: bool = True) -> None:
+        super().__init__()
+        weights = ResNet50_Weights.DEFAULT if pretrained else None
+        backbone = resnet50(weights=weights)
+        self.backbone = nn.Sequential(*list(backbone.children())[:-2])
+        self.pool = nn.AdaptiveAvgPool2d((encoded_size, encoded_size))
+        for p in self.parameters():
+            p.requires_grad = False
+        self.eval()
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        x = self.pool(self.backbone(images))
+        return x.flatten(2).permute(0, 2, 1)
+
+    def fine_tune(self, blocks: tuple[int, ...] = ()) -> None:
+        for i in blocks:
+            for p in self.backbone[i].parameters():
+                p.requires_grad = True
+
+    def train(self, mode: bool = True) -> "Encoder":
+        super().train(mode)
+        self.backbone.eval()
+        return self

pipeline/src/capit/overfit.py

@@ -0,0 +1,65 @@
+"""Overfit-one-batch harness (Stage 2.3, killer gate #1).
+
+Trains the decoder on ONE batch of N distinct subsample images (encoder frozen, features
+cached once) and reports whether it memorizes them. Trains on the full caption loss
+(CE + doubly-stochastic regularizer); the meaningful signal is the CE term, since the
+regularizer floors near ~0.9 for short captions so a sub-0.5 total is unreachable.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch.optim import Adam
+
+from capit.config import config
+from capit.data.dataset import CaptionDataset, build_transform, collate_fn
+from capit.data.vocab import END, START, Vocab
+from capit.losses import caption_loss, word_cross_entropy
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+
+
+def _distinct_image_batch(ds: CaptionDataset, n: int):
+    seen: set[str] = set()
+    picks = []
+    for idx, (rec, _) in enumerate(ds.pairs):
+        if rec["filename"] not in seen:
+            seen.add(rec["filename"])
+            picks.append(idx)
+        if len(picks) == n:
+            break
+    if len(picks) < n:
+        raise ValueError(f"requested {n} distinct images but dataset has only {len(picks)}")
+    return collate_fn([ds[i] for i in picks])
+
+
+def run_overfit(steps: int = 400, n_images: int = 4) -> dict:
+    torch.manual_seed(config.seed)
+    vocab = Vocab.load(config.vocab_path)
+    ds = CaptionDataset(config.subsample_root, "train", vocab, build_transform())
+    images, captions, lengths = _distinct_image_batch(ds, n_images)
+
+    encoder = Encoder(pretrained=True)
+    with torch.no_grad():
+        features = encoder(images)  # frozen → cache once, never recompute
+
+    decoder = Decoder(vocab_size=len(vocab))
+    decoder.train()
+    opt = Adam(decoder.parameters(), lr=config.decoder_lr)
+
+    ce_curve = []
+    for step in range(steps):
+        logits, alphas, _ = decoder(features, captions, lengths)
+        ce = word_cross_entropy(logits, captions)
+        loss = caption_loss(logits, alphas, captions)
+        if not torch.isfinite(loss):
+            raise RuntimeError(f"loss diverged to {loss.item()} at step {step}")
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        ce_curve.append(ce.item())
+
+    decoded = decoder.greedy(features, vocab.word2id[START], vocab.word2id[END])
+    decoded_words = [vocab.decode(ids) for ids in decoded]
+    target_words = [vocab.decode(captions[i, 1 : lengths[i] - 1].tolist()) for i in range(n_images)]
+    return {"ce_curve": ce_curve, "final_ce": ce_curve[-1], "decoded": decoded_words, "targets": target_words}

pipeline/src/capit/serving.py

@@ -0,0 +1,97 @@
+"""Serving artifact: build / load the self-contained inference contract (capit-sat.pt).
+
+The artifact carries everything inference needs (encoder + decoder weights, model dims, the
+exact preprocess spec, and the vocab hash) and nothing training-only. The backend rebuilds the
+model from this alone — it never imports training code or touches a checkpoint.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+from torchvision import transforms
+
+from capit.config import config
+from capit.data.vocab import END, START, Vocab
+from capit.decode import beam_search
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+
+
+def build_artifact(encoder: Encoder, decoder: Decoder, vocab: Vocab) -> dict[str, Any]:
+    return {
+        "encoder_state": encoder.state_dict(),
+        "decoder_state": decoder.state_dict(),
+        "config": {
+            "vocab_size": len(vocab),
+            "encoder_dim": config.encoder_dim,
+            "decoder_dim": config.decoder_dim,
+            "attention_dim": config.attention_dim,
+            "embed_dim": config.embed_dim,
+            "encoded_size": config.encoded_size,
+            "dropout": config.dropout,
+            "beam_k": config.beam_k,
+            "beam_alpha": config.beam_alpha,
+            "max_decode_len": config.max_decode_len,
+        },
+        "preprocess": {
+            "resize": config.resize,
+            "crop": config.crop,
+            "mean": list(config.imagenet_mean),
+            "std": list(config.imagenet_std),
+        },
+        "vocab_sha256": vocab.sha256(),
+    }
+
+
+def center_crop_box(width: int, height: int, preprocess: dict[str, Any]) -> dict[str, float]:
+    """The Resize(short=resize)+CenterCrop(crop) region, as fractions of the original image.
+
+    Resize preserves aspect, so resized-image fractions equal original-image fractions. The
+    frontend positions the attention overlay over this box (non-square uploads misalign otherwise).
+    """
+    resize, crop = preprocess["resize"], preprocess["crop"]
+    scale = resize / min(width, height)
+    rw, rh = width * scale, height * scale
+    return {"x": (rw - crop) / 2 / rw, "y": (rh - crop) / 2 / rh, "w": crop / rw, "h": crop / rh}
+
+
+def make_transform(preprocess: dict[str, Any]):
+    return transforms.Compose(
+        [
+            transforms.Resize(preprocess["resize"]),
+            transforms.CenterCrop(preprocess["crop"]),
+            transforms.ToTensor(),
+            transforms.Normalize(preprocess["mean"], preprocess["std"]),
+        ]
+    )
+
+
+def load_artifact(artifact_path: Path | str, vocab_path: Path | str):
+    blob = torch.load(artifact_path, map_location="cpu", weights_only=True)
+    vocab = Vocab.load(vocab_path)
+    if vocab.sha256() != blob["vocab_sha256"]:
+        raise ValueError(f"vocab mismatch: vocab {vocab.sha256()[:8]} != artifact {blob['vocab_sha256'][:8]}")
+    cfg = blob["config"]
+    encoder = Encoder(encoded_size=cfg["encoded_size"], pretrained=False)
+    encoder.load_state_dict(blob["encoder_state"])
+    encoder.eval()
+    decoder = Decoder(
+        vocab_size=cfg["vocab_size"],
+        embed_dim=cfg["embed_dim"],
+        decoder_dim=cfg["decoder_dim"],
+        attention_dim=cfg["attention_dim"],
+        dropout=cfg["dropout"],
+    )
+    decoder.load_state_dict(blob["decoder_state"])
+    decoder.eval()
+    return encoder, decoder, vocab, blob["preprocess"]
+
+
+@torch.no_grad()
+def caption(encoder: Encoder, decoder: Decoder, vocab: Vocab, image: torch.Tensor, k: int = config.beam_k):
+    features = encoder(image.unsqueeze(0))
+    tokens, alphas, beams = beam_search(decoder, features, vocab.word2id[START], vocab.word2id[END], k=k)
+    return vocab.decode(tokens), alphas, beams

pipeline/src/capit/train.py

@@ -0,0 +1,146 @@
+"""Train the capit model.
+
+CLI: python -m capit.train --data-root <root> --ckpt-dir <dir> --resume auto
+
+Checkpoints the decoder + optimizer every epoch (latest.pt) and on val-BLEU improvement (best.pt); the frozen encoder is reconstructed from ImageNet weights, not stored. Early stops on val BLEU-4 (nltk, a cheap relative signal — reported scores come from Stage 4.2).
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import cast
+
+import torch
+from nltk.translate.bleu_score import SmoothingFunction, corpus_bleu
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+
+from capit.checkpoint import load, save
+from capit.config import config
+from capit.data.dataset import CaptionDataset, build_transform, collate_fn
+from capit.data.vocab import END, START, Vocab
+from capit.losses import caption_loss
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+
+
+def corpus_bleu4(references: dict[int, list[list[str]]], candidates: dict[int, list[str]]) -> float:
+    order = list(references)
+    if not order:
+        raise ValueError("corpus_bleu4: no references (val split empty?)")
+    refs = [references[i] for i in order]
+    hyps = [candidates[i] for i in order]
+    empty = sum(not h for h in hyps)
+    if empty:
+        print(f"warning: {empty}/{len(hyps)} val captions are empty (BLEU will read low)")
+    return cast(float, corpus_bleu(refs, hyps, smoothing_function=SmoothingFunction().method1)) * 100
+
+
+def _train_one_epoch(encoder, decoder, loader, opt, device) -> float:
+    decoder.train()
+    total = 0.0
+    for images, captions, lengths in loader:
+        images, captions = images.to(device), captions.to(device)
+        with torch.no_grad():
+            features = encoder(images)
+        logits, alphas, _ = decoder(features, captions, lengths)
+        loss = caption_loss(logits, alphas, captions)
+        if not torch.isfinite(loss):
+            raise FloatingPointError(f"non-finite loss {loss.item()} — aborting before it corrupts latest.pt")
+        opt.zero_grad()
+        loss.backward()
+        clip_grad_norm_(decoder.parameters(), config.grad_clip)
+        opt.step()
+        total += loss.item()
+    return total / len(loader)
+
+
+@torch.no_grad()
+def _evaluate(encoder, decoder, ds, vocab, device) -> tuple[float, dict[int, list[str]]]:
+    decoder.eval()
+    candidates: dict[int, list[str]] = {}
+    for image_id, image in ds.iter_images():
+        features = encoder(image.unsqueeze(0).to(device))
+        ids = decoder.greedy(features, vocab.word2id[START], vocab.word2id[END])[0]
+        candidates[image_id] = vocab.decode(ids)
+    return corpus_bleu4(ds.references(), candidates), candidates
+
+
+def train(
+    data_root: str | Path,
+    ckpt_dir: str | Path,
+    resume: str = "auto",
+    vocab_path: str | Path | None = None,
+    max_epochs: int | None = None,
+    num_workers: int | None = None,
+) -> None:
+    max_epochs = config.max_epochs if max_epochs is None else max_epochs
+    num_workers = config.num_workers if num_workers is None else num_workers
+    torch.manual_seed(config.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt_dir = Path(ckpt_dir)
+    if "drive/MyDrive" in str(ckpt_dir) and not Path("/content/drive/MyDrive").is_dir():
+        raise RuntimeError(f"--ckpt-dir {ckpt_dir} is on Drive but it isn't mounted — run the mount cell first")
+    latest_path, best_path = ckpt_dir / "latest.pt", ckpt_dir / "best.pt"
+
+    vocab = Vocab.load(vocab_path or config.vocab_path)
+    transform = build_transform()
+    train_ds = CaptionDataset(data_root, "train", vocab, transform)
+    val_ds = CaptionDataset(data_root, "val", vocab, transform)
+    generator = torch.Generator().manual_seed(config.seed)
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=config.batch_size,
+        shuffle=True,
+        generator=generator,
+        collate_fn=collate_fn,
+        num_workers=num_workers,
+    )
+
+    encoder = Encoder(pretrained=True).to(device)
+    decoder = Decoder(vocab_size=len(vocab)).to(device)
+    opt = Adam(decoder.parameters(), lr=config.decoder_lr)
+
+    start_epoch, best_bleu4, best_epoch = 0, -1.0, 0
+    if resume == "auto" and latest_path.is_file():
+        state = load(latest_path)
+        if state.vocab_sha256 != vocab.sha256():
+            raise ValueError("vocab mismatch: checkpoint vocab_sha256 != current vocab.json")
+        decoder.load_state_dict(state.model_state)
+        opt.load_state_dict(state.optim_state)
+        start_epoch, best_bleu4, best_epoch = state.epoch + 1, state.best_bleu4, state.best_epoch
+        print(f"resumed at epoch {start_epoch} (best_bleu4 {best_bleu4:.2f} at epoch {best_epoch})")
+
+    for epoch in range(start_epoch, max_epochs):
+        train_loss = _train_one_epoch(encoder, decoder, train_loader, opt, device)
+        val_bleu4, candidates = _evaluate(encoder, decoder, val_ds, vocab, device)
+
+        improved = val_bleu4 > best_bleu4
+        if improved:
+            best_bleu4, best_epoch = val_bleu4, epoch
+        save(latest_path, decoder, opt, epoch, best_bleu4, best_epoch, vocab.sha256())
+        if improved:
+            save(best_path, decoder, opt, epoch, best_bleu4, best_epoch, vocab.sha256())
+
+        sample = " ".join(next(iter(candidates.values())))
+        flag = " *" if improved else ""
+        print(f"epoch {epoch}  loss {train_loss:.4f}  val_bleu4 {val_bleu4:.2f}  best {best_bleu4:.2f}{flag}  | {sample}")
+        if epoch - best_epoch >= config.patience:
+            print(f"early stop: no val BLEU-4 improvement in {config.patience} epochs (last loss {train_loss:.4f})")
+            break
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-root", default=str(config.subsample_root))
+    parser.add_argument("--ckpt-dir", default=str(config.ckpt_dir))
+    parser.add_argument("--vocab-path", default=str(config.vocab_path))
+    parser.add_argument("--resume", default="auto", choices=["auto", "none"])
+    args = parser.parse_args()
+    train(args.data_root, args.ckpt_dir, args.resume, vocab_path=args.vocab_path)
+
+
+if __name__ == "__main__":
+    main()

pipeline/tests/test_attention.py

@@ -0,0 +1,68 @@
+"""Stage 2.2 — BahdanauAttention. Pure tensor tests; no data needed."""
+
+import pytest
+import torch
+
+from capit.config import config
+from capit.models.attention import BahdanauAttention
+
+B = 3
+L = config.encoded_size**2  # 196 locations
+
+
+@pytest.fixture(autouse=True)
+def _seed():
+    torch.manual_seed(config.seed)
+
+
+def _inputs(requires_grad: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    features = torch.randn(B, L, config.encoder_dim, requires_grad=requires_grad)
+    h = torch.randn(B, config.decoder_dim, requires_grad=requires_grad)
+    return features, h
+
+
+def test_output_shapes():
+    context, alpha = BahdanauAttention()(*_inputs())
+    assert context.shape == (B, config.encoder_dim)
+    assert alpha.shape == (B, L)
+
+
+def test_alpha_sums_to_one():
+    _, alpha = BahdanauAttention()(*_inputs())
+    assert torch.allclose(alpha.sum(dim=1), torch.ones(B), atol=1e-5)
+
+
+def test_alpha_is_a_distribution():
+    _, alpha = BahdanauAttention()(*_inputs())
+    assert (alpha >= 0).all()
+    assert not torch.isnan(alpha).any()
+
+
+def test_no_nan_in_context():
+    context, _ = BahdanauAttention()(*_inputs())
+    assert not torch.isnan(context).any()
+
+
+def test_alpha_depends_on_hidden_state():
+    att = BahdanauAttention()
+    features, _ = _inputs()
+    _, a1 = att(features, torch.randn(B, config.decoder_dim))
+    _, a2 = att(features, torch.randn(B, config.decoder_dim))
+    assert not torch.allclose(a1, a2)
+
+
+def test_batch_independence():
+    att = BahdanauAttention()
+    features, h = _inputs()
+    ctx_all, alpha_all = att(features, h)
+    ctx_row, alpha_row = att(features[:1], h[:1])
+    assert torch.allclose(alpha_all[:1], alpha_row, atol=1e-6)
+    assert torch.allclose(ctx_all[:1], ctx_row, atol=1e-6)
+
+
+def test_gradient_flows_to_both_inputs():
+    features, h = _inputs(requires_grad=True)
+    context, _ = BahdanauAttention()(features, h)
+    context.sum().backward()
+    assert features.grad is not None and features.grad.abs().sum() > 0
+    assert h.grad is not None and h.grad.abs().sum() > 0

pipeline/tests/test_checkpoint.py

@@ -0,0 +1,63 @@
+"""Stage 3.1 — checkpoint round-trip including optimizer state. Hermetic, no data."""
+
+import torch
+from torch import nn
+from torch.optim import Adam
+
+from capit.checkpoint import load, save
+
+
+def _model() -> nn.Module:
+    torch.manual_seed(13)
+    return nn.Linear(4, 4)
+
+
+def _step(model: nn.Module, opt: Adam) -> float:
+    loss = model(torch.ones(2, 4)).pow(2).mean()
+    opt.zero_grad()
+    loss.backward()
+    opt.step()
+    return loss.item()
+
+
+def test_round_trip_identity(tmp_path):
+    model = _model()
+    opt = Adam(model.parameters())
+    _step(model, opt)
+    path = tmp_path / "ckpt.pt"
+    save(path, model, opt, epoch=3, best_bleu4=12.5, best_epoch=2, vocab_sha256="abc")
+
+    ts = load(path)
+    assert (ts.epoch, ts.best_bleu4, ts.best_epoch, ts.vocab_sha256) == (3, 12.5, 2, "abc")
+    fresh = nn.Linear(4, 4)
+    fresh.load_state_dict(ts.model_state)
+    for a, b in zip(fresh.state_dict().values(), model.state_dict().values()):
+        assert torch.equal(a, b)
+
+
+def test_optimizer_state_round_trips(tmp_path):
+    model = _model()
+    opt = Adam(model.parameters())
+    _step(model, opt)
+    _step(model, opt)
+    path = tmp_path / "ckpt.pt"
+    save(path, model, opt, epoch=2, best_bleu4=0.0, best_epoch=0, vocab_sha256="x")
+    live = [_step(model, opt), _step(model, opt)]
+
+    model2 = nn.Linear(4, 4)
+    opt2 = Adam(model2.parameters())
+    ts = load(path)
+    model2.load_state_dict(ts.model_state)
+    opt2.load_state_dict(ts.optim_state)
+    resumed = [_step(model2, opt2), _step(model2, opt2)]
+
+    assert live == resumed
+
+
+def test_save_leaves_no_tmp(tmp_path):
+    model = _model()
+    opt = Adam(model.parameters())
+    path = tmp_path / "ckpt.pt"
+    save(path, model, opt, epoch=0, best_bleu4=0.0, best_epoch=0, vocab_sha256="x")
+    assert path.exists()
+    assert not (tmp_path / "ckpt.pt.tmp").exists()

pipeline/tests/test_dataset.py

@@ -0,0 +1,109 @@
+"""CaptionDataset.
+
+Runs against the dev subsample + built vocab; skips when either is absent.
+"""
+
+from collections import Counter
+import pytest
+import torch
+from torch.utils.data import DataLoader
+
+from capit.config import config
+from capit.data.dataset import CaptionDataset, build_transform, collate_fn
+from capit.data.vocab import END, PAD, START, Vocab
+
+pytestmark = pytest.mark.skipif(
+    not (config.subsample_json.is_file() and config.vocab_path.is_file()),
+    reason="dev subsample or vocab not built",
+)
+
+
+@pytest.fixture(scope="module")
+def vocab() -> Vocab:
+    return Vocab.load(config.vocab_path)
+
+
+@pytest.fixture(scope="module")
+def train_ds(vocab: Vocab) -> CaptionDataset:
+    return CaptionDataset(config.subsample_root, "train", vocab, build_transform())
+
+
+def test_len_is_five_per_image(train_ds: CaptionDataset):
+    n_images = len({r["filename"] for r, _ in train_ds.pairs})
+    assert len(train_ds) == 5 * n_images
+    assert len(train_ds) == 5 * config.subsample_train
+
+
+def test_each_image_has_exactly_five_pairs(train_ds: CaptionDataset):
+    per_image = Counter(r["filename"] for r, _ in train_ds.pairs)
+    assert set(per_image.values()) == {5}
+
+
+def test_item_shapes_and_dtypes(train_ds: CaptionDataset):
+    img, ids, clen = train_ds[0]
+    assert img.shape == (3, config.crop, config.crop)
+    assert img.dtype == torch.float32
+    assert ids.shape == (train_ds.max_len,)
+    assert ids.dtype == torch.long
+    assert isinstance(clen, int)
+
+
+def test_start_end_and_padding(train_ds: CaptionDataset, vocab: Vocab):
+    _, ids, clen = train_ds[0]
+    assert ids[0].item() == vocab.word2id[START]
+    assert ids[clen - 1].item() == vocab.word2id[END]
+    assert all(ids[j].item() == vocab.word2id[PAD] for j in range(clen, train_ds.max_len))
+
+
+def test_references_five_per_image(vocab: Vocab):
+    ds = CaptionDataset(config.subsample_root, "test", vocab, build_transform())
+    refs = ds.references()
+    assert len(refs) == config.subsample_test
+    assert all(len(r) == 5 for r in refs.values())
+
+
+def test_references_are_raw_str_tokens(vocab: Vocab):
+    ds = CaptionDataset(config.subsample_root, "test", vocab, build_transform())
+    refs = ds.references()
+    assert all(isinstance(tok, str) for caps in refs.values() for caption in caps for tok in caption)
+
+
+def test_iter_images_once_per_image(vocab: Vocab):
+    ds = CaptionDataset(config.subsample_root, "test", vocab, build_transform())
+    ids = [iid for iid, _ in ds.iter_images()]
+    assert len(ids) == config.subsample_test
+    assert len(set(ids)) == config.subsample_test
+
+
+def test_image_id_consistent_across_eval_views(vocab: Vocab):
+    ds = CaptionDataset(config.subsample_root, "test", vocab, build_transform())
+    ref_ids = set(ds.references().keys())
+    img_ids = {iid for iid, _ in ds.iter_images()}
+    assert img_ids == ref_ids == set(range(config.subsample_test))
+
+
+def test_unknown_split_raises(vocab: Vocab):
+    with pytest.raises(ValueError):
+        CaptionDataset(config.subsample_root, "nope", vocab, build_transform())
+
+
+def test_collate_sorts_descending_with_alignment(train_ds: CaptionDataset, vocab: Vocab):
+    items = sorted((train_ds[i] for i in range(8)), key=lambda x: x[2])  # ascending → force re-sort
+    images, captions, lengths = collate_fn(items)
+    lens = lengths.tolist()
+    assert images.shape == (8, 3, config.crop, config.crop)
+    assert captions.shape == (8, train_ds.max_len)
+    assert lens == sorted(lens, reverse=True)
+    assert all(n <= train_ds.max_len for n in lens)
+    pad = vocab.word2id[PAD]
+    for row, n in enumerate(lens):
+        assert int((captions[row] != pad).sum()) == n
+
+
+def test_determinism_seeded(train_ds: CaptionDataset):
+    def first_batch() -> torch.Tensor:
+        g = torch.Generator().manual_seed(config.seed)
+        loader = DataLoader(train_ds, batch_size=4, shuffle=True, generator=g, collate_fn=collate_fn)
+        return next(iter(loader))[1]
+
+    assert torch.equal(first_batch(), first_batch())

pipeline/tests/test_decode.py

@@ -0,0 +1,45 @@
+"""Stage 4.1 — greedy + beam decoding. Hermetic: random decoder + random features."""
+
+import pytest
+import torch
+
+from capit.config import config
+from capit.decode import beam_search, greedy
+from capit.models.decoder import Decoder
+
+V = 30
+L = config.encoded_size**2
+START, END = 1, 2
+
+
+@pytest.fixture(autouse=True)
+def _seed():
+    torch.manual_seed(config.seed)
+
+
+def _features() -> torch.Tensor:
+    return torch.randn(1, L, config.encoder_dim)
+
+
+def test_beam_k1_equals_greedy():
+    dec = Decoder(vocab_size=V)
+    feats = _features()
+    g_ids, _ = greedy(dec, feats, START, END)
+    b_ids, _, _ = beam_search(dec, feats, START, END, k=1)
+    assert b_ids == g_ids
+
+
+def test_greedy_alphas_align_with_ids():
+    dec = Decoder(vocab_size=V)
+    ids, alphas = greedy(dec, _features(), START, END)
+    assert alphas.shape == (len(ids), L)
+    assert not torch.isnan(alphas).any()
+
+
+def test_beam_k3_caption_and_road_not_taken():
+    dec = Decoder(vocab_size=V)
+    ids, alphas, beams = beam_search(dec, _features(), START, END, k=3)
+    assert alphas.shape == (len(ids), L)
+    assert beams[0][0] == ids  # winner first
+    assert all(isinstance(s, float) for _, s in beams)
+    assert END not in ids  # specials stripped

pipeline/tests/test_decoder.py

@@ -0,0 +1,84 @@
+"""Stage 2.3 — Decoder. Hermetic tests on random features + synthetic captions."""
+
+import pytest
+import torch
+
+from capit.config import config
+from capit.losses import caption_loss
+from capit.models.decoder import Decoder
+
+V = 12
+B = 3
+L = config.encoded_size**2
+
+# captions sorted by length descending; PAD=0 START=1 END=2
+CAPTIONS = torch.tensor(
+    [
+        [1, 5, 6, 7, 8, 2],
+        [1, 5, 6, 2, 0, 0],
+        [1, 7, 2, 0, 0, 0],
+    ]
+)
+LENGTHS = torch.tensor([6, 4, 3])
+
+
+@pytest.fixture(autouse=True)
+def _seed():
+    torch.manual_seed(config.seed)
+
+
+def _features() -> torch.Tensor:
+    return torch.randn(B, L, config.encoder_dim)
+
+
+def test_forward_shapes():
+    logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
+    assert decode_lengths == [5, 3, 2]
+    assert logits.shape == (B, 5, V)
+    assert alphas.shape == (B, 5, L)
+
+
+def test_no_nan():
+    logits, alphas, _ = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
+    assert not torch.isnan(logits).any()
+    assert not torch.isnan(alphas).any()
+
+
+def test_shrinking_leaves_finished_rows_zero():
+    logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
+    for i, dl in enumerate(decode_lengths):
+        assert torch.all(logits[i, dl:] == 0)
+        assert torch.all(alphas[i, dl:] == 0)
+
+
+def test_padded_positions_contribute_zero_loss():
+    logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS)
+    base = caption_loss(logits, alphas, CAPTIONS)
+    corrupted = logits.clone()
+    for i, dl in enumerate(decode_lengths):
+        corrupted[i, dl:] = 999.0
+    assert torch.isclose(base, caption_loss(corrupted, alphas, CAPTIONS))
+
+
+def test_forward_rejects_unsorted_batch():
+    captions = torch.tensor([[1, 7, 2, 0, 0, 0], [1, 5, 6, 7, 8, 2]])  # ascending lengths
+    lengths = torch.tensor([3, 6])
+    with pytest.raises(ValueError):
+        Decoder(vocab_size=V)(torch.randn(2, L, config.encoder_dim), captions, lengths)
+
+
+def test_gradient_reaches_decoder_params():
+    dec = Decoder(vocab_size=V)
+    logits, alphas, _ = dec(_features(), CAPTIONS, LENGTHS)
+    caption_loss(logits, alphas, CAPTIONS).backward()
+    for p in (dec.embedding.weight, dec.fc.weight, dec.f_beta.weight):
+        assert p.grad is not None and p.grad.abs().sum() > 0
+
+
+def test_greedy_returns_tokens_and_restores_mode():
+    dec = Decoder(vocab_size=V)
+    dec.train()
+    out = dec.greedy(_features(), start_id=1, end_id=2, max_len=10)
+    assert len(out) == B
+    assert all(isinstance(seq, list) and 2 not in seq for seq in out)
+    assert dec.training

pipeline/tests/test_download.py

@@ -0,0 +1,62 @@
+"""Stage 0.2 exit gate — dataset integrity.
+
+The whole module is skipped when the data isn't on disk, so CI never needs the dataset.
+When present, it verifies the counts the rest of the pipeline relies on.
+"""
+
+import json
+from collections import Counter
+import pytest
+from PIL import Image
+from capit.config import config
+
+EXPECTED_JPGS = 8091
+EXPECTED_IMAGE_RECORDS = 8000
+EXPECTED_CAPTIONS_PER_IMAGE = 5
+EXPECTED_SPLITS = {"train": 6000, "val": 1000, "test": 1000}
+
+
+def _data_present() -> bool:
+    return config.images_dir.is_dir() and config.karpathy_json.is_file()
+
+
+pytestmark = pytest.mark.skipif(
+    not _data_present(),
+    reason="Flickr8k data not on disk (run: python -m capit.data.download)",
+)
+
+
+@pytest.fixture(scope="module")
+def karpathy() -> dict:
+    with open(config.karpathy_json) as f:
+        return json.load(f)
+
+
+def test_dataset_integrity(karpathy: dict) -> None:
+    # 8,091 images on disk (Kaggle mirror).
+    jpgs = list(config.images_dir.glob("*.jpg"))
+    assert len(jpgs) == EXPECTED_JPGS, f"expected {EXPECTED_JPGS} jpgs, found {len(jpgs)}"
+
+    # The JSON is the source of truth: 8,000 records, exactly 5 captions each.
+    images = karpathy["images"]
+    assert len(images) == EXPECTED_IMAGE_RECORDS
+    assert all(len(img["sentences"]) == EXPECTED_CAPTIONS_PER_IMAGE for img in images)
+
+    # Canonical Karpathy split (flickr8k has only train/val/test — no restval).
+    splits = dict(Counter(img["split"] for img in images))
+    assert splits == EXPECTED_SPLITS, f"split tally {splits} != {EXPECTED_SPLITS}"
+
+    # Every referenced image exists on disk; the ~91 unreferenced Kaggle images may not.
+    disk = {p.name for p in jpgs}
+    missing = [img["filename"] for img in images if img["filename"] not in disk]
+    assert not missing, f"{len(missing)} referenced images absent, e.g. {missing[:3]}"
+
+
+def test_spot_check_images_open(karpathy: dict) -> None:
+    # 3 images open with PIL and their captions read sensibly.
+    for img in karpathy["images"][:3]:
+        path = config.images_dir / img["filename"]
+        with Image.open(path) as im:
+            im.load()
+        caption = " ".join(img["sentences"][0]["tokens"])
+        assert len(caption.split()) >= 3, f"caption too short: {caption!r}"

pipeline/tests/test_encoder.py

@@ -0,0 +1,87 @@
+"""Encoder.
+
+Structural tests run offline (random weights). The non-degenerate-features test needs
+ImageNet weights + a real subsample image and skips when either is absent.
+"""
+
+import pytest
+import torch
+from capit.config import config
+from capit.models.encoder import Encoder
+
+
+@pytest.fixture(scope="module")
+def encoder() -> Encoder:
+    return Encoder(pretrained=False)
+
+
+def test_output_shape(encoder: Encoder):
+    feats = encoder(torch.randn(2, 3, config.crop, config.crop))
+    assert feats.shape == (2, config.encoded_size**2, config.encoder_dim)
+
+
+def test_all_params_frozen(encoder: Encoder):
+    total = sum(p.numel() for p in encoder.parameters())
+    frozen = sum(p.numel() for p in encoder.parameters() if not p.requires_grad)
+    assert total > 0
+    assert frozen == total
+
+
+def test_no_nan_on_random_input(encoder: Encoder):
+    feats = encoder(torch.randn(2, 3, config.crop, config.crop))
+    assert not torch.isnan(feats).any()
+
+
+def test_eval_mode_determinism(encoder: Encoder):
+    encoder.eval()
+    x = torch.randn(1, 3, config.crop, config.crop)
+    with torch.no_grad():
+        assert torch.equal(encoder(x), encoder(x))
+
+
+def test_eval_uses_running_stats_not_batch(encoder: Encoder):
+    encoder.eval()
+    x = torch.randn(1, 3, config.crop, config.crop)
+    others = torch.randn(3, 3, config.crop, config.crop)
+    with torch.no_grad():
+        alone = encoder(x)
+        batched = encoder(torch.cat([x, others]))[:1]
+    assert torch.allclose(alone, batched, atol=1e-5)
+
+
+def test_train_keeps_backbone_in_eval(encoder: Encoder):
+    encoder.train()
+    assert not encoder.backbone.training
+    encoder.eval()
+
+
+def test_fine_tune_default_keeps_all_frozen():
+    enc = Encoder(pretrained=False)
+    enc.fine_tune()
+    assert all(not p.requires_grad for p in enc.parameters())
+
+
+def test_fine_tune_unfreezes_exactly_that_block():
+    enc = Encoder(pretrained=False)
+    enc.fine_tune((7,))  # index 7 = layer4
+    trainable = sum(p.numel() for p in enc.parameters() if p.requires_grad)
+    layer4 = sum(p.numel() for p in enc.backbone[7].parameters())
+    assert trainable == layer4 > 0
+
+
+@pytest.mark.skipif(
+    not (config.subsample_json.is_file() and config.vocab_path.is_file()),
+    reason="dev subsample or vocab not built",
+)
+def test_real_image_features_non_degenerate():
+    from capit.data.dataset import CaptionDataset, build_transform
+    from capit.data.vocab import Vocab
+
+    vocab = Vocab.load(config.vocab_path)
+    ds = CaptionDataset(config.subsample_root, "test", vocab, build_transform())
+    _, image = next(ds.iter_images())
+    encoder = Encoder(pretrained=True)
+    with torch.no_grad():
+        feats = encoder(image.unsqueeze(0))
+    assert not torch.isnan(feats).any()
+    assert feats.std().item() > 0

pipeline/tests/test_evaluate.py

@@ -0,0 +1,37 @@
+"""Stage 4.2 — official BLEU/CIDEr harness. Hermetic scorer sanity + opt-in eval-on-test gate."""
+
+import os
+
+import pytest
+
+from capit.config import config
+from capit.evaluate import evaluate, score
+
+
+def test_score_perfect_match():
+    refs = {
+        0: [["a", "brown", "dog", "runs", "fast"], ["the", "dog", "runs", "very", "fast"]],
+        1: [["a", "small", "cat", "sleeps", "quietly"]],
+    }
+    cands = {0: ["a", "brown", "dog", "runs", "fast"], 1: ["a", "small", "cat", "sleeps", "quietly"]}
+    s = score(refs, cands)
+    assert s["BLEU-4"] == pytest.approx(100.0, abs=1e-6)
+    assert s["BLEU-1"] == pytest.approx(100.0, abs=1e-6)
+
+
+def test_score_disjoint():
+    refs = {0: [["a", "brown", "dog", "runs", "fast"]]}
+    cands = {0: ["xyz", "qpr", "lmn", "uvw", "abc"]}
+    s = score(refs, cands)
+    assert s["BLEU-1"] == pytest.approx(0.0, abs=1e-6)
+    assert s["BLEU-4"] == pytest.approx(0.0, abs=1e-6)
+
+
+@pytest.mark.skipif(
+    os.environ.get("RUN_EVAL") != "1" or not (config.ckpt_dir / "best.pt").exists(),
+    reason="set RUN_EVAL=1 and provide data/checkpoints/best.pt to run the test-set gate",
+)
+def test_eval_on_test_meets_gate():
+    # pycocoevalcap corpus BLEU (unsmoothed) != train.corpus_bleu4 (NLTK smoothed); not directly comparable to best_bleu4.
+    results = evaluate(config.flickr8k_dir, config.ckpt_dir / "best.pt", beams=(3,))
+    assert results[3]["BLEU-4"] >= 16.0

pipeline/tests/test_losses.py

@@ -0,0 +1,39 @@
+"""Stage 2.3 — caption loss. Hand-built logits pin the shift, padding, and regularizer."""
+
+import torch
+
+from capit.config import config
+from capit.losses import caption_loss, word_cross_entropy
+
+V = 12
+
+
+def _peaked(rows: list[list[int]]) -> torch.Tensor:
+    logits = torch.full((1, len(rows[0]), V), -10.0)
+    for t, tok in enumerate(rows[0]):
+        logits[0, t, tok] = 10.0
+    return logits
+
+
+def test_cross_entropy_targets_are_shifted_by_one():
+    captions = torch.tensor([[1, 5, 7, 2]])  # T=3, shifted targets = [5, 7, 2]
+    assert word_cross_entropy(_peaked([[5, 7, 2]]), captions).item() < 0.01
+    # predicting the inputs (captions[:, :-1]) instead of the shifted targets must score badly
+    assert word_cross_entropy(_peaked([[1, 5, 7]]), captions).item() > 1.0
+
+
+def test_pad_targets_are_ignored():
+    captions = torch.tensor([[1, 5, 2, 0]])  # T=3, targets = [5, 2, 0]; position 2 is PAD
+    logits = _peaked([[5, 2, 0]])
+    base = word_cross_entropy(logits, captions)
+    corrupted = logits.clone()
+    corrupted[0, 2] = torch.randn(V) * 50  # corrupt the PAD-target position
+    assert torch.isclose(base, word_cross_entropy(corrupted, captions))
+
+
+def test_regularizer_is_alpha_c_on_zero_attention():
+    captions = torch.tensor([[1, 5, 2, 0], [1, 6, 2, 0]])
+    logits = torch.zeros(2, 3, V)
+    alphas = torch.zeros(2, 3, 196)  # time-sum 0 per location → (1-0)^2 = 1
+    reg = caption_loss(logits, alphas, captions) - word_cross_entropy(logits, captions)
+    assert torch.isclose(reg, torch.tensor(config.alpha_c))

pipeline/tests/test_overfit.py

@@ -0,0 +1,25 @@
+"""Stage 2.3 exit gate — overfit one batch (killer gate #1).
+
+Slow (~3 min): opt-in via RUN_OVERFIT=1, and needs the subsample + vocab. Otherwise the
+gate lives in scripts/overfit_one_batch.py (run + eyeball the decoded captions).
+"""
+
+import os
+
+import pytest
+
+from capit.config import config
+from capit.overfit import run_overfit
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_OVERFIT") != "1"
+    or not (config.subsample_json.is_file() and config.vocab_path.is_file()),
+    reason="slow gate — set RUN_OVERFIT=1 (needs subsample + vocab)",
+)
+
+
+def test_overfit_one_batch():
+    result = run_overfit(steps=200)
+    assert result["final_ce"] < 0.5, f"ce {result['final_ce']:.3f}: model failed to memorize"
+    matches = sum(d == t for d, t in zip(result["decoded"], result["targets"]))
+    assert matches == len(result["targets"]), f"{matches}/{len(result['targets'])} captions reproduced"

pipeline/tests/test_serving.py

@@ -0,0 +1,53 @@
+"""Stage 4.4 — serving artifact round-trip. Hermetic: random encoder/decoder, no checkpoint/network."""
+
+import json
+
+import pytest
+import torch
+from PIL import Image
+
+from capit.data.vocab import SPECIALS, Vocab
+from capit.models.decoder import Decoder
+from capit.models.encoder import Encoder
+from capit.serving import build_artifact, caption, load_artifact, make_transform
+
+
+@pytest.fixture
+def artifact(tmp_path):
+    torch.manual_seed(0)
+    vocab = Vocab(SPECIALS + ["a", "dog", "runs", "fast"])
+    vocab_path = tmp_path / "vocab.json"
+    vocab.save(vocab_path)
+    encoder = Encoder(pretrained=False).eval()
+    decoder = Decoder(vocab_size=len(vocab))
+    blob = build_artifact(encoder, decoder, vocab)
+    path = tmp_path / "capit-sat.pt"
+    torch.save(blob, path)
+    return path, vocab_path, encoder, decoder, vocab
+
+
+def test_round_trip_reproduces_caption(artifact):
+    path, vocab_path, encoder, decoder, vocab = artifact
+    enc2, dec2, vocab2, preprocess = load_artifact(path, vocab_path)
+    assert vocab2.id2word == vocab.id2word
+    assert {"resize", "crop", "mean", "std"} <= preprocess.keys()
+    torch.manual_seed(1)
+    image = torch.rand(3, preprocess["crop"], preprocess["crop"])
+    words, alphas, _ = caption(enc2, dec2, vocab2, image)
+    ref_words, ref_alphas, _ = caption(encoder, decoder, vocab, image)
+    assert words == ref_words
+    assert torch.equal(alphas, ref_alphas)
+
+
+def test_vocab_guard_raises_on_mismatch(artifact, tmp_path):
+    path, vocab_path, *_ = artifact
+    wrong = tmp_path / "wrong.json"
+    wrong.write_text(json.dumps(SPECIALS + ["different", "words", "entirely"]))
+    with pytest.raises(ValueError, match="vocab mismatch"):
+        load_artifact(path, wrong)
+
+
+def test_make_transform_shapes_image():
+    t = make_transform({"resize": 256, "crop": 224, "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]})
+    out = t(Image.new("RGB", (400, 300)))
+    assert out.shape == (3, 224, 224)

pipeline/tests/test_subsample.py

@@ -0,0 +1,73 @@
+"""Dev subsample.
+
+Pure selection tests always run; structural checks on the on-disk artifact skip when
+it hasn't been generated.
+"""
+
+import json
+from collections import Counter
+
+import pytest
+
+from capit.config import config
+from capit.data.records import select_records
+
+
+def _records(split: str, n: int) -> list[dict]:
+    return [{"filename": f"{split}_{i:03d}.jpg", "split": split, "imgid": i} for i in range(n)]
+
+
+SYNTH = _records("train", 20) + _records("val", 8) + _records("test", 8)
+COUNTS = {"train": 5, "val": 3, "test": 2}
+
+
+def test_select_counts_and_splits():
+    sel = select_records(SYNTH, COUNTS, seed=13)
+    assert len(sel) == 10
+    assert Counter(r["split"] for r in sel) == {"train": 5, "val": 3, "test": 2}
+
+
+def test_select_is_deterministic():
+    assert select_records(SYNTH, COUNTS, 13) == select_records(SYNTH, COUNTS, 13)
+    assert select_records(SYNTH, COUNTS, 13) != select_records(SYNTH, COUNTS, 14)
+
+
+def test_select_records_are_verbatim():
+    sel = select_records(SYNTH, COUNTS, 13)
+    assert all(r in SYNTH for r in sel)
+
+
+def test_select_raises_when_pool_too_small():
+    with pytest.raises(ValueError):
+        select_records(SYNTH, {"train": 999}, 13)
+
+
+def _artifact_present() -> bool:
+    return config.subsample_json.is_file() and config.subsample_images_dir.is_dir()
+
+
+pytestmark_structural = pytest.mark.skipif(
+    not _artifact_present(),
+    reason="dev subsample not generated (run: python pipeline/scripts/make_subsample.py)",
+)
+
+
+@pytestmark_structural
+def test_artifact_structure():
+    sub = json.loads(config.subsample_json.read_text())
+    images = sub["images"]
+    assert len(images) == sum(config.subsample_counts.values())
+    assert Counter(r["split"] for r in images) == config.subsample_counts
+
+    jpgs = {p.name for p in config.subsample_images_dir.glob("*.jpg")}
+    assert len(jpgs) == len(images)
+    assert all(r["filename"] in jpgs for r in images)
+
+    assert sub["dataset"] == json.loads(config.karpathy_json.read_text())["dataset"]
+
+
+@pytestmark_structural
+def test_artifact_records_are_verbatim():
+    sub = json.loads(config.subsample_json.read_text())
+    full = {r["filename"]: r for r in json.loads(config.karpathy_json.read_text())["images"]}
+    assert all(r == full[r["filename"]] for r in sub["images"])