f

data/mimic_cxr_resized_builder.py

@@ -289,7 +289,37 @@ def build_mimic_cxr_resized_instruct_json(
                                         "structured_findings": structured})
 
     # ── Pass 2: optional VQA attach ─────────────────────────────────────────
+    # Match policy: VQA `image_path` and manifest `image_relpath` should be the
+    # same string, but historical subset-build scripts differ in whether the
+    # `files/` prefix is included. Build a normalized image_index that maps
+    # BOTH the original key and its files/-stripped variant back to the
+    # canonical manifest path. On a lookup miss, also try stripping/adding
+    # `files/` from the VQA path. This handles the three known formats:
+    #   manifest "files/p10/...", vqa "files/p10/..."    → exact match
+    #   manifest "files/p10/...", vqa "p10/..."          → +files/ on VQA
+    #   manifest "p10/...",       vqa "files/p10/..."    → -files/ on VQA
+    def _strip_files(p: str) -> str:
+        return p[len("files/"):] if p.startswith("files/") else p
+
+    norm_index: Dict[str, str] = {}
+    for k in image_index:
+        norm_index[k] = k                  # exact
+        norm_index[_strip_files(k)] = k    # without files/
+
+    def _resolve_vqa_path(p: str) -> Optional[str]:
+        if p in norm_index:
+            return norm_index[p]
+        stripped = _strip_files(p)
+        if stripped in norm_index:
+            return norm_index[stripped]
+        with_prefix = f"files/{p}" if not p.startswith("files/") else p
+        if with_prefix in norm_index:
+            return norm_index[with_prefix]
+        return None
+
     n_vqa = n_vqa_dropped = 0
+    _logged_diag = False
+    _dropped_samples: List[str] = []     # for diagnostic on total drop
     if vqa_dir:
         vqa_dir = Path(vqa_dir)
         for fname, split_label in _VQA_FILES:
