f

configs/train_config.yaml

@@ -222,7 +222,7 @@ training:
 # ── Stage 1: Train projection only ───────────
 stage1:
   enabled: true
-  num_epochs: 2
+  num_epochs: 3
   learning_rate: 1.0e-3
   freeze_llm: true
   freeze_encoder: true
@@ -252,7 +252,7 @@ stage1:
 # ── Stage 2: Full instruction tuning ─────────
 stage2:
   enabled: true
-  num_epochs: 10
+  num_epochs: 5
   learning_rate: 2.0e-4
   freeze_encoder: true
   freeze_llm: false                    # train via LoRA

data/dataset.py

@@ -389,7 +389,13 @@ class CXRInstructDataset(Dataset):
             input_ids: (L,)    L ≤ cutoff_len
             labels:    (L,)    with -100 for prompt positions
         """
-        full_text      = prompt + " " + target
+        # Append EOS to target so model learns to terminate generation.
+        # LlamaTokenizer.encode(add_special_tokens=True) prepends <s> (BOS)
+        # but does NOT append </s> (EOS) by default. Without an EOS in the
+        # labels, the model never learns "report is done → stop" and at
+        # inference time it runs until max_new_tokens, looping on high-prob
+        # template phrases. Adding it once here costs ~1 token per sample.
+        full_text      = prompt + " " + target + self.tokenizer.eos_token
         prompt_encoded = self.tokenizer.encode(prompt, add_special_tokens=True)
         full_encoded   = self.tokenizer.encode(
             full_text,

model/cxr_vlm.py

@@ -429,6 +429,7 @@ class CXRVisionLanguageModel(nn.Module):
         temperature:    float = 0.1,
         do_sample:      bool = False,
         num_beams:      int = 1,
+        **gen_kwargs,
     ) -> List[str]:
         """
         Generate text for a batch of images and prompts.
@@ -467,6 +468,12 @@ class CXRVisionLanguageModel(nn.Module):
         )
 
         # Generate
+        # Extra HF generate kwargs (repetition_penalty, no_repeat_ngram_size,
+        # top_p, top_k, length_penalty, early_stopping, …) flow through
+        # gen_kwargs so we don't have to add them one by one. Anti-repetition
+        # knobs in particular fix the greedy-decoding loop where the model
+        # latches onto a high-prob template phrase (e.g. "The vertebral body
+        # heights are preserved.") and emits it until max_new_tokens.
         output_ids = self.llm.generate(
             inputs_embeds  = inputs_embeds,
             attention_mask = attention_mask,
@@ -476,6 +483,7 @@ class CXRVisionLanguageModel(nn.Module):
             num_beams      = num_beams,
             pad_token_id   = self.tokenizer.pad_token_id,
             eos_token_id   = self.tokenizer.eos_token_id,
+            **gen_kwargs,
         )
 
         # Decode — strip the prompt part

scripts/_build_eos_test_notebook.py

