d

scripts/gcp_entrypoint.py

@@ -20,9 +20,15 @@ Optional env vars (defaults shown):
   S1_EPOCHS           = 2
   S2_EPOCHS           = 7
   MODE                = resume        # 'fresh' | 'resume'
-  EXPLICIT_RUN_ID     = ''            # only matters when MODE=resume
+  EXPLICIT_RUN_ID     = ''            # passes through as --run_id for both modes
   HF_RUNS_REPO        = hieu3636/cxr-vlm-runs
   WORK                = /workspace
+  USE_ITC             = true          # BLIP-2-style Stage 1 contrastive alignment.
+                                      # When true, downloads CXR-BERT text-embed cache from
+                                      # {HF_USER}/cxr-vlm-data/cxr_bert_text_embeddings/.
+                                      # Must have run scripts/precompute_cxrbert_embeddings.ipynb
+                                      # first. Set 'false' to use original LM-loss Stage 1.
+  ITC_TEMPERATURE     = 0.07
 """
 
 from __future__ import annotations
@@ -63,6 +69,8 @@ MODE            = env("MODE", "resume")
 EXPLICIT_RUN_ID = env("EXPLICIT_RUN_ID", "")
 HF_RUNS_REPO    = env("HF_RUNS_REPO", "hieu3636/cxr-vlm-runs")
 WORK            = Path(env("WORK", "/workspace"))
+USE_ITC         = env("USE_ITC", "true").lower() in ("1", "true", "yes")
+ITC_TEMPERATURE = float(env("ITC_TEMPERATURE", "0.07"))
 
 assert DATASET_NAME in ("IU-Xray", "MIMIC-CXR", "MIMIC-CXR_resized"), DATASET_NAME
 assert MODE in ("fresh", "resume"), MODE
@@ -297,6 +305,46 @@ 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")
 
+# ── Stage 1 ITC toggle (mirrors notebook cell-itc) ────────────────────────────
+# Repo defaults to itc.enabled=true with text_embed_cache=null which causes
+# train.py to raise at Stage 1 start. Either supply the precomputed cache OR
+# explicitly disable ITC (falls back to RaDialog-style LM-loss alignment).
+_itc_cache_by_ds = {
+    "IU-Xray":           "cxrbert_text_embeds_iu_xray.pt",
+    "MIMIC-CXR_resized": "cxrbert_text_embeds_mimic_resized.pt",
+    "MIMIC-CXR":         "cxrbert_text_embeds_mimic.pt",
+}
+
+if USE_ITC:
+    cache_filename = _itc_cache_by_ds[DATASET_NAME]
+    print(f"[gcp_entrypoint] downloading ITC cache {cache_filename} …")
+    cache_path = hf_hub_download(
+        repo_id=f"{HF_USER}/cxr-vlm-data",
+        repo_type="dataset",
+        filename=f"cxr_bert_text_embeddings/{cache_filename}",
+        token=HF_TOKEN,
+    )
+    print(f"[gcp_entrypoint] ITC cache → {cache_path}")
+
+    # Stage-1 contrastive batch (Vicuna not loaded → can go much bigger)
+    if   _vram_gb >= 35: itc_bs = 256
+    elif _vram_gb >= 22: itc_bs = 192
+    elif _vram_gb >= 14: itc_bs = 96
+    else:                itc_bs = 32
+
+    train_cfg.stage1.itc.enabled                     = True
+    train_cfg.stage1.itc.text_embed_cache            = cache_path
+    train_cfg.stage1.itc.temperature                 = ITC_TEMPERATURE
+    train_cfg.stage1.itc.fallback_to_impression      = True
+    train_cfg.stage1.itc.per_device_train_batch_size = itc_bs
+    train_cfg.stage1.itc.per_device_eval_batch_size  = itc_bs
+    train_cfg.stage1.itc.gradient_accumulation_steps = 1
+    print(f"[gcp_entrypoint] ITC ENABLED  batch={itc_bs}  temp={ITC_TEMPERATURE}")
+else:
+    # Force-disable so train.py uses original LM-loss Stage 1.
+    train_cfg.stage1.itc.enabled = False
+    print("[gcp_entrypoint] ITC DISABLED — Stage 1 uses LM-loss alignment.")
+
 OmegaConf.save(train_cfg, train_cfg_path)
 OmegaConf.save(model_cfg, model_cfg_path)
 print("[gcp_entrypoint] configs patched.")