@@ -297,10 +327,27 @@ def build_mimic_cxr_resized_instruct_json(
             if not vqa_file.is_file():
                 print(f"[mimic_cxr_resized_builder] VQA missing: {vqa_file} — skipping {split_label}")
                 continue
-            for row in json.load(open(vqa_file, encoding="utf-8")):
+            rows = json.load(open(vqa_file, encoding="utf-8"))
+
+            # One-shot diagnostic — print sample paths from both sides so any
+            # remaining mismatch is immediately visible in the log.
+            if not _logged_diag:
+                _logged_diag = True
+                print(f"[mimic_cxr_resized_builder] VQA path-format diagnostic "
+                      f"(file: {fname}):")
+                _idx_samples = list(image_index.keys())[:3]
+                _vqa_samples = [str(r.get("image_path", "")).lstrip("/")
+                                for r in rows[:3]]
+                print(f"  manifest image_relpath samples : {_idx_samples}")
+                print(f"  vqa      image_path     samples : {_vqa_samples}")
+
+            for row in rows:
                 img_path = str(row.get("image_path", "")).lstrip("/")
-                if img_path not in image_index:
+                canonical = _resolve_vqa_path(img_path)
+                if canonical is None:
                     n_vqa_dropped += 1
+                    if len(_dropped_samples) < 10:
+                        _dropped_samples.append(img_path)
                     continue
                 ans = row.get("answer", [])
                 answer = (", ".join(map(str, ans)) if isinstance(ans, list)
@@ -315,7 +362,7 @@ def build_mimic_cxr_resized_instruct_json(
                 subj_str  = f"p{subj}" if subj and not subj.startswith("p") else subj
                 study_str = f"s{study}" if study and not study.startswith("s") else study
                 samples.append({
-                    "image_path": img_path, "image_paths": None,
+                    "image_path": canonical, "image_paths": None,
                     "task": "vqa", "target": answer,
                     "question": row["question"],
                     "structured_findings": structured,
@@ -325,6 +372,21 @@ def build_mimic_cxr_resized_instruct_json(
                 })
                 n_vqa += 1
 
+    # Loud warning if ALL vqa rows were dropped — this almost always means a
+    # path-format mismatch the heuristic above didn't cover, e.g. different
+    # file extension (.dcm vs .jpg) or different study/dicom id encoding.
+    if vqa_dir and n_vqa == 0 and n_vqa_dropped > 0:
+        print(f"[mimic_cxr_resized_builder] !! ALL {n_vqa_dropped:,} VQA rows "
+              f"dropped — path mismatch not handled by the heuristic.")
+        print(f"  First 10 unresolved VQA image_paths:")
+        for p in _dropped_samples:
+            print(f"    {p!r}")
+        print(f"  First 5 manifest image_relpath keys:")
+        for k in list(image_index.keys())[:5]:
+            print(f"    {k!r}")
+        print(f"  Inspect both sets and adjust _resolve_vqa_path() in "
+              f"mimic_cxr_resized_builder.py.")
+
     # ── Write ───────────────────���───────────────────────────────────────────
     output_path.parent.mkdir(parents=True, exist_ok=True)
     with open(output_path, "w", encoding="utf-8") as f:

scripts/_build_eval_notebook.py

@@ -732,8 +732,67 @@ print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))
 """))
 
 
+CELLS.append(md("eval-verify-md", """\
+## 7. Verify dataset before running eval
+
+Quick pre-flight check: triggers the instruct-JSON builder (if not cached), then prints **per-split × per-task** sample counts. Catches issues like "VQA = 0 samples" (path-format mismatch in the builder) **before** spending 2h on an eval that has nothing to evaluate.
+
+If `vqa` column shows 0 in the test split:
+- the model was likely **not trained on VQA** either (same JSON cache used both ways)
+- options: skip VQA via `TASK='findings'` then run a second job with `TASK='impression'`, OR fix the builder + retrain.
+"""))
+
+
+CELLS.append(code("eval-verify", """\
+import json as _json
+from collections import Counter
+from utils.dataset_resolver import resolve_dataset_spec
+
+# Reload the patched config snapshot the eval subprocess will see.
+train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
+spec = resolve_dataset_spec(train_cfg)
+print(f'Dataset      : {spec.dataset_name}')
+print(f'Instruct JSON: {spec.instruct_json}')
+print(f'Image root   : {spec.image_root}')
+print(f'Tasks (cfg)  : {spec.tasks}')
+print()
+
+# Load the JSON the dataset module will read.
+samples = _json.loads(open(spec.instruct_json, encoding='utf-8').read())
+print(f'Total samples in JSON: {len(samples):,}')
+
+# Cross-tab: (split, task) -> count
+ctab = Counter((s['split'], s['task']) for s in samples)
+splits = sorted({k[0] for k in ctab})
+tasks  = sorted({k[1] for k in ctab})
+
+# Pretty table
+col_w = max(10, max(len(t) for t in tasks) + 2)
+hdr = f'{\"split\":<10} | ' + ' | '.join(f'{t:>{col_w}}' for t in tasks) + ' |     total'
+print(hdr); print('-' * len(hdr))
+for sp in splits:
+    row_vals = [ctab.get((sp, t), 0) for t in tasks]
+    total = sum(row_vals)
+    print(f'{sp:<10} | ' + ' | '.join(f'{v:>{col_w},}' for v in row_vals)
+          + f' | {total:>9,}')
+
+# Loud warning if VQA is expected but missing.
+test_vqa = ctab.get(('test', 'vqa'), 0)
+if 'vqa' in spec.tasks and test_vqa == 0:
+    print()
+    print('!! WARNING: vqa task is configured but TEST split has 0 vqa samples.')
+    print('   This usually means the dataset builder dropped all VQA rows due to')
+    print('   path-format mismatch between vqa/*.json and manifest_*.csv.')
+    print('   Check the builder log above for the line:')
+    print('     [mimic_cxr_resized_builder] vqa added/dropped : N / M')
+    print('   If N=0 the model was likely NOT trained on VQA either — same JSON cache.')
+elif 'vqa' in spec.tasks:
+    print(f'\\nVQA in test split: {test_vqa:,} samples — OK')
+"""))
+
+
 CELLS.append(md("eval-run-md", """\
-## 7. Run evaluation
+## 8. 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.
@@ -795,7 +854,7 @@ python -u -m evaluation.evaluate \\
 
 
 CELLS.append(md("eval-summary-md", """\
-## 8. Show the metrics summary
+## 9. Show the metrics summary
 """))
 
 

scripts/cxrvlm_colab_eval.ipynb

@@ -763,12 +763,79 @@
     "print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "eval-verify-md",
+   "metadata": {},
+   "source": [
+    "## 7. Verify dataset before running eval\n",
+    "\n",
+    "Quick pre-flight check: triggers the instruct-JSON builder (if not cached), then prints **per-split × per-task** sample counts. Catches issues like \"VQA = 0 samples\" (path-format mismatch in the builder) **before** spending 2h on an eval that has nothing to evaluate.\n",
+    "\n",
+    "If `vqa` column shows 0 in the test split:\n",
+    "- the model was likely **not trained on VQA** either (same JSON cache used both ways)\n",
+    "- options: skip VQA via `TASK='findings'` then run a second job with `TASK='impression'`, OR fix the builder + retrain.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "eval-verify",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import json as _json\n",
+    "from collections import Counter\n",
+    "from utils.dataset_resolver import resolve_dataset_spec\n",
+    "\n",
+    "# Reload the patched config snapshot the eval subprocess will see.\n",
+    "train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\n",
+    "spec = resolve_dataset_spec(train_cfg)\n",
+    "print(f'Dataset      : {spec.dataset_name}')\n",
+    "print(f'Instruct JSON: {spec.instruct_json}')\n",
+    "print(f'Image root   : {spec.image_root}')\n",
+    "print(f'Tasks (cfg)  : {spec.tasks}')\n",
+    "print()\n",
+    "\n",
+    "# Load the JSON the dataset module will read.\n",
+    "samples = _json.loads(open(spec.instruct_json, encoding='utf-8').read())\n",
+    "print(f'Total samples in JSON: {len(samples):,}')\n",
+    "\n",
+    "# Cross-tab: (split, task) -> count\n",
+    "ctab = Counter((s['split'], s['task']) for s in samples)\n",
+    "splits = sorted({k[0] for k in ctab})\n",
+    "tasks  = sorted({k[1] for k in ctab})\n",
+    "\n",
+    "# Pretty table\n",
+    "col_w = max(10, max(len(t) for t in tasks) + 2)\n",
+    "hdr = f'{\"split\":<10} | ' + ' | '.join(f'{t:>{col_w}}' for t in tasks) + ' |     total'\n",
+    "print(hdr); print('-' * len(hdr))\n",
+    "for sp in splits:\n",
+    "    row_vals = [ctab.get((sp, t), 0) for t in tasks]\n",
+    "    total = sum(row_vals)\n",
+    "    print(f'{sp:<10} | ' + ' | '.join(f'{v:>{col_w},}' for v in row_vals)\n",
+    "          + f' | {total:>9,}')\n",
+    "\n",
+    "# Loud warning if VQA is expected but missing.\n",
+    "test_vqa = ctab.get(('test', 'vqa'), 0)\n",
+    "if 'vqa' in spec.tasks and test_vqa == 0:\n",
+    "    print()\n",
+    "    print('!! WARNING: vqa task is configured but TEST split has 0 vqa samples.')\n",
+    "    print('   This usually means the dataset builder dropped all VQA rows due to')\n",
+    "    print('   path-format mismatch between vqa/*.json and manifest_*.csv.')\n",
+    "    print('   Check the builder log above for the line:')\n",
+    "    print('     [mimic_cxr_resized_builder] vqa added/dropped : N / M')\n",
+    "    print('   If N=0 the model was likely NOT trained on VQA either — same JSON cache.')\n",
+    "elif 'vqa' in spec.tasks:\n",
+    "    print(f'\\nVQA in test split: {test_vqa:,} samples — OK')\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "eval-run-md",
    "metadata": {},
    "source": [
-    "## 7. Run evaluation\n",
+    "## 8. 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",
@@ -838,7 +905,7 @@
    "id": "eval-summary-md",
    "metadata": {},
    "source": [
-    "## 8. Show the metrics summary\n"
+    "## 9. Show the metrics summary\n"
    ]
   },
   {

scripts/gcp_finetune_vqa_entrypoint.py

@@ -0,0 +1,455 @@
+"""GCP Vertex AI Custom Job entrypoint — **VQA finetune** from an existing run.
+
+Take a fully-trained run (stage 2 done) and continue training Stage 2 with
+**VQA-heavy task weights** + low LR + few epochs, into a NEW run_id (so the
+original checkpoint stays untouched on HF Hub).
+
+Continuation strategy:
+  - run_1/stage2/best/   ← already-trained projection + LoRA on HF
+       checkpoint_projection.pt
+       checkpoint_lora/
+       checkpoint_chexpert_classifier.pt        (if present)
+
+  We download those and place them at:
+       run_2/stage1_projection/stage1_final_projection.pt   ← renamed
+       run_2/stage1_projection/stage1_final_lora/           ← renamed
+       run_2/stage1_projection/stage1_final_chexpert_classifier.pt
+
+  Then train.py:
+    1. detect_resume_point sees stage1_final_projection.pt → ("stage2", None)
+       → skips Stage 1 entirely
+    2. Builds the model, calls load_checkpoint(stage1_final.pt) which loads
+       BOTH the projection AND the LoRA from the renamed files (load_checkpoint
+       derives the LoRA dir from the .pt stem).
+    3. Runs Stage 2 fresh with the new task weights + LR + epochs, starting
+       from the loaded weights. Optimizer state is reset — that's deliberate
+       (different task mix; old momentum would point the wrong way).
+
+Required env vars:
+  HF_TOKEN              — HuggingFace token (read+write)
+  DATASET_NAME          — 'IU-Xray' | 'MIMIC-CXR' | 'MIMIC-CXR_resized'
+  SOURCE_RUN_ID         — run on HF_RUNS_REPO whose stage2/best is the seed
+                          e.g. 'MIMIC-CXR_resized_run_1'
+  TARGET_RUN_ID         — NEW run id to write into, e.g.
+                          'MIMIC-CXR_resized_run_2'
+
+Optional env vars (defaults shown — tuned for "vqa-heavy with rehearsal"):
+  HF_USER               = hieu3636
+  HF_RUNS_REPO          = hieu3636/cxr-vlm-runs
+  SOURCE_CKPT_PICK      = best                  # 'best' | 'last'
+  REPORT_MODE           =                       # blank → from source run's snapshot
+  IMAGE_MODE            =                       # blank → from source run's snapshot
+  W_FINDINGS            = 0.15                  # rehearsal weight
+  W_IMPRESSION          = 0.10                  # rehearsal weight
+  W_VQA                 = 0.75                  # focus
+  S2_EPOCHS             = 3
+  S2_LR                 = 5e-5                  # 4× lower than original 2e-4
+  STRICT_VQA_REQUIRED   = 1                     # fail fast if VQA samples = 0
+  WORK                  = /workspace
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import zipfile
+from pathlib import Path
+
+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")
+
+
+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_finetune_vqa] 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)
+SOURCE_RUN_ID        = env("SOURCE_RUN_ID", required=True)
+TARGET_RUN_ID        = env("TARGET_RUN_ID", required=True)
+HF_USER              = env("HF_USER", "hieu3636")
+HF_RUNS_REPO         = env("HF_RUNS_REPO", "hieu3636/cxr-vlm-runs")
+SOURCE_CKPT_PICK     = env("SOURCE_CKPT_PICK", "best")
+REPORT_MODE_OVERRIDE = env("REPORT_MODE", "")
+IMAGE_MODE_OVERRIDE  = env("IMAGE_MODE", "")
+W_FINDINGS           = float(env("W_FINDINGS",   "0.15"))
+W_IMPRESSION         = float(env("W_IMPRESSION", "0.10"))
+W_VQA                = float(env("W_VQA",        "0.75"))
+S2_EPOCHS            = int(env("S2_EPOCHS", "3"))
+S2_LR                = float(env("S2_LR",   "5e-5"))
+STRICT_VQA_REQUIRED  = env("STRICT_VQA_REQUIRED", "1") in ("1", "true", "True")
+WORK                 = Path(env("WORK", "/workspace"))
+
+assert DATASET_NAME in ("IU-Xray", "MIMIC-CXR", "MIMIC-CXR_resized"), DATASET_NAME
+assert SOURCE_CKPT_PICK in ("best", "last"), SOURCE_CKPT_PICK
+assert TARGET_RUN_ID != SOURCE_RUN_ID, \
+    f"TARGET_RUN_ID must differ from SOURCE_RUN_ID — refusing to overwrite the source run."
+
+PROJECT       = Path(__file__).resolve().parent.parent
+DATA_SRC      = WORK / "data"
+RUN_PULL_ROOT = WORK / "run_pull"
+CKPT_ROOT     = WORK / "ckpt"
+for d in (DATA_SRC, RUN_PULL_ROOT, CKPT_ROOT):
+    d.mkdir(parents=True, exist_ok=True)
+
+print(f"[gcp_finetune_vqa] PROJECT       = {PROJECT}")
+print(f"[gcp_finetune_vqa] WORK          = {WORK}")
+print(f"[gcp_finetune_vqa] DATASET_NAME  = {DATASET_NAME}")
+print(f"[gcp_finetune_vqa] SOURCE_RUN_ID = {SOURCE_RUN_ID}  (ckpt: stage2/{SOURCE_CKPT_PICK})")
+print(f"[gcp_finetune_vqa] TARGET_RUN_ID = {TARGET_RUN_ID}  (NEW — original stays untouched)")
+print(f"[gcp_finetune_vqa] Task mix      = findings:{W_FINDINGS}  impression:{W_IMPRESSION}  vqa:{W_VQA}")
+print(f"[gcp_finetune_vqa] S2_EPOCHS={S2_EPOCHS}  S2_LR={S2_LR}")
+
+# ── 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_finetune_vqa] {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"
+        )
+        tar_files = sorted(
+            f for f in all_files
+            if f.startswith("MIMIC-CXR_resized/") and f.endswith(".tar")
+        )
+        print(f"[gcp_finetune_vqa] {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_finetune_vqa]  [{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_finetune_vqa] {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_finetune_vqa] 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_finetune_vqa] {marker} already present — skipping download.")
+
+print(f"[gcp_finetune_vqa] DATA_SRC contents: {sorted(os.listdir(DATA_SRC))}")
+
+# ── 3) Pull SOURCE run's stage2/best + config snapshot from HF runs repo ─────
+print(f"[gcp_finetune_vqa] pulling {SOURCE_RUN_ID}/{{configs,stage2/{SOURCE_CKPT_PICK}}} "
+      f"from {HF_RUNS_REPO} …")
+snapshot_download(
+    repo_id=HF_RUNS_REPO,
+    repo_type="model",
+    token=HF_TOKEN,
+    allow_patterns=[
+        f"{SOURCE_RUN_ID}/configs/**",
+        f"{SOURCE_RUN_ID}/run_meta.json",
+        f"{SOURCE_RUN_ID}/stage2/{SOURCE_CKPT_PICK}/**",
+    ],
+    local_dir=str(RUN_PULL_ROOT),
+)
+SRC_DIR     = RUN_PULL_ROOT / SOURCE_RUN_ID
+SRC_S2_DIR  = SRC_DIR / "stage2" / SOURCE_CKPT_PICK
+SRC_PROJ    = SRC_S2_DIR / "checkpoint_projection.pt"
+SRC_LORA    = SRC_S2_DIR / "checkpoint_lora"
+SRC_CHEXPRT = SRC_S2_DIR / "checkpoint_chexpert_classifier.pt"
+
+assert SRC_PROJ.is_file(), f"source projection missing: {SRC_PROJ}"
+assert (SRC_LORA / "adapter_config.json").is_file(), \
+    f"source LoRA dir missing: {SRC_LORA}"
+
+# ── 4) Seed TARGET_RUN_ID dir: place source ckpt as stage1_final_* so that
+#    detect_resume_point() returns ("stage2", None) and load_checkpoint() picks
+#    up BOTH projection AND lora from the renamed files. ─────────────────────
+TARGET_DIR  = CKPT_ROOT / TARGET_RUN_ID
+TGT_S1_DIR  = TARGET_DIR / "stage1_projection"
+TGT_S1_DIR.mkdir(parents=True, exist_ok=True)
+
+shutil.copy2(SRC_PROJ, TGT_S1_DIR / "stage1_final_projection.pt")
+print(f"[gcp_finetune_vqa] seeded projection -> {TGT_S1_DIR / 'stage1_final_projection.pt'}")
+
+tgt_lora = TGT_S1_DIR / "stage1_final_lora"
+if tgt_lora.exists():
+    shutil.rmtree(tgt_lora)
+shutil.copytree(SRC_LORA, tgt_lora)
+print(f"[gcp_finetune_vqa] seeded LoRA       -> {tgt_lora}")
+
+if SRC_CHEXPRT.is_file():
+    shutil.copy2(SRC_CHEXPRT, TGT_S1_DIR / "stage1_final_chexpert_classifier.pt")
+    print(f"[gcp_finetune_vqa] seeded chexpert  -> "
+          f"{TGT_S1_DIR / 'stage1_final_chexpert_classifier.pt'}")
+
+# ── 5) Build configs (start from source snapshot to preserve architecture) ───
+import torch  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+
+SAVED_CFG_DIR   = SRC_DIR / "configs"
+SAVED_TRAIN_CFG = SAVED_CFG_DIR / "train_config.yaml"
+SAVED_MODEL_CFG = SAVED_CFG_DIR / "model_config.yaml"
+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_finetune_vqa] train_cfg <- {SAVED_TRAIN_CFG}")
+else:
+    train_cfg = OmegaConf.load(repo_train_cfg_path)
+    print(f"[gcp_finetune_vqa] train_cfg <- repo default")
+
+if SAVED_MODEL_CFG.is_file():
+    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
+    print(f"[gcp_finetune_vqa] model_cfg <- {SAVED_MODEL_CFG}")
+else:
+    model_cfg = OmegaConf.load(repo_model_cfg_path)
+
+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_finetune_vqa] 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_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)
+
+# ── Task mix: VQA-heavy with rehearsal ───────────────────────────────────────
+train_cfg.tasks.findings_generation.enabled    = W_FINDINGS > 0
+train_cfg.tasks.findings_generation.weight     = W_FINDINGS
+train_cfg.tasks.impression_generation.enabled  = W_IMPRESSION > 0
+train_cfg.tasks.impression_generation.weight   = W_IMPRESSION
+train_cfg.tasks.vqa.enabled                    = W_VQA > 0
+train_cfg.tasks.vqa.weight                     = W_VQA
+
+# ── Stage 1 disabled (we seeded its outputs from source). Stage 2 = short
+#    finetune with low LR. Stage1 ITC explicitly off — irrelevant here. ─────
+train_cfg.stage1.enabled = False
+if "itc" in train_cfg.stage1:
+    train_cfg.stage1.itc.enabled = False
+train_cfg.stage2.enabled       = True
+train_cfg.stage2.num_epochs    = S2_EPOCHS
+train_cfg.stage2.learning_rate = S2_LR
+
+# ── GPU auto-profile (mirrors gcp_entrypoint.py) ────────────────────────────
+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:
+    _profile = dict(label="A100/H100 80GB",
+                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
+                    gradient_accumulation_steps=2, dataloader_num_workers=16,
+                    gradient_checkpointing=False)
+elif _vram_gb >= 35:
+    _profile = dict(label="A100 40GB",
+                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
+                    gradient_accumulation_steps=2, dataloader_num_workers=12,
+                    gradient_checkpointing=False)
+elif _vram_gb >= 22:
+    _profile = dict(label="3090 / L4 / A10 (24GB)",
+                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
+                    gradient_accumulation_steps=2, dataloader_num_workers=8,
+                    gradient_checkpointing=True)
+elif _vram_gb >= 14:
+    _profile = dict(label="T4 / V100 (15-16GB)",
+                    per_device_train_batch_size=1, per_device_eval_batch_size=1,
+                    gradient_accumulation_steps=16, dataloader_num_workers=2,
+                    gradient_checkpointing=True)
+else:
+    _profile = dict(label=f"unknown ({_vram_gb:.0f}GB)",
+                    per_device_train_batch_size=1, per_device_eval_batch_size=1,
+                    gradient_accumulation_steps=16, dataloader_num_workers=2,
+                    gradient_checkpointing=True)
+
+_profile["bf16"] = bool(_bf16_ok)
+_profile["fp16"] = not _bf16_ok
+_profile["attn_implementation"] = (
+    "flash_attention_2" if (_fa2_ok and _flash_attn_installed) else "sdpa"
+)
+_profile["optim"] = "paged_adamw_8bit" if _cap >= (8, 0) else "adamw_torch"
+_profile["bnb_4bit_compute_dtype"] = "bfloat16" if _bf16_ok else "float16"
+_profile["torch_dtype"]            = "bfloat16" if _bf16_ok else "float16"
+
+print(f"[gcp_finetune_vqa] 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_finetune_vqa] → {_profile['label']}")
+
+train_cfg.training.per_device_train_batch_size   = _profile["per_device_train_batch_size"]
+train_cfg.training.per_device_eval_batch_size    = _profile["per_device_eval_batch_size"]
+train_cfg.training.gradient_accumulation_steps   = _profile["gradient_accumulation_steps"]
+train_cfg.training.dataloader_num_workers        = _profile["dataloader_num_workers"]
+train_cfg.training.fp16                          = _profile["fp16"]
+train_cfg.training.bf16                          = _profile["bf16"]
+train_cfg.training.dataloader_pin_memory         = True
+train_cfg.training.dataloader_persistent_workers = True
+train_cfg.training.optim                         = _profile["optim"]
+
+model_cfg.llm.attn_implementation       = _profile["attn_implementation"]
+model_cfg.llm.gradient_checkpointing    = _profile["gradient_checkpointing"]
+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.load_in_8bit              = False
+model_cfg.llm.load_in_4bit              = True
+
+# CheXpert classifier — enable iff its checkpoint came along with source run.
+if (TGT_S1_DIR / "stage1_final_chexpert_classifier.pt").is_file():
+    model_cfg.chexpert_classifier.enabled = True
+else:
+    model_cfg.chexpert_classifier.enabled = False
+
+# HF Hub uploads → TARGET_RUN_ID (so source run_1 stays untouched).
+train_cfg.wandb.enabled         = False
+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
+train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / "run_id.txt")
+
+# Pin run_id = target so resolve_run_id picks it up.
+(CKPT_ROOT / "run_id.txt").write_text(TARGET_RUN_ID)
+print(f"[gcp_finetune_vqa] pinned run_id = {TARGET_RUN_ID}")
+
+OmegaConf.save(train_cfg, repo_train_cfg_path)
+OmegaConf.save(model_cfg, repo_model_cfg_path)
+print("[gcp_finetune_vqa] configs patched.")
+
+# ── 6) Pre-flight check: VQA samples > 0 ─────────────────────────────────────
+# Trigger the builder, then verify the train split actually has vqa samples.
+# Bail fast — saves hours of compute if the path-fix didn't work.
+import json as _json
+from utils.dataset_resolver import resolve_dataset_spec  # noqa: E402
+
+print("[gcp_finetune_vqa] pre-flight: triggering dataset builder + verifying VQA …")
+spec = resolve_dataset_spec(train_cfg)
+samples = _json.load(open(spec.instruct_json, encoding="utf-8"))
+train_counts = {}
+for s in samples:
+    if s.get("split") == "train":
+        train_counts[s["task"]] = train_counts.get(s["task"], 0) + 1
+print(f"[gcp_finetune_vqa] train-split task counts: {train_counts}")
+
+if STRICT_VQA_REQUIRED and train_counts.get("vqa", 0) == 0:
+    print(f"[gcp_finetune_vqa] !! FATAL: train-split has 0 VQA samples.")
+    print(f"   The dataset builder dropped all VQA rows — path-format mismatch.")
+    print(f"   Inspect the builder log above for the line:")
+    print(f"     [mimic_cxr_resized_builder] vqa added/dropped : N / M")
+    print(f"   Fix the builder, push to HF, then re-submit. Set "
+          f"STRICT_VQA_REQUIRED=0 to bypass this check (not recommended).")
+    sys.exit(2)
+
+# ── 7) Launch training (Stage 2 only — Stage 1 seeded from source) ──────────
+cmd = [
+    "python", "-u", "-m", "training.train",
+    "--model_config", str(repo_model_cfg_path),
+    "--train_config", str(repo_train_cfg_path),
+    "--mode",         "resume",
+    "--run_id",       TARGET_RUN_ID,
+]
+print(f"[gcp_finetune_vqa] launching: {' '.join(cmd)}", flush=True)
+os.chdir(PROJECT)
+sys.exit(subprocess.call(cmd))

scripts/vertex_finetune_vqa_job.yaml

@@ -0,0 +1,87 @@
+# Vertex AI Custom Job — CXR-VLM **VQA finetune** on L4
+#
+# Takes an existing trained run (SOURCE_RUN_ID) and continues Stage 2 with
+# VQA-heavy task weights into a NEW run (TARGET_RUN_ID). Source run on HF is
+# untouched.
+#
+# Submit:
+#   gcloud ai custom-jobs create \
+#     --region=us-central1 \
+#     --display-name=cxr-vlm-vqa-finetune \
+#     --config=scripts/vertex_finetune_vqa_job.yaml
+#
+# Stream logs:
+#   gcloud ai custom-jobs stream-logs <JOB_ID> --region=us-central1
+#
+# Output:
+#   - Stage 2 checkpoints land at hieu3636/cxr-vlm-runs/<TARGET_RUN_ID>/stage2/{last,best}/
+#   - Original run_1 untouched.
+#
+# Pre-flight: the entrypoint refuses to launch training if the train split has
+# 0 VQA samples after the builder runs (i.e. path-fix didn't work). You see
+# the failure in ~30 min instead of after 3h of wasted GPU.
+
+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
+  containerSpec:
+    imageUri: us-central1-docker.pkg.dev/cxr-vlm-thesis/cxr-vlm/cxr-vlm-env:cu128
+    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_finetune_vqa_entrypoint.py …"
+      exec python scripts/gcp_finetune_vqa_entrypoint.py
+    env:
+    # ── Required ────────────────────────────────────────────────────────────
+    - name: HF_TOKEN
+      value: 
+    - name: DATASET_NAME
+      value: MIMIC-CXR_resized
+    - name: SOURCE_RUN_ID
+      value: MIMIC-CXR_resized_run_1     # the trained run to fine-tune from
+    - name: TARGET_RUN_ID
+      value: MIMIC-CXR_resized_run_2     # NEW — won't collide with run_1
+    # ── Optional ────────────────────────────────────────────────────────────
+    - name: HF_RUNS_REPO
+      value: hieu3636/cxr-vlm-runs
+    - name: SOURCE_CKPT_PICK
+      value: best                        # 'best' | 'last'
+    # ── Task mix (VQA-heavy with rehearsal) ─────────────────────────────────
+    - name: W_FINDINGS
+      value: "0.15"                      # rehearsal — keeps findings skill
+    - name: W_IMPRESSION
+      value: "0.10"                      # rehearsal — keeps impression skill
+    - name: W_VQA
+      value: "0.75"                      # focus on VQA
+    # ── Schedule ────────────────────────────────────────────────────────────
+    - name: S2_EPOCHS
+      value: "3"                         # short finetune (was 7 originally)
+    - name: S2_LR
+      value: "5e-5"                      # 4× lower than original 2e-4
+    # Safety: refuse to train if VQA build still drops everything.
+    - name: STRICT_VQA_REQUIRED
+      value: "1"
+
+scheduling:
+  strategy: STANDARD
+  # Stage 2 finetune (3 epochs, mixed batch ~80k samples) on L4 ≈ 3-4h.
+  # 6h ceiling leaves slack for data pull + builder.
+  timeout: 21600s