@@ -0,0 +1,673 @@
+"""
+Builds scripts/cxrvlm_eos_test.ipynb — A/B test for the EOS-in-labels fix.
+
+Hypothesis (from dataset.py:392): training labels don't end with </s>, so the
+model never learned to emit EOS, so generation runs to max_new_tokens and
+loops on high-prob template phrases.
+
+Test design (cheap — ~30-45 min on L4):
+  1. Load existing trained checkpoint (run_id picked in selectors).
+  2. Pick 5 test images + their GT findings.
+  3. PHASE A — generate on the 5 images with several settings, record:
+       avg_gen_tokens, hit_max_rate, distinct_sentence_ratio,
+       sample outputs.
+  4. Mini-FT for 1 epoch on 100 train samples using a dataset SUBCLASS that
+     appends `tokenizer.eos_token` to every target (the only change).
+  5. PHASE B — re-generate on the SAME 5 images with the SAME settings.
+  6. Print before/after comparison table + side-by-side text.
+
+Interpretation:
+  - If hit_max_rate drops sharply and avg_gen_tokens decreases → EOS hypothesis
+    confirmed dominant (the model CAN learn EOS in ~100 LoRA steps).
+  - If barely changes → another factor (greedy bias, prompt mismatch, quant
+    noise) is more dominant than EOS-in-labels.
+
+Run from repo root:
+  python scripts/_build_eos_test_notebook.py
+"""
+
+from pathlib import Path
+import nbformat as nbf
+
+
+def md(text: str):
+    return nbf.v4.new_markdown_cell(text)
+
+
+def code(text: str):
+    return nbf.v4.new_code_cell(text)
+
+
+cells = []
+
+# ─── Title ──────────────────────────────────────────────────────────────────
+cells.append(md("""# CXR-VLM — EOS Fix A/B Test
+
+**Hypothesis:** `data/dataset.py` tokenizes `prompt + " " + target` without appending `</s>` to the target. `LlamaTokenizer.encode(add_special_tokens=True)` adds BOS but NOT EOS, so labels never include EOS → model never learns "report finished → stop" → at inference it runs to `max_new_tokens` and loops on high-probability template phrases.
+
+**Test (~30-45 min on L4, ~60-90 min on T4):**
+1. Load the existing trained checkpoint.
+2. Pick 5 test images + their GT findings text.
+3. **PHASE A (BEFORE)** — generate with current model. Record `avg_gen_tokens`, `hit_max_rate` (% of samples that ran out of token budget), `distinct_sentence_ratio` (1.0 = no repeats, <0.5 = heavy loop).
+4. **Mini fine-tune** for 1 epoch on 100 training samples using a dataset subclass that appends `tokenizer.eos_token` — this is the ONLY change.
+5. **PHASE B (AFTER)** — re-generate the same 5 images with the same settings. Recompute metrics.
+6. Compare. If `hit_max_rate` drops sharply → EOS hypothesis confirmed dominant. If barely changes → another factor (greedy bias / quantization / prompt mismatch) is more important than missing-EOS.
+
+Run this notebook standalone; it pulls code + checkpoint + a small dataset slice from HF.
+"""))
+
+# ─── Selectors ──────────────────────────────────────────────────────────────
+cells.append(md("## 0. Selectors"))
+cells.append(code("""# ── Platform ─────────────────────────────────────────────────────
+PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'
+
+# ── Source repos ─────────────────────────────────────────────────
+HF_USER       = 'hieu3636'
+HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'
+HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'
+HF_DATA_REPO  = f'{HF_USER}/cxr-vlm-data'
+
+# ── Which trained run to start from ──────────────────────────────
+RUN_ID     = 'MIMIC-CXR_resized_run_1'
+CKPT_PICK  = 'best'                    # 'best' | 'last'
+
+# ── Test setup ───────────────────────────────────────────────────
+NUM_TEST_IMAGES   = 5                  # generate before+after on these
+NUM_TRAIN_SAMPLES = 100                # mini-FT budget
+TASK              = 'findings'         # which task to test on
+MAX_NEW_TOKENS    = 300                # generation cap (proxy for "didn't EOS")
+
+# ── Mini fine-tune hparams ───────────────────────────────────────
+FT_LR              = 2e-5
+FT_BATCH_SIZE      = 2                 # keep tiny — L4 has 24GB but model is 4-bit + LoRA
+FT_EPOCHS          = 1
+FT_GRAD_ACCUM      = 4                 # effective batch = 8
+
+# ── Generation settings (used IDENTICALLY in BEFORE and AFTER) ───
+GEN_DO_SAMPLE = False                  # greedy → makes EOS effect very visible
+GEN_NUM_BEAMS = 1                      # NO beam search — isolates the EOS variable
+
+assert PLATFORM in ('kaggle', 'colab', 'lightning', 'gcp', 'local')
+print(f'PLATFORM={PLATFORM}  RUN_ID={RUN_ID}/{CKPT_PICK}')
+print(f'Test: {NUM_TEST_IMAGES} images   FT: {NUM_TRAIN_SAMPLES} samples × {FT_EPOCHS} epoch(s)')
+"""))
+
+# ─── Env + pip ──────────────────────────────────────────────────────────────
+cells.append(md("## 1. Env + pip"))
+cells.append(code("""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, json, time, random, re
+from pathlib import Path
+"""))
+
+cells.append(code("""!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
+
+import torch as _t
+if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):
+    print('[pip] Ampere+/Ada -> flash-attn install (may take 5-10 min)')
+    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5
+else:
+    print('[pip] T4/V100 -> skipping flash-attn')
+"""))
+
+cells.append(code("""import torch, transformers, peft, huggingface_hub, httpx
+print('torch       :', torch.__version__, '| cuda:', torch.cuda.is_available())
+print('transformers:', transformers.__version__)
+print('peft        :', peft.__version__)
+
+# httpx 0.28+ shim
+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 required'
+_p = torch.cuda.get_device_properties(0)
+print(f'GPU: {_p.name}  ({_p.total_memory/1e9:.1f} GB)')
+"""))
+
+# ─── Paths + code/checkpoint/data pull ──────────────────────────────────────
+cells.append(md("## 2. Paths + pull code + checkpoint + data slice"))
+cells.append(code("""# WORK + HF_TOKEN
+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':
+    for c in (Path('/home/jupyter'), Path('/workspace')):
+        if c.exists() or os.access(c.parent, os.W_OK):
+            WORK = c; break
+    else:
+        WORK = Path.home() / 'cxr-vlm-work'
+else:
+    WORK = Path.home() / 'cxr-vlm-work'
+WORK.mkdir(parents=True, exist_ok=True)
+assert os.environ.get('HF_TOKEN'), 'HF_TOKEN missing in platform secrets'
+
+from huggingface_hub import snapshot_download, hf_hub_download
+
+print('Pulling code …')
+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)
+"""))
+
+cells.append(code("""# Pull checkpoint (configs + stage2/{CKPT_PICK})
+RUN_PULL_ROOT = WORK / 'run_pull'
+RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)
+
+print(f'Pulling {RUN_ID}/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
+assert (CKPT_DIR_PULLED / 'checkpoint_projection.pt').is_file()
+assert (CKPT_DIR_PULLED / 'checkpoint_lora' / 'adapter_config.json').is_file()
+print('checkpoint OK')
+
+SAVED_MODEL_CFG = RUN_DIR_PULLED / 'configs' / 'model_config.yaml'
+SAVED_TRAIN_CFG = RUN_DIR_PULLED / 'configs' / 'train_config.yaml'
+"""))
+
+cells.append(code("""# Pull a thin dataset slice: manifests + instruct JSON + 1 image tar shard.
+# 1 shard ≈ a few thousand resized JPGs, far more than we need.
+import tarfile
+
+DATA_SRC = WORK / 'data_src'
+DATA_DIR = DATA_SRC / 'MIMIC-CXR_resized'
+DATA_DIR.mkdir(parents=True, exist_ok=True)
+
+# Metadata (CSV manifests + instruct JSONs)
+print('Pulling manifests + instruct JSONs …')
+snapshot_download(
+    repo_id=HF_DATA_REPO, repo_type='dataset',
+    token=os.environ['HF_TOKEN'],
+    allow_patterns=[
+        'MIMIC-CXR_resized/*.csv',
+        'MIMIC-CXR_resized/*.json',
+        'MIMIC-CXR_resized/*.txt',
+    ],
+    local_dir=str(DATA_SRC),
+)
+
+# List tar shards on HF and pull the smallest one (usually train shard 0)
+from huggingface_hub import HfApi
+api = HfApi(token=os.environ['HF_TOKEN'])
+all_files = api.list_repo_files(repo_id=HF_DATA_REPO, repo_type='dataset')
+shards = sorted(f for f in all_files
+                if f.startswith('MIMIC-CXR_resized/') and f.endswith('.tar'))
+assert shards, 'No tar shards found on HF data repo.'
+
+# Prefer a train shard for FT samples
+train_shard = next((s for s in shards if 'train' in s.lower()), shards[0])
+print(f'Pulling 1 tar shard: {train_shard}')
+shard_path = Path(hf_hub_download(
+    repo_id=HF_DATA_REPO, repo_type='dataset',
+    filename=train_shard, token=os.environ['HF_TOKEN'],
+    local_dir=str(DATA_SRC),
+))
+with tarfile.open(shard_path) as t:
+    t.extractall(DATA_DIR)
+shard_path.unlink(missing_ok=True)
+print(f'Data ready under {DATA_DIR}')
+print('Top-level entries:', sorted(p.name for p in DATA_DIR.iterdir())[:10])
+"""))
+
+# ─── Build configs + load model ─────────────────────────────────────────────
+cells.append(md("## 3. GPU profile + build configs + load model"))
+cells.append(code("""import torch
+_p = torch.cuda.get_device_properties(0)
+_cap = (_p.major, _p.minor)
+_bf16_ok = torch.cuda.is_bf16_supported()
+_fa2_ok = _cap >= (8, 0)
+_fa2_installed = False
+if _fa2_ok:
+    try:
+        import flash_attn; _fa2_installed = True
+    except Exception:
+        pass
+
+PROFILE = dict(
+    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 _fa2_installed) else 'sdpa',
+)
+print(f'GPU={_p.name}  cap=sm_{_cap[0]}{_cap[1]}  bf16={_bf16_ok}  FA2={_fa2_installed}')
+print(f'-> dtype={PROFILE["torch_dtype"]}  attn={PROFILE["attn_implementation"]}')
+"""))
+
+cells.append(code("""from omegaconf import OmegaConf
+
+model_cfg = OmegaConf.load(SAVED_MODEL_CFG) if SAVED_MODEL_CFG.is_file() \\
+            else OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')
+train_cfg = OmegaConf.load(SAVED_TRAIN_CFG) if SAVED_TRAIN_CFG.is_file() \\
+            else OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
+
+# 4-bit Vicuna + profile
+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
+# Keep grad checkpointing ON during the mini-FT — saves VRAM, irrelevant for gen.
+model_cfg.llm.gradient_checkpointing    = True
+
+# CheXpert classifier off — not needed for this test
+model_cfg.chexpert_classifier.enabled = False
+print('configs ready')
+"""))
+
+cells.append(code("""import time
+from model import CXRVisionLanguageModel
+from model.rad_dino import BioViLTEncoder
+from utils.checkpoint import load_checkpoint
+
+print('[1/3] Building model … (cold cache: 5-9 min)')
+t0 = time.time()
+model = CXRVisionLanguageModel(model_cfg)
+print(f'      built in {time.time()-t0:.1f}s')
+
+print(f'[2/3] Loading checkpoint from {CKPT_DIR_PULLED} …')
+t0 = time.time()
+# CRITICAL: pass the DIRECTORY, not the .pt file (load_checkpoint splits suffix
+# off the stem; passing checkpoint_projection.pt silently skips both).
+load_checkpoint(model, str(CKPT_DIR_PULLED))
+print(f'      loaded in {time.time()-t0:.1f}s')
+
+print('[3/3] cuda + eval()')
+model = model.to('cuda')
+model.eval()
+TRANSFORM = BioViLTEncoder.get_transform('val')
+print(f'VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB')
+"""))
+
+# ─── Build test set ─────────────────────────────────────────────────────────
+cells.append(md("""## 4. Build the test set (5 images with GT findings)
+
+We sample test images from the instruct JSON that was used at training time. We need:
+1. An `image_path` that resolves under `DATA_DIR` (so the file is on disk after our 1-shard pull).
+2. A non-empty `target` to compare against qualitatively.
+3. `task == TASK` (default `'findings'`)."""))
+
+cells.append(code("""from utils.dataset_resolver import resolve_dataset_spec
+
+# Point the config at our pulled data dir, then let the resolver pick the right
+# instruct JSON (it auto-builds if missing, matching the report/image mode that
+# the trained model expects).
+train_cfg.data.dataset_name = 'MIMIC-CXR_resized'
+train_cfg.data.mimic_cxr_resized.root = str(DATA_DIR)
+
+spec = resolve_dataset_spec(train_cfg)
+INSTRUCT_JSON = spec.instruct_json
+IMAGE_ROOT    = Path(spec.image_root)
+print(f'report_mode={spec.report_mode}  image_mode={spec.image_mode}')
+print('instruct JSON:', INSTRUCT_JSON)
+print('image_root   :', IMAGE_ROOT)
+
+all_entries = json.load(open(INSTRUCT_JSON))
+print(f'{len(all_entries):,} total entries')
+"""))
+
+cells.append(code("""# Filter to entries whose image is actually on disk (we only pulled 1 shard).
+def _img_present(entry):
+    p = entry.get('image_path') or (entry.get('image_paths') or [None])[0]
+    return p and (IMAGE_ROOT / p).is_file()
+
+train_pool = [e for e in all_entries
+              if e.get('split') == 'train' and e.get('task') == TASK
+                 and e.get('target') and _img_present(e)]
+test_pool  = [e for e in all_entries
+              if e.get('split') in ('test', 'validate')
+                 and e.get('task') == TASK
+                 and e.get('target') and _img_present(e)]
+# If test set isn't covered by our single shard, fall back to held-out train.
+if len(test_pool) < NUM_TEST_IMAGES:
+    print(f'(only {len(test_pool)} test/val entries in this shard; '
+          f'using held-out train entries for the test set)')
+    held_out  = train_pool[-NUM_TEST_IMAGES:]
+    train_pool = train_pool[:-NUM_TEST_IMAGES]
+    test_pool = held_out
+
+random.seed(0)
+random.shuffle(train_pool)
+TRAIN_ENTRIES = train_pool[:NUM_TRAIN_SAMPLES]
+TEST_ENTRIES  = test_pool[:NUM_TEST_IMAGES]
+
+print(f'TRAIN_ENTRIES: {len(TRAIN_ENTRIES)}  TEST_ENTRIES: {len(TEST_ENTRIES)}')
+assert len(TRAIN_ENTRIES) >= 10 and len(TEST_ENTRIES) >= 3, \\
+    'Not enough samples in this shard — pull a second shard.'
+
+for i, e in enumerate(TEST_ENTRIES):
+    print(f'\\n[{i}] {e["image_path"]}')
+    print(f'    GT ({len(e["target"].split())} words): {e["target"][:160]}…')
+"""))
+
+# ─── Metrics ────────────────────────────────────────────────────────────────
+cells.append(md("""## 5. Metrics helper
+
+Three signals to track:
+
+- **`avg_gen_tokens`** — mean output length in tokens. If model never emits EOS, this saturates near `MAX_NEW_TOKENS`.
+- **`hit_max_rate`** — fraction of samples where output length ≥ `MAX_NEW_TOKENS - 5`. Direct proxy for "didn't emit EOS".
+- **`distinct_sentence_ratio`** — (# distinct sentences) / (total sentences). 1.0 = no repeats, < 0.5 = heavy loop."""))
+
+cells.append(code("""def split_sentences(text: str):
+    return [s.strip() for s in re.split(r'(?<=[.!?])\\s+', text.strip()) if s.strip()]
+
+def measure_outputs(outputs, max_new_tokens, tokenizer):
+    n = len(outputs)
+    if n == 0:
+        return {}
+    lengths = [len(tokenizer.encode(o, add_special_tokens=False)) for o in outputs]
+    sent_counts, distinct_ratios = [], []
+    for o in outputs:
+        ss = split_sentences(o)
+        if not ss:
+            sent_counts.append(0); distinct_ratios.append(1.0); continue
+        # Normalize whitespace + de-id tokens for dedup
+        norm = [re.sub(r'_+|\\s+', ' ', s.lower()) for s in ss]
+        sent_counts.append(len(ss))
+        distinct_ratios.append(len(set(norm)) / len(norm))
+    return dict(
+        n                       = n,
+        avg_gen_tokens          = sum(lengths) / n,
+        max_gen_tokens          = max(lengths),
+        hit_max_rate            = sum(1 for L in lengths if L >= max_new_tokens - 5) / n,
+        avg_sentences           = sum(sent_counts) / n,
+        distinct_sentence_ratio = sum(distinct_ratios) / n,
+    )
+
+def fmt_metrics(label, m, mnt):
+    print(f'{label:<10s}'
+          f'  avg_tok={m["avg_gen_tokens"]:6.1f}'
+          f'  max_tok={m["max_gen_tokens"]:4d}'
+          f'  hit_max%={m["hit_max_rate"]*100:5.1f}'
+          f'  sentences={m["avg_sentences"]:4.1f}'
+          f'  distinct_sent%={m["distinct_sentence_ratio"]*100:5.1f}'
+          f'  (cap={mnt})')
+"""))
+
+# ─── PHASE A — BEFORE ───────────────────────────────────────────────────────
+cells.append(md("""## 6. PHASE A — BEFORE fine-tune
+
+Generate on the 5 test images with the model as-is. Greedy + `num_beams=1` is intentional — it makes the EOS effect visible. Beam search would mask it."""))
+
+cells.append(code("""from PIL import Image
+from data.prompt_templates import (
+    build_findings_prompt, build_impression_prompt,
+    build_report_prompt,   build_vqa_prompt,
+)
+
+def _build_prompt(task, structured_findings=None, question=None):
+    return {
+        'findings':   lambda: build_findings_prompt(structured_findings, randomize=False),
+        'impression': lambda: build_impression_prompt(structured_findings, randomize=False),
+        'report':     lambda: build_report_prompt(structured_findings, randomize=False),
+        'vqa':        lambda: build_vqa_prompt(question, structured_findings),
+    }[task]()
+
+@torch.no_grad()
+def generate_on_test_set(entries, label, max_new_tokens=MAX_NEW_TOKENS):
+    model.eval()
+    outs = []
+    for e in entries:
+        img = Image.open(IMAGE_ROOT / e['image_path']).convert('RGB')
+        img_t = TRANSFORM(img).unsqueeze(0).to('cuda')
+        prompt = _build_prompt(e['task'])
+        out = model.generate(
+            images         = img_t,
+            prompts        = [prompt],
+            max_new_tokens = max_new_tokens,
+            temperature    = 1.0,
+            do_sample      = GEN_DO_SAMPLE,
+            num_beams      = GEN_NUM_BEAMS,
+        )[0]
+        outs.append(out)
+    metrics = measure_outputs(outs, max_new_tokens, model.tokenizer)
+    fmt_metrics(label, metrics, max_new_tokens)
+    return outs, metrics
+
+print('Generating BEFORE …')
+BEFORE_OUTS, BEFORE_METRICS = generate_on_test_set(TEST_ENTRIES, 'BEFORE')
+"""))
+
+cells.append(code("""# Show 2 example outputs side-by-side with GT
+for i in range(min(2, len(TEST_ENTRIES))):
+    print('═' * 80)
+    print(f'Image: {TEST_ENTRIES[i]["image_path"]}')
+    print('-' * 80, '\\nGT:'); print(TEST_ENTRIES[i]['target'])
+    print('-' * 80, '\\nBEFORE generation:'); print(BEFORE_OUTS[i])
+    print()
+"""))
+
+# ─── Mini fine-tune with EOS fix ────────────────────────────────────────────
+cells.append(md("""## 7. Mini fine-tune — the only change is appending EOS to targets
+
+We subclass `CXRInstructDataset` and override `_tokenize_with_labels` to append `tokenizer.eos_token` before encoding. Everything else (prompt format, LR, optimizer, model architecture) is identical to the original training. So any behavior change must come from the EOS.
+
+100 samples × 1 epoch with grad_accum=4, batch=2 → ~12-13 optimizer steps. ~5-10 minutes on L4."""))
+
+cells.append(code("""from data.dataset import CXRInstructDataset
+from data.collator import CXRDataCollator
+from torch.utils.data import Subset, DataLoader
+
+class CXRInstructDataset_EOS(CXRInstructDataset):
+    '''Same dataset as production, but targets get </s> appended.'''
+    def _tokenize_with_labels(self, prompt: str, target: str):
+        full_text      = prompt + ' ' + target + self.tokenizer.eos_token
+        prompt_encoded = self.tokenizer.encode(prompt, add_special_tokens=True)
+        full_encoded   = self.tokenizer.encode(
+            full_text,
+            add_special_tokens = True,
+            max_length         = self.cutoff_len,
+            truncation         = True,
+        )
+        input_ids = torch.tensor(full_encoded, dtype=torch.long)
+        labels    = input_ids.clone()
+        labels[: min(len(prompt_encoded), self.cutoff_len)] = -100
+        return input_ids, labels
+
+# Sanity-check: confirm the override actually appends EOS
+_eos_id = model.tokenizer.eos_token_id
+print('EOS token id =', _eos_id, '   token =', repr(model.tokenizer.eos_token))
+"""))
+
+cells.append(code("""# Build train dataset using the SAME instruct JSON; restrict to our 100 entries.
+# Trick: write a temp JSON with just TRAIN_ENTRIES so we don't change CXRInstructDataset filter logic.
+import tempfile
+
+tmp_json = WORK / 'tmp_train_subset.json'
+tmp_json.write_text(json.dumps(TRAIN_ENTRIES))
+
+ft_dataset = CXRInstructDataset_EOS(
+    data_path  = str(tmp_json),
+    image_root = str(IMAGE_ROOT),
+    tokenizer  = model.tokenizer,
+    transform  = BioViLTEncoder.get_transform('train'),
+    task       = TASK,
+    split      = 'train',
+    cutoff_len = 512,
+)
+print(f'FT dataset: {len(ft_dataset)} samples')
+
+collator = CXRDataCollator(model.tokenizer.pad_token_id)
+loader = DataLoader(
+    ft_dataset,
+    batch_size  = FT_BATCH_SIZE,
+    shuffle     = True,
+    collate_fn  = collator,
+    num_workers = 0,
+)
+print(f'Steps per epoch: {len(loader)} (× {FT_EPOCHS} epoch(s), grad_accum={FT_GRAD_ACCUM})')
+"""))
+
+cells.append(code("""# Mini fine-tune loop. Trains projection + LoRA (everything that has requires_grad).
+from torch.optim import AdamW
+
+trainable = [p for p in model.parameters() if p.requires_grad]
+n_trainable = sum(p.numel() for p in trainable)
+print(f'Trainable params: {n_trainable/1e6:.1f}M')
+
+optimizer = AdamW(trainable, lr=FT_LR)
+
+model.train()
+# QLoRA needs grad checkpointing kwarg
+if hasattr(model.llm, 'gradient_checkpointing_enable'):
+    model.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})
+
+step = 0
+optimizer.zero_grad()
+t0 = time.time()
+for epoch in range(FT_EPOCHS):
+    for batch_idx, batch in enumerate(loader):
+        batch = {k: (v.cuda(non_blocking=True) if torch.is_tensor(v) else v)
+                 for k, v in batch.items()}
+        out = model(**{k: batch[k] for k in ('images', 'input_ids', 'attention_mask', 'labels')
+                       if k in batch})
+        loss = out['loss'] if isinstance(out, dict) else out.loss
+        (loss / FT_GRAD_ACCUM).backward()
+
+        if (batch_idx + 1) % FT_GRAD_ACCUM == 0 or batch_idx + 1 == len(loader):
+            optimizer.step()
+            optimizer.zero_grad()
+            step += 1
+
+        if batch_idx % 4 == 0:
+            elapsed = time.time() - t0
+            print(f'epoch {epoch+1}  batch {batch_idx+1}/{len(loader)}  '
+                  f'loss={loss.item():.3f}  ({elapsed:.0f}s elapsed)')
+
+print(f'\\n✔ Mini-FT done in {time.time()-t0:.0f}s, {step} optimizer steps')
+"""))
+
+# ─── PHASE B — AFTER ────────────────────────────────────────────────────────
+cells.append(md("## 8. PHASE B — AFTER fine-tune (same images, same settings)"))
+
+cells.append(code("""# Disable grad checkpointing for cleaner generate
+if hasattr(model.llm, 'gradient_checkpointing_disable'):
+    model.llm.gradient_checkpointing_disable()
+
+print('Generating AFTER …')
+AFTER_OUTS, AFTER_METRICS = generate_on_test_set(TEST_ENTRIES, 'AFTER')
+"""))
+
+cells.append(code("""# Show same 2 examples side-by-side: GT vs BEFORE vs AFTER
+for i in range(min(2, len(TEST_ENTRIES))):
+    print('═' * 80)
+    print(f'Image: {TEST_ENTRIES[i]["image_path"]}')
+    print('-' * 80, '\\nGT:');     print(TEST_ENTRIES[i]['target'])
+    print('-' * 80, '\\nBEFORE:'); print(BEFORE_OUTS[i])
+    print('-' * 80, '\\nAFTER :'); print(AFTER_OUTS[i])
+    print()
+"""))
+
+# ─── Compare ────────────────────────────────────────────────────────────────
+cells.append(md("""## 9. Verdict
+
+| Signal | Hypothesis-confirmed direction |
+|---|---|
+| `avg_gen_tokens` | **down** (model stops earlier) |
+| `hit_max_rate` | **down sharply** — fewer samples truncated at cap |
+| `distinct_sentence_ratio` | **up** (less looping) |
+
+If at least 2 of these 3 move strongly in the predicted direction after just 100 samples of fine-tune, that's solid evidence the missing EOS in training labels was the dominant cause. If they don't budge, the looping comes from another factor (quantization noise / greedy bias / data overfitting on boilerplate) that this fix alone won't solve."""))
+
+cells.append(code("""print('═' * 80)
+print('SUMMARY')
+print('═' * 80)
+fmt_metrics('BEFORE', BEFORE_METRICS, MAX_NEW_TOKENS)
+fmt_metrics('AFTER ', AFTER_METRICS,  MAX_NEW_TOKENS)
+print()
+
+deltas = {
+    'avg_gen_tokens         ': AFTER_METRICS['avg_gen_tokens']          - BEFORE_METRICS['avg_gen_tokens'],
+    'hit_max_rate           ': AFTER_METRICS['hit_max_rate']            - BEFORE_METRICS['hit_max_rate'],
+    'distinct_sentence_ratio': AFTER_METRICS['distinct_sentence_ratio'] - BEFORE_METRICS['distinct_sentence_ratio'],
+}
+print('Δ AFTER − BEFORE:')
+for k, v in deltas.items():
+    arrow = '↓' if v < 0 else ('↑' if v > 0 else '·')
+    print(f'  {k}: {v:+.3f}  {arrow}')
+
+print()
+# Heuristic verdict
+ok_len   = AFTER_METRICS['avg_gen_tokens']  < BEFORE_METRICS['avg_gen_tokens'] - 20
+ok_hit   = AFTER_METRICS['hit_max_rate']    < BEFORE_METRICS['hit_max_rate']   - 0.15
+ok_dist  = AFTER_METRICS['distinct_sentence_ratio'] > BEFORE_METRICS['distinct_sentence_ratio'] + 0.10
+n_signals = sum([ok_len, ok_hit, ok_dist])
+
+if n_signals >= 2:
+    print(f'✔ {n_signals}/3 signals confirm — EOS-in-labels appears to be the dominant cause.')
+    print('  Recommendation: apply the fix in dataset.py and retrain (full run).')
+elif n_signals == 1:
+    print(f'~ {n_signals}/3 signals — EOS contributes but is not alone.')
+    print('  Likely co-factors: greedy decoding bias, quantization noise, boilerplate overfit.')
+    print('  Recommendation: retrain with the fix AND switch eval to beam search (num_beams=4).')
+else:
+    print('✘ 0/3 signals — EOS alone is NOT the dominant cause.')
+    print('  Try: (a) more FT samples or epochs, (b) bigger LR, (c) compare beam vs greedy on')
+    print('  the BEFORE model — if beam fixes it, the issue is decoding not training.')
+"""))
+
+# ─── Build + write ──────────────────────────────────────────────────────────
+nb = nbf.v4.new_notebook()
+nb.cells = cells
+nb.metadata = {
+    'kernelspec': {'name': 'python3', 'display_name': 'Python 3'},
+    'language_info': {'name': 'python'},
+}
+
+OUT = Path(__file__).resolve().parent / 'cxrvlm_eos_test.ipynb'
+nbf.write(nb, OUT)
+print(f'Wrote {OUT}  ({len(cells)} cells)')

