"""
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)')