update vertex jobs and remove hardcoded secrets

evaluation/evaluate.py

@@ -125,8 +125,14 @@ def run_inference(
 
     hypotheses, references, questions = [], [], []
 
+    # mininterval=1.0 keeps the bar visibly updating in Colab subprocess
+    # output (default 10s feels like the kernel froze on T4). dynamic_ncols
+    # lets it adapt to whatever width the notebook renders.
     for i in tqdm(range(0, len(task_samples), batch_size),
-                  desc=f"Evaluating {task}"):
+                  desc=f"Evaluating {task}",
+                  mininterval=1.0,
+                  dynamic_ncols=True,
+                  file=sys.stdout):
         batch_samples = task_samples[i:i + batch_size]
 
         images, prompts = [], []
@@ -229,14 +235,32 @@ def main():
             explicit_run_id = run_id,
         )
 
-    # Build and load model
-    logger.info(f"Loading model from checkpoint: {args.checkpoint}")
+    # Build and load model — these phases are SILENT for several minutes on
+    # T4 (4-bit Vicuna shard download + quantization). Print breadcrumbs with
+    # flush=True so notebook users can see where the kernel is sitting.
+    import time as _time
+    print(f"[evaluate] === phase: build model ===", flush=True)
+    print(f"[evaluate] (T4 cold cache: ~5-8 min — encoder load + Vicuna-7B "
+          f"shard download + 4-bit NF4 quantization)", flush=True)
+    _t0 = _time.time()
     model = CXRVisionLanguageModel(model_cfg)
+    print(f"[evaluate] model built in {_time.time()-_t0:.1f}s", flush=True)
+
+    print(f"[evaluate] === phase: load checkpoint ===", flush=True)
+    print(f"[evaluate] checkpoint: {args.checkpoint}", flush=True)
+    _t0 = _time.time()
     load_checkpoint(model, args.checkpoint)
+    print(f"[evaluate] checkpoint loaded in {_time.time()-_t0:.1f}s", flush=True)
+
+    print(f"[evaluate] === phase: move model to {args.device} + eval() ===", flush=True)
+    _t0 = _time.time()
     model = model.to(args.device)
     model.eval()
+    print(f"[evaluate] model ready in {_time.time()-_t0:.1f}s", flush=True)
 
     # Load test dataset (for the chosen dataset)
+    print(f"[evaluate] === phase: build dataset (split={args.split}) ===", flush=True)
+    _t0 = _time.time()
     dataset = CXRInstructDataset(
         data_path    = spec.instruct_json,
         image_root   = spec.image_root,
@@ -249,6 +273,8 @@ def main():
         max_images   = spec.max_images,
         feature_cache_dir = getattr(train_cfg.data, "feature_cache_dir", None) or None,
     )
+    print(f"[evaluate] dataset built in {_time.time()-_t0:.1f}s "
+          f"({len(dataset.samples):,} total samples in split={args.split!r})", flush=True)
 
     # Build task list, intersected with what's available for this dataset.
     if args.task == "all":
@@ -263,10 +289,19 @@ def main():
         else:
             tasks_to_eval = [args.task]
 
+    # Print per-task sample counts up-front so the user can estimate runtime
+    # before tqdm starts.
+    print(f"[evaluate] tasks to evaluate: {tasks_to_eval}", flush=True)
+    for _t in tasks_to_eval:
+        _n = sum(1 for s in dataset.samples if s["task"] == _t)
+        print(f"[evaluate]   {_t:<11s}: {_n:,} samples "
+              f"(~{_n // max(args.batch_size, 1):,} batches)", flush=True)
+
     all_results = {}
 
     for task in tasks_to_eval:
         logger.info(f"\nEvaluating task: {task.upper()}")