scripts/cxrvlm_colab_inference.ipynb

@@ -55,10 +55,15 @@
     "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",
+    "MAX_NEW_TOKENS       = 200       # report-length ceiling — caps loops too\n",
+    "TEMPERATURE          = 0.1\n",
+    "DO_SAMPLE            = False     # greedy by default — deterministic\n",
+    "NUM_BEAMS            = 1\n",
+    "# Anti-repetition (greedy decoding on a clinical LM tends to loop on template\n",
+    "# phrases — '___ at 11:30 a.m.', 'vertebral body heights preserved' — bumping\n",
+    "# these two breaks the loop without losing determinism).\n",
+    "REPETITION_PENALTY   = 1.3\n",
+    "NO_REPEAT_NGRAM_SIZE = 4\n",
     "\n",
     "assert PLATFORM  in ('kaggle', 'colab', 'lightning', 'gcp', 'local')\n",
     "assert CKPT_PICK in ('best', 'last')\n",
@@ -491,10 +496,12 @@
     "    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",
+    "    max_new_tokens:        int   = None,\n",
+    "    temperature:           float = None,\n",
+    "    do_sample:             bool  = None,\n",
+    "    num_beams:             int   = None,\n",
+    "    repetition_penalty:    float = None,\n",
+    "    no_repeat_ngram_size:  int   = None,\n",
     ") -> str:\n",
     "    '''Run the model on one image.\n",
     "\n",
@@ -516,12 +523,14 @@
     "        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",
+    "        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",
+    "        repetition_penalty   = repetition_penalty   if repetition_penalty   is not None else REPETITION_PENALTY,\n",
+    "        no_repeat_ngram_size = no_repeat_ngram_size if no_repeat_ngram_size is not None else NO_REPEAT_NGRAM_SIZE,\n",
     "    )\n",
     "    return out[0]\n",
     "\n",

scripts/cxrvlm_eos_test.ipynb

@@ -0,0 +1,840 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ad40c0fb",
+   "metadata": {},
+   "source": [
+    "# CXR-VLM — EOS Fix A/B Test\n",
+    "\n",
+    "**Hypothesis:** `data/dataset.py` tokenizes `prompt + \" \" + target` without appending `</s>` to the target. `LlamaTokenizer.encode(add_special_tokens=True)` adds BOS but NOT EOS, so labels never include EOS → model never learns \"report finished → stop\" → at inference it runs to `max_new_tokens` and loops on high-probability template phrases.\n",
+    "\n",
+    "**Test (~30-45 min on L4, ~60-90 min on T4):**\n",
+    "1. Load the existing trained checkpoint.\n",
+    "2. Pick 5 test images + their GT findings text.\n",
+    "3. **PHASE A (BEFORE)** — generate with current model. Record `avg_gen_tokens`, `hit_max_rate` (% of samples that ran out of token budget), `distinct_sentence_ratio` (1.0 = no repeats, <0.5 = heavy loop).\n",
+    "4. **Mini fine-tune** for 1 epoch on 100 training samples using a dataset subclass that appends `tokenizer.eos_token` — this is the ONLY change.\n",
+    "5. **PHASE B (AFTER)** — re-generate the same 5 images with the same settings. Recompute metrics.\n",
+    "6. Compare. If `hit_max_rate` drops sharply → EOS hypothesis confirmed dominant. If barely changes → another factor (greedy bias / quantization / prompt mismatch) is more important than missing-EOS.\n",
+    "\n",
+    "Run this notebook standalone; it pulls code + checkpoint + a small dataset slice from HF.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6a1882eb",
+   "metadata": {},
+   "source": [
+    "## 0. Selectors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47da9294",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ── Platform ─────────────────────────────────────────────────────\n",
+    "PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'\n",
+    "\n",
+    "# ── Source repos ─────────────────────────────────────────────────\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",
+    "HF_DATA_REPO  = f'{HF_USER}/cxr-vlm-data'\n",
+    "\n",
+    "# ── Which trained run to start from ──────────────────────────────\n",
+    "RUN_ID     = 'MIMIC-CXR_resized_run_1'\n",
+    "CKPT_PICK  = 'best'                    # 'best' | 'last'\n",
+    "\n",
+    "# ── Test setup ───────────────────────────────────────────────────\n",
+    "NUM_TEST_IMAGES   = 5                  # generate before+after on these\n",
+    "NUM_TRAIN_SAMPLES = 100                # mini-FT budget\n",
+    "TASK              = 'findings'         # which task to test on\n",
+    "MAX_NEW_TOKENS    = 300                # generation cap (proxy for \"didn't EOS\")\n",
+    "\n",
+    "# ── Mini fine-tune hparams ───────────────────────────────────────\n",
+    "FT_LR              = 2e-5\n",
+    "FT_BATCH_SIZE      = 2                 # keep tiny — L4 has 24GB but model is 4-bit + LoRA\n",
+    "FT_EPOCHS          = 1\n",
+    "FT_GRAD_ACCUM      = 4                 # effective batch = 8\n",
+    "\n",
+    "# ── Generation settings (used IDENTICALLY in BEFORE and AFTER) ───\n",
+    "GEN_DO_SAMPLE = False                  # greedy → makes EOS effect very visible\n",
+    "GEN_NUM_BEAMS = 1                      # NO beam search — isolates the EOS variable\n",
+    "\n",
+    "assert PLATFORM in ('kaggle', 'colab', 'lightning', 'gcp', 'local')\n",
+    "print(f'PLATFORM={PLATFORM}  RUN_ID={RUN_ID}/{CKPT_PICK}')\n",
+    "print(f'Test: {NUM_TEST_IMAGES} images   FT: {NUM_TRAIN_SAMPLES} samples × {FT_EPOCHS} epoch(s)')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1c9e3253",
+   "metadata": {},
+   "source": [
+    "## 1. Env + pip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93eee293",
+   "metadata": {},
+   "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, json, time, random, re\n",
+    "from pathlib import Path\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7d23c73",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!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\n",
+    "\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 -> flash-attn install (may take 5-10 min)')\n",
+    "    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5\n",
+    "else:\n",
+    "    print('[pip] T4/V100 -> skipping flash-attn')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d4fc533",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch, transformers, peft, huggingface_hub, httpx\n",
+    "print('torch       :', torch.__version__, '| cuda:', torch.cuda.is_available())\n",
+    "print('transformers:', transformers.__version__)\n",
+    "print('peft        :', peft.__version__)\n",
+    "\n",
+    "# httpx 0.28+ shim\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', '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",
+    "_patch_httpx()\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA required'\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": "370928ae",
+   "metadata": {},
+   "source": [
+    "## 2. Paths + pull code + checkpoint + data slice"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "788cc5c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# WORK + HF_TOKEN\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",
+    "    for c in (Path('/home/jupyter'), Path('/workspace')):\n",
+    "        if c.exists() or os.access(c.parent, os.W_OK):\n",
+    "            WORK = c; break\n",
+    "    else:\n",
+    "        WORK = Path.home() / 'cxr-vlm-work'\n",
+    "else:\n",
+    "    WORK = Path.home() / 'cxr-vlm-work'\n",
+    "WORK.mkdir(parents=True, exist_ok=True)\n",
+    "assert os.environ.get('HF_TOKEN'), 'HF_TOKEN missing in platform secrets'\n",
+    "\n",
+    "from huggingface_hub import snapshot_download, hf_hub_download\n",
+    "\n",
+    "print('Pulling code …')\n",
+    "CODE_SRC = Path(snapshot_download(\n",
+    "    repo_id=HF_CODE_REPO, repo_type='model',\n",
+    "    token=os.environ['HF_TOKEN'],\n",
+    "    local_dir=str(WORK / 'cxr-vlm-code'),\n",
+    "))\n",
+    "PROJECT = WORK / 'cxr_vlm'\n",
+    "if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():\n",
+    "    shutil.copytree(CODE_SRC, PROJECT)\n",
+    "os.chdir(PROJECT)\n",
+    "sys.path.insert(0, str(PROJECT))\n",
+    "print('PROJECT =', PROJECT)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a12ff7f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Pull checkpoint (configs + stage2/{CKPT_PICK})\n",
+    "RUN_PULL_ROOT = WORK / 'run_pull'\n",
+    "RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "print(f'Pulling {RUN_ID}/stage2/{CKPT_PICK} from {HF_RUNS_REPO} …')\n",
+    "snapshot_download(\n",
+    "    repo_id=HF_RUNS_REPO, 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",
+    "RUN_DIR_PULLED  = RUN_PULL_ROOT / RUN_ID\n",
+    "CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK\n",
+    "assert (CKPT_DIR_PULLED / 'checkpoint_projection.pt').is_file()\n",
+    "assert (CKPT_DIR_PULLED / 'checkpoint_lora' / 'adapter_config.json').is_file()\n",
+    "print('checkpoint OK')\n",
+    "\n",
+    "SAVED_MODEL_CFG = RUN_DIR_PULLED / 'configs' / 'model_config.yaml'\n",
+    "SAVED_TRAIN_CFG = RUN_DIR_PULLED / 'configs' / 'train_config.yaml'\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad8e2aa3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Pull a thin dataset slice: manifests + instruct JSON + 1 image tar shard.\n",
+    "# 1 shard ≈ a few thousand resized JPGs, far more than we need.\n",
+    "import tarfile\n",
+    "\n",
+    "DATA_SRC = WORK / 'data_src'\n",
+    "DATA_DIR = DATA_SRC / 'MIMIC-CXR_resized'\n",
+    "DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "# Metadata (CSV manifests + instruct JSONs)\n",
+    "print('Pulling manifests + instruct JSONs …')\n",
+    "snapshot_download(\n",
+    "    repo_id=HF_DATA_REPO, repo_type='dataset',\n",
+    "    token=os.environ['HF_TOKEN'],\n",
+    "    allow_patterns=[\n",
+    "        'MIMIC-CXR_resized/*.csv',\n",
+    "        'MIMIC-CXR_resized/*.json',\n",
+    "        'MIMIC-CXR_resized/*.txt',\n",
+    "    ],\n",
+    "    local_dir=str(DATA_SRC),\n",
+    ")\n",
+    "\n",
+    "# List tar shards on HF and pull the smallest one (usually train shard 0)\n",
+    "from huggingface_hub import HfApi\n",
+    "api = HfApi(token=os.environ['HF_TOKEN'])\n",
+    "all_files = api.list_repo_files(repo_id=HF_DATA_REPO, repo_type='dataset')\n",
+    "shards = sorted(f for f in all_files\n",
+    "                if f.startswith('MIMIC-CXR_resized/') and f.endswith('.tar'))\n",
+    "assert shards, 'No tar shards found on HF data repo.'\n",
+    "\n",
+    "# Prefer a train shard for FT samples\n",
+    "train_shard = next((s for s in shards if 'train' in s.lower()), shards[0])\n",
+    "print(f'Pulling 1 tar shard: {train_shard}')\n",
+    "shard_path = Path(hf_hub_download(\n",
+    "    repo_id=HF_DATA_REPO, repo_type='dataset',\n",
+    "    filename=train_shard, token=os.environ['HF_TOKEN'],\n",
+    "    local_dir=str(DATA_SRC),\n",
+    "))\n",
+    "with tarfile.open(shard_path) as t:\n",
+    "    t.extractall(DATA_DIR)\n",
+    "shard_path.unlink(missing_ok=True)\n",
+    "print(f'Data ready under {DATA_DIR}')\n",
+    "print('Top-level entries:', sorted(p.name for p in DATA_DIR.iterdir())[:10])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c3666fd",
+   "metadata": {},
+   "source": [
+    "## 3. GPU profile + build configs + load model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d6ac34e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "_p = torch.cuda.get_device_properties(0)\n",
+    "_cap = (_p.major, _p.minor)\n",
+    "_bf16_ok = torch.cuda.is_bf16_supported()\n",
+    "_fa2_ok = _cap >= (8, 0)\n",
+    "_fa2_installed = False\n",
+    "if _fa2_ok:\n",
+    "    try:\n",
+    "        import flash_attn; _fa2_installed = True\n",
+    "    except Exception:\n",
+    "        pass\n",
+    "\n",
+    "PROFILE = dict(\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 _fa2_installed) else 'sdpa',\n",
+    ")\n",
+    "print(f'GPU={_p.name}  cap=sm_{_cap[0]}{_cap[1]}  bf16={_bf16_ok}  FA2={_fa2_installed}')\n",
+    "print(f'-> dtype={PROFILE[\"torch_dtype\"]}  attn={PROFILE[\"attn_implementation\"]}')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7072c0a6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from omegaconf import OmegaConf\n",
+    "\n",
+    "model_cfg = OmegaConf.load(SAVED_MODEL_CFG) if SAVED_MODEL_CFG.is_file() \\\n",
+    "            else OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')\n",
+    "train_cfg = OmegaConf.load(SAVED_TRAIN_CFG) if SAVED_TRAIN_CFG.is_file() \\\n",
+    "            else OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\n",
+    "\n",
+    "# 4-bit Vicuna + profile\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",
+    "# Keep grad checkpointing ON during the mini-FT — saves VRAM, irrelevant for gen.\n",
+    "model_cfg.llm.gradient_checkpointing    = True\n",
+    "\n",
+    "# CheXpert classifier off — not needed for this test\n",
+    "model_cfg.chexpert_classifier.enabled = False\n",
+    "print('configs ready')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38a9fc47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\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 … (cold cache: 5-9 min)')\n",
+    "t0 = time.time()\n",
+    "model = CXRVisionLanguageModel(model_cfg)\n",
+    "print(f'      built in {time.time()-t0:.1f}s')\n",
+    "\n",
+    "print(f'[2/3] Loading checkpoint from {CKPT_DIR_PULLED} …')\n",
+    "t0 = time.time()\n",
+    "# CRITICAL: pass the DIRECTORY, not the .pt file (load_checkpoint splits suffix\n",
+    "# off the stem; passing checkpoint_projection.pt silently skips both).\n",
+    "load_checkpoint(model, str(CKPT_DIR_PULLED))\n",
+    "print(f'      loaded in {time.time()-t0:.1f}s')\n",
+    "\n",
+    "print('[3/3] cuda + eval()')\n",
+    "model = model.to('cuda')\n",
+    "model.eval()\n",
+    "TRANSFORM = BioViLTEncoder.get_transform('val')\n",
+    "print(f'VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ca325a5a",
+   "metadata": {},
+   "source": [
+    "## 4. Build the test set (5 images with GT findings)\n",
+    "\n",
+    "We sample test images from the instruct JSON that was used at training time. We need:\n",
+    "1. An `image_path` that resolves under `DATA_DIR` (so the file is on disk after our 1-shard pull).\n",
+    "2. A non-empty `target` to compare against qualitatively.\n",
+    "3. `task == TASK` (default `'findings'`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "783d8217",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils.dataset_resolver import resolve_dataset_spec\n",
+    "\n",
+    "# Point the config at our pulled data dir, then let the resolver pick the right\n",
+    "# instruct JSON (it auto-builds if missing, matching the report/image mode that\n",
+    "# the trained model expects).\n",
+    "train_cfg.data.dataset_name = 'MIMIC-CXR_resized'\n",
+    "train_cfg.data.mimic_cxr_resized.root = str(DATA_DIR)\n",
+    "\n",
+    "spec = resolve_dataset_spec(train_cfg)\n",
+    "INSTRUCT_JSON = spec.instruct_json\n",
+    "IMAGE_ROOT    = Path(spec.image_root)\n",
+    "print(f'report_mode={spec.report_mode}  image_mode={spec.image_mode}')\n",
+    "print('instruct JSON:', INSTRUCT_JSON)\n",
+    "print('image_root   :', IMAGE_ROOT)\n",
+    "\n",
+    "all_entries = json.load(open(INSTRUCT_JSON))\n",
+    "print(f'{len(all_entries):,} total entries')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c2b8f2c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Filter to entries whose image is actually on disk (we only pulled 1 shard).\n",
+    "def _img_present(entry):\n",
+    "    p = entry.get('image_path') or (entry.get('image_paths') or [None])[0]\n",
+    "    return p and (IMAGE_ROOT / p).is_file()\n",
+    "\n",
+    "train_pool = [e for e in all_entries\n",
+    "              if e.get('split') == 'train' and e.get('task') == TASK\n",
+    "                 and e.get('target') and _img_present(e)]\n",
+    "test_pool  = [e for e in all_entries\n",
+    "              if e.get('split') in ('test', 'validate')\n",
+    "                 and e.get('task') == TASK\n",
+    "                 and e.get('target') and _img_present(e)]\n",
+    "# If test set isn't covered by our single shard, fall back to held-out train.\n",
+    "if len(test_pool) < NUM_TEST_IMAGES:\n",
+    "    print(f'(only {len(test_pool)} test/val entries in this shard; '\n",
+    "          f'using held-out train entries for the test set)')\n",
+    "    held_out  = train_pool[-NUM_TEST_IMAGES:]\n",
+    "    train_pool = train_pool[:-NUM_TEST_IMAGES]\n",
+    "    test_pool = held_out\n",
+    "\n",
+    "random.seed(0)\n",
+    "random.shuffle(train_pool)\n",
+    "TRAIN_ENTRIES = train_pool[:NUM_TRAIN_SAMPLES]\n",
+    "TEST_ENTRIES  = test_pool[:NUM_TEST_IMAGES]\n",
+    "\n",
+    "print(f'TRAIN_ENTRIES: {len(TRAIN_ENTRIES)}  TEST_ENTRIES: {len(TEST_ENTRIES)}')\n",
+    "assert len(TRAIN_ENTRIES) >= 10 and len(TEST_ENTRIES) >= 3, \\\n",
+    "    'Not enough samples in this shard — pull a second shard.'\n",
+    "\n",
+    "for i, e in enumerate(TEST_ENTRIES):\n",
+    "    print(f'\\n[{i}] {e[\"image_path\"]}')\n",
+    "    print(f'    GT ({len(e[\"target\"].split())} words): {e[\"target\"][:160]}…')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea68dfb1",
+   "metadata": {},
+   "source": [
+    "## 5. Metrics helper\n",
+    "\n",
+    "Three signals to track:\n",
+    "\n",
+    "- **`avg_gen_tokens`** — mean output length in tokens. If model never emits EOS, this saturates near `MAX_NEW_TOKENS`.\n",
+    "- **`hit_max_rate`** — fraction of samples where output length ≥ `MAX_NEW_TOKENS - 5`. Direct proxy for \"didn't emit EOS\".\n",
+    "- **`distinct_sentence_ratio`** — (# distinct sentences) / (total sentences). 1.0 = no repeats, < 0.5 = heavy loop."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b6479cf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def split_sentences(text: str):\n",
+    "    return [s.strip() for s in re.split(r'(?<=[.!?])\\s+', text.strip()) if s.strip()]\n",
+    "\n",
+    "def measure_outputs(outputs, max_new_tokens, tokenizer):\n",
+    "    n = len(outputs)\n",
+    "    if n == 0:\n",
+    "        return {}\n",
+    "    lengths = [len(tokenizer.encode(o, add_special_tokens=False)) for o in outputs]\n",
+    "    sent_counts, distinct_ratios = [], []\n",
+    "    for o in outputs:\n",
+    "        ss = split_sentences(o)\n",
+    "        if not ss:\n",
+    "            sent_counts.append(0); distinct_ratios.append(1.0); continue\n",
+    "        # Normalize whitespace + de-id tokens for dedup\n",
+    "        norm = [re.sub(r'_+|\\s+', ' ', s.lower()) for s in ss]\n",
+    "        sent_counts.append(len(ss))\n",
+    "        distinct_ratios.append(len(set(norm)) / len(norm))\n",
+    "    return dict(\n",
+    "        n                       = n,\n",
+    "        avg_gen_tokens          = sum(lengths) / n,\n",
+    "        max_gen_tokens          = max(lengths),\n",
+    "        hit_max_rate            = sum(1 for L in lengths if L >= max_new_tokens - 5) / n,\n",
+    "        avg_sentences           = sum(sent_counts) / n,\n",
+    "        distinct_sentence_ratio = sum(distinct_ratios) / n,\n",
+    "    )\n",
+    "\n",
+    "def fmt_metrics(label, m, mnt):\n",
+    "    print(f'{label:<10s}'\n",
+    "          f'  avg_tok={m[\"avg_gen_tokens\"]:6.1f}'\n",
+    "          f'  max_tok={m[\"max_gen_tokens\"]:4d}'\n",
+    "          f'  hit_max%={m[\"hit_max_rate\"]*100:5.1f}'\n",
+    "          f'  sentences={m[\"avg_sentences\"]:4.1f}'\n",
+    "          f'  distinct_sent%={m[\"distinct_sentence_ratio\"]*100:5.1f}'\n",
+    "          f'  (cap={mnt})')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b8f16ed1",
+   "metadata": {},
+   "source": [
+    "## 6. PHASE A — BEFORE fine-tune\n",
+    "\n",
+    "Generate on the 5 test images with the model as-is. Greedy + `num_beams=1` is intentional — it makes the EOS effect visible. Beam search would mask it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "80e533c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "from data.prompt_templates import (\n",
+    "    build_findings_prompt, build_impression_prompt,\n",
+    "    build_report_prompt,   build_vqa_prompt,\n",
+    ")\n",
+    "\n",
+    "def _build_prompt(task, structured_findings=None, question=None):\n",
+    "    return {\n",
+    "        'findings':   lambda: build_findings_prompt(structured_findings, randomize=False),\n",
+    "        'impression': lambda: build_impression_prompt(structured_findings, randomize=False),\n",
+    "        'report':     lambda: build_report_prompt(structured_findings, randomize=False),\n",
+    "        'vqa':        lambda: build_vqa_prompt(question, structured_findings),\n",
+    "    }[task]()\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def generate_on_test_set(entries, label, max_new_tokens=MAX_NEW_TOKENS):\n",
+    "    model.eval()\n",
+    "    outs = []\n",
+    "    for e in entries:\n",
+    "        img = Image.open(IMAGE_ROOT / e['image_path']).convert('RGB')\n",
+    "        img_t = TRANSFORM(img).unsqueeze(0).to('cuda')\n",
+    "        prompt = _build_prompt(e['task'])\n",
+    "        out = model.generate(\n",
+    "            images         = img_t,\n",
+    "            prompts        = [prompt],\n",
+    "            max_new_tokens = max_new_tokens,\n",
+    "            temperature    = 1.0,\n",
+    "            do_sample      = GEN_DO_SAMPLE,\n",
+    "            num_beams      = GEN_NUM_BEAMS,\n",
+    "        )[0]\n",
+    "        outs.append(out)\n",
+    "    metrics = measure_outputs(outs, max_new_tokens, model.tokenizer)\n",
+    "    fmt_metrics(label, metrics, max_new_tokens)\n",
+    "    return outs, metrics\n",
+    "\n",
+    "print('Generating BEFORE …')\n",
+    "BEFORE_OUTS, BEFORE_METRICS = generate_on_test_set(TEST_ENTRIES, 'BEFORE')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16f21302",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Show 2 example outputs side-by-side with GT\n",
+    "for i in range(min(2, len(TEST_ENTRIES))):\n",
+    "    print('═' * 80)\n",
+    "    print(f'Image: {TEST_ENTRIES[i][\"image_path\"]}')\n",
+    "    print('-' * 80, '\\nGT:'); print(TEST_ENTRIES[i]['target'])\n",
+    "    print('-' * 80, '\\nBEFORE generation:'); print(BEFORE_OUTS[i])\n",
+    "    print()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cf56b845",
+   "metadata": {},
+   "source": [
+    "## 7. Mini fine-tune — the only change is appending EOS to targets\n",
+    "\n",
+    "We subclass `CXRInstructDataset` and override `_tokenize_with_labels` to append `tokenizer.eos_token` before encoding. Everything else (prompt format, LR, optimizer, model architecture) is identical to the original training. So any behavior change must come from the EOS.\n",
+    "\n",
+    "100 samples × 1 epoch with grad_accum=4, batch=2 → ~12-13 optimizer steps. ~5-10 minutes on L4."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e34db1b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from data.dataset import CXRInstructDataset\n",
+    "from data.collator import CXRDataCollator\n",
+    "from torch.utils.data import Subset, DataLoader\n",
+    "\n",
+    "class CXRInstructDataset_EOS(CXRInstructDataset):\n",
+    "    '''Same dataset as production, but targets get </s> appended.'''\n",
+    "    def _tokenize_with_labels(self, prompt: str, target: str):\n",
+    "        full_text      = prompt + ' ' + target + self.tokenizer.eos_token\n",
+    "        prompt_encoded = self.tokenizer.encode(prompt, add_special_tokens=True)\n",
+    "        full_encoded   = self.tokenizer.encode(\n",
+    "            full_text,\n",
+    "            add_special_tokens = True,\n",
+    "            max_length         = self.cutoff_len,\n",
+    "            truncation         = True,\n",
+    "        )\n",
+    "        input_ids = torch.tensor(full_encoded, dtype=torch.long)\n",
+    "        labels    = input_ids.clone()\n",
+    "        labels[: min(len(prompt_encoded), self.cutoff_len)] = -100\n",
+    "        return input_ids, labels\n",
+    "\n",
+    "# Sanity-check: confirm the override actually appends EOS\n",
+    "_eos_id = model.tokenizer.eos_token_id\n",
+    "print('EOS token id =', _eos_id, '   token =', repr(model.tokenizer.eos_token))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28d87bdb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build train dataset using the SAME instruct JSON; restrict to our 100 entries.\n",
+    "# Trick: write a temp JSON with just TRAIN_ENTRIES so we don't change CXRInstructDataset filter logic.\n",
+    "import tempfile\n",
+    "\n",
+    "tmp_json = WORK / 'tmp_train_subset.json'\n",
+    "tmp_json.write_text(json.dumps(TRAIN_ENTRIES))\n",
+    "\n",
+    "ft_dataset = CXRInstructDataset_EOS(\n",
+    "    data_path  = str(tmp_json),\n",
+    "    image_root = str(IMAGE_ROOT),\n",
+    "    tokenizer  = model.tokenizer,\n",
+    "    transform  = BioViLTEncoder.get_transform('train'),\n",
+    "    task       = TASK,\n",
+    "    split      = 'train',\n",
+    "    cutoff_len = 512,\n",
+    ")\n",
+    "print(f'FT dataset: {len(ft_dataset)} samples')\n",
+    "\n",
+    "collator = CXRDataCollator(model.tokenizer.pad_token_id)\n",
+    "loader = DataLoader(\n",
+    "    ft_dataset,\n",
+    "    batch_size  = FT_BATCH_SIZE,\n",
+    "    shuffle     = True,\n",
+    "    collate_fn  = collator,\n",
+    "    num_workers = 0,\n",
+    ")\n",
+    "print(f'Steps per epoch: {len(loader)} (× {FT_EPOCHS} epoch(s), grad_accum={FT_GRAD_ACCUM})')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7175df4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Mini fine-tune loop. Trains projection + LoRA (everything that has requires_grad).\n",
+    "from torch.optim import AdamW\n",
+    "\n",
+    "trainable = [p for p in model.parameters() if p.requires_grad]\n",
+    "n_trainable = sum(p.numel() for p in trainable)\n",
+    "print(f'Trainable params: {n_trainable/1e6:.1f}M')\n",
+    "\n",
+    "optimizer = AdamW(trainable, lr=FT_LR)\n",
+    "\n",
+    "model.train()\n",
+    "# QLoRA needs grad checkpointing kwarg\n",
+    "if hasattr(model.llm, 'gradient_checkpointing_enable'):\n",
+    "    model.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})\n",
+    "\n",
+    "step = 0\n",
+    "optimizer.zero_grad()\n",
+    "t0 = time.time()\n",
+    "for epoch in range(FT_EPOCHS):\n",
+    "    for batch_idx, batch in enumerate(loader):\n",
+    "        batch = {k: (v.cuda(non_blocking=True) if torch.is_tensor(v) else v)\n",
+    "                 for k, v in batch.items()}\n",
+    "        out = model(**{k: batch[k] for k in ('images', 'input_ids', 'attention_mask', 'labels')\n",
+    "                       if k in batch})\n",
+    "        loss = out['loss'] if isinstance(out, dict) else out.loss\n",
+    "        (loss / FT_GRAD_ACCUM).backward()\n",
+    "\n",
+    "        if (batch_idx + 1) % FT_GRAD_ACCUM == 0 or batch_idx + 1 == len(loader):\n",
+    "            optimizer.step()\n",
+    "            optimizer.zero_grad()\n",
+    "            step += 1\n",
+    "\n",
+    "        if batch_idx % 4 == 0:\n",
+    "            elapsed = time.time() - t0\n",
+    "            print(f'epoch {epoch+1}  batch {batch_idx+1}/{len(loader)}  '\n",
+    "                  f'loss={loss.item():.3f}  ({elapsed:.0f}s elapsed)')\n",
+    "\n",
+    "print(f'\\n✔ Mini-FT done in {time.time()-t0:.0f}s, {step} optimizer steps')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7e18a80",
+   "metadata": {},
+   "source": [
+    "## 8. PHASE B — AFTER fine-tune (same images, same settings)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43e2a5b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Disable grad checkpointing for cleaner generate\n",
+    "if hasattr(model.llm, 'gradient_checkpointing_disable'):\n",
+    "    model.llm.gradient_checkpointing_disable()\n",
+    "\n",
+    "print('Generating AFTER …')\n",
+    "AFTER_OUTS, AFTER_METRICS = generate_on_test_set(TEST_ENTRIES, 'AFTER')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0212f1de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Show same 2 examples side-by-side: GT vs BEFORE vs AFTER\n",
+    "for i in range(min(2, len(TEST_ENTRIES))):\n",
+    "    print('═' * 80)\n",
+    "    print(f'Image: {TEST_ENTRIES[i][\"image_path\"]}')\n",
+    "    print('-' * 80, '\\nGT:');     print(TEST_ENTRIES[i]['target'])\n",
+    "    print('-' * 80, '\\nBEFORE:'); print(BEFORE_OUTS[i])\n",
+    "    print('-' * 80, '\\nAFTER :'); print(AFTER_OUTS[i])\n",
+    "    print()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3601f44b",
+   "metadata": {},
+   "source": [
+    "## 9. Verdict\n",
+    "\n",
+    "| Signal | Hypothesis-confirmed direction |\n",
+    "|---|---|\n",
+    "| `avg_gen_tokens` | **down** (model stops earlier) |\n",
+    "| `hit_max_rate` | **down sharply** — fewer samples truncated at cap |\n",
+    "| `distinct_sentence_ratio` | **up** (less looping) |\n",
+    "\n",
+    "If at least 2 of these 3 move strongly in the predicted direction after just 100 samples of fine-tune, that's solid evidence the missing EOS in training labels was the dominant cause. If they don't budge, the looping comes from another factor (quantization noise / greedy bias / data overfitting on boilerplate) that this fix alone won't solve."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6cc53af9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print('═' * 80)\n",
+    "print('SUMMARY')\n",
+    "print('═' * 80)\n",
+    "fmt_metrics('BEFORE', BEFORE_METRICS, MAX_NEW_TOKENS)\n",
+    "fmt_metrics('AFTER ', AFTER_METRICS,  MAX_NEW_TOKENS)\n",
+    "print()\n",
+    "\n",
+    "deltas = {\n",
+    "    'avg_gen_tokens         ': AFTER_METRICS['avg_gen_tokens']          - BEFORE_METRICS['avg_gen_tokens'],\n",
+    "    'hit_max_rate           ': AFTER_METRICS['hit_max_rate']            - BEFORE_METRICS['hit_max_rate'],\n",
+    "    'distinct_sentence_ratio': AFTER_METRICS['distinct_sentence_ratio'] - BEFORE_METRICS['distinct_sentence_ratio'],\n",
+    "}\n",
+    "print('Δ AFTER − BEFORE:')\n",
+    "for k, v in deltas.items():\n",
+    "    arrow = '↓' if v < 0 else ('↑' if v > 0 else '·')\n",
+    "    print(f'  {k}: {v:+.3f}  {arrow}')\n",
+    "\n",
+    "print()\n",
+    "# Heuristic verdict\n",
+    "ok_len   = AFTER_METRICS['avg_gen_tokens']  < BEFORE_METRICS['avg_gen_tokens'] - 20\n",
+    "ok_hit   = AFTER_METRICS['hit_max_rate']    < BEFORE_METRICS['hit_max_rate']   - 0.15\n",
+    "ok_dist  = AFTER_METRICS['distinct_sentence_ratio'] > BEFORE_METRICS['distinct_sentence_ratio'] + 0.10\n",
+    "n_signals = sum([ok_len, ok_hit, ok_dist])\n",
+    "\n",
+    "if n_signals >= 2:\n",
+    "    print(f'✔ {n_signals}/3 signals confirm — EOS-in-labels appears to be the dominant cause.')\n",
+    "    print('  Recommendation: apply the fix in dataset.py and retrain (full run).')\n",
+    "elif n_signals == 1:\n",
+    "    print(f'~ {n_signals}/3 signals — EOS contributes but is not alone.')\n",
+    "    print('  Likely co-factors: greedy decoding bias, quantization noise, boilerplate overfit.')\n",
+    "    print('  Recommendation: retrain with the fix AND switch eval to beam search (num_beams=4).')\n",
+    "else:\n",
+    "    print('✘ 0/3 signals — EOS alone is NOT the dominant cause.')\n",
+    "    print('  Try: (a) more FT samples or epochs, (b) bigger LR, (c) compare beam vs greedy on')\n",
+    "    print('  the BEFORE model — if beam fixes it, the issue is decoding not training.')\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/vertex_job.yaml

@@ -52,15 +52,15 @@ workerPoolSpecs:
     - name: IMAGE_MODE
       value: all_views_split
     - name: MODE
-      value: resume
+      value: fresh
     - name: EXPLICIT_RUN_ID
-      value: MIMIC-CXR_resized_run_1
+      value: MIMIC-CXR_resized_run_3
     - name: HF_RUNS_REPO
       value: hieu3636/cxr-vlm-runs
     - name: S1_EPOCHS
-      value: "2"
+      value: "3"
     - name: S2_EPOCHS
-      value: "7"
+      value: "5"
 
 # Scheduling: SPOT = ~70% cheaper, may be preempted. Repo has auto-resume so safe.
 # Use STANDARD if you want guaranteed uninterrupted run (3.4× cost).