+        print(f"[evaluate] === phase: inference (task={task}) ===", flush=True)
 
         predictions = run_inference(
             model          = model,

scripts/_build_eval_notebook.py

@@ -0,0 +1,875 @@
+"""
+Build scripts/cxrvlm_colab_eval.ipynb from cell sources defined here.
+
+Run:
+    python scripts/_build_eval_notebook.py
+"""
+
+import json
+from pathlib import Path
+
+NB_PATH = Path(__file__).parent / "cxrvlm_colab_eval.ipynb"
+
+
+def md(cell_id, src):
+    return {
+        "cell_type": "markdown",
+        "id": cell_id,
+        "metadata": {},
+        "source": src.splitlines(keepends=True),
+    }
+
+
+def code(cell_id, src):
+    return {
+        "cell_type": "code",
+        "id": cell_id,
+        "metadata": {},
+        "execution_count": None,
+        "outputs": [],
+        "source": src.splitlines(keepends=True),
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Cell sources
+# ─────────────────────────────────────────────────────────────────────
+
+CELLS = []
+
+
+CELLS.append(md("eval-0", """\
+# CXR-VLM — Evaluation Notebook (Colab T4)
+
+Standalone evaluation + inference for a trained CXR-VLM run.
+
+What it does:
+1. Pulls project code from `<HF_USER>/cxr-vlm-code`.
+2. Pulls the chosen dataset from `<HF_USER>/cxr-vlm-data` (same layout the trainer used).
+3. Pulls a trained run checkpoint from `<HF_USER>/cxr-vlm-runs/<RUN_ID>/stage2/{best|last}/`.
+4. Pulls the run's config snapshot (`<RUN_ID>/configs/{model,train}_config.yaml`) so the model is rebuilt **with the exact same architecture / report_mode / image_mode** as training.
+5. **Auto-detects the GPU** (T4 / L4 / 3090 / A10 / A100 / H100) and patches the configs accordingly — fp16+SDPA on Turing, bf16+FA2 on Ampere/Ada, batch size scaled to VRAM. Also patches local dataset paths.
+6. Runs `python -m evaluation.evaluate` on the **test split** for all available tasks.
+7. Saves predictions + metrics under `RESULTS_DIR/<RUN_ID>/` (and optionally uploads them to `<HF_USER>/cxr-vlm-runs/<RUN_ID>/results/`).
+
+Set the variables in **section 0** and run all cells top-to-bottom.
+
+**Want interactive inference on individual images?** Use the companion notebook **`cxrvlm_colab_inference.ipynb`** — same model pull, but with image preview + free-form prompts.
+"""))
+
+
+CELLS.append(md("eval-select-md", """\
+## 0. Select run + dataset + options
+
+Change the variables in the cell below. `RUN_ID` decides which trained model is pulled; everything else (dataset name, report/image mode) defaults to whatever was used at training time (read from the run's config snapshot on HF), but you can override.
+"""))
+
+
+CELLS.append(code("eval-select", """\
+# ── Platform ─────────────────────────────────────────────────────
+PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'
+
+# ── Source repos on HuggingFace ──────────────────────────────────
+HF_USER       = 'hieu3636'              # owner of cxr-vlm-{code,data,runs}
+HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'
+HF_DATA_REPO  = f'{HF_USER}/cxr-vlm-data'
+HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'
+
+# ── Which trained run to evaluate ────────────────────────────────
+# This MUST be an existing folder on HF_RUNS_REPO.
+# Example: 'MIMIC-CXR_resized_run_1' | 'IU-Xray_run_2' | 'MIMIC-CXR_run_3'
+RUN_ID        = 'MIMIC-CXR_resized_run_1'
+
+# Which stage-2 checkpoint to load.
+#   'best' → {RUN_ID}/stage2/best/  (final / best eval_loss)
+#   'last' → {RUN_ID}/stage2/last/  (most recent intermediate save)
+CKPT_PICK     = 'best'
+
+# ── Dataset (auto-derived from RUN_ID prefix; override if needed) ──
+# Supported: 'MIMIC-CXR' | 'MIMIC-CXR_resized' | 'IU-Xray'
+DATASET_NAME  = None        # None → auto-detect from RUN_ID prefix
+REPORT_MODE   = None        # None → read from run's saved train_config.yaml
+IMAGE_MODE    = None        # None → read from run's saved train_config.yaml
+
+# ── What to evaluate ─────────────────────────────────────────────
+TASK          = 'all'       # 'all' | 'findings' | 'impression' | 'report' | 'vqa'
+SPLIT         = 'test'      # always 'test' for this notebook; here for visibility
+BATCH_SIZE    = None        # None → auto from GPU profile (T4:1, L4:4, A100:8, H100:16)
+MAX_NEW_TOKENS = 300
+
+# ── LLM-as-judge for VQA (optional; needs OPENAI_API_KEY) ────────
+LLM_JUDGE             = False
+LLM_JUDGE_MODEL       = 'gpt-4o-mini'
+LLM_JUDGE_BASE_URL    = None          # e.g. Gemini OpenAI-compat endpoint
+LLM_JUDGE_MAX_SAMPLES = None          # cap cost; None → all VQA samples
+
+# ── Output ───────────────────────────────────────────────────────
+# Local folder where predictions_*.json + metrics_summary.json land.
+# Files end up at  {LOCAL_RESULTS_DIR}/{RUN_ID}/...
+LOCAL_RESULTS_DIR = 'results'
+
+# Push the results folder back to HF_RUNS_REPO under
+#   {RUN_ID}/results/predictions_*.json + metrics_summary.json
+UPLOAD_RESULTS_TO_HF = True
+
+# ── Auto-derive DATASET_NAME from RUN_ID prefix if not set ───────
+if DATASET_NAME is None:
+    for cand in ('MIMIC-CXR_resized', 'MIMIC-CXR', 'IU-Xray'):
+        if RUN_ID.startswith(cand + '_run_'):
+            DATASET_NAME = cand
+            break
+    assert DATASET_NAME is not None, \\
+        f"Cannot auto-derive DATASET_NAME from RUN_ID={RUN_ID!r}. " \\
+        "Expected prefix one of: MIMIC-CXR_resized / MIMIC-CXR / IU-Xray. Set DATASET_NAME explicitly."
+
+assert PLATFORM     in ('kaggle', 'colab', 'lightning', 'gcp', 'local')
+assert DATASET_NAME in ('MIMIC-CXR', 'MIMIC-CXR_resized', 'IU-Xray')
+assert CKPT_PICK    in ('best', 'last')
+assert TASK         in ('all', 'findings', 'impression', 'report', 'vqa')
+
+print(f'PLATFORM      = {PLATFORM}')
+print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')
+print(f'DATASET_NAME  = {DATASET_NAME}')
+print(f'TASK          = {TASK}   SPLIT = {SPLIT}')
+print(f'LOCAL_RESULTS_DIR = {LOCAL_RESULTS_DIR}   (upload→HF: {UPLOAD_RESULTS_TO_HF})')
+"""))
+
+
+CELLS.append(md("eval-env-md", """\
+## 1. Environment + pip install (matches the training notebook)
+"""))
+
+
+CELLS.append(code("eval-env", """\
+import os
+os.environ['CUDA_VISIBLE_DEVICES']            = '0'      # single-GPU
+os.environ['TOKENIZERS_PARALLELISM']          = 'false'
+os.environ['BITSANDBYTES_NOWELCOME']          = '1'
+os.environ['HF_HUB_DISABLE_PROGRESS_BARS']    = '1'
+os.environ['TRANSFORMERS_VERBOSITY']          = 'warning'
+os.environ['PYTHONUNBUFFERED']                = '1'
+
+import sys, shutil, subprocess
+from pathlib import Path
+"""))
+
+
+CELLS.append(code("eval-pip", """\
+import os as _os
+
+# Same dependency set as the training notebook — keeps load_checkpoint /
+# evaluate.py running against the exact stack the model was trained with.
+!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate
+
+# Let pip pick latest bnb that matches Colab's CUDA + triton.
+!pip install -q -U bitsandbytes
+
+# Pin transformers / peft to the same window the trainer uses.
+!pip install -q \\
+    'transformers>=4.46,<4.50' \\
+    'peft>=0.13,<0.15' \\
+    'accelerate>=1.0' \\
+    'huggingface_hub>=0.27,<1.0' \\
+    omegaconf sentencepiece 'protobuf>=3.20' \\
+    nltk rouge-score bert-score sacrebleu
+
+# Best-effort flash-attn install for Ampere+/Ada GPUs (L4, 3090, A10, A100, H100).
+# Silent fail is OK — cxr_vlm.py auto-falls-back to SDPA if FA2 isn't importable.
+# Skipped entirely on Turing (T4/V100) since FA2 requires sm_80+.
+import torch as _t
+if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):
+    print('[pip] Ampere+/Ada detected -> trying flash-attn (5-10 min if building from source)')
+    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5
+else:
+    print('[pip] Pre-Ampere GPU (or no CUDA) -> skipping flash-attn install')
+"""))
+
+
+CELLS.append(code("eval-versions", """\
+import torch, transformers, bitsandbytes, peft, accelerate, huggingface_hub, httpx
+
+print('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())
+print('transformers   :', transformers.__version__)
+print('bitsandbytes   :', bitsandbytes.__version__)
+print('peft           :', peft.__version__)
+print('accelerate     :', accelerate.__version__)
+print('huggingface_hub:', huggingface_hub.__version__)
+print('httpx          :', httpx.__version__)
+
+# httpx 0.28+ compat shim — same patch as the training notebook.
+# transformers <=4.50 calls httpx.Client.head(..., allow_redirects=True)
+# which httpx 0.28 removed; translate the kwarg at the call site.
+def _patch_httpx():
+    if tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28):
+        return
+    if getattr(httpx.Client, '_cxr_vlm_compat_patched', False):
+        return
+    def _make(orig):
+        def patched(self, *args, **kwargs):
+            if 'allow_redirects' in kwargs:
+                kwargs['follow_redirects'] = kwargs.pop('allow_redirects')
+            kwargs.pop('proxies', None)
+            return orig(self, *args, **kwargs)
+        return patched
+    for cls in (httpx.Client, httpx.AsyncClient):
+        for m in ('request', 'get', 'head', 'post', 'put',
+                  'patch', 'delete', 'options'):
+            if hasattr(cls, m):
+                setattr(cls, m, _make(getattr(cls, m)))
+    httpx.Client._cxr_vlm_compat_patched = True
+    print(f'httpx {httpx.__version__}: monkey-patched allow_redirects -> follow_redirects')
+
+_patch_httpx()
+
+assert torch.cuda.is_available(), 'CUDA not available — refusing to evaluate on CPU.'
+_p = torch.cuda.get_device_properties(0)
+print(f'\\nGPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')
+print(f'Compute cap : sm_{_p.major}{_p.minor}  (BF16 ok: {torch.cuda.is_bf16_supported()})')
+"""))
+
+
+CELLS.append(md("eval-paths-md", """\
+## 2. Paths + pull code + pull dataset
+
+Identical to the training notebook for sections that overlap (HF code + per-dataset payload).
+"""))
+
+
+CELLS.append(code("eval-paths", """\
+# ── 1) WORK dir + HF_TOKEN bootstrap (platform-specific) ───────────
+if PLATFORM == 'kaggle':
+    from kaggle_secrets import UserSecretsClient
+    os.environ['HF_TOKEN'] = UserSecretsClient().get_secret('HF_TOKEN')
+    WORK = Path('/kaggle/working')
+elif PLATFORM == 'colab':
+    from google.colab import userdata
+    os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')
+    WORK = Path('/content')
+elif PLATFORM == 'lightning':
+    WORK = Path('/teamspace/studios/this_studio')
+elif PLATFORM == 'gcp':
+    WORK = Path('/workspace')
+else:  # 'local'
+    WORK = Path.home() / 'cxr-vlm-work'
+WORK.mkdir(parents=True, exist_ok=True)
+
+assert os.environ.get('HF_TOKEN'), \\
+    'HF_TOKEN missing — set it via the platform secrets UI before re-running.'
+
+from huggingface_hub import snapshot_download, hf_hub_download, HfApi
+
+# ── 2) Code: flat folder, snapshot_download ──
+print(f'Pulling code from {HF_CODE_REPO} …')
+CODE_SRC = Path(snapshot_download(
+    repo_id   = HF_CODE_REPO,
+    repo_type = 'model',
+    token     = os.environ['HF_TOKEN'],
+    local_dir = str(WORK / 'cxr-vlm-code'),
+))
+
+# ── 3) Data: layout depends on DATASET_NAME (same logic as train notebook) ──
+DATA_SRC = WORK / 'data'
+DATA_SRC.mkdir(parents=True, exist_ok=True)
+
+if DATASET_NAME == 'MIMIC-CXR_resized':
+    import tarfile
+    mr_dir = DATA_SRC / 'MIMIC-CXR_resized'
+    mr_dir.mkdir(parents=True, exist_ok=True)
+    files_dir = mr_dir / 'files'
+
+    manifests_present = all(
+        (mr_dir / f).is_file() for f in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv')
+    )
+    if manifests_present and files_dir.is_dir() and any(files_dir.glob('p*')):
+        print(f'{mr_dir} already populated — skipping download.')
+    else:
+        api = HfApi(token=os.environ['HF_TOKEN'])
+        all_files = api.list_repo_files(repo_id=HF_DATA_REPO, repo_type='dataset')
+        mr_files  = [f for f in all_files if f.startswith('MIMIC-CXR_resized/')]
+        tar_files = sorted(f for f in mr_files if f.endswith('.tar'))
+        meta_files = [f for f in mr_files if not f.endswith('.tar')]
+        print(f'MIMIC-CXR_resized on HF: {len(tar_files)} tar shards + {len(meta_files)} metadata files')
+
+        # Manifests / vqa / SHARDS.txt
+        snapshot_download(
+            repo_id        = HF_DATA_REPO,
+            repo_type      = 'dataset',
+            allow_patterns = ['MIMIC-CXR_resized/*.csv',
+                              'MIMIC-CXR_resized/*.json',
+                              'MIMIC-CXR_resized/*.txt',
+                              'MIMIC-CXR_resized/vqa/**'],
+            token          = os.environ['HF_TOKEN'],
+            local_dir      = str(DATA_SRC),
+        )
+
+        # Tar shards — sequential extract + delete to keep peak disk low.
+        for i, tf in enumerate(tar_files, 1):
+            print(f'  [{i}/{len(tar_files)}] {tf}')
+            tar_path = Path(hf_hub_download(
+                repo_id=HF_DATA_REPO, repo_type='dataset',
+                filename=tf, token=os.environ['HF_TOKEN'],
+                local_dir=str(DATA_SRC),
+            ))
+            with tarfile.open(tar_path) as t:
+                t.extractall(mr_dir)
+            tar_path.unlink(missing_ok=True)
+        print(f'  done. {mr_dir} ready.')
+
+else:
+    # MIMIC-CXR / IU-Xray — single zip per dataset.
+    import zipfile
+    zip_name = f'{DATASET_NAME}.zip'
+    marker   = DATA_SRC / DATASET_NAME
+
+    if not marker.exists():
+        print(f'Pulling {zip_name} from HF …')
+        zpath = hf_hub_download(
+            repo_id   = HF_DATA_REPO,
+            filename  = zip_name,
+            repo_type = 'dataset',
+            token     = os.environ['HF_TOKEN'],
+            local_dir = str(DATA_SRC),
+        )
+        print(f'  unzipping -> {DATA_SRC}')
+        with zipfile.ZipFile(zpath) as zf:
+            zf.extractall(DATA_SRC)
+        try:
+            os.remove(zpath)
+        except OSError:
+            pass
+    else:
+        print(f'{marker} already present — skipping download.')
+
+print(f'Contents of {DATA_SRC}: {sorted(os.listdir(DATA_SRC))}')
+
+# ── 4) Copy code into writable PROJECT dir + chdir ─────────────────
+PROJECT = WORK / 'cxr_vlm'
+if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():
+    shutil.copytree(CODE_SRC, PROJECT)
+
+os.chdir(PROJECT)
+sys.path.insert(0, str(PROJECT))
+print('CODE_SRC :', CODE_SRC)
+print('DATA_SRC :', DATA_SRC)
+print('PROJECT  :', PROJECT)
+print('WORK     :', WORK)
+"""))
+
+
+CELLS.append(md("eval-find-md", """\
+## 3. Locate the dataset payload on disk
+"""))
+
+
+CELLS.append(code("eval-find", """\
+# Same dataset-payload finders as the training notebook. Filling in the
+# variables the config cell below will consume.
+
+def find_split_parent(root: Path) -> Path:
+    for cand in [root, root / 'MIMIC-CXR', root / 'data' / 'MIMIC-CXR']:
+        if (cand / 'train').exists() and (cand / 'valid').exists() and (cand / 'test').exists():
+            return cand
+    for p in root.rglob('train'):
+        if p.is_dir() and (p.parent / 'valid').exists() and (p.parent / 'test').exists():
+            return p.parent
+    raise FileNotFoundError('Could not find train/ valid/ test/ under ' + str(root))
+
+
+def find_mimic_resized_root(root: Path) -> Path:
+    for cand in [root / 'MIMIC-CXR_resized', root, *root.rglob('MIMIC-CXR_resized')]:
+        if (cand / 'manifest_train.csv').is_file():
+            return cand
+    raise FileNotFoundError(
+        f'Could not find MIMIC-CXR_resized payload under {root}. '
+        'Expected manifest_train.csv (alongside manifest_val.csv / manifest_test.csv).'
+    )
+
+
+def find_iu_dirs(root: Path):
+    for cand in [root / 'IU-Xray', *root.rglob('IU-Xray')]:
+        if not cand.is_dir():
+            continue
+        imgs = cand / 'images'
+        lbls = cand / 'labels'
+        if imgs.is_dir() and lbls.is_dir() and any(lbls.glob('*.xml')):
+            return imgs, lbls
+        legacy = lbls / 'ecgen-radiology'
+        if imgs.is_dir() and legacy.is_dir() and any(legacy.glob('*.xml')):
+            return imgs, legacy
+    img_dir = lbl_dir = None
+    for cand in [root / 'images', *root.rglob('images')]:
+        if cand.is_dir() and any(cand.glob('CXR*.png')):
+            img_dir = cand; break
+    for cand in [root / 'labels', *root.rglob('labels')]:
+        if cand.is_dir() and any(cand.glob('*.xml')):
+            lbl_dir = cand; break
+    if lbl_dir is None:
+        for cand in root.rglob('ecgen-radiology'):
+            if cand.is_dir() and any(cand.glob('*.xml')):
+                lbl_dir = cand; break
+    return img_dir, lbl_dir
+
+
+CXR_ROOT      = None
+VQA_ROOT      = None
+MR_ROOT       = None
+IU_IMAGES_DIR = None
+IU_LABELS_DIR = None
+
+if DATASET_NAME == 'MIMIC-CXR':
+    CXR_ROOT = find_split_parent(DATA_SRC)
+    print('MIMIC-CXR root:', CXR_ROOT)
+    for s in ('train', 'valid', 'test'):
+        d = CXR_ROOT / s
+        assert d.exists(), f'Missing split dir: {d}'
+        print(f'  {s:<6s} -> {d}')
+
+    for p in DATA_SRC.rglob('MIMIC-Ext-MIMIC-CXR-VQA'):
+        cand = p / 'dataset'
+        if cand.exists() and (cand / 'train.json').exists():
+            VQA_ROOT = cand; break
+    if VQA_ROOT is None:
+        print('VQA root: NOT FOUND -> VQA task will be skipped')
+    else:
+        print('VQA root:', VQA_ROOT)
+
+elif DATASET_NAME == 'MIMIC-CXR_resized':
+    MR_ROOT = find_mimic_resized_root(DATA_SRC)
+    print('MIMIC-CXR_resized root:', MR_ROOT)
+    for cf in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv'):
+        f = MR_ROOT / cf
+        print(f'  {cf}: {"OK" if f.is_file() else "MISSING"}')
+    for sub in ('files', 'vqa'):
+        d = MR_ROOT / sub
+        print(f'  {sub:<5s}: {"OK" if d.is_dir() else "MISSING"}  ({d})')
+
+else:  # IU-Xray
+    IU_IMAGES_DIR, IU_LABELS_DIR = find_iu_dirs(DATA_SRC)
+    assert IU_IMAGES_DIR is not None, f'IU images/ not found under {DATA_SRC}'
+    assert IU_LABELS_DIR is not None, f'IU labels/ (with *.xml) not found under {DATA_SRC}'
+    print('IU images dir:', IU_IMAGES_DIR, '->', len(list(IU_IMAGES_DIR.glob('*.png'))), 'PNGs')
+    print('IU labels dir:', IU_LABELS_DIR, '->', len(list(IU_LABELS_DIR.glob('*.xml'))), 'XMLs')
+"""))
+
+
+CELLS.append(md("eval-pull-run-md", """\
+## 4. Pull trained checkpoint + run config snapshot from HF Runs repo
+
+Layout on `HF_RUNS_REPO`:
+
+```
+{RUN_ID}/
+  configs/{model,train}_config.yaml   ← snapshot taken at training time
+  stage2/best/   checkpoint_projection.pt + checkpoint_lora/ [+ checkpoint_chexpert_classifier.pt]
+  stage2/last/   same shape, intermediate save
+```
+
+We pull `{RUN_ID}/configs/` and `{RUN_ID}/stage2/{CKPT_PICK}/` into `{WORK}/run_pull/{RUN_ID}/`.
+"""))
+
+
+CELLS.append(code("eval-pull-run", """\
+from huggingface_hub import snapshot_download
+
+RUN_PULL_ROOT = WORK / 'run_pull'
+RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)
+
+print(f'Pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …')
+snapshot_download(
+    repo_id   = HF_RUNS_REPO,
+    repo_type = 'model',
+    token     = os.environ['HF_TOKEN'],
+    allow_patterns = [
+        f'{RUN_ID}/configs/**',
+        f'{RUN_ID}/run_meta.json',
+        f'{RUN_ID}/stage2/{CKPT_PICK}/**',
+    ],
+    local_dir = str(RUN_PULL_ROOT),
+)
+
+RUN_DIR_PULLED = RUN_PULL_ROOT / RUN_ID
+assert RUN_DIR_PULLED.is_dir(), f'Pull failed — {RUN_DIR_PULLED} missing.'
+
+CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK
+PROJ_PT         = CKPT_DIR_PULLED / 'checkpoint_projection.pt'
+LORA_DIR        = CKPT_DIR_PULLED / 'checkpoint_lora'
+CHEXPERT_PT     = CKPT_DIR_PULLED / 'checkpoint_chexpert_classifier.pt'
+
+assert PROJ_PT.is_file(), \\
+    f'Projection weights not found at {PROJ_PT}. ' \\
+    f'Check that {RUN_ID}/stage2/{CKPT_PICK}/ exists on {HF_RUNS_REPO}.'
+assert (LORA_DIR / 'adapter_config.json').is_file(), \\
+    f'LoRA adapter_config.json not found in {LORA_DIR}. Stage-2 checkpoint partial?'
+
+print()
+print(f'  projection : {PROJ_PT}      ({PROJ_PT.stat().st_size/1e6:.1f} MB)')
+print(f'  lora       : {LORA_DIR}/    ({sum(p.stat().st_size for p in LORA_DIR.rglob("*") if p.is_file())/1e6:.1f} MB)')
+print(f'  chexpert   : {CHEXPERT_PT}  (exists: {CHEXPERT_PT.is_file()})')
+
+# Saved configs (may or may not exist on older runs)
+SAVED_CFG_DIR = RUN_DIR_PULLED / 'configs'
+SAVED_TRAIN_CFG = SAVED_CFG_DIR / 'train_config.yaml'
+SAVED_MODEL_CFG = SAVED_CFG_DIR / 'model_config.yaml'
+print()
+print(f'  saved train_cfg : {SAVED_TRAIN_CFG}  (exists: {SAVED_TRAIN_CFG.is_file()})')
+print(f'  saved model_cfg : {SAVED_MODEL_CFG}  (exists: {SAVED_MODEL_CFG.is_file()})')
+"""))
+
+
+CELLS.append(md("eval-gpu-md", """\
+## 5. Auto-detect GPU profile
+
+Mirrors the training notebook: picks precision (bf16/fp16), attention backend (FA2/SDPA), and an eval batch size based on the actual GPU's compute capability + VRAM.
+
+| Bucket | Examples | Precision | Attn | Eval batch |
+|---|---|---|---|---|
+| 70+ GB | A100/H100 80GB | bf16 | FA2 (if installed) | 16 |
+| 35–69 GB | A100 40GB | bf16 | FA2 | 8 |
+| 22–34 GB | 3090 / L4 / A10 24GB | bf16 | FA2 | 4 |
+| 14–21 GB | T4 / V100 16GB | fp16 | SDPA | 1 |
+
+The eval batch sizes are smaller than training's because generation builds a KV cache that scales with `batch × max_new_tokens`. Override by setting `BATCH_SIZE` in section 0 to a concrete number.
+"""))
+
+
+CELLS.append(code("eval-gpu", """\
+import torch
+
+assert torch.cuda.is_available(), 'CUDA not available — refusing to write a CPU profile.'
+_props   = torch.cuda.get_device_properties(0)
+_cap     = (_props.major, _props.minor)
+_vram_gb = _props.total_memory / 1e9
+_bf16_ok = torch.cuda.is_bf16_supported()
+_fa2_ok  = _cap >= (8, 0)        # FA2 needs Ampere+ (sm_80 or newer)
+
+# Detect whether flash-attn package is actually importable. FA2 falls back to
+# SDPA inside cxr_vlm.py if missing, but knowing here lets us print clearly.
+_flash_attn_installed = False
+if _fa2_ok:
+    try:
+        import flash_attn  # noqa: F401
+        _flash_attn_installed = True
+    except Exception:
+        _flash_attn_installed = False
+
+print(f'GPU         : {_props.name}  ({_vram_gb:.1f} GB)')
+print(f'Compute cap : sm_{_cap[0]}{_cap[1]}')
+print(f'BF16 native : {_bf16_ok}')
+print(f'FA2 capable : {_fa2_ok}   flash-attn installed: {_flash_attn_installed}')
+
+# Eval batch size — smaller than training because generation builds a KV cache
+# that scales with batch × max_new_tokens.
+if _vram_gb >= 70:
+    GPU_LABEL, _EVAL_BS, _NW = 'A100/H100 80GB',         16, 16
+elif _vram_gb >= 35:
+    GPU_LABEL, _EVAL_BS, _NW = 'A100 40GB',               8, 12
+elif _vram_gb >= 22:
+    GPU_LABEL, _EVAL_BS, _NW = 'RTX 3090 / L4 / A10 24GB', 4, 8
+elif _vram_gb >= 14:
+    GPU_LABEL, _EVAL_BS, _NW = 'T4 / V100 (15-16GB)',     1, 2
+else:
+    GPU_LABEL, _EVAL_BS, _NW = f'unknown ({_vram_gb:.0f}GB) - conservative', 1, 2
+
+PROFILE = dict(
+    label                       = GPU_LABEL,
+    bf16                        = bool(_bf16_ok),
+    fp16                        = not _bf16_ok,
+    torch_dtype                 = 'bfloat16' if _bf16_ok else 'float16',
+    bnb_4bit_compute_dtype      = 'bfloat16' if _bf16_ok else 'float16',
+    attn_implementation         = 'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa',
+    per_device_eval_batch_size  = _EVAL_BS,
+    dataloader_num_workers      = _NW,
+)
+
+# Allow the section-0 selector to override (e.g. BATCH_SIZE=8 to push harder).
+if BATCH_SIZE is not None:
+    PROFILE['per_device_eval_batch_size'] = int(BATCH_SIZE)
+    print(f'BATCH_SIZE override: section 0 forced batch={BATCH_SIZE}')
+
+# Final BATCH_SIZE the rest of the notebook (eval-run cell) will use.
+BATCH_SIZE = PROFILE['per_device_eval_batch_size']
+
+print(f'\\n-> Profile  : {PROFILE["label"]}')
+for k, v in PROFILE.items():
+    if k == 'label': continue
+    print(f'    {k:<32s} = {v}')
+"""))
+
+
+CELLS.append(md("eval-cfg-md", """\
+## 6. Build configs
+
+Strategy: start from the run's saved config snapshot if present (so `report_mode`, `image_mode`, `lora.r`, `num_image_tokens`, etc. match training). If absent, fall back to the repo defaults. Then patch:
+- dataset paths to the local download
+- compute (precision + attn backend + batch) from the auto-detected `PROFILE`
+- HF Hub tracker → uses `HF_RUNS_REPO` so the results upload lands under `{RUN_ID}/results/`
+- pin `run_id` to `RUN_ID` so `evaluate.py` writes under `{LOCAL_RESULTS_DIR}/{RUN_ID}/`
+"""))
+
+
+CELLS.append(code("eval-cfg", """\
+from omegaconf import OmegaConf
+import torch
+
+# ── 1) Load base configs: prefer the run's saved snapshot ─────────
+if SAVED_TRAIN_CFG.is_file():
+    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
+    print(f'train_cfg <- {SAVED_TRAIN_CFG}')
+else:
+    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
+    print(f'train_cfg <- repo default (no snapshot on HF)')
+
+if SAVED_MODEL_CFG.is_file():
+    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
+    print(f'model_cfg <- {SAVED_MODEL_CFG}')
+else:
+    model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')
+    print(f'model_cfg <- repo default (no snapshot on HF)')
+
+# ── 2) Allow notebook overrides for report/image mode ─────────────
+if REPORT_MODE is not None:
+    train_cfg.data.report_mode = REPORT_MODE
+if IMAGE_MODE is not None:
+    train_cfg.data.image_mode = IMAGE_MODE
+print(f'report_mode = {train_cfg.data.report_mode}   image_mode = {train_cfg.data.image_mode}')
+
+# ── 3) Patch dataset paths for the local download ────────────────
+train_cfg.data.dataset_name = DATASET_NAME
+
+if DATASET_NAME == 'MIMIC-CXR':
+    train_cfg.data.mimic_cxr_root   = str(CXR_ROOT)
+    train_cfg.data.mimic_auto_build = True
+
+    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))
+           or sorted(DATA_SRC.rglob('*chexbert*.csv')))
+    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None
+    train_cfg.data.mimic_vqa_root     = str(VQA_ROOT) if VQA_ROOT is not None else None
+
+    out_dir = PROJECT / 'data' / 'data_files'
+    out_dir.mkdir(parents=True, exist_ok=True)
+    train_cfg.data.instruct_json = str(out_dir / 'mimic_cxr_instruct_unified.json')
+
+elif DATASET_NAME == 'MIMIC-CXR_resized':
+    train_cfg.data.mimic_cxr_resized.root          = str(MR_ROOT)
+    train_cfg.data.mimic_cxr_resized.manifest_dir  = None
+    train_cfg.data.mimic_cxr_resized.vqa_dir       = None
+    train_cfg.data.mimic_cxr_resized.reports_root  = None
+    train_cfg.data.mimic_cxr_resized.auto_build    = True
+    out_dir = PROJECT / 'data' / 'data_files'
+    out_dir.mkdir(parents=True, exist_ok=True)
+    train_cfg.data.mimic_cxr_resized.instruct_json = str(out_dir / 'mimic_cxr_resized_instruct.json')
+
+else:  # IU-Xray
+    train_cfg.data.iu_xray.images_dir = str(IU_IMAGES_DIR)
+    train_cfg.data.iu_xray.labels_dir = str(IU_LABELS_DIR)
+    train_cfg.data.iu_xray.auto_build = True
+    out_dir = PROJECT / 'data' / 'data_files'
+    out_dir.mkdir(parents=True, exist_ok=True)
+    train_cfg.data.iu_xray.instruct_json = str(out_dir / 'iu_xray_instruct.json')
+
+train_cfg.data.train_split = 'train'
+train_cfg.data.val_split   = 'validate'
+train_cfg.data.test_split  = 'test'
+
+# ── 4) Apply auto-detected compute profile (overrides saved config) ──
+train_cfg.training.fp16 = PROFILE['fp16']
+train_cfg.training.bf16 = PROFILE['bf16']
+train_cfg.training.per_device_train_batch_size = PROFILE['per_device_eval_batch_size']
+train_cfg.training.per_device_eval_batch_size  = PROFILE['per_device_eval_batch_size']
+train_cfg.training.dataloader_num_workers      = PROFILE['dataloader_num_workers']
+train_cfg.training.dataloader_pin_memory       = True
+
+# Disable feature cache for eval (test images haven't been precomputed).
+train_cfg.data.feature_cache_dir = None
+
+# 4-bit QLoRA — must match how the trainer set it up for the saved LoRA to load.
+model_cfg.llm.load_in_4bit              = True
+model_cfg.llm.load_in_8bit              = False
+model_cfg.llm.attn_implementation       = PROFILE['attn_implementation']
+model_cfg.llm.torch_dtype               = PROFILE['torch_dtype']
+model_cfg.llm.bnb_4bit_compute_dtype    = PROFILE['bnb_4bit_compute_dtype']
+model_cfg.llm.bnb_4bit_quant_type       = 'nf4'
+model_cfg.llm.bnb_4bit_use_double_quant = True
+model_cfg.llm.gradient_checkpointing    = False     # eval-only - no backward pass
+
+# ── 5) CheXpert classifier ───────────────────────────────────────
+# The training notebook keeps it disabled (oracle PNU from CSV/manifest).
+# Only enable here if the trained run actually has a learned classifier
+# checkpoint saved alongside.
+if CHEXPERT_PT.is_file():
+    model_cfg.chexpert_classifier.enabled = True
+    print(f'CheXpert classifier checkpoint found at {CHEXPERT_PT} -> enabled')
+else:
+    model_cfg.chexpert_classifier.enabled = False
+    print('No CheXpert classifier checkpoint -> disabled (oracle PNU from CSV/manifest)')
+
+# ── 6) HF Hub tracker — points results uploads at HF_RUNS_REPO ──
+CKPT_ROOT = WORK / 'ckpt_eval'
+CKPT_ROOT.mkdir(parents=True, exist_ok=True)
+train_cfg.training.output_root = str(CKPT_ROOT)
+
+if UPLOAD_RESULTS_TO_HF:
+    train_cfg.hf_hub.enabled        = True
+    train_cfg.hf_hub.repo_id        = HF_RUNS_REPO
+    train_cfg.hf_hub.token_env      = 'HF_TOKEN'
+    train_cfg.hf_hub.private        = True
+else:
+    train_cfg.hf_hub.enabled = False
+
+train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')
+
+# Pin RUN_ID so resolve_run_id picks it up exactly.
+Path(train_cfg.hf_hub.run_state_file).write_text(RUN_ID)
+
+# ── 7) Save patched configs into the project so the subprocess sees them ──
+OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')
+OmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')
+
+print('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))
+print('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))
+print('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))
+print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))
+"""))
+
+
+CELLS.append(md("eval-run-md", """\
+## 7. Run evaluation
+
+Calls `python -m evaluation.evaluate` as a subprocess. It will:
+1. Build the unified instruct JSON for the chosen `(report_mode, image_mode)` if missing.
+2. Build the model from the patched configs.
+3. Load the projection + LoRA from `CKPT_DIR_PULLED`.
+4. Iterate the **test** split, generate predictions, score every task.
+5. Write `{LOCAL_RESULTS_DIR}/{RUN_ID}/predictions_*.json` + `metrics_summary.json`.
+6. If `UPLOAD_RESULTS_TO_HF=True`, upload the folder to `{RUN_ID}/results/` on `HF_RUNS_REPO`.
+"""))
+
+
+CELLS.append(code("eval-run", """\
+import shlex
+
+RESULTS_DIR_LOCAL = WORK / LOCAL_RESULTS_DIR
+RESULTS_DIR_LOCAL.mkdir(parents=True, exist_ok=True)
+
+# evaluate.py expects --checkpoint to be a `<dir>/<name>_projection.pt` path;
+# it derives the LoRA dir + CheXpert classifier path from the same folder.
+CKPT_ARG = str(PROJ_PT)
+
+extra = ''
+if LLM_JUDGE:
+    extra += ' --llm_judge'
+    extra += f' --llm_judge_model {shlex.quote(LLM_JUDGE_MODEL)}'
+    if LLM_JUDGE_BASE_URL:
+        extra += f' --llm_judge_base_url {shlex.quote(LLM_JUDGE_BASE_URL)}'
+    if LLM_JUDGE_MAX_SAMPLES:
+        extra += f' --llm_judge_max_samples {int(LLM_JUDGE_MAX_SAMPLES)}'
+
+if not UPLOAD_RESULTS_TO_HF:
+    extra += ' --no_hf_upload'
+
+print(f'Evaluating run_id    : {RUN_ID}')
+print(f'Checkpoint           : {CKPT_ARG}')
+print(f'Task                 : {TASK}   (split={SPLIT})')
+print(f'Local results dir    : {RESULTS_DIR_LOCAL}/{RUN_ID}/')
+print(f'Upload to HF         : {UPLOAD_RESULTS_TO_HF}')
+print()
+
+# NOTE: do NOT set HF_HUB_DISABLE_PROGRESS_BARS=1 here. On a cold cache the
+# 4-bit Vicuna shard download is ~13GB and takes minutes on Colab T4 — hiding
+# the bar makes the cell look frozen. TQDM_MININTERVAL=1.0 forces the per-task
+# tqdm in run_inference to refresh every 1s.
+!TRANSFORMERS_VERBOSITY=warning TOKENIZERS_PARALLELISM=false BITSANDBYTES_NOWELCOME=1 \\
+PYTHONUNBUFFERED=1 TQDM_MININTERVAL=1.0 \\
+python -u -m evaluation.evaluate \\
+    --model_config    configs/model_config.yaml \\
+    --train_config    configs/train_config.yaml \\
+    --checkpoint      "{CKPT_ARG}" \\
+    --run_id          "{RUN_ID}" \\
+    --task            {TASK} \\
+    --split           {SPLIT} \\
+    --batch_size      {BATCH_SIZE} \\
+    --max_new_tokens  {MAX_NEW_TOKENS} \\
+    --output_dir      "{RESULTS_DIR_LOCAL}" \\
+    --device          cuda{extra}
+"""))
+
+
+CELLS.append(md("eval-summary-md", """\
+## 8. Show the metrics summary
+"""))
+
+
+CELLS.append(code("eval-summary", """\
+import json as _json
+
+_summary_path = RESULTS_DIR_LOCAL / RUN_ID / 'metrics_summary.json'
+if not _summary_path.is_file():
+    print(f'No metrics_summary.json at {_summary_path}. Did evaluate.py error out?')
+else:
+    summary = _json.loads(_summary_path.read_text())
+    print(f'Dataset    : {summary.get("dataset_name")}')
+    print(f'Run ID     : {summary.get("run_id")}')
+    print(f'Split      : {summary.get("split")}')
+    print()
+    for task, metrics in (summary.get('metrics') or {}).items():
+        print(f'─── {task.upper()} ───')
+        for k, v in metrics.items():
+            if isinstance(v, float):
+                print(f'  {k:<22s}: {v:.4f}')
+            else:
+                print(f'  {k:<22s}: {v}')
+        print()
+
+    # List per-task prediction files for convenience
+    print('Per-task prediction files:')
+    for f in sorted((RESULTS_DIR_LOCAL / RUN_ID).glob('predictions_*.json')):
+        n = len(_json.loads(f.read_text()))
+        print(f'  {f.name}  ({n} samples)')
+"""))
+
+
+CELLS.append(md("eval-cleanup-md", """\
+### Done
+
+Final artifacts (also pushed to HF if `UPLOAD_RESULTS_TO_HF=True`):
+
+```
+{LOCAL_RESULTS_DIR}/{RUN_ID}/
+  predictions_findings.json
+  predictions_impression.json
+  predictions_vqa.json       (MIMIC datasets only)
+  metrics_summary.json
+```
+
+For free-form inference on individual images (with image preview, custom VQA questions, etc.) use the separate **`cxrvlm_colab_inference.ipynb`** notebook.
+"""))
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Write notebook
+# ─────────────────────────────────────────────────────────────────────
+
+nb = {
+    "cells": CELLS,
+    "metadata": {
+        "accelerator": "GPU",
+        "colab": {
+            "gpuType": "T4",
+            "provenance": [],
+            "machine_shape": "hm",
+        },
+        "kernelspec": {
+            "display_name": "Python 3",
+            "name": "python3",
+        },
+        "language_info": {
+            "name": "python",
+            "version": "3.10",
+        },
+    },
+    "nbformat": 4,
+    "nbformat_minor": 5,
+}
+
+NB_PATH.write_text(json.dumps(nb, indent=1, ensure_ascii=False), encoding="utf-8")
+print(f"wrote {NB_PATH}  ({len(CELLS)} cells)")

scripts/_build_inference_notebook.py

@@ -0,0 +1,676 @@
+"""
+Build scripts/cxrvlm_colab_inference.ipynb from cell sources defined here.
+
+Run:
+    python scripts/_build_inference_notebook.py
+"""
+
+import json
+from pathlib import Path
+
+NB_PATH = Path(__file__).parent / "cxrvlm_colab_inference.ipynb"
+
+
+def md(cell_id, src):
+    return {
+        "cell_type": "markdown",
+        "id": cell_id,
+        "metadata": {},
+        "source": src.splitlines(keepends=True),
+    }
+
+
+def code(cell_id, src):
+    return {
+        "cell_type": "code",
+        "id": cell_id,
+        "metadata": {},
+        "execution_count": None,
+        "outputs": [],
+        "source": src.splitlines(keepends=True),
+    }
+
+
+CELLS = []
+
+# ─────────────────────────────────────────────────────────────────────
+
+CELLS.append(md("inf-0", """\
+# CXR-VLM — Inference Notebook (Colab T4 / GCP)
+
+Free-form inference on individual chest X-ray images. Use this when you want to:
+- Run the trained model on a **new image** (not in the test split).
+- Ask **custom VQA questions**.
+- See the image + generated text side-by-side for qualitative checks.
+- Iterate on prompts without re-running the heavy `evaluate.py` pipeline.
+
+For aggregate test-set metrics use the companion notebook **`cxrvlm_colab_eval.ipynb`**.
+
+### Sections
+1. Selectors (which trained model to load + compute device)
+2. Env + pip
+3. Pull code from HF
+4. Pull trained checkpoint + config snapshot from HF
+5. Auto-detect GPU profile (bf16+FA2 on Ampere+/Ada, fp16+SDPA on Turing)
+6. Build configs (4-bit Vicuna with the profile)
+7. **Load model in-kernel** (one-time, ~5–8 min on T4 cold cache, ~3 min on L4)
+8. **Inference helpers** (single image → findings / impression / report / VQA / full cascade)
+9. **Examples** — single image, VQA, batch folder
+"""))
+
+
+CELLS.append(md("inf-select-md", """\
+## 0. Select trained run + device
+"""))
+
+
+CELLS.append(code("inf-select", """\
+# ── Platform ─────────────────────────────────────────────────────
+PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'
+
+# ── Source repos on HuggingFace ──────────────────────────────────
+HF_USER       = 'hieu3636'
+HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'
+HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'
+
+# ── Which trained run to use ─────────────────────────────────────
+RUN_ID        = 'MIMIC-CXR_resized_run_1'
+CKPT_PICK     = 'best'                 # 'best' | 'last'
+
+# ── Generation defaults (override per-call later if desired) ─────
+MAX_NEW_TOKENS  = 300
+TEMPERATURE     = 0.1
+DO_SAMPLE       = False                # greedy by default — deterministic
+NUM_BEAMS       = 1
+
+assert PLATFORM  in ('kaggle', 'colab', 'lightning', 'gcp', 'local')
+assert CKPT_PICK in ('best', 'last')
+
+print(f'PLATFORM      = {PLATFORM}')
+print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')
+"""))
+
+
+CELLS.append(md("inf-env-md", "## 1. Environment + pip install"))
+
+
+CELLS.append(code("inf-env", """\
+import os
+os.environ['CUDA_VISIBLE_DEVICES']            = '0'
+os.environ['TOKENIZERS_PARALLELISM']          = 'false'
+os.environ['BITSANDBYTES_NOWELCOME']          = '1'
+os.environ['TRANSFORMERS_VERBOSITY']          = 'warning'
+os.environ['PYTHONUNBUFFERED']                = '1'
+
+import sys, shutil, subprocess
+from pathlib import Path
+"""))
+
+
+CELLS.append(code("inf-pip", """\
+# Same dependency window as the training/eval notebooks.
+!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate
+!pip install -q -U bitsandbytes
+!pip install -q \\
+    'transformers>=4.46,<4.50' \\
+    'peft>=0.13,<0.15' \\
+    'accelerate>=1.0' \\
+    'huggingface_hub>=0.27,<1.0' \\
+    omegaconf sentencepiece 'protobuf>=3.20' \\
+    pillow matplotlib
+
+# Best-effort flash-attn install for Ampere+/Ada GPUs (L4, 3090, A10, A100, H100).
+# Silent fail is OK — cxr_vlm.py auto-falls-back to SDPA if FA2 isn't importable.
+import torch as _t
+if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):
+    print('[pip] Ampere+/Ada detected -> trying flash-attn (5-10 min if building from source)')
+    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5
+else:
+    print('[pip] Pre-Ampere GPU (or no CUDA) -> skipping flash-attn install')
+"""))
+
+
+CELLS.append(code("inf-versions", """\
+import torch, transformers, peft, huggingface_hub, httpx
+
+print('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())
+print('transformers   :', transformers.__version__)
+print('peft           :', peft.__version__)
+print('huggingface_hub:', huggingface_hub.__version__)
+print('httpx          :', httpx.__version__)
+
+# httpx 0.28+ shim — same as the other notebooks.
+def _patch_httpx():
+    if tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28):
+        return
+    if getattr(httpx.Client, '_cxr_vlm_compat_patched', False):
+        return
+    def _make(orig):
+        def patched(self, *args, **kwargs):
+            if 'allow_redirects' in kwargs:
+                kwargs['follow_redirects'] = kwargs.pop('allow_redirects')
+            kwargs.pop('proxies', None)
+            return orig(self, *args, **kwargs)
+        return patched
+    for cls in (httpx.Client, httpx.AsyncClient):
+        for m in ('request', 'get', 'head', 'post', 'put',
+                  'patch', 'delete', 'options'):
+            if hasattr(cls, m):
+                setattr(cls, m, _make(getattr(cls, m)))
+    httpx.Client._cxr_vlm_compat_patched = True
+
+_patch_httpx()
+
+assert torch.cuda.is_available(), 'CUDA not available.'
+_p = torch.cuda.get_device_properties(0)
+print(f'GPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')
+"""))
+
+
+CELLS.append(md("inf-paths-md", "## 2. Paths + pull code"))
+
+
+CELLS.append(code("inf-paths", """\
+# WORK dir + HF_TOKEN bootstrap (platform-specific).
+if PLATFORM == 'kaggle':
+    from kaggle_secrets import UserSecretsClient
+    os.environ['HF_TOKEN'] = UserSecretsClient().get_secret('HF_TOKEN')
+    WORK = Path('/kaggle/working')
+elif PLATFORM == 'colab':
+    from google.colab import userdata
+    os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')
+    WORK = Path('/content')
+elif PLATFORM == 'lightning':
+    WORK = Path('/teamspace/studios/this_studio')
+elif PLATFORM == 'gcp':
+    # Vertex AI Workbench → /home/jupyter. Generic GCE VM → /workspace.
+    # Pick whichever exists; create the latter if neither is writable.
+    for _cand in (Path('/home/jupyter'), Path('/workspace')):
+        if _cand.exists() or os.access(_cand.parent, os.W_OK):
+            WORK = _cand
+            break
+    else:
+        WORK = Path.home() / 'cxr-vlm-work'
+else:  # 'local'
+    WORK = Path.home() / 'cxr-vlm-work'
+WORK.mkdir(parents=True, exist_ok=True)
+
+assert os.environ.get('HF_TOKEN'), \\
+    'HF_TOKEN missing — set it via the platform secrets UI before re-running.'
+
+from huggingface_hub import snapshot_download
+
+print(f'Pulling code from {HF_CODE_REPO} …')
+CODE_SRC = Path(snapshot_download(
+    repo_id   = HF_CODE_REPO,
+    repo_type = 'model',
+    token     = os.environ['HF_TOKEN'],
+    local_dir = str(WORK / 'cxr-vlm-code'),
+))
+
+PROJECT = WORK / 'cxr_vlm'
+if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():
+    shutil.copytree(CODE_SRC, PROJECT)
+
+os.chdir(PROJECT)
+sys.path.insert(0, str(PROJECT))
+print('PROJECT :', PROJECT)
+print('WORK    :', WORK)
+"""))
+
+
+CELLS.append(md("inf-pull-run-md", """\
+## 3. Pull trained checkpoint + config snapshot
+
+Same as the eval notebook: pulls `{RUN_ID}/configs/` + `{RUN_ID}/stage2/{best|last}/` from `HF_RUNS_REPO`. No dataset payload needed for free-form inference.
+"""))
+
+
+CELLS.append(code("inf-pull-run", """\
+from huggingface_hub import snapshot_download
+
+RUN_PULL_ROOT = WORK / 'run_pull'
+RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)
+
+print(f'Pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …')
+snapshot_download(
+    repo_id   = HF_RUNS_REPO,
+    repo_type = 'model',
+    token     = os.environ['HF_TOKEN'],
+    allow_patterns = [
+        f'{RUN_ID}/configs/**',
+        f'{RUN_ID}/run_meta.json',
+        f'{RUN_ID}/stage2/{CKPT_PICK}/**',
+    ],
+    local_dir = str(RUN_PULL_ROOT),
+)
+
+RUN_DIR_PULLED = RUN_PULL_ROOT / RUN_ID
+CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK
+PROJ_PT         = CKPT_DIR_PULLED / 'checkpoint_projection.pt'
+LORA_DIR        = CKPT_DIR_PULLED / 'checkpoint_lora'
+CHEXPERT_PT     = CKPT_DIR_PULLED / 'checkpoint_chexpert_classifier.pt'
+
+assert PROJ_PT.is_file(), f'Projection weights not found at {PROJ_PT}.'
+assert (LORA_DIR / 'adapter_config.json').is_file(), \\
+    f'LoRA adapter_config.json not found in {LORA_DIR}.'
+
+print(f'  projection : {PROJ_PT}      ({PROJ_PT.stat().st_size/1e6:.1f} MB)')
+print(f'  lora       : {LORA_DIR}/    ({sum(p.stat().st_size for p in LORA_DIR.rglob("*") if p.is_file())/1e6:.1f} MB)')
+print(f'  chexpert   : {CHEXPERT_PT}  (exists: {CHEXPERT_PT.is_file()})')
+
+SAVED_CFG_DIR   = RUN_DIR_PULLED / 'configs'
+SAVED_TRAIN_CFG = SAVED_CFG_DIR / 'train_config.yaml'
+SAVED_MODEL_CFG = SAVED_CFG_DIR / 'model_config.yaml'
+"""))
+
+
+CELLS.append(md("inf-gpu-md", """\
+## 4. Auto-detect GPU profile
+
+Picks precision + attention backend based on the actual GPU:
+- **Turing (T4, V100)** → fp16 + SDPA
+- **Ampere+ / Ada (3090, L4, A10, A100, H100)** → bf16 + FA2 (if flash-attn installed; falls back to SDPA otherwise)
+"""))
+
+
+CELLS.append(code("inf-gpu", """\
+import torch
+
+assert torch.cuda.is_available(), 'CUDA not available.'
+_p       = torch.cuda.get_device_properties(0)
+_cap     = (_p.major, _p.minor)
+_vram_gb = _p.total_memory / 1e9
+_bf16_ok = torch.cuda.is_bf16_supported()
+_fa2_ok  = _cap >= (8, 0)
+
+_flash_attn_installed = False
+if _fa2_ok:
+    try:
+        import flash_attn  # noqa: F401
+        _flash_attn_installed = True
+    except Exception:
+        _flash_attn_installed = False
+
+PROFILE = dict(
+    label                  = _p.name,
+    torch_dtype            = 'bfloat16' if _bf16_ok else 'float16',
+    bnb_4bit_compute_dtype = 'bfloat16' if _bf16_ok else 'float16',
+    attn_implementation    = 'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa',
+)
+
+print(f'GPU         : {_p.name}  ({_vram_gb:.1f} GB)')
+print(f'Compute cap : sm_{_cap[0]}{_cap[1]}   bf16 ok: {_bf16_ok}   FA2 wheel: {_flash_attn_installed}')
+print(f'-> precision: {PROFILE["torch_dtype"]}   attn: {PROFILE["attn_implementation"]}')
+"""))
+
+
+CELLS.append(md("inf-cfg-md", "## 5. Build configs"))
+
+
+CELLS.append(code("inf-cfg", """\
+from omegaconf import OmegaConf
+
+if SAVED_MODEL_CFG.is_file():
+    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
+    print(f'model_cfg <- {SAVED_MODEL_CFG}')
+else:
+    model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')
+    print(f'model_cfg <- repo default')
+
+if SAVED_TRAIN_CFG.is_file():
+    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
+else:
+    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
+
+# Compute from auto-detected PROFILE (previous cell).
+model_cfg.llm.load_in_4bit              = True
+model_cfg.llm.load_in_8bit              = False
+model_cfg.llm.attn_implementation       = PROFILE['attn_implementation']
+model_cfg.llm.torch_dtype               = PROFILE['torch_dtype']
+model_cfg.llm.bnb_4bit_compute_dtype    = PROFILE['bnb_4bit_compute_dtype']
+model_cfg.llm.bnb_4bit_quant_type       = 'nf4'
+model_cfg.llm.bnb_4bit_use_double_quant = True
+model_cfg.llm.gradient_checkpointing    = False    # inference only
+
+# CheXpert classifier — enable iff its checkpoint is pulled.
+if CHEXPERT_PT.is_file():
+    model_cfg.chexpert_classifier.enabled = True
+    print(f'CheXpert classifier checkpoint found -> ENABLED')
+else:
+    model_cfg.chexpert_classifier.enabled = False
+    print('No CheXpert classifier in this run -> DISABLED '
+          '(structured_findings will be None unless you pass it explicitly)')
+
+print('\\nmodel_cfg.llm:')
+print(OmegaConf.to_yaml(model_cfg.llm))
+"""))
+
+
+CELLS.append(md("inf-load-md", """\
+## 6. Load model into kernel
+
+This is the slow phase on cold cache: encoder + 4-bit Vicuna shard download + quantization. Run **once per session** and reuse the `model` object for all subsequent inferences.
+
+Approx timing: ~8 min on T4 cold, ~3 min on L4 cold, seconds when cache is warm.
+"""))
+
+
+CELLS.append(code("inf-load", """\
+import time, torch
+from model import CXRVisionLanguageModel
+from model.rad_dino import BioViLTEncoder
+from utils.checkpoint import load_checkpoint
+
+print('[1/3] Building model (Vicuna-7B 4-bit + RAD-DINO + LoRA)…')
+print('      On T4 cold cache this takes ~5–8 min for the shard download + quantization.')
+_t0 = time.time()
+model = CXRVisionLanguageModel(model_cfg)
+print(f'      built in {time.time()-_t0:.1f}s', flush=True)
+
+print(f'[2/3] Loading checkpoint from {PROJ_PT.name} …')
+_t0 = time.time()
+load_checkpoint(model, str(PROJ_PT))
+print(f'      loaded in {time.time()-_t0:.1f}s', flush=True)
+
+print('[3/3] Moving to cuda + eval()')
+_t0 = time.time()
+model = model.to('cuda').eval()
+print(f'      ready in {time.time()-_t0:.1f}s', flush=True)
+
+# Cache transform + chexpert helpers for reuse.
+TRANSFORM = BioViLTEncoder.get_transform('val')
+print('\\nModel ready. VRAM used:',
+      f'{torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB')
+"""))
+
+
+CELLS.append(md("inf-helpers-md", """\
+## 7. Inference helpers
+
+Two convenience functions:
+
+- `predict(image, task, question=None, structured_findings=None, ...)` — single image, single task.
+- `predict_report(image, ...)` — cascade: generates findings first, then feeds the model's own findings as context for impression. Returns both sections plus the merged report.
+
+`image` can be a filesystem path (str/Path) or a `PIL.Image.Image`.
+
+### About `structured_findings` (the PNU prompt block)
+
+Your training run baked **oracle CheXpert labels from the CSV/manifest** into the prompt as the `Positive / Negative / Uncertain Abnormalities` block. At inference time you don't have GT labels — you have three choices:
+
+1. **`None` (default)** — no PNU block prepended. The model sees only the image + the task instruction. This is the realistic deployment mode.
+2. **Pass a PNU string yourself** — for oracle ablation or testing how much PNU helps:
+   ```
+   structured_findings='Positive Abnormalities: Cardiomegaly\\nNegative Abnormalities: No Finding, ...\\nUncertain Abnormalities: None'
+   ```
+3. **Train a CheXpert classifier** (Stage 0 in the training pipeline) and load it via `model_cfg.chexpert_classifier.checkpoint` — your current runs do not have one, so this path is disabled in section 4.
+"""))
+
+
+CELLS.append(code("inf-helpers", """\
+import torch
+from pathlib import Path
+from typing import Optional, Union
+from PIL import Image
+
+from data.prompt_templates import (
+    build_findings_prompt, build_impression_prompt,
+    build_report_prompt,   build_vqa_prompt,
+)
+
+ImageLike = Union[str, Path, Image.Image]
+
+
+def _to_tensor(image: ImageLike) -> torch.Tensor:
+    '''Load + transform a CXR into a (1, C, H, W) tensor on cuda.'''
+    if isinstance(image, (str, Path)):
+        img = Image.open(image).convert('RGB')
+    else:
+        img = image.convert('RGB') if image.mode != 'RGB' else image
+    t = TRANSFORM(img)                                # (C, H, W)
+    return t.unsqueeze(0).to('cuda')                  # (1, C, H, W)
+
+
+@torch.no_grad()
+def predict(
+    image:               ImageLike,
+    task:                str,
+    question:            Optional[str] = None,
+    structured_findings: Optional[str] = None,
+    max_new_tokens:      int = None,
+    temperature:         float = None,
+    do_sample:           bool = None,
+    num_beams:           int = None,
+) -> str:
+    '''Run the model on one image.
+
+    task ∈ {'findings', 'impression', 'report', 'vqa'}.
+    For task='vqa', `question` is required.
+    '''
+    assert task in ('findings', 'impression', 'report', 'vqa'), f'bad task: {task}'
+    if task == 'vqa':
+        assert question, 'VQA requires `question`.'
+
+    sf = structured_findings
+    if task == 'findings':
+        prompt = build_findings_prompt(sf, randomize=False)
+    elif task == 'impression':
+        prompt = build_impression_prompt(sf, randomize=False)
+    elif task == 'report':
+        prompt = build_report_prompt(sf, randomize=False)
+    else:
+        prompt = build_vqa_prompt(question, sf)
+
+    out = model.generate(
+        images         = _to_tensor(image),
+        prompts        = [prompt],
+        max_new_tokens = max_new_tokens if max_new_tokens is not None else MAX_NEW_TOKENS,
+        temperature    = temperature    if temperature    is not None else TEMPERATURE,
+        do_sample      = do_sample      if do_sample      is not None else DO_SAMPLE,
+        num_beams      = num_beams      if num_beams      is not None else NUM_BEAMS,
+    )
+    return out[0]
+
+
+@torch.no_grad()
+def predict_report(
+    image:               ImageLike,
+    structured_findings: Optional[str] = None,
+    max_new_tokens:      int = None,
+):
+    '''Cascade: generate findings first, then feed them as context for impression.
+
+    Matches the `split_cascade` training recipe but with model-generated (not
+    GT) findings — i.e. a real end-to-end cascade (no GT leakage).
+
+    Returns dict {findings, impression, report}.
+    '''
+    # 1) Findings (uses PNU context only — None by default)
+    findings = predict(image, task='findings',
+                       structured_findings=structured_findings,
+                       max_new_tokens=max_new_tokens)
+
+    # 2) Impression: feed the generated findings as the structured_findings
+    #    context — same shape as the cascade trainer saw at training time
+    #    but with the model's own findings instead of GT.
+    impression = predict(image, task='impression',
+                         structured_findings=findings,
+                         max_new_tokens=max_new_tokens)
+
+    return {
+        'findings':   findings,
+        'impression': impression,
+        'report':     f'Findings: {findings}\\n\\nImpression: {impression}',
+    }
+"""))
+
+
+CELLS.append(md("inf-show-md", "## 8. Display helper (optional, pretty-prints image + outputs)"))
+
+
+CELLS.append(code("inf-show", """\
+import matplotlib.pyplot as plt
+from PIL import Image as _PILImage
+
+def show(image, title=None, max_size=512):
+    '''Render an image inline with optional title.'''
+    if isinstance(image, (str, Path)):
+        img = _PILImage.open(image).convert('RGB')
+    else:
+        img = image
+    img.thumbnail((max_size, max_size))
+    plt.figure(figsize=(6, 6))
+    plt.imshow(img)
+    plt.axis('off')
+    if title:
+        plt.title(title)
+    plt.show()
+
+def pretty(d: dict, header: str = None):
+    '''Pretty-print a dict of {label: text}.'''
+    if header:
+        print('=' * 80)
+        print(header)
+    for k, v in d.items():
+        print('-' * 80)
+        print(f'[{k.upper()}]')
+        print(v)
+    print('=' * 80)
+"""))
+
+
+CELLS.append(md("inf-ex1-md", """\
+## 9. Examples
+
+### Example A — single image, all tasks
+
+Edit `IMAGE_PATH` and `VQA_QUESTION` below, then run.
+"""))
+
+
+CELLS.append(code("inf-ex1", """\
+# Pick an image — anywhere on disk. Colab: drag-and-drop into the Files pane
+# (left sidebar) and use '/content/<filename>'. Local: any path that exists.
+IMAGE_PATH = '/content/sample_cxr.jpg'
+
+# Optional PNU override (see helpers section for format). None → no PNU block.
+STRUCTURED_FINDINGS = None
+
+# A clinical question for the VQA call below.
+VQA_QUESTION = 'Is there any pleural effusion visible in this chest X-ray?'
+
+if not Path(IMAGE_PATH).is_file():
+    print(f'!! {IMAGE_PATH} does not exist. Upload an image or change IMAGE_PATH.')
+else:
+    show(IMAGE_PATH, title=Path(IMAGE_PATH).name)
+
+    if STRUCTURED_FINDINGS:
+        print('Structured findings (PNU) override:')
+        print(STRUCTURED_FINDINGS)
+        print()
+
+    findings   = predict(IMAGE_PATH, task='findings',
+                         structured_findings=STRUCTURED_FINDINGS)
+    impression = predict(IMAGE_PATH, task='impression',
+                         structured_findings=STRUCTURED_FINDINGS)
+    vqa_ans    = predict(IMAGE_PATH, task='vqa', question=VQA_QUESTION,
+                         structured_findings=STRUCTURED_FINDINGS)
+
+    pretty({
+        'findings':       findings,
+        'impression':     impression,
+        f'vqa ({VQA_QUESTION!r})': vqa_ans,
+    }, header=Path(IMAGE_PATH).name)
+"""))
+
+
+CELLS.append(md("inf-ex2-md", """\
+### Example B — cascade report (findings → impression on model output)
+"""))
+
+
+CELLS.append(code("inf-ex2", """\
+if Path(IMAGE_PATH).is_file():
+    out = predict_report(IMAGE_PATH)
+    show(IMAGE_PATH, title='Cascade report')
+    pretty(out, header=f'Cascade: {Path(IMAGE_PATH).name}')
+"""))
+
+
+CELLS.append(md("inf-ex3-md", """\
+### Example C — batch a folder of images
+
+Drops all results into `BATCH_OUT_JSON` for downstream analysis.
+"""))
+
+
+CELLS.append(code("inf-ex3", """\
+import json
+from tqdm.auto import tqdm
+
+BATCH_FOLDER   = '/content/cxr_inbox'            # folder of *.jpg / *.png
+BATCH_TASK     = 'report'                        # 'findings' | 'impression' | 'report' | 'vqa' | 'cascade'
+BATCH_QUESTION = None                            # only used when BATCH_TASK == 'vqa'
+BATCH_OUT_JSON = WORK / f'inference_{RUN_ID}_{BATCH_TASK}.json'
+
+EXTS = {'.jpg', '.jpeg', '.png'}
+folder = Path(BATCH_FOLDER)
+imgs = sorted(p for p in folder.rglob('*') if p.suffix.lower() in EXTS) \\
+       if folder.is_dir() else []
+print(f'{len(imgs)} image(s) under {folder}')
+
+results = []
+for p in tqdm(imgs, desc=f'Inference [{BATCH_TASK}]'):
+    try:
+        if BATCH_TASK == 'cascade':
+            out = predict_report(p)
+        else:
+            out = {BATCH_TASK: predict(p, task=BATCH_TASK,
+                                       question=BATCH_QUESTION)}
+        results.append({'image': str(p), **out})
+    except Exception as e:
+        results.append({'image': str(p), 'error': f'{type(e).__name__}: {e}'})
+
+if results:
+    BATCH_OUT_JSON.write_text(json.dumps(results, indent=2))
+    print(f'Wrote {len(results)} entries -> {BATCH_OUT_JSON}')
+"""))
+
+
+CELLS.append(md("inf-done-md", """\
+### Tips
+
+- Loading the model takes 5–8 min on T4 cold cache. After that, each generation is ~3–8 s.
+- To switch to a different trained run: change `RUN_ID` in section 0, re-run sections 3–5 (skip the env / pip sections).
+- For deterministic outputs keep `DO_SAMPLE=False`. For diversity, set `DO_SAMPLE=True, TEMPERATURE=0.7`.
+- `predict_report()` is the realistic end-to-end pipeline (no GT leakage). `predict(task='report')` only works on runs trained with `report_mode='merged'`.
+"""))
+
+
+# ─────────────────────────────────────────────────────────────────────
+
+nb = {
+    "cells": CELLS,
+    "metadata": {
+        "accelerator": "GPU",
+        "colab": {
+            "gpuType": "T4",
+            "provenance": [],
+            "machine_shape": "hm",
+        },
+        "kernelspec": {
+            "display_name": "Python 3",
+            "name": "python3",
+        },
+        "language_info": {
+            "name": "python",
+            "version": "3.10",
+        },
+    },
+    "nbformat": 4,
+    "nbformat_minor": 5,
+}
+
+NB_PATH.write_text(json.dumps(nb, indent=1, ensure_ascii=False), encoding="utf-8")
+print(f"wrote {NB_PATH}  ({len(CELLS)} cells)")

scripts/cxrvlm_colab_eval.ipynb

@@ -0,0 +1,917 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "eval-0",
+   "metadata": {},
+   "source": [
+    "# CXR-VLM — Evaluation Notebook (Colab T4)\n",
+    "\n",
+    "Standalone evaluation + inference for a trained CXR-VLM run.\n",
+    "\n",
+    "What it does:\n",
+    "1. Pulls project code from `<HF_USER>/cxr-vlm-code`.\n",
+    "2. Pulls the chosen dataset from `<HF_USER>/cxr-vlm-data` (same layout the trainer used).\n",
+    "3. Pulls a trained run checkpoint from `<HF_USER>/cxr-vlm-runs/<RUN_ID>/stage2/{best|last}/`.\n",
+    "4. Pulls the run's config snapshot (`<RUN_ID>/configs/{model,train}_config.yaml`) so the model is rebuilt **with the exact same architecture / report_mode / image_mode** as training.\n",
+    "5. **Auto-detects the GPU** (T4 / L4 / 3090 / A10 / A100 / H100) and patches the configs accordingly — fp16+SDPA on Turing, bf16+FA2 on Ampere/Ada, batch size scaled to VRAM. Also patches local dataset paths.\n",
+    "6. Runs `python -m evaluation.evaluate` on the **test split** for all available tasks.\n",
+    "7. Saves predictions + metrics under `RESULTS_DIR/<RUN_ID>/` (and optionally uploads them to `<HF_USER>/cxr-vlm-runs/<RUN_ID>/results/`).\n",
+    "\n",
+    "Set the variables in **section 0** and run all cells top-to-bottom.\n",
+    "\n",
+    "**Want interactive inference on individual images?** Use the companion notebook **`cxrvlm_colab_inference.ipynb`** — same model pull, but with image preview + free-form prompts.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-select-md",
+   "metadata": {},
+   "source": [
+    "## 0. Select run + dataset + options\n",
+    "\n",
+    "Change the variables in the cell below. `RUN_ID` decides which trained model is pulled; everything else (dataset name, report/image mode) defaults to whatever was used at training time (read from the run's config snapshot on HF), but you can override.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-select",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# ── Platform ─────────────────────────────────────────────────────\n",
+    "PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'\n",
+    "\n",
+    "# ── Source repos on HuggingFace ──────────────────────────────────\n",
+    "HF_USER       = 'hieu3636'              # owner of cxr-vlm-{code,data,runs}\n",
+    "HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'\n",
+    "HF_DATA_REPO  = f'{HF_USER}/cxr-vlm-data'\n",
+    "HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'\n",
+    "\n",
+    "# ── Which trained run to evaluate ────────────────────────────────\n",
+    "# This MUST be an existing folder on HF_RUNS_REPO.\n",
+    "# Example: 'MIMIC-CXR_resized_run_1' | 'IU-Xray_run_2' | 'MIMIC-CXR_run_3'\n",
+    "RUN_ID        = 'MIMIC-CXR_resized_run_1'\n",
+    "\n",
+    "# Which stage-2 checkpoint to load.\n",
+    "#   'best' → {RUN_ID}/stage2/best/  (final / best eval_loss)\n",
+    "#   'last' → {RUN_ID}/stage2/last/  (most recent intermediate save)\n",
+    "CKPT_PICK     = 'best'\n",
+    "\n",
+    "# ── Dataset (auto-derived from RUN_ID prefix; override if needed) ──\n",
+    "# Supported: 'MIMIC-CXR' | 'MIMIC-CXR_resized' | 'IU-Xray'\n",
+    "DATASET_NAME  = None        # None → auto-detect from RUN_ID prefix\n",
+    "REPORT_MODE   = None        # None → read from run's saved train_config.yaml\n",
+    "IMAGE_MODE    = None        # None → read from run's saved train_config.yaml\n",
+    "\n",
+    "# ── What to evaluate ─────────────────────────────────────────────\n",
+    "TASK          = 'all'       # 'all' | 'findings' | 'impression' | 'report' | 'vqa'\n",
+    "SPLIT         = 'test'      # always 'test' for this notebook; here for visibility\n",
+    "BATCH_SIZE    = None        # None → auto from GPU profile (T4:1, L4:4, A100:8, H100:16)\n",
+    "MAX_NEW_TOKENS = 300\n",
+    "\n",
+    "# ── LLM-as-judge for VQA (optional; needs OPENAI_API_KEY) ────────\n",
+    "LLM_JUDGE             = False\n",
+    "LLM_JUDGE_MODEL       = 'gpt-4o-mini'\n",
+    "LLM_JUDGE_BASE_URL    = None          # e.g. Gemini OpenAI-compat endpoint\n",
+    "LLM_JUDGE_MAX_SAMPLES = None          # cap cost; None → all VQA samples\n",
+    "\n",
+    "# ── Output ───────────────────────────────────────────────────────\n",
+    "# Local folder where predictions_*.json + metrics_summary.json land.\n",
+    "# Files end up at  {LOCAL_RESULTS_DIR}/{RUN_ID}/...\n",
+    "LOCAL_RESULTS_DIR = 'results'\n",
+    "\n",
+    "# Push the results folder back to HF_RUNS_REPO under\n",
+    "#   {RUN_ID}/results/predictions_*.json + metrics_summary.json\n",
+    "UPLOAD_RESULTS_TO_HF = True\n",
+    "\n",
+    "# ── Auto-derive DATASET_NAME from RUN_ID prefix if not set ───────\n",
+    "if DATASET_NAME is None:\n",
+    "    for cand in ('MIMIC-CXR_resized', 'MIMIC-CXR', 'IU-Xray'):\n",
+    "        if RUN_ID.startswith(cand + '_run_'):\n",
+    "            DATASET_NAME = cand\n",
+    "            break\n",
+    "    assert DATASET_NAME is not None, \\\n",
+    "        f\"Cannot auto-derive DATASET_NAME from RUN_ID={RUN_ID!r}. \" \\\n",
+    "        \"Expected prefix one of: MIMIC-CXR_resized / MIMIC-CXR / IU-Xray. Set DATASET_NAME explicitly.\"\n",
+    "\n",
+    "assert PLATFORM     in ('kaggle', 'colab', 'lightning', 'gcp', 'local')\n",
+    "assert DATASET_NAME in ('MIMIC-CXR', 'MIMIC-CXR_resized', 'IU-Xray')\n",
+    "assert CKPT_PICK    in ('best', 'last')\n",
+    "assert TASK         in ('all', 'findings', 'impression', 'report', 'vqa')\n",
+    "\n",
+    "print(f'PLATFORM      = {PLATFORM}')\n",
+    "print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')\n",
+    "print(f'DATASET_NAME  = {DATASET_NAME}')\n",
+    "print(f'TASK          = {TASK}   SPLIT = {SPLIT}')\n",
+    "print(f'LOCAL_RESULTS_DIR = {LOCAL_RESULTS_DIR}   (upload→HF: {UPLOAD_RESULTS_TO_HF})')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-env-md",
+   "metadata": {},
+   "source": [
+    "## 1. Environment + pip install (matches the training notebook)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-env",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES']            = '0'      # single-GPU\n",
+    "os.environ['TOKENIZERS_PARALLELISM']          = 'false'\n",
+    "os.environ['BITSANDBYTES_NOWELCOME']          = '1'\n",
+    "os.environ['HF_HUB_DISABLE_PROGRESS_BARS']    = '1'\n",
+    "os.environ['TRANSFORMERS_VERBOSITY']          = 'warning'\n",
+    "os.environ['PYTHONUNBUFFERED']                = '1'\n",
+    "\n",
+    "import sys, shutil, subprocess\n",
+    "from pathlib import Path\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-pip",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import os as _os\n",
+    "\n",
+    "# Same dependency set as the training notebook — keeps load_checkpoint /\n",
+    "# evaluate.py running against the exact stack the model was trained with.\n",
+    "!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate\n",
+    "\n",
+    "# Let pip pick latest bnb that matches Colab's CUDA + triton.\n",
+    "!pip install -q -U bitsandbytes\n",
+    "\n",
+    "# Pin transformers / peft to the same window the trainer uses.\n",
+    "!pip install -q \\\n",
+    "    'transformers>=4.46,<4.50' \\\n",
+    "    'peft>=0.13,<0.15' \\\n",
+    "    'accelerate>=1.0' \\\n",
+    "    'huggingface_hub>=0.27,<1.0' \\\n",
+    "    omegaconf sentencepiece 'protobuf>=3.20' \\\n",
+    "    nltk rouge-score bert-score sacrebleu\n",
+    "\n",
+    "# Best-effort flash-attn install for Ampere+/Ada GPUs (L4, 3090, A10, A100, H100).\n",
+    "# Silent fail is OK — cxr_vlm.py auto-falls-back to SDPA if FA2 isn't importable.\n",
+    "# Skipped entirely on Turing (T4/V100) since FA2 requires sm_80+.\n",
+    "import torch as _t\n",
+    "if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):\n",
+    "    print('[pip] Ampere+/Ada detected -> trying flash-attn (5-10 min if building from source)')\n",
+    "    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5\n",
+    "else:\n",
+    "    print('[pip] Pre-Ampere GPU (or no CUDA) -> skipping flash-attn install')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-versions",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch, transformers, bitsandbytes, peft, accelerate, huggingface_hub, httpx\n",
+    "\n",
+    "print('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())\n",
+    "print('transformers   :', transformers.__version__)\n",
+    "print('bitsandbytes   :', bitsandbytes.__version__)\n",
+    "print('peft           :', peft.__version__)\n",
+    "print('accelerate     :', accelerate.__version__)\n",
+    "print('huggingface_hub:', huggingface_hub.__version__)\n",
+    "print('httpx          :', httpx.__version__)\n",
+    "\n",
+    "# httpx 0.28+ compat shim — same patch as the training notebook.\n",
+    "# transformers <=4.50 calls httpx.Client.head(..., allow_redirects=True)\n",
+    "# which httpx 0.28 removed; translate the kwarg at the call site.\n",
+    "def _patch_httpx():\n",
+    "    if tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28):\n",
+    "        return\n",
+    "    if getattr(httpx.Client, '_cxr_vlm_compat_patched', False):\n",
+    "        return\n",
+    "    def _make(orig):\n",
+    "        def patched(self, *args, **kwargs):\n",
+    "            if 'allow_redirects' in kwargs:\n",
+    "                kwargs['follow_redirects'] = kwargs.pop('allow_redirects')\n",
+    "            kwargs.pop('proxies', None)\n",
+    "            return orig(self, *args, **kwargs)\n",
+    "        return patched\n",
+    "    for cls in (httpx.Client, httpx.AsyncClient):\n",
+    "        for m in ('request', 'get', 'head', 'post', 'put',\n",
+    "                  'patch', 'delete', 'options'):\n",
+    "            if hasattr(cls, m):\n",
+    "                setattr(cls, m, _make(getattr(cls, m)))\n",
+    "    httpx.Client._cxr_vlm_compat_patched = True\n",
+    "    print(f'httpx {httpx.__version__}: monkey-patched allow_redirects -> follow_redirects')\n",
+    "\n",
+    "_patch_httpx()\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA not available — refusing to evaluate on CPU.'\n",
+    "_p = torch.cuda.get_device_properties(0)\n",
+    "print(f'\\nGPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')\n",
+    "print(f'Compute cap : sm_{_p.major}{_p.minor}  (BF16 ok: {torch.cuda.is_bf16_supported()})')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-paths-md",
+   "metadata": {},
+   "source": [
+    "## 2. Paths + pull code + pull dataset\n",
+    "\n",
+    "Identical to the training notebook for sections that overlap (HF code + per-dataset payload).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-paths",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# ── 1) WORK dir + HF_TOKEN bootstrap (platform-specific) ───────────\n",
+    "if PLATFORM == 'kaggle':\n",
+    "    from kaggle_secrets import UserSecretsClient\n",
+    "    os.environ['HF_TOKEN'] = UserSecretsClient().get_secret('HF_TOKEN')\n",
+    "    WORK = Path('/kaggle/working')\n",
+    "elif PLATFORM == 'colab':\n",
+    "    from google.colab import userdata\n",
+    "    os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')\n",
+    "    WORK = Path('/content')\n",
+    "elif PLATFORM == 'lightning':\n",
+    "    WORK = Path('/teamspace/studios/this_studio')\n",
+    "elif PLATFORM == 'gcp':\n",
+    "    WORK = Path('/workspace')\n",
+    "else:  # 'local'\n",
+    "    WORK = Path.home() / 'cxr-vlm-work'\n",
+    "WORK.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "assert os.environ.get('HF_TOKEN'), \\\n",
+    "    'HF_TOKEN missing — set it via the platform secrets UI before re-running.'\n",
+    "\n",
+    "from huggingface_hub import snapshot_download, hf_hub_download, HfApi\n",
+    "\n",
+    "# ── 2) Code: flat folder, snapshot_download ──\n",
+    "print(f'Pulling code from {HF_CODE_REPO} …')\n",
+    "CODE_SRC = Path(snapshot_download(\n",
+    "    repo_id   = HF_CODE_REPO,\n",
+    "    repo_type = 'model',\n",
+    "    token     = os.environ['HF_TOKEN'],\n",
+    "    local_dir = str(WORK / 'cxr-vlm-code'),\n",
+    "))\n",
+    "\n",
+    "# ── 3) Data: layout depends on DATASET_NAME (same logic as train notebook) ──\n",
+    "DATA_SRC = WORK / 'data'\n",
+    "DATA_SRC.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "if DATASET_NAME == 'MIMIC-CXR_resized':\n",
+    "    import tarfile\n",
+    "    mr_dir = DATA_SRC / 'MIMIC-CXR_resized'\n",
+    "    mr_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    files_dir = mr_dir / 'files'\n",
+    "\n",
+    "    manifests_present = all(\n",
+    "        (mr_dir / f).is_file() for f in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv')\n",
+    "    )\n",
+    "    if manifests_present and files_dir.is_dir() and any(files_dir.glob('p*')):\n",
+    "        print(f'{mr_dir} already populated — skipping download.')\n",
+    "    else:\n",
+    "        api = HfApi(token=os.environ['HF_TOKEN'])\n",
+    "        all_files = api.list_repo_files(repo_id=HF_DATA_REPO, repo_type='dataset')\n",
+    "        mr_files  = [f for f in all_files if f.startswith('MIMIC-CXR_resized/')]\n",
+    "        tar_files = sorted(f for f in mr_files if f.endswith('.tar'))\n",
+    "        meta_files = [f for f in mr_files if not f.endswith('.tar')]\n",
+    "        print(f'MIMIC-CXR_resized on HF: {len(tar_files)} tar shards + {len(meta_files)} metadata files')\n",
+    "\n",
+    "        # Manifests / vqa / SHARDS.txt\n",
+    "        snapshot_download(\n",
+    "            repo_id        = HF_DATA_REPO,\n",
+    "            repo_type      = 'dataset',\n",
+    "            allow_patterns = ['MIMIC-CXR_resized/*.csv',\n",
+    "                              'MIMIC-CXR_resized/*.json',\n",
+    "                              'MIMIC-CXR_resized/*.txt',\n",
+    "                              'MIMIC-CXR_resized/vqa/**'],\n",
+    "            token          = os.environ['HF_TOKEN'],\n",
+    "            local_dir      = str(DATA_SRC),\n",
+    "        )\n",
+    "\n",
+    "        # Tar shards — sequential extract + delete to keep peak disk low.\n",
+    "        for i, tf in enumerate(tar_files, 1):\n",
+    "            print(f'  [{i}/{len(tar_files)}] {tf}')\n",
+    "            tar_path = Path(hf_hub_download(\n",
+    "                repo_id=HF_DATA_REPO, repo_type='dataset',\n",
+    "                filename=tf, token=os.environ['HF_TOKEN'],\n",
+    "                local_dir=str(DATA_SRC),\n",
+    "            ))\n",
+    "            with tarfile.open(tar_path) as t:\n",
+    "                t.extractall(mr_dir)\n",
+    "            tar_path.unlink(missing_ok=True)\n",
+    "        print(f'  done. {mr_dir} ready.')\n",
+    "\n",
+    "else:\n",
+    "    # MIMIC-CXR / IU-Xray — single zip per dataset.\n",
+    "    import zipfile\n",
+    "    zip_name = f'{DATASET_NAME}.zip'\n",
+    "    marker   = DATA_SRC / DATASET_NAME\n",
+    "\n",
+    "    if not marker.exists():\n",
+    "        print(f'Pulling {zip_name} from HF …')\n",
+    "        zpath = hf_hub_download(\n",
+    "            repo_id   = HF_DATA_REPO,\n",
+    "            filename  = zip_name,\n",
+    "            repo_type = 'dataset',\n",
+    "            token     = os.environ['HF_TOKEN'],\n",
+    "            local_dir = str(DATA_SRC),\n",
+    "        )\n",
+    "        print(f'  unzipping -> {DATA_SRC}')\n",
+    "        with zipfile.ZipFile(zpath) as zf:\n",
+    "            zf.extractall(DATA_SRC)\n",
+    "        try:\n",
+    "            os.remove(zpath)\n",
+    "        except OSError:\n",
+    "            pass\n",
+    "    else:\n",
+    "        print(f'{marker} already present — skipping download.')\n",
+    "\n",
+    "print(f'Contents of {DATA_SRC}: {sorted(os.listdir(DATA_SRC))}')\n",
+    "\n",
+    "# ── 4) Copy code into writable PROJECT dir + chdir ─────────────────\n",
+    "PROJECT = WORK / 'cxr_vlm'\n",
+    "if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():\n",
+    "    shutil.copytree(CODE_SRC, PROJECT)\n",
+    "\n",
+    "os.chdir(PROJECT)\n",
+    "sys.path.insert(0, str(PROJECT))\n",
+    "print('CODE_SRC :', CODE_SRC)\n",
+    "print('DATA_SRC :', DATA_SRC)\n",
+    "print('PROJECT  :', PROJECT)\n",
+    "print('WORK     :', WORK)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-find-md",
+   "metadata": {},
+   "source": [
+    "## 3. Locate the dataset payload on disk\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-find",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Same dataset-payload finders as the training notebook. Filling in the\n",
+    "# variables the config cell below will consume.\n",
+    "\n",
+    "def find_split_parent(root: Path) -> Path:\n",
+    "    for cand in [root, root / 'MIMIC-CXR', root / 'data' / 'MIMIC-CXR']:\n",
+    "        if (cand / 'train').exists() and (cand / 'valid').exists() and (cand / 'test').exists():\n",
+    "            return cand\n",
+    "    for p in root.rglob('train'):\n",
+    "        if p.is_dir() and (p.parent / 'valid').exists() and (p.parent / 'test').exists():\n",
+    "            return p.parent\n",
+    "    raise FileNotFoundError('Could not find train/ valid/ test/ under ' + str(root))\n",
+    "\n",
+    "\n",
+    "def find_mimic_resized_root(root: Path) -> Path:\n",
+    "    for cand in [root / 'MIMIC-CXR_resized', root, *root.rglob('MIMIC-CXR_resized')]:\n",
+    "        if (cand / 'manifest_train.csv').is_file():\n",
+    "            return cand\n",
+    "    raise FileNotFoundError(\n",
+    "        f'Could not find MIMIC-CXR_resized payload under {root}. '\n",
+    "        'Expected manifest_train.csv (alongside manifest_val.csv / manifest_test.csv).'\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def find_iu_dirs(root: Path):\n",
+    "    for cand in [root / 'IU-Xray', *root.rglob('IU-Xray')]:\n",
+    "        if not cand.is_dir():\n",
+    "            continue\n",
+    "        imgs = cand / 'images'\n",
+    "        lbls = cand / 'labels'\n",
+    "        if imgs.is_dir() and lbls.is_dir() and any(lbls.glob('*.xml')):\n",
+    "            return imgs, lbls\n",
+    "        legacy = lbls / 'ecgen-radiology'\n",
+    "        if imgs.is_dir() and legacy.is_dir() and any(legacy.glob('*.xml')):\n",
+    "            return imgs, legacy\n",
+    "    img_dir = lbl_dir = None\n",
+    "    for cand in [root / 'images', *root.rglob('images')]:\n",
+    "        if cand.is_dir() and any(cand.glob('CXR*.png')):\n",
+    "            img_dir = cand; break\n",
+    "    for cand in [root / 'labels', *root.rglob('labels')]:\n",
+    "        if cand.is_dir() and any(cand.glob('*.xml')):\n",
+    "            lbl_dir = cand; break\n",
+    "    if lbl_dir is None:\n",
+    "        for cand in root.rglob('ecgen-radiology'):\n",
+    "            if cand.is_dir() and any(cand.glob('*.xml')):\n",
+    "                lbl_dir = cand; break\n",
+    "    return img_dir, lbl_dir\n",
+    "\n",
+    "\n",
+    "CXR_ROOT      = None\n",
+    "VQA_ROOT      = None\n",
+    "MR_ROOT       = None\n",
+    "IU_IMAGES_DIR = None\n",
+    "IU_LABELS_DIR = None\n",
+    "\n",
+    "if DATASET_NAME == 'MIMIC-CXR':\n",
+    "    CXR_ROOT = find_split_parent(DATA_SRC)\n",
+    "    print('MIMIC-CXR root:', CXR_ROOT)\n",
+    "    for s in ('train', 'valid', 'test'):\n",
+    "        d = CXR_ROOT / s\n",
+    "        assert d.exists(), f'Missing split dir: {d}'\n",
+    "        print(f'  {s:<6s} -> {d}')\n",
+    "\n",
+    "    for p in DATA_SRC.rglob('MIMIC-Ext-MIMIC-CXR-VQA'):\n",
+    "        cand = p / 'dataset'\n",
+    "        if cand.exists() and (cand / 'train.json').exists():\n",
+    "            VQA_ROOT = cand; break\n",
+    "    if VQA_ROOT is None:\n",
+    "        print('VQA root: NOT FOUND -> VQA task will be skipped')\n",
+    "    else:\n",
+    "        print('VQA root:', VQA_ROOT)\n",
+    "\n",
+    "elif DATASET_NAME == 'MIMIC-CXR_resized':\n",
+    "    MR_ROOT = find_mimic_resized_root(DATA_SRC)\n",
+    "    print('MIMIC-CXR_resized root:', MR_ROOT)\n",
+    "    for cf in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv'):\n",
+    "        f = MR_ROOT / cf\n",
+    "        print(f'  {cf}: {\"OK\" if f.is_file() else \"MISSING\"}')\n",
+    "    for sub in ('files', 'vqa'):\n",
+    "        d = MR_ROOT / sub\n",
+    "        print(f'  {sub:<5s}: {\"OK\" if d.is_dir() else \"MISSING\"}  ({d})')\n",
+    "\n",
+    "else:  # IU-Xray\n",
+    "    IU_IMAGES_DIR, IU_LABELS_DIR = find_iu_dirs(DATA_SRC)\n",
+    "    assert IU_IMAGES_DIR is not None, f'IU images/ not found under {DATA_SRC}'\n",
+    "    assert IU_LABELS_DIR is not None, f'IU labels/ (with *.xml) not found under {DATA_SRC}'\n",
+    "    print('IU images dir:', IU_IMAGES_DIR, '->', len(list(IU_IMAGES_DIR.glob('*.png'))), 'PNGs')\n",
+    "    print('IU labels dir:', IU_LABELS_DIR, '->', len(list(IU_LABELS_DIR.glob('*.xml'))), 'XMLs')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-pull-run-md",
+   "metadata": {},
+   "source": [
+    "## 4. Pull trained checkpoint + run config snapshot from HF Runs repo\n",
+    "\n",
+    "Layout on `HF_RUNS_REPO`:\n",
+    "\n",
+    "```\n",
+    "{RUN_ID}/\n",
+    "  configs/{model,train}_config.yaml   ← snapshot taken at training time\n",
+    "  stage2/best/   checkpoint_projection.pt + checkpoint_lora/ [+ checkpoint_chexpert_classifier.pt]\n",
+    "  stage2/last/   same shape, intermediate save\n",
+    "```\n",
+    "\n",
+    "We pull `{RUN_ID}/configs/` and `{RUN_ID}/stage2/{CKPT_PICK}/` into `{WORK}/run_pull/{RUN_ID}/`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-pull-run",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "RUN_PULL_ROOT = WORK / 'run_pull'\n",
+    "RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "print(f'Pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …')\n",
+    "snapshot_download(\n",
+    "    repo_id   = HF_RUNS_REPO,\n",
+    "    repo_type = 'model',\n",
+    "    token     = os.environ['HF_TOKEN'],\n",
+    "    allow_patterns = [\n",
+    "        f'{RUN_ID}/configs/**',\n",
+    "        f'{RUN_ID}/run_meta.json',\n",
+    "        f'{RUN_ID}/stage2/{CKPT_PICK}/**',\n",
+    "    ],\n",
+    "    local_dir = str(RUN_PULL_ROOT),\n",
+    ")\n",
+    "\n",
+    "RUN_DIR_PULLED = RUN_PULL_ROOT / RUN_ID\n",
+    "assert RUN_DIR_PULLED.is_dir(), f'Pull failed — {RUN_DIR_PULLED} missing.'\n",
+    "\n",
+    "CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK\n",
+    "PROJ_PT         = CKPT_DIR_PULLED / 'checkpoint_projection.pt'\n",
+    "LORA_DIR        = CKPT_DIR_PULLED / 'checkpoint_lora'\n",
+    "CHEXPERT_PT     = CKPT_DIR_PULLED / 'checkpoint_chexpert_classifier.pt'\n",
+    "\n",
+    "assert PROJ_PT.is_file(), \\\n",
+    "    f'Projection weights not found at {PROJ_PT}. ' \\\n",
+    "    f'Check that {RUN_ID}/stage2/{CKPT_PICK}/ exists on {HF_RUNS_REPO}.'\n",
+    "assert (LORA_DIR / 'adapter_config.json').is_file(), \\\n",
+    "    f'LoRA adapter_config.json not found in {LORA_DIR}. Stage-2 checkpoint partial?'\n",
+    "\n",
+    "print()\n",
+    "print(f'  projection : {PROJ_PT}      ({PROJ_PT.stat().st_size/1e6:.1f} MB)')\n",
+    "print(f'  lora       : {LORA_DIR}/    ({sum(p.stat().st_size for p in LORA_DIR.rglob(\"*\") if p.is_file())/1e6:.1f} MB)')\n",
+    "print(f'  chexpert   : {CHEXPERT_PT}  (exists: {CHEXPERT_PT.is_file()})')\n",
+    "\n",
+    "# Saved configs (may or may not exist on older runs)\n",
+    "SAVED_CFG_DIR = RUN_DIR_PULLED / 'configs'\n",
+    "SAVED_TRAIN_CFG = SAVED_CFG_DIR / 'train_config.yaml'\n",
+    "SAVED_MODEL_CFG = SAVED_CFG_DIR / 'model_config.yaml'\n",
+    "print()\n",
+    "print(f'  saved train_cfg : {SAVED_TRAIN_CFG}  (exists: {SAVED_TRAIN_CFG.is_file()})')\n",
+    "print(f'  saved model_cfg : {SAVED_MODEL_CFG}  (exists: {SAVED_MODEL_CFG.is_file()})')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-gpu-md",
+   "metadata": {},
+   "source": [
+    "## 5. Auto-detect GPU profile\n",
+    "\n",
+    "Mirrors the training notebook: picks precision (bf16/fp16), attention backend (FA2/SDPA), and an eval batch size based on the actual GPU's compute capability + VRAM.\n",
+    "\n",
+    "| Bucket | Examples | Precision | Attn | Eval batch |\n",
+    "|---|---|---|---|---|\n",
+    "| 70+ GB | A100/H100 80GB | bf16 | FA2 (if installed) | 16 |\n",
+    "| 35–69 GB | A100 40GB | bf16 | FA2 | 8 |\n",
+    "| 22–34 GB | 3090 / L4 / A10 24GB | bf16 | FA2 | 4 |\n",
+    "| 14–21 GB | T4 / V100 16GB | fp16 | SDPA | 1 |\n",
+    "\n",
+    "The eval batch sizes are smaller than training's because generation builds a KV cache that scales with `batch × max_new_tokens`. Override by setting `BATCH_SIZE` in section 0 to a concrete number.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-gpu",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA not available — refusing to write a CPU profile.'\n",
+    "_props   = torch.cuda.get_device_properties(0)\n",
+    "_cap     = (_props.major, _props.minor)\n",
+    "_vram_gb = _props.total_memory / 1e9\n",
+    "_bf16_ok = torch.cuda.is_bf16_supported()\n",
+    "_fa2_ok  = _cap >= (8, 0)        # FA2 needs Ampere+ (sm_80 or newer)\n",
+    "\n",
+    "# Detect whether flash-attn package is actually importable. FA2 falls back to\n",
+    "# SDPA inside cxr_vlm.py if missing, but knowing here lets us print clearly.\n",
+    "_flash_attn_installed = False\n",
+    "if _fa2_ok:\n",
+    "    try:\n",
+    "        import flash_attn  # noqa: F401\n",
+    "        _flash_attn_installed = True\n",
+    "    except Exception:\n",
+    "        _flash_attn_installed = False\n",
+    "\n",
+    "print(f'GPU         : {_props.name}  ({_vram_gb:.1f} GB)')\n",
+    "print(f'Compute cap : sm_{_cap[0]}{_cap[1]}')\n",
+    "print(f'BF16 native : {_bf16_ok}')\n",
+    "print(f'FA2 capable : {_fa2_ok}   flash-attn installed: {_flash_attn_installed}')\n",
+    "\n",
+    "# Eval batch size — smaller than training because generation builds a KV cache\n",
+    "# that scales with batch × max_new_tokens.\n",
+    "if _vram_gb >= 70:\n",
+    "    GPU_LABEL, _EVAL_BS, _NW = 'A100/H100 80GB',         16, 16\n",
+    "elif _vram_gb >= 35:\n",
+    "    GPU_LABEL, _EVAL_BS, _NW = 'A100 40GB',               8, 12\n",
+    "elif _vram_gb >= 22:\n",
+    "    GPU_LABEL, _EVAL_BS, _NW = 'RTX 3090 / L4 / A10 24GB', 4, 8\n",
+    "elif _vram_gb >= 14:\n",
+    "    GPU_LABEL, _EVAL_BS, _NW = 'T4 / V100 (15-16GB)',     1, 2\n",
+    "else:\n",
+    "    GPU_LABEL, _EVAL_BS, _NW = f'unknown ({_vram_gb:.0f}GB) - conservative', 1, 2\n",
+    "\n",
+    "PROFILE = dict(\n",
+    "    label                       = GPU_LABEL,\n",
+    "    bf16                        = bool(_bf16_ok),\n",
+    "    fp16                        = not _bf16_ok,\n",
+    "    torch_dtype                 = 'bfloat16' if _bf16_ok else 'float16',\n",
+    "    bnb_4bit_compute_dtype      = 'bfloat16' if _bf16_ok else 'float16',\n",
+    "    attn_implementation         = 'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa',\n",
+    "    per_device_eval_batch_size  = _EVAL_BS,\n",
+    "    dataloader_num_workers      = _NW,\n",
+    ")\n",
+    "\n",
+    "# Allow the section-0 selector to override (e.g. BATCH_SIZE=8 to push harder).\n",
+    "if BATCH_SIZE is not None:\n",
+    "    PROFILE['per_device_eval_batch_size'] = int(BATCH_SIZE)\n",
+    "    print(f'BATCH_SIZE override: section 0 forced batch={BATCH_SIZE}')\n",
+    "\n",
+    "# Final BATCH_SIZE the rest of the notebook (eval-run cell) will use.\n",
+    "BATCH_SIZE = PROFILE['per_device_eval_batch_size']\n",
+    "\n",
+    "print(f'\\n-> Profile  : {PROFILE[\"label\"]}')\n",
+    "for k, v in PROFILE.items():\n",
+    "    if k == 'label': continue\n",
+    "    print(f'    {k:<32s} = {v}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-cfg-md",
+   "metadata": {},
+   "source": [
+    "## 6. Build configs\n",
+    "\n",
+    "Strategy: start from the run's saved config snapshot if present (so `report_mode`, `image_mode`, `lora.r`, `num_image_tokens`, etc. match training). If absent, fall back to the repo defaults. Then patch:\n",
+    "- dataset paths to the local download\n",
+    "- compute (precision + attn backend + batch) from the auto-detected `PROFILE`\n",
+    "- HF Hub tracker → uses `HF_RUNS_REPO` so the results upload lands under `{RUN_ID}/results/`\n",
+    "- pin `run_id` to `RUN_ID` so `evaluate.py` writes under `{LOCAL_RESULTS_DIR}/{RUN_ID}/`\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-cfg",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from omegaconf import OmegaConf\n",
+    "import torch\n",
+    "\n",
+    "# ── 1) Load base configs: prefer the run's saved snapshot ─────────\n",
+    "if SAVED_TRAIN_CFG.is_file():\n",
+    "    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)\n",
+    "    print(f'train_cfg <- {SAVED_TRAIN_CFG}')\n",
+    "else:\n",
+    "    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\n",
+    "    print(f'train_cfg <- repo default (no snapshot on HF)')\n",
+    "\n",
+    "if SAVED_MODEL_CFG.is_file():\n",
+    "    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)\n",
+    "    print(f'model_cfg <- {SAVED_MODEL_CFG}')\n",
+    "else:\n",
+    "    model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')\n",
+    "    print(f'model_cfg <- repo default (no snapshot on HF)')\n",
+    "\n",
+    "# ── 2) Allow notebook overrides for report/image mode ─────────────\n",
+    "if REPORT_MODE is not None:\n",
+    "    train_cfg.data.report_mode = REPORT_MODE\n",
+    "if IMAGE_MODE is not None:\n",
+    "    train_cfg.data.image_mode = IMAGE_MODE\n",
+    "print(f'report_mode = {train_cfg.data.report_mode}   image_mode = {train_cfg.data.image_mode}')\n",
+    "\n",
+    "# ── 3) Patch dataset paths for the local download ────────────────\n",
+    "train_cfg.data.dataset_name = DATASET_NAME\n",
+    "\n",
+    "if DATASET_NAME == 'MIMIC-CXR':\n",
+    "    train_cfg.data.mimic_cxr_root   = str(CXR_ROOT)\n",
+    "    train_cfg.data.mimic_auto_build = True\n",
+    "\n",
+    "    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))\n",
+    "           or sorted(DATA_SRC.rglob('*chexbert*.csv')))\n",
+    "    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None\n",
+    "    train_cfg.data.mimic_vqa_root     = str(VQA_ROOT) if VQA_ROOT is not None else None\n",
+    "\n",
+    "    out_dir = PROJECT / 'data' / 'data_files'\n",
+    "    out_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    train_cfg.data.instruct_json = str(out_dir / 'mimic_cxr_instruct_unified.json')\n",
+    "\n",
+    "elif DATASET_NAME == 'MIMIC-CXR_resized':\n",
+    "    train_cfg.data.mimic_cxr_resized.root          = str(MR_ROOT)\n",
+    "    train_cfg.data.mimic_cxr_resized.manifest_dir  = None\n",
+    "    train_cfg.data.mimic_cxr_resized.vqa_dir       = None\n",
+    "    train_cfg.data.mimic_cxr_resized.reports_root  = None\n",
+    "    train_cfg.data.mimic_cxr_resized.auto_build    = True\n",
+    "    out_dir = PROJECT / 'data' / 'data_files'\n",
+    "    out_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    train_cfg.data.mimic_cxr_resized.instruct_json = str(out_dir / 'mimic_cxr_resized_instruct.json')\n",
+    "\n",
+    "else:  # IU-Xray\n",
+    "    train_cfg.data.iu_xray.images_dir = str(IU_IMAGES_DIR)\n",
+    "    train_cfg.data.iu_xray.labels_dir = str(IU_LABELS_DIR)\n",
+    "    train_cfg.data.iu_xray.auto_build = True\n",
+    "    out_dir = PROJECT / 'data' / 'data_files'\n",
+    "    out_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    train_cfg.data.iu_xray.instruct_json = str(out_dir / 'iu_xray_instruct.json')\n",
+    "\n",
+    "train_cfg.data.train_split = 'train'\n",
+    "train_cfg.data.val_split   = 'validate'\n",
+    "train_cfg.data.test_split  = 'test'\n",
+    "\n",
+    "# ── 4) Apply auto-detected compute profile (overrides saved config) ──\n",
+    "train_cfg.training.fp16 = PROFILE['fp16']\n",
+    "train_cfg.training.bf16 = PROFILE['bf16']\n",
+    "train_cfg.training.per_device_train_batch_size = PROFILE['per_device_eval_batch_size']\n",
+    "train_cfg.training.per_device_eval_batch_size  = PROFILE['per_device_eval_batch_size']\n",
+    "train_cfg.training.dataloader_num_workers      = PROFILE['dataloader_num_workers']\n",
+    "train_cfg.training.dataloader_pin_memory       = True\n",
+    "\n",
+    "# Disable feature cache for eval (test images haven't been precomputed).\n",
+    "train_cfg.data.feature_cache_dir = None\n",
+    "\n",
+    "# 4-bit QLoRA — must match how the trainer set it up for the saved LoRA to load.\n",
+    "model_cfg.llm.load_in_4bit              = True\n",
+    "model_cfg.llm.load_in_8bit              = False\n",
+    "model_cfg.llm.attn_implementation       = PROFILE['attn_implementation']\n",
+    "model_cfg.llm.torch_dtype               = PROFILE['torch_dtype']\n",
+    "model_cfg.llm.bnb_4bit_compute_dtype    = PROFILE['bnb_4bit_compute_dtype']\n",
+    "model_cfg.llm.bnb_4bit_quant_type       = 'nf4'\n",
+    "model_cfg.llm.bnb_4bit_use_double_quant = True\n",
+    "model_cfg.llm.gradient_checkpointing    = False     # eval-only - no backward pass\n",
+    "\n",
+    "# ── 5) CheXpert classifier ───────────────────────────────────────\n",
+    "# The training notebook keeps it disabled (oracle PNU from CSV/manifest).\n",
+    "# Only enable here if the trained run actually has a learned classifier\n",
+    "# checkpoint saved alongside.\n",
+    "if CHEXPERT_PT.is_file():\n",
+    "    model_cfg.chexpert_classifier.enabled = True\n",
+    "    print(f'CheXpert classifier checkpoint found at {CHEXPERT_PT} -> enabled')\n",
+    "else:\n",
+    "    model_cfg.chexpert_classifier.enabled = False\n",
+    "    print('No CheXpert classifier checkpoint -> disabled (oracle PNU from CSV/manifest)')\n",
+    "\n",
+    "# ── 6) HF Hub tracker — points results uploads at HF_RUNS_REPO ──\n",
+    "CKPT_ROOT = WORK / 'ckpt_eval'\n",
+    "CKPT_ROOT.mkdir(parents=True, exist_ok=True)\n",
+    "train_cfg.training.output_root = str(CKPT_ROOT)\n",
+    "\n",
+    "if UPLOAD_RESULTS_TO_HF:\n",
+    "    train_cfg.hf_hub.enabled        = True\n",
+    "    train_cfg.hf_hub.repo_id        = HF_RUNS_REPO\n",
+    "    train_cfg.hf_hub.token_env      = 'HF_TOKEN'\n",
+    "    train_cfg.hf_hub.private        = True\n",
+    "else:\n",
+    "    train_cfg.hf_hub.enabled = False\n",
+    "\n",
+    "train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')\n",
+    "\n",
+    "# Pin RUN_ID so resolve_run_id picks it up exactly.\n",
+    "Path(train_cfg.hf_hub.run_state_file).write_text(RUN_ID)\n",
+    "\n",
+    "# ── 7) Save patched configs into the project so the subprocess sees them ──\n",
+    "OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')\n",
+    "OmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')\n",
+    "\n",
+    "print('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))\n",
+    "print('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))\n",
+    "print('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))\n",
+    "print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-run-md",
+   "metadata": {},
+   "source": [
+    "## 7. Run evaluation\n",
+    "\n",
+    "Calls `python -m evaluation.evaluate` as a subprocess. It will:\n",
+    "1. Build the unified instruct JSON for the chosen `(report_mode, image_mode)` if missing.\n",
+    "2. Build the model from the patched configs.\n",
+    "3. Load the projection + LoRA from `CKPT_DIR_PULLED`.\n",
+    "4. Iterate the **test** split, generate predictions, score every task.\n",
+    "5. Write `{LOCAL_RESULTS_DIR}/{RUN_ID}/predictions_*.json` + `metrics_summary.json`.\n",
+    "6. If `UPLOAD_RESULTS_TO_HF=True`, upload the folder to `{RUN_ID}/results/` on `HF_RUNS_REPO`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-run",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import shlex\n",
+    "\n",
+    "RESULTS_DIR_LOCAL = WORK / LOCAL_RESULTS_DIR\n",
+    "RESULTS_DIR_LOCAL.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "# evaluate.py expects --checkpoint to be a `<dir>/<name>_projection.pt` path;\n",
+    "# it derives the LoRA dir + CheXpert classifier path from the same folder.\n",
+    "CKPT_ARG = str(PROJ_PT)\n",
+    "\n",
+    "extra = ''\n",
+    "if LLM_JUDGE:\n",
+    "    extra += ' --llm_judge'\n",
+    "    extra += f' --llm_judge_model {shlex.quote(LLM_JUDGE_MODEL)}'\n",
+    "    if LLM_JUDGE_BASE_URL:\n",
+    "        extra += f' --llm_judge_base_url {shlex.quote(LLM_JUDGE_BASE_URL)}'\n",
+    "    if LLM_JUDGE_MAX_SAMPLES:\n",
+    "        extra += f' --llm_judge_max_samples {int(LLM_JUDGE_MAX_SAMPLES)}'\n",
+    "\n",
+    "if not UPLOAD_RESULTS_TO_HF:\n",
+    "    extra += ' --no_hf_upload'\n",
+    "\n",
+    "print(f'Evaluating run_id    : {RUN_ID}')\n",
+    "print(f'Checkpoint           : {CKPT_ARG}')\n",
+    "print(f'Task                 : {TASK}   (split={SPLIT})')\n",
+    "print(f'Local results dir    : {RESULTS_DIR_LOCAL}/{RUN_ID}/')\n",
+    "print(f'Upload to HF         : {UPLOAD_RESULTS_TO_HF}')\n",
+    "print()\n",
+    "\n",
+    "# NOTE: do NOT set HF_HUB_DISABLE_PROGRESS_BARS=1 here. On a cold cache the\n",
+    "# 4-bit Vicuna shard download is ~13GB and takes minutes on Colab T4 — hiding\n",
+    "# the bar makes the cell look frozen. TQDM_MININTERVAL=1.0 forces the per-task\n",
+    "# tqdm in run_inference to refresh every 1s.\n",
+    "!TRANSFORMERS_VERBOSITY=warning TOKENIZERS_PARALLELISM=false BITSANDBYTES_NOWELCOME=1 \\\n",
+    "PYTHONUNBUFFERED=1 TQDM_MININTERVAL=1.0 \\\n",
+    "python -u -m evaluation.evaluate \\\n",
+    "    --model_config    configs/model_config.yaml \\\n",
+    "    --train_config    configs/train_config.yaml \\\n",
+    "    --checkpoint      \"{CKPT_ARG}\" \\\n",
+    "    --run_id          \"{RUN_ID}\" \\\n",
+    "    --task            {TASK} \\\n",
+    "    --split           {SPLIT} \\\n",
+    "    --batch_size      {BATCH_SIZE} \\\n",
+    "    --max_new_tokens  {MAX_NEW_TOKENS} \\\n",
+    "    --output_dir      \"{RESULTS_DIR_LOCAL}\" \\\n",
+    "    --device          cuda{extra}\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-summary-md",
+   "metadata": {},
+   "source": [
+    "## 8. Show the metrics summary\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-summary",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import json as _json\n",
+    "\n",
+    "_summary_path = RESULTS_DIR_LOCAL / RUN_ID / 'metrics_summary.json'\n",
+    "if not _summary_path.is_file():\n",
+    "    print(f'No metrics_summary.json at {_summary_path}. Did evaluate.py error out?')\n",
+    "else:\n",
+    "    summary = _json.loads(_summary_path.read_text())\n",
+    "    print(f'Dataset    : {summary.get(\"dataset_name\")}')\n",
+    "    print(f'Run ID     : {summary.get(\"run_id\")}')\n",
+    "    print(f'Split      : {summary.get(\"split\")}')\n",
+    "    print()\n",
+    "    for task, metrics in (summary.get('metrics') or {}).items():\n",
+    "        print(f'─── {task.upper()} ───')\n",
+    "        for k, v in metrics.items():\n",
+    "            if isinstance(v, float):\n",
+    "                print(f'  {k:<22s}: {v:.4f}')\n",
+    "            else:\n",
+    "                print(f'  {k:<22s}: {v}')\n",
+    "        print()\n",
+    "\n",
+    "    # List per-task prediction files for convenience\n",
+    "    print('Per-task prediction files:')\n",
+    "    for f in sorted((RESULTS_DIR_LOCAL / RUN_ID).glob('predictions_*.json')):\n",
+    "        n = len(_json.loads(f.read_text()))\n",
+    "        print(f'  {f.name}  ({n} samples)')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eval-cleanup-md",
+   "metadata": {},
+   "source": [
+    "### Done\n",
+    "\n",
+    "Final artifacts (also pushed to HF if `UPLOAD_RESULTS_TO_HF=True`):\n",
+    "\n",
+    "```\n",
+    "{LOCAL_RESULTS_DIR}/{RUN_ID}/\n",
+    "  predictions_findings.json\n",
+    "  predictions_impression.json\n",
+    "  predictions_vqa.json       (MIMIC datasets only)\n",
+    "  metrics_summary.json\n",
+    "```\n",
+    "\n",
+    "For free-form inference on individual images (with image preview, custom VQA questions, etc.) use the separate **`cxrvlm_colab_inference.ipynb`** notebook.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": [],
+   "machine_shape": "hm"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/cxrvlm_colab_inference.ipynb

@@ -0,0 +1,755 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "inf-0",
+   "metadata": {},
+   "source": [
+    "# CXR-VLM — Inference Notebook (Colab T4 / GCP)\n",
+    "\n",
+    "Free-form inference on individual chest X-ray images. Use this when you want to:\n",
+    "- Run the trained model on a **new image** (not in the test split).\n",
+    "- Ask **custom VQA questions**.\n",
+    "- See the image + generated text side-by-side for qualitative checks.\n",
+    "- Iterate on prompts without re-running the heavy `evaluate.py` pipeline.\n",
+    "\n",
+    "For aggregate test-set metrics use the companion notebook **`cxrvlm_colab_eval.ipynb`**.\n",
+    "\n",
+    "### Sections\n",
+    "1. Selectors (which trained model to load + compute device)\n",
+    "2. Env + pip\n",
+    "3. Pull code from HF\n",
+    "4. Pull trained checkpoint + config snapshot from HF\n",
+    "5. Auto-detect GPU profile (bf16+FA2 on Ampere+/Ada, fp16+SDPA on Turing)\n",
+    "6. Build configs (4-bit Vicuna with the profile)\n",
+    "7. **Load model in-kernel** (one-time, ~5–8 min on T4 cold cache, ~3 min on L4)\n",
+    "8. **Inference helpers** (single image → findings / impression / report / VQA / full cascade)\n",
+    "9. **Examples** — single image, VQA, batch folder\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-select-md",
+   "metadata": {},
+   "source": [
+    "## 0. Select trained run + device\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-select",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# ── Platform ─────────────────────────────────────────────────────\n",
+    "PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'\n",
+    "\n",
+    "# ── Source repos on HuggingFace ──────────────────────────────────\n",
+    "HF_USER       = 'hieu3636'\n",
+    "HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'\n",
+    "HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'\n",
+    "\n",
+    "# ── Which trained run to use ─────────────────────────────────────\n",
+    "RUN_ID        = 'MIMIC-CXR_resized_run_1'\n",
+    "CKPT_PICK     = 'best'                 # 'best' | 'last'\n",
+    "\n",
+    "# ── Generation defaults (override per-call later if desired) ─────\n",
+    "MAX_NEW_TOKENS  = 300\n",
+    "TEMPERATURE     = 0.1\n",
+    "DO_SAMPLE       = False                # greedy by default — deterministic\n",
+    "NUM_BEAMS       = 1\n",
+    "\n",
+    "assert PLATFORM  in ('kaggle', 'colab', 'lightning', 'gcp', 'local')\n",
+    "assert CKPT_PICK in ('best', 'last')\n",
+    "\n",
+    "print(f'PLATFORM      = {PLATFORM}')\n",
+    "print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-env-md",
+   "metadata": {},
+   "source": [
+    "## 1. Environment + pip install"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-env",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES']            = '0'\n",
+    "os.environ['TOKENIZERS_PARALLELISM']          = 'false'\n",
+    "os.environ['BITSANDBYTES_NOWELCOME']          = '1'\n",
+    "os.environ['TRANSFORMERS_VERBOSITY']          = 'warning'\n",
+    "os.environ['PYTHONUNBUFFERED']                = '1'\n",
+    "\n",
+    "import sys, shutil, subprocess\n",
+    "from pathlib import Path\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-pip",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Same dependency window as the training/eval notebooks.\n",
+    "!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate\n",
+    "!pip install -q -U bitsandbytes\n",
+    "!pip install -q \\\n",
+    "    'transformers>=4.46,<4.50' \\\n",
+    "    'peft>=0.13,<0.15' \\\n",
+    "    'accelerate>=1.0' \\\n",
+    "    'huggingface_hub>=0.27,<1.0' \\\n",
+    "    omegaconf sentencepiece 'protobuf>=3.20' \\\n",
+    "    pillow matplotlib\n",
+    "\n",
+    "# Best-effort flash-attn install for Ampere+/Ada GPUs (L4, 3090, A10, A100, H100).\n",
+    "# Silent fail is OK — cxr_vlm.py auto-falls-back to SDPA if FA2 isn't importable.\n",
+    "import torch as _t\n",
+    "if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):\n",
+    "    print('[pip] Ampere+/Ada detected -> trying flash-attn (5-10 min if building from source)')\n",
+    "    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5\n",
+    "else:\n",
+    "    print('[pip] Pre-Ampere GPU (or no CUDA) -> skipping flash-attn install')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-versions",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch, transformers, peft, huggingface_hub, httpx\n",
+    "\n",
+    "print('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())\n",
+    "print('transformers   :', transformers.__version__)\n",
+    "print('peft           :', peft.__version__)\n",
+    "print('huggingface_hub:', huggingface_hub.__version__)\n",
+    "print('httpx          :', httpx.__version__)\n",
+    "\n",
+    "# httpx 0.28+ shim — same as the other notebooks.\n",
+    "def _patch_httpx():\n",
+    "    if tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28):\n",
+    "        return\n",
+    "    if getattr(httpx.Client, '_cxr_vlm_compat_patched', False):\n",
+    "        return\n",
+    "    def _make(orig):\n",
+    "        def patched(self, *args, **kwargs):\n",
+    "            if 'allow_redirects' in kwargs:\n",
+    "                kwargs['follow_redirects'] = kwargs.pop('allow_redirects')\n",
+    "            kwargs.pop('proxies', None)\n",
+    "            return orig(self, *args, **kwargs)\n",
+    "        return patched\n",
+    "    for cls in (httpx.Client, httpx.AsyncClient):\n",
+    "        for m in ('request', 'get', 'head', 'post', 'put',\n",
+    "                  'patch', 'delete', 'options'):\n",
+    "            if hasattr(cls, m):\n",
+    "                setattr(cls, m, _make(getattr(cls, m)))\n",
+    "    httpx.Client._cxr_vlm_compat_patched = True\n",
+    "\n",
+    "_patch_httpx()\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA not available.'\n",
+    "_p = torch.cuda.get_device_properties(0)\n",
+    "print(f'GPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-paths-md",
+   "metadata": {},
+   "source": [
+    "## 2. Paths + pull code"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-paths",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# WORK dir + HF_TOKEN bootstrap (platform-specific).\n",
+    "if PLATFORM == 'kaggle':\n",
+    "    from kaggle_secrets import UserSecretsClient\n",
+    "    os.environ['HF_TOKEN'] = UserSecretsClient().get_secret('HF_TOKEN')\n",
+    "    WORK = Path('/kaggle/working')\n",
+    "elif PLATFORM == 'colab':\n",
+    "    from google.colab import userdata\n",
+    "    os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')\n",
+    "    WORK = Path('/content')\n",
+    "elif PLATFORM == 'lightning':\n",
+    "    WORK = Path('/teamspace/studios/this_studio')\n",
+    "elif PLATFORM == 'gcp':\n",
+    "    # Vertex AI Workbench → /home/jupyter. Generic GCE VM → /workspace.\n",
+    "    # Pick whichever exists; create the latter if neither is writable.\n",
+    "    for _cand in (Path('/home/jupyter'), Path('/workspace')):\n",
+    "        if _cand.exists() or os.access(_cand.parent, os.W_OK):\n",
+    "            WORK = _cand\n",
+    "            break\n",
+    "    else:\n",
+    "        WORK = Path.home() / 'cxr-vlm-work'\n",
+    "else:  # 'local'\n",
+    "    WORK = Path.home() / 'cxr-vlm-work'\n",
+    "WORK.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "assert os.environ.get('HF_TOKEN'), \\\n",
+    "    'HF_TOKEN missing — set it via the platform secrets UI before re-running.'\n",
+    "\n",
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "print(f'Pulling code from {HF_CODE_REPO} …')\n",
+    "CODE_SRC = Path(snapshot_download(\n",
+    "    repo_id   = HF_CODE_REPO,\n",
+    "    repo_type = 'model',\n",
+    "    token     = os.environ['HF_TOKEN'],\n",
+    "    local_dir = str(WORK / 'cxr-vlm-code'),\n",
+    "))\n",
+    "\n",
+    "PROJECT = WORK / 'cxr_vlm'\n",
+    "if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():\n",
+    "    shutil.copytree(CODE_SRC, PROJECT)\n",
+    "\n",
+    "os.chdir(PROJECT)\n",
+    "sys.path.insert(0, str(PROJECT))\n",
+    "print('PROJECT :', PROJECT)\n",
+    "print('WORK    :', WORK)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-pull-run-md",
+   "metadata": {},
+   "source": [
+    "## 3. Pull trained checkpoint + config snapshot\n",
+    "\n",
+    "Same as the eval notebook: pulls `{RUN_ID}/configs/` + `{RUN_ID}/stage2/{best|last}/` from `HF_RUNS_REPO`. No dataset payload needed for free-form inference.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-pull-run",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "RUN_PULL_ROOT = WORK / 'run_pull'\n",
+    "RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "print(f'Pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …')\n",
+    "snapshot_download(\n",
+    "    repo_id   = HF_RUNS_REPO,\n",
+    "    repo_type = 'model',\n",
+    "    token     = os.environ['HF_TOKEN'],\n",
+    "    allow_patterns = [\n",
+    "        f'{RUN_ID}/configs/**',\n",
+    "        f'{RUN_ID}/run_meta.json',\n",
+    "        f'{RUN_ID}/stage2/{CKPT_PICK}/**',\n",
+    "    ],\n",
+    "    local_dir = str(RUN_PULL_ROOT),\n",
+    ")\n",
+    "\n",
+    "RUN_DIR_PULLED = RUN_PULL_ROOT / RUN_ID\n",
+    "CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK\n",
+    "PROJ_PT         = CKPT_DIR_PULLED / 'checkpoint_projection.pt'\n",
+    "LORA_DIR        = CKPT_DIR_PULLED / 'checkpoint_lora'\n",
+    "CHEXPERT_PT     = CKPT_DIR_PULLED / 'checkpoint_chexpert_classifier.pt'\n",
+    "\n",
+    "assert PROJ_PT.is_file(), f'Projection weights not found at {PROJ_PT}.'\n",
+    "assert (LORA_DIR / 'adapter_config.json').is_file(), \\\n",
+    "    f'LoRA adapter_config.json not found in {LORA_DIR}.'\n",
+    "\n",
+    "print(f'  projection : {PROJ_PT}      ({PROJ_PT.stat().st_size/1e6:.1f} MB)')\n",
+    "print(f'  lora       : {LORA_DIR}/    ({sum(p.stat().st_size for p in LORA_DIR.rglob(\"*\") if p.is_file())/1e6:.1f} MB)')\n",
+    "print(f'  chexpert   : {CHEXPERT_PT}  (exists: {CHEXPERT_PT.is_file()})')\n",
+    "\n",
+    "SAVED_CFG_DIR   = RUN_DIR_PULLED / 'configs'\n",
+    "SAVED_TRAIN_CFG = SAVED_CFG_DIR / 'train_config.yaml'\n",
+    "SAVED_MODEL_CFG = SAVED_CFG_DIR / 'model_config.yaml'\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-gpu-md",
+   "metadata": {},
+   "source": [
+    "## 4. Auto-detect GPU profile\n",
+    "\n",
+    "Picks precision + attention backend based on the actual GPU:\n",
+    "- **Turing (T4, V100)** → fp16 + SDPA\n",
+    "- **Ampere+ / Ada (3090, L4, A10, A100, H100)** → bf16 + FA2 (if flash-attn installed; falls back to SDPA otherwise)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-gpu",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA not available.'\n",
+    "_p       = torch.cuda.get_device_properties(0)\n",
+    "_cap     = (_p.major, _p.minor)\n",
+    "_vram_gb = _p.total_memory / 1e9\n",
+    "_bf16_ok = torch.cuda.is_bf16_supported()\n",
+    "_fa2_ok  = _cap >= (8, 0)\n",
+    "\n",
+    "_flash_attn_installed = False\n",
+    "if _fa2_ok:\n",
+    "    try:\n",
+    "        import flash_attn  # noqa: F401\n",
+    "        _flash_attn_installed = True\n",
+    "    except Exception:\n",
+    "        _flash_attn_installed = False\n",
+    "\n",
+    "PROFILE = dict(\n",
+    "    label                  = _p.name,\n",
+    "    torch_dtype            = 'bfloat16' if _bf16_ok else 'float16',\n",
+    "    bnb_4bit_compute_dtype = 'bfloat16' if _bf16_ok else 'float16',\n",
+    "    attn_implementation    = 'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa',\n",
+    ")\n",
+    "\n",
+    "print(f'GPU         : {_p.name}  ({_vram_gb:.1f} GB)')\n",
+    "print(f'Compute cap : sm_{_cap[0]}{_cap[1]}   bf16 ok: {_bf16_ok}   FA2 wheel: {_flash_attn_installed}')\n",
+    "print(f'-> precision: {PROFILE[\"torch_dtype\"]}   attn: {PROFILE[\"attn_implementation\"]}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-cfg-md",
+   "metadata": {},
+   "source": [
+    "## 5. Build configs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-cfg",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from omegaconf import OmegaConf\n",
+    "\n",
+    "if SAVED_MODEL_CFG.is_file():\n",
+    "    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)\n",
+    "    print(f'model_cfg <- {SAVED_MODEL_CFG}')\n",
+    "else:\n",
+    "    model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')\n",
+    "    print(f'model_cfg <- repo default')\n",
+    "\n",
+    "if SAVED_TRAIN_CFG.is_file():\n",
+    "    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)\n",
+    "else:\n",
+    "    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\n",
+    "\n",
+    "# Compute from auto-detected PROFILE (previous cell).\n",
+    "model_cfg.llm.load_in_4bit              = True\n",
+    "model_cfg.llm.load_in_8bit              = False\n",
+    "model_cfg.llm.attn_implementation       = PROFILE['attn_implementation']\n",
+    "model_cfg.llm.torch_dtype               = PROFILE['torch_dtype']\n",
+    "model_cfg.llm.bnb_4bit_compute_dtype    = PROFILE['bnb_4bit_compute_dtype']\n",
+    "model_cfg.llm.bnb_4bit_quant_type       = 'nf4'\n",
+    "model_cfg.llm.bnb_4bit_use_double_quant = True\n",
+    "model_cfg.llm.gradient_checkpointing    = False    # inference only\n",
+    "\n",
+    "# CheXpert classifier — enable iff its checkpoint is pulled.\n",
+    "if CHEXPERT_PT.is_file():\n",
+    "    model_cfg.chexpert_classifier.enabled = True\n",
+    "    print(f'CheXpert classifier checkpoint found -> ENABLED')\n",
+    "else:\n",
+    "    model_cfg.chexpert_classifier.enabled = False\n",
+    "    print('No CheXpert classifier in this run -> DISABLED '\n",
+    "          '(structured_findings will be None unless you pass it explicitly)')\n",
+    "\n",
+    "print('\\nmodel_cfg.llm:')\n",
+    "print(OmegaConf.to_yaml(model_cfg.llm))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-load-md",
+   "metadata": {},
+   "source": [
+    "## 6. Load model into kernel\n",
+    "\n",
+    "This is the slow phase on cold cache: encoder + 4-bit Vicuna shard download + quantization. Run **once per session** and reuse the `model` object for all subsequent inferences.\n",
+    "\n",
+    "Approx timing: ~8 min on T4 cold, ~3 min on L4 cold, seconds when cache is warm.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-load",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import time, torch\n",
+    "from model import CXRVisionLanguageModel\n",
+    "from model.rad_dino import BioViLTEncoder\n",
+    "from utils.checkpoint import load_checkpoint\n",
+    "\n",
+    "print('[1/3] Building model (Vicuna-7B 4-bit + RAD-DINO + LoRA)…')\n",
+    "print('      On T4 cold cache this takes ~5–8 min for the shard download + quantization.')\n",
+    "_t0 = time.time()\n",
+    "model = CXRVisionLanguageModel(model_cfg)\n",
+    "print(f'      built in {time.time()-_t0:.1f}s', flush=True)\n",
+    "\n",
+    "print(f'[2/3] Loading checkpoint from {PROJ_PT.name} …')\n",
+    "_t0 = time.time()\n",
+    "load_checkpoint(model, str(PROJ_PT))\n",
+    "print(f'      loaded in {time.time()-_t0:.1f}s', flush=True)\n",
+    "\n",
+    "print('[3/3] Moving to cuda + eval()')\n",
+    "_t0 = time.time()\n",
+    "model = model.to('cuda').eval()\n",
+    "print(f'      ready in {time.time()-_t0:.1f}s', flush=True)\n",
+    "\n",
+    "# Cache transform + chexpert helpers for reuse.\n",
+    "TRANSFORM = BioViLTEncoder.get_transform('val')\n",
+    "print('\\nModel ready. VRAM used:',\n",
+    "      f'{torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-helpers-md",
+   "metadata": {},
+   "source": [
+    "## 7. Inference helpers\n",
+    "\n",
+    "Two convenience functions:\n",
+    "\n",
+    "- `predict(image, task, question=None, structured_findings=None, ...)` — single image, single task.\n",
+    "- `predict_report(image, ...)` — cascade: generates findings first, then feeds the model's own findings as context for impression. Returns both sections plus the merged report.\n",
+    "\n",
+    "`image` can be a filesystem path (str/Path) or a `PIL.Image.Image`.\n",
+    "\n",
+    "### About `structured_findings` (the PNU prompt block)\n",
+    "\n",
+    "Your training run baked **oracle CheXpert labels from the CSV/manifest** into the prompt as the `Positive / Negative / Uncertain Abnormalities` block. At inference time you don't have GT labels — you have three choices:\n",
+    "\n",
+    "1. **`None` (default)** — no PNU block prepended. The model sees only the image + the task instruction. This is the realistic deployment mode.\n",
+    "2. **Pass a PNU string yourself** — for oracle ablation or testing how much PNU helps:\n",
+    "   ```\n",
+    "   structured_findings='Positive Abnormalities: Cardiomegaly\\nNegative Abnormalities: No Finding, ...\\nUncertain Abnormalities: None'\n",
+    "   ```\n",
+    "3. **Train a CheXpert classifier** (Stage 0 in the training pipeline) and load it via `model_cfg.chexpert_classifier.checkpoint` — your current runs do not have one, so this path is disabled in section 4.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-helpers",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from pathlib import Path\n",
+    "from typing import Optional, Union\n",
+    "from PIL import Image\n",
+    "\n",
+    "from data.prompt_templates import (\n",
+    "    build_findings_prompt, build_impression_prompt,\n",
+    "    build_report_prompt,   build_vqa_prompt,\n",
+    ")\n",
+    "\n",
+    "ImageLike = Union[str, Path, Image.Image]\n",
+    "\n",
+    "\n",
+    "def _to_tensor(image: ImageLike) -> torch.Tensor:\n",
+    "    '''Load + transform a CXR into a (1, C, H, W) tensor on cuda.'''\n",
+    "    if isinstance(image, (str, Path)):\n",
+    "        img = Image.open(image).convert('RGB')\n",
+    "    else:\n",
+    "        img = image.convert('RGB') if image.mode != 'RGB' else image\n",
+    "    t = TRANSFORM(img)                                # (C, H, W)\n",
+    "    return t.unsqueeze(0).to('cuda')                  # (1, C, H, W)\n",
+    "\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def predict(\n",
+    "    image:               ImageLike,\n",
+    "    task:                str,\n",
+    "    question:            Optional[str] = None,\n",
+    "    structured_findings: Optional[str] = None,\n",
+    "    max_new_tokens:      int = None,\n",
+    "    temperature:         float = None,\n",
+    "    do_sample:           bool = None,\n",
+    "    num_beams:           int = None,\n",
+    ") -> str:\n",
+    "    '''Run the model on one image.\n",
+    "\n",
+    "    task ∈ {'findings', 'impression', 'report', 'vqa'}.\n",
+    "    For task='vqa', `question` is required.\n",
+    "    '''\n",
+    "    assert task in ('findings', 'impression', 'report', 'vqa'), f'bad task: {task}'\n",
+    "    if task == 'vqa':\n",
+    "        assert question, 'VQA requires `question`.'\n",
+    "\n",
+    "    sf = structured_findings\n",
+    "    if task == 'findings':\n",
+    "        prompt = build_findings_prompt(sf, randomize=False)\n",
+    "    elif task == 'impression':\n",
+    "        prompt = build_impression_prompt(sf, randomize=False)\n",
+    "    elif task == 'report':\n",
+    "        prompt = build_report_prompt(sf, randomize=False)\n",
+    "    else:\n",
+    "        prompt = build_vqa_prompt(question, sf)\n",
+    "\n",
+    "    out = model.generate(\n",
+    "        images         = _to_tensor(image),\n",
+    "        prompts        = [prompt],\n",
+    "        max_new_tokens = max_new_tokens if max_new_tokens is not None else MAX_NEW_TOKENS,\n",
+    "        temperature    = temperature    if temperature    is not None else TEMPERATURE,\n",
+    "        do_sample      = do_sample      if do_sample      is not None else DO_SAMPLE,\n",
+    "        num_beams      = num_beams      if num_beams      is not None else NUM_BEAMS,\n",
+    "    )\n",
+    "    return out[0]\n",
+    "\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def predict_report(\n",
+    "    image:               ImageLike,\n",
+    "    structured_findings: Optional[str] = None,\n",
+    "    max_new_tokens:      int = None,\n",
+    "):\n",
+    "    '''Cascade: generate findings first, then feed them as context for impression.\n",
+    "\n",
+    "    Matches the `split_cascade` training recipe but with model-generated (not\n",
+    "    GT) findings — i.e. a real end-to-end cascade (no GT leakage).\n",
+    "\n",
+    "    Returns dict {findings, impression, report}.\n",
+    "    '''\n",
+    "    # 1) Findings (uses PNU context only — None by default)\n",
+    "    findings = predict(image, task='findings',\n",
+    "                       structured_findings=structured_findings,\n",
+    "                       max_new_tokens=max_new_tokens)\n",
+    "\n",
+    "    # 2) Impression: feed the generated findings as the structured_findings\n",
+    "    #    context — same shape as the cascade trainer saw at training time\n",
+    "    #    but with the model's own findings instead of GT.\n",
+    "    impression = predict(image, task='impression',\n",
+    "                         structured_findings=findings,\n",
+    "                         max_new_tokens=max_new_tokens)\n",
+    "\n",
+    "    return {\n",
+    "        'findings':   findings,\n",
+    "        'impression': impression,\n",
+    "        'report':     f'Findings: {findings}\\n\\nImpression: {impression}',\n",
+    "    }\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-show-md",
+   "metadata": {},
+   "source": [
+    "## 8. Display helper (optional, pretty-prints image + outputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-show",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import Image as _PILImage\n",
+    "\n",
+    "def show(image, title=None, max_size=512):\n",
+    "    '''Render an image inline with optional title.'''\n",
+    "    if isinstance(image, (str, Path)):\n",
+    "        img = _PILImage.open(image).convert('RGB')\n",
+    "    else:\n",
+    "        img = image\n",
+    "    img.thumbnail((max_size, max_size))\n",
+    "    plt.figure(figsize=(6, 6))\n",
+    "    plt.imshow(img)\n",
+    "    plt.axis('off')\n",
+    "    if title:\n",
+    "        plt.title(title)\n",
+    "    plt.show()\n",
+    "\n",
+    "def pretty(d: dict, header: str = None):\n",
+    "    '''Pretty-print a dict of {label: text}.'''\n",
+    "    if header:\n",
+    "        print('=' * 80)\n",
+    "        print(header)\n",
+    "    for k, v in d.items():\n",
+    "        print('-' * 80)\n",
+    "        print(f'[{k.upper()}]')\n",
+    "        print(v)\n",
+    "    print('=' * 80)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-ex1-md",
+   "metadata": {},
+   "source": [
+    "## 9. Examples\n",
+    "\n",
+    "### Example A — single image, all tasks\n",
+    "\n",
+    "Edit `IMAGE_PATH` and `VQA_QUESTION` below, then run.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-ex1",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Pick an image — anywhere on disk. Colab: drag-and-drop into the Files pane\n",
+    "# (left sidebar) and use '/content/<filename>'. Local: any path that exists.\n",
+    "IMAGE_PATH = '/content/sample_cxr.jpg'\n",
+    "\n",
+    "# Optional PNU override (see helpers section for format). None → no PNU block.\n",
+    "STRUCTURED_FINDINGS = None\n",
+    "\n",
+    "# A clinical question for the VQA call below.\n",
+    "VQA_QUESTION = 'Is there any pleural effusion visible in this chest X-ray?'\n",
+    "\n",
+    "if not Path(IMAGE_PATH).is_file():\n",
+    "    print(f'!! {IMAGE_PATH} does not exist. Upload an image or change IMAGE_PATH.')\n",
+    "else:\n",
+    "    show(IMAGE_PATH, title=Path(IMAGE_PATH).name)\n",
+    "\n",
+    "    if STRUCTURED_FINDINGS:\n",
+    "        print('Structured findings (PNU) override:')\n",
+    "        print(STRUCTURED_FINDINGS)\n",
+    "        print()\n",
+    "\n",
+    "    findings   = predict(IMAGE_PATH, task='findings',\n",
+    "                         structured_findings=STRUCTURED_FINDINGS)\n",
+    "    impression = predict(IMAGE_PATH, task='impression',\n",
+    "                         structured_findings=STRUCTURED_FINDINGS)\n",
+    "    vqa_ans    = predict(IMAGE_PATH, task='vqa', question=VQA_QUESTION,\n",
+    "                         structured_findings=STRUCTURED_FINDINGS)\n",
+    "\n",
+    "    pretty({\n",
+    "        'findings':       findings,\n",
+    "        'impression':     impression,\n",
+    "        f'vqa ({VQA_QUESTION!r})': vqa_ans,\n",
+    "    }, header=Path(IMAGE_PATH).name)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-ex2-md",
+   "metadata": {},
+   "source": [
+    "### Example B — cascade report (findings → impression on model output)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-ex2",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "if Path(IMAGE_PATH).is_file():\n",
+    "    out = predict_report(IMAGE_PATH)\n",
+    "    show(IMAGE_PATH, title='Cascade report')\n",
+    "    pretty(out, header=f'Cascade: {Path(IMAGE_PATH).name}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-ex3-md",
+   "metadata": {},
+   "source": [
+    "### Example C — batch a folder of images\n",
+    "\n",
+    "Drops all results into `BATCH_OUT_JSON` for downstream analysis.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "inf-ex3",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "BATCH_FOLDER   = '/content/cxr_inbox'            # folder of *.jpg / *.png\n",
+    "BATCH_TASK     = 'report'                        # 'findings' | 'impression' | 'report' | 'vqa' | 'cascade'\n",
+    "BATCH_QUESTION = None                            # only used when BATCH_TASK == 'vqa'\n",
+    "BATCH_OUT_JSON = WORK / f'inference_{RUN_ID}_{BATCH_TASK}.json'\n",
+    "\n",
+    "EXTS = {'.jpg', '.jpeg', '.png'}\n",
+    "folder = Path(BATCH_FOLDER)\n",
+    "imgs = sorted(p for p in folder.rglob('*') if p.suffix.lower() in EXTS) \\\n",
+    "       if folder.is_dir() else []\n",
+    "print(f'{len(imgs)} image(s) under {folder}')\n",
+    "\n",
+    "results = []\n",
+    "for p in tqdm(imgs, desc=f'Inference [{BATCH_TASK}]'):\n",
+    "    try:\n",
+    "        if BATCH_TASK == 'cascade':\n",
+    "            out = predict_report(p)\n",
+    "        else:\n",
+    "            out = {BATCH_TASK: predict(p, task=BATCH_TASK,\n",
+    "                                       question=BATCH_QUESTION)}\n",
+    "        results.append({'image': str(p), **out})\n",
+    "    except Exception as e:\n",
+    "        results.append({'image': str(p), 'error': f'{type(e).__name__}: {e}'})\n",
+    "\n",
+    "if results:\n",
+    "    BATCH_OUT_JSON.write_text(json.dumps(results, indent=2))\n",
+    "    print(f'Wrote {len(results)} entries -> {BATCH_OUT_JSON}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inf-done-md",
+   "metadata": {},
+   "source": [
+    "### Tips\n",
+    "\n",
+    "- Loading the model takes 5–8 min on T4 cold cache. After that, each generation is ~3–8 s.\n",
+    "- To switch to a different trained run: change `RUN_ID` in section 0, re-run sections 3–5 (skip the env / pip sections).\n",
+    "- For deterministic outputs keep `DO_SAMPLE=False`. For diversity, set `DO_SAMPLE=True, TEMPERATURE=0.7`.\n",
+    "- `predict_report()` is the realistic end-to-end pipeline (no GT leakage). `predict(task='report')` only works on runs trained with `report_mode='merged'`.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": [],
+   "machine_shape": "hm"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/gcp_eval_entrypoint.py

@@ -0,0 +1,390 @@
+"""GCP Vertex AI Custom Job entrypoint — **evaluation** (not training).
+
+Mirrors gcp_entrypoint.py but:
+  1. Pulls dataset from HF Hub (same logic).
+  2. Pulls a trained checkpoint from <HF_RUNS_REPO>/<RUN_ID>/stage2/{best|last}/
+     into /workspace/run_pull/<RUN_ID>/stage2/<CKPT_PICK>/.
+  3. Pulls <RUN_ID>/configs/{model,train}_config.yaml (if present) so the model
+     is rebuilt with the EXACT architecture used at training time. Falls back
+     to the repo defaults when the snapshot is absent.
+  4. Patches the configs for the detected GPU (T4 → fp16+SDPA; L4/A100 →
+     bf16+FA2 if flash-attn installed) and pins paths.
+  5. Pins run_id.txt to RUN_ID so evaluate.py writes under {output_dir}/{RUN_ID}/.
+  6. Execs `python -m evaluation.evaluate --checkpoint ... --run_id ...
+     --task all --split test --output_dir /workspace/results`.
+
+Required env vars:
+  HF_TOKEN         — HuggingFace token (read code+data+runs, write runs)
+  DATASET_NAME     — 'IU-Xray' | 'MIMIC-CXR' | 'MIMIC-CXR_resized'
+  RUN_ID           — folder name on HF_RUNS_REPO to evaluate, e.g.
+                     'MIMIC-CXR_resized_run_1'
+
+Optional env vars (defaults shown):
+  HF_USER             = hieu3636
+  HF_RUNS_REPO        = hieu3636/cxr-vlm-runs
+  CKPT_PICK           = best                # 'best' | 'last' (stage2 sub-folder)
+  REPORT_MODE         =                     # blank → use the run's saved snapshot
+  IMAGE_MODE          =                     # blank → use the run's saved snapshot
+  TASK                = all                 # 'all'|'findings'|'impression'|'report'|'vqa'
+  SPLIT               = test
+  BATCH_SIZE          =                     # blank → from GPU profile
+  MAX_NEW_TOKENS      = 300
+  UPLOAD_RESULTS_TO_HF= 1                   # set '0' to skip upload
+  LLM_JUDGE           = 0                   # set '1' to enable VQA judge
+  LLM_JUDGE_MODEL     = gpt-4o-mini
+  LLM_JUDGE_BASE_URL  =                     # for OpenAI-compat endpoints
+  LLM_JUDGE_MAX_SAMPLES=                    # cap samples for cost control
+  WORK                = /workspace
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import zipfile
+from pathlib import Path
+
+# Tame HF/transformers chatter so logs are readable in Cloud Logging.
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
+os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+os.environ.setdefault("TRANSFORMERS_VERBOSITY", "warning")
+os.environ.setdefault("PYTHONUNBUFFERED", "1")
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
+# tqdm flushes ≥ every 1s — keeps Cloud Logging looking alive during inference.
+os.environ.setdefault("TQDM_MININTERVAL", "1.0")
+
+
+def env(name: str, default: str | None = None, *, required: bool = False) -> str:
+    val = os.environ.get(name, default)
+    if required and not val:
+        sys.exit(f"[gcp_eval_entrypoint] ERROR: required env var {name} not set")
+    return val or ""
+
+
+# ── 1) Resolve config from env ────────────────────────────────────────────────
+HF_TOKEN              = env("HF_TOKEN", required=True)
+DATASET_NAME          = env("DATASET_NAME", required=True)
+RUN_ID                = env("RUN_ID", required=True)
+HF_USER               = env("HF_USER", "hieu3636")
+HF_RUNS_REPO          = env("HF_RUNS_REPO", "hieu3636/cxr-vlm-runs")
+CKPT_PICK             = env("CKPT_PICK", "best")
+REPORT_MODE_OVERRIDE  = env("REPORT_MODE", "")
+IMAGE_MODE_OVERRIDE   = env("IMAGE_MODE", "")
+TASK                  = env("TASK", "all")
+SPLIT                 = env("SPLIT", "test")
+BATCH_SIZE_OVERRIDE   = env("BATCH_SIZE", "")
+MAX_NEW_TOKENS        = int(env("MAX_NEW_TOKENS", "300"))
+UPLOAD_RESULTS_TO_HF  = env("UPLOAD_RESULTS_TO_HF", "1") not in ("0", "false", "False", "")
+LLM_JUDGE             = env("LLM_JUDGE", "0") in ("1", "true", "True")
+LLM_JUDGE_MODEL       = env("LLM_JUDGE_MODEL", "gpt-4o-mini")
+LLM_JUDGE_BASE_URL    = env("LLM_JUDGE_BASE_URL", "")
+LLM_JUDGE_MAX_SAMPLES = env("LLM_JUDGE_MAX_SAMPLES", "")
+WORK                  = Path(env("WORK", "/workspace"))
+
+assert DATASET_NAME in ("IU-Xray", "MIMIC-CXR", "MIMIC-CXR_resized"), DATASET_NAME
+assert CKPT_PICK in ("best", "last"), CKPT_PICK
+assert TASK in ("all", "findings", "impression", "report", "vqa"), TASK
+
+PROJECT       = Path(__file__).resolve().parent.parent         # /workspace/code
+DATA_SRC      = WORK / "data"
+RUN_PULL_ROOT = WORK / "run_pull"
+RESULTS_DIR   = WORK / "results"
+CKPT_ROOT     = WORK / "ckpt_eval"
+for d in (DATA_SRC, RUN_PULL_ROOT, RESULTS_DIR, CKPT_ROOT):
+    d.mkdir(parents=True, exist_ok=True)
+
+print(f"[gcp_eval] PROJECT      = {PROJECT}")
+print(f"[gcp_eval] WORK         = {WORK}")
+print(f"[gcp_eval] DATASET_NAME = {DATASET_NAME}")
+print(f"[gcp_eval] RUN_ID       = {RUN_ID}  (ckpt: stage2/{CKPT_PICK})")
+print(f"[gcp_eval] TASK         = {TASK}    SPLIT = {SPLIT}")
+print(f"[gcp_eval] UPLOAD_TO_HF = {UPLOAD_RESULTS_TO_HF}")
+
+# ── 2) Download dataset payload from HF Hub ───────────────────────────────────
+from huggingface_hub import HfApi, hf_hub_download, snapshot_download  # noqa: E402
+
+if DATASET_NAME == "MIMIC-CXR_resized":
+    mr_dir = DATA_SRC / "MIMIC-CXR_resized"
+    mr_dir.mkdir(parents=True, exist_ok=True)
+    files_dir = mr_dir / "files"
+    manifests_present = all(
+        (mr_dir / f).is_file()
+        for f in ("manifest_train.csv", "manifest_val.csv", "manifest_test.csv")
+    )
+    if manifests_present and files_dir.is_dir() and any(files_dir.glob("p*")):
+        print(f"[gcp_eval] {mr_dir} already populated — skipping download.")
+    else:
+        api = HfApi(token=HF_TOKEN)
+        all_files = api.list_repo_files(
+            repo_id=f"{HF_USER}/cxr-vlm-data", repo_type="dataset"
+        )
+        mr_files = [f for f in all_files if f.startswith("MIMIC-CXR_resized/")]
+        tar_files = sorted(f for f in mr_files if f.endswith(".tar"))
+        print(f"[gcp_eval] {len(tar_files)} tar shards on HF")
+
+        snapshot_download(
+            repo_id=f"{HF_USER}/cxr-vlm-data",
+            repo_type="dataset",
+            allow_patterns=[
+                "MIMIC-CXR_resized/*.csv",
+                "MIMIC-CXR_resized/*.json",
+                "MIMIC-CXR_resized/*.txt",
+                "MIMIC-CXR_resized/vqa/**",
+            ],
+            token=HF_TOKEN,
+            local_dir=str(DATA_SRC),
+        )
+
+        for i, tf in enumerate(tar_files, 1):
+            print(f"[gcp_eval]  [{i}/{len(tar_files)}] {tf}", flush=True)
+            tp = Path(hf_hub_download(
+                repo_id=f"{HF_USER}/cxr-vlm-data",
+                repo_type="dataset",
+                filename=tf,
+                token=HF_TOKEN,
+                local_dir=str(DATA_SRC),
+            ))
+            with tarfile.open(tp) as t:
+                t.extractall(mr_dir)
+            tp.unlink(missing_ok=True)
+        print(f"[gcp_eval] {mr_dir} ready.")
+
+    DATA_ROOT_RESIZED = mr_dir
+
+else:
+    zip_name = f"{DATASET_NAME}.zip"
+    marker = DATA_SRC / DATASET_NAME
+    if not marker.exists():
+        print(f"[gcp_eval] downloading {zip_name} ...")
+        zpath = hf_hub_download(
+            repo_id=f"{HF_USER}/cxr-vlm-data",
+            filename=zip_name,
+            repo_type="dataset",
+            token=HF_TOKEN,
+            local_dir=str(DATA_SRC),
+        )
+        with zipfile.ZipFile(zpath) as zf:
+            zf.extractall(DATA_SRC)
+        try:
+            os.remove(zpath)
+        except OSError:
+            pass
+    else:
+        print(f"[gcp_eval] {marker} already present — skipping download.")
+
+print(f"[gcp_eval] DATA_SRC contents: {sorted(os.listdir(DATA_SRC))}")
+
+# ── 3) Pull checkpoint + saved config snapshot from runs repo ────────────────
+print(f"[gcp_eval] pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …")
+snapshot_download(
+    repo_id=HF_RUNS_REPO,
+    repo_type="model",
+    token=HF_TOKEN,
+    allow_patterns=[
+        f"{RUN_ID}/configs/**",
+        f"{RUN_ID}/run_meta.json",
+        f"{RUN_ID}/stage2/{CKPT_PICK}/**",
+    ],
+    local_dir=str(RUN_PULL_ROOT),
+)
+RUN_DIR_PULLED  = RUN_PULL_ROOT / RUN_ID
+CKPT_DIR_PULLED = RUN_DIR_PULLED / "stage2" / CKPT_PICK
+PROJ_PT         = CKPT_DIR_PULLED / "checkpoint_projection.pt"
+LORA_DIR        = CKPT_DIR_PULLED / "checkpoint_lora"
+CHEXPERT_PT     = CKPT_DIR_PULLED / "checkpoint_chexpert_classifier.pt"
+SAVED_CFG_DIR   = RUN_DIR_PULLED / "configs"
+SAVED_TRAIN_CFG = SAVED_CFG_DIR / "train_config.yaml"
+SAVED_MODEL_CFG = SAVED_CFG_DIR / "model_config.yaml"
+
+assert PROJ_PT.is_file(), f"projection weights not found: {PROJ_PT}"
+assert (LORA_DIR / "adapter_config.json").is_file(), \
+    f"LoRA adapter_config.json missing in {LORA_DIR}"
+print(f"[gcp_eval]   projection : {PROJ_PT}  ({PROJ_PT.stat().st_size/1e6:.1f} MB)")
+print(f"[gcp_eval]   lora       : {LORA_DIR}/")
+print(f"[gcp_eval]   chexpert   : exists={CHEXPERT_PT.is_file()}")
+
+# ── 4) Build configs ─────────────────────────────────────────────────────────
+import torch  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+
+repo_train_cfg_path = PROJECT / "configs" / "train_config.yaml"
+repo_model_cfg_path = PROJECT / "configs" / "model_config.yaml"
+
+if SAVED_TRAIN_CFG.is_file():
+    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
+    print(f"[gcp_eval] train_cfg <- {SAVED_TRAIN_CFG}")
+else:
+    train_cfg = OmegaConf.load(repo_train_cfg_path)
+    print(f"[gcp_eval] train_cfg <- repo default (no snapshot)")
+
+if SAVED_MODEL_CFG.is_file():
+    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
+    print(f"[gcp_eval] model_cfg <- {SAVED_MODEL_CFG}")
+else:
+    model_cfg = OmegaConf.load(repo_model_cfg_path)
+    print(f"[gcp_eval] model_cfg <- repo default (no snapshot)")
+
+# Optional report/image mode overrides — leave blank to use saved values.
+if REPORT_MODE_OVERRIDE:
+    train_cfg.data.report_mode = REPORT_MODE_OVERRIDE
+if IMAGE_MODE_OVERRIDE:
+    train_cfg.data.image_mode  = IMAGE_MODE_OVERRIDE
+print(f"[gcp_eval] report_mode = {train_cfg.data.report_mode}  image_mode = {train_cfg.data.image_mode}")
+
+# Dataset paths (mirror gcp_entrypoint.py)
+train_cfg.data.dataset_name           = DATASET_NAME
+train_cfg.data.max_images_per_sample  = int(getattr(train_cfg.data, "max_images_per_sample", 2))
+
+out_dir = PROJECT / "data" / "data_files"
+out_dir.mkdir(parents=True, exist_ok=True)
+
+if DATASET_NAME == "MIMIC-CXR_resized":
+    train_cfg.data.mimic_cxr_resized.root          = str(DATA_ROOT_RESIZED)
+    train_cfg.data.mimic_cxr_resized.manifest_dir  = None
+    train_cfg.data.mimic_cxr_resized.vqa_dir       = None
+    train_cfg.data.mimic_cxr_resized.reports_root  = None
+    train_cfg.data.mimic_cxr_resized.auto_build    = True
+    train_cfg.data.mimic_cxr_resized.instruct_json = str(
+        out_dir / "mimic_cxr_resized_instruct.json")
+elif DATASET_NAME == "MIMIC-CXR":
+    def _find_mimic_root(root: Path) -> Path:
+        for cand in [root / "MIMIC-CXR", root]:
+            if (cand / "train").exists() and (cand / "valid").exists() and (cand / "test").exists():
+                return cand
+        for p in root.rglob("train"):
+            if p.is_dir() and (p.parent / "valid").exists() and (p.parent / "test").exists():
+                return p.parent
+        raise FileNotFoundError(f"MIMIC-CXR train/valid/test not found under {root}")
+    cxr_root = _find_mimic_root(DATA_SRC)
+    train_cfg.data.mimic_cxr_root   = str(cxr_root)
+    train_cfg.data.instruct_json    = str(out_dir / "mimic_cxr_instruct_unified.json")
+    train_cfg.data.mimic_auto_build = True
+    _cx = sorted(DATA_SRC.rglob("*chexpert*.csv")) or sorted(DATA_SRC.rglob("*chexbert*.csv"))
+    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None
+    _vqa = list(DATA_SRC.rglob("vqa"))
+    train_cfg.data.mimic_vqa_root = str(_vqa[0]) if _vqa else None
+else:  # IU-Xray
+    iu_root = DATA_SRC / "IU-Xray"
+    train_cfg.data.iu_xray.images_dir    = str(iu_root / "images")
+    train_cfg.data.iu_xray.labels_dir    = str(iu_root / "labels")
+    train_cfg.data.iu_xray.instruct_json = str(out_dir / "iu_xray_instruct.json")
+    train_cfg.data.iu_xray.auto_build    = True
+
+train_cfg.data.train_split = "train"
+train_cfg.data.val_split   = "validate"
+train_cfg.data.test_split  = "test"
+train_cfg.data.feature_cache_dir = None
+train_cfg.training.output_root = str(CKPT_ROOT)
+
+# ── GPU auto-profile (eval batch sizes — smaller than training because the KV
+#    cache built during generation scales with batch × max_new_tokens). ──
+assert torch.cuda.is_available(), "CUDA not available in container"
+_props   = torch.cuda.get_device_properties(0)
+_cap     = (_props.major, _props.minor)
+_vram_gb = _props.total_memory / 1e9
+_bf16_ok = torch.cuda.is_bf16_supported()
+_fa2_ok  = _cap >= (8, 0)
+
+_flash_attn_installed = False
+if _fa2_ok:
+    try:
+        import flash_attn  # noqa: F401
+        _flash_attn_installed = True
+    except Exception:
+        _flash_attn_installed = False
+
+if _vram_gb >= 70:
+    _label, _eval_bs, _nw = "A100/H100 80GB", 16, 16
+elif _vram_gb >= 35:
+    _label, _eval_bs, _nw = "A100 40GB", 8, 12
+elif _vram_gb >= 22:
+    _label, _eval_bs, _nw = "3090 / L4 / A10 (24GB)", 4, 8
+elif _vram_gb >= 14:
+    _label, _eval_bs, _nw = "T4 / V100 (15-16GB)", 1, 2
+else:
+    _label, _eval_bs, _nw = f"unknown ({_vram_gb:.0f}GB)", 1, 2
+
+if BATCH_SIZE_OVERRIDE:
+    _eval_bs = int(BATCH_SIZE_OVERRIDE)
+    print(f"[gcp_eval] BATCH_SIZE override -> {_eval_bs}")
+
+print(f"[gcp_eval] GPU: {_props.name}  {_vram_gb:.1f}GB  sm_{_cap[0]}{_cap[1]}  "
+      f"bf16={_bf16_ok}  fa2={_fa2_ok}  fa2_wheel={_flash_attn_installed}")
+print(f"[gcp_eval] -> profile {_label}  eval_batch={_eval_bs}")
+
+train_cfg.training.per_device_train_batch_size   = _eval_bs
+train_cfg.training.per_device_eval_batch_size    = _eval_bs
+train_cfg.training.dataloader_num_workers        = _nw
+train_cfg.training.fp16                          = not _bf16_ok
+train_cfg.training.bf16                          = bool(_bf16_ok)
+train_cfg.training.dataloader_pin_memory         = True
+train_cfg.training.dataloader_persistent_workers = True
+
+model_cfg.llm.attn_implementation       = (
+    "flash_attention_2" if (_fa2_ok and _flash_attn_installed) else "sdpa"
+)
+model_cfg.llm.gradient_checkpointing    = False     # eval — no backward
+model_cfg.llm.torch_dtype               = "bfloat16" if _bf16_ok else "float16"
+model_cfg.llm.bnb_4bit_compute_dtype    = "bfloat16" if _bf16_ok else "float16"
+model_cfg.llm.bnb_4bit_quant_type       = "nf4"
+model_cfg.llm.bnb_4bit_use_double_quant = True
+model_cfg.llm.load_in_8bit              = False
+model_cfg.llm.load_in_4bit              = True
+
+# CheXpert classifier — enable iff its checkpoint is pulled with the run.
+if CHEXPERT_PT.is_file():
+    model_cfg.chexpert_classifier.enabled = True
+    print("[gcp_eval] CheXpert classifier checkpoint found -> enabled")
+else:
+    model_cfg.chexpert_classifier.enabled = False
+
+# HF Hub tracker
+train_cfg.wandb.enabled         = False
+if UPLOAD_RESULTS_TO_HF:
+    train_cfg.hf_hub.enabled    = True
+    train_cfg.hf_hub.repo_id    = HF_RUNS_REPO
+    train_cfg.hf_hub.token_env  = "HF_TOKEN"
+    train_cfg.hf_hub.private    = True
+else:
+    train_cfg.hf_hub.enabled    = False
+train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / "run_id.txt")
+
+# Pin run_id so resolve_run_id picks up exactly the one we're evaluating.
+(CKPT_ROOT / "run_id.txt").write_text(RUN_ID)
+print(f"[gcp_eval] pinned run_id = {RUN_ID}")
+
+# Save patched configs back into the project so evaluate.py reads them.
+OmegaConf.save(train_cfg, repo_train_cfg_path)
+OmegaConf.save(model_cfg, repo_model_cfg_path)
+print("[gcp_eval] configs patched.")
+
+# ── 5) Launch evaluation ─────────────────────────────────────────────────────
+cmd = [
+    "python", "-u", "-m", "evaluation.evaluate",
+    "--model_config",   str(repo_model_cfg_path),
+    "--train_config",   str(repo_train_cfg_path),
+    "--checkpoint",     str(PROJ_PT),
+    "--run_id",         RUN_ID,
+    "--task",           TASK,
+    "--split",          SPLIT,
+    "--batch_size",     str(_eval_bs),
+    "--max_new_tokens", str(MAX_NEW_TOKENS),
+    "--output_dir",     str(RESULTS_DIR),
+    "--device",         "cuda",
+]
+if not UPLOAD_RESULTS_TO_HF:
+    cmd.append("--no_hf_upload")
+if LLM_JUDGE:
+    cmd += ["--llm_judge", "--llm_judge_model", LLM_JUDGE_MODEL]
+    if LLM_JUDGE_BASE_URL:
+        cmd += ["--llm_judge_base_url", LLM_JUDGE_BASE_URL]
+    if LLM_JUDGE_MAX_SAMPLES:
+        cmd += ["--llm_judge_max_samples", LLM_JUDGE_MAX_SAMPLES]
+
+print(f"[gcp_eval] launching: {' '.join(cmd)}", flush=True)
+os.chdir(PROJECT)
+sys.exit(subprocess.call(cmd))

scripts/vertex_eval_job.yaml

@@ -0,0 +1,102 @@
+# Vertex AI Custom Job — CXR-VLM EVALUATION on L4
+#
+# Submit with:
+#   gcloud ai custom-jobs create \
+#     --region=us-central1 \
+#     --display-name=cxr-vlm-eval-run1 \
+#     --config=scripts/vertex_eval_job.yaml
+#
+# Watch logs (streaming):
+#   gcloud ai custom-jobs stream-logs <JOB_ID> --region=us-central1
+#
+# This reuses the SAME Docker image you trained with — no rebuild needed.
+# Quota: uses CustomModelTrainingNvidiaL4Gpus (NOT global GPU quota), so it
+# works even when Workbench / Compute Engine GPU creation is blocked.
+#
+# Output:
+#   - Predictions + metrics_summary.json land under
+#     /workspace/results/<RUN_ID>/ inside the container
+#   - With UPLOAD_RESULTS_TO_HF=1 they're also uploaded to
+#     hieu3636/cxr-vlm-runs/<RUN_ID>/results/
+# (After the job ends the container is destroyed, so HF upload is what you'll
+#  actually retrieve afterward.)
+#
+# Typical wall-clock on L4 for MIMIC-CXR_resized (10k test samples, all tasks):
+#   ~1.5-2.5 hours total (3-5x faster than T4).
+
+workerPoolSpecs:
+- machineSpec:
+    machineType: g2-standard-8           # 8 vCPU, 32GB RAM, 1×L4 24GB
+    acceleratorType: NVIDIA_L4
+    acceleratorCount: 1
+  replicaCount: 1
+  diskSpec:
+    bootDiskType: pd-ssd
+    bootDiskSizeGb: 200                  # enough for MIMIC-CXR_resized (~40GB extracted) + Vicuna cache
+  containerSpec:
+    imageUri: us-central1-docker.pkg.dev/cxr-vlm-thesis/cxr-vlm/cxr-vlm-env:cu128
+    # Same bootstrap pattern as the training job: pull code from HF Hub,
+    # then exec the eval entrypoint.
+    command:
+    - bash
+    - -c
+    - |
+      set -e
+      echo "[bootstrap] downloading code from HF Hub …"
+      python -c "
+      from huggingface_hub import snapshot_download
+      import os
+      snapshot_download('hieu3636/cxr-vlm-code',
+                        repo_type='model',
+                        token=os.environ['HF_TOKEN'],
+                        local_dir='/workspace/code')
+      "
+      cd /workspace/code
+      echo "[bootstrap] exec scripts/gcp_eval_entrypoint.py …"
+      exec python scripts/gcp_eval_entrypoint.py
+    env:
+    # ── Required ────────────────────────────────────────────────────────────
+    - name: HF_TOKEN
+      value: 
+    - name: DATASET_NAME
+      value: MIMIC-CXR_resized           # 'IU-Xray' | 'MIMIC-CXR' | 'MIMIC-CXR_resized'
+    - name: RUN_ID
+      value: MIMIC-CXR_resized_run_1     # which run on HF_RUNS_REPO to evaluate
+    # ── Optional ────────────────────────────────────────────────────────────
+    - name: HF_RUNS_REPO
+      value: hieu3636/cxr-vlm-runs
+    - name: CKPT_PICK
+      value: best                        # 'best' | 'last'
+    - name: TASK
+      value: all                         # 'all'|'findings'|'impression'|'report'|'vqa'
+    - name: SPLIT
+      value: test
+    - name: MAX_NEW_TOKENS
+      value: "300"
+    - name: UPLOAD_RESULTS_TO_HF
+      value: "1"                         # '0' to disable HF upload
+    # Leave REPORT_MODE / IMAGE_MODE blank → use the run's saved snapshot.
+    # Set explicitly only if you need to override (rare).
+    # - name: REPORT_MODE
+    #   value: split_cascade
+    # - name: IMAGE_MODE
+    #   value: all_views_split
+    # Leave BATCH_SIZE blank → auto from GPU profile (L4 -> 4).
+    # - name: BATCH_SIZE
+    #   value: "8"
+    # LLM-as-judge for VQA (extra cost; needs OPENAI_API_KEY).
+    # - name: LLM_JUDGE
+    #   value: "1"
+    # - name: OPENAI_API_KEY
+    #   value: sk-...
+    # - name: LLM_JUDGE_MODEL
+    #   value: gpt-4o-mini
+
+scheduling:
+  # STANDARD on-demand. Eval is short (~2h) so spot interruption isn't worth
+  # the risk — evaluate.py doesn't checkpoint mid-task. If preempted you'd
+  # restart from the beginning.
+  strategy: STANDARD
+  # 6h ceiling. MIMIC-CXR_resized on L4 finishes in ~2h; raise this if
+  # you're evaluating the full MIMIC-CXR with VQA.
+  timeout: 21600s

scripts/vertex_job.yaml

@@ -19,7 +19,7 @@ workerPoolSpecs:
     acceleratorCount: 1
   replicaCount: 1
   diskSpec:
-    bootDiskType: pd-balanced
+    bootDiskType: pd-ssd
     bootDiskSizeGb: 200
   containerSpec:
     imageUri: us-central1-docker.pkg.dev/cxr-vlm-thesis/cxr-vlm/cxr-vlm-env:cu128
@@ -44,7 +44,7 @@ workerPoolSpecs:
       exec python scripts/gcp_entrypoint.py
     env:
     - name: HF_TOKEN
-      value: REPLACE_WITH_YOUR_HF_TOKEN
+      value: 
     - name: DATASET_NAME
       value: MIMIC-CXR_resized
     - name: REPORT_MODE
@@ -65,10 +65,10 @@ workerPoolSpecs:
 # Scheduling: SPOT = ~70% cheaper, may be preempted. Repo has auto-resume so safe.
 # Use STANDARD if you want guaranteed uninterrupted run (3.4× cost).
 scheduling:
-  strategy: SPOT
-  # Max time job can wait in queue for resources before failing.
-  # 1 hour is reasonable for L4 in us-central1 (usually available immediately).
-  maxWaitDuration: 3600s
+  # STANDARD = on-demand (~$0.71/h for g2-standard-8). Switch to SPOT (~$0.21/h)
+  # after requesting `custom_model_training_preemptible_nvidia_l4_gpus` quota —
+  # it's a SEPARATE quota from the regular L4 one.
+  strategy: STANDARD
   # Max wall-clock time the job is allowed to run. If you hit this, it's killed.
   # 24h is a sane ceiling — refuses to silently run forever.
-  timeout: 86400s
+  timeout: 200000s