train_hf.py

"""BrainAge training script for Hugging Face Spaces (GPU).

Uses a mounted Volume for the preprocessed dataset, runs two-phase training,
pushes checkpoints + metrics to a HF Model repo, and serves a live progress
page on port 7860.

Environment variables (set as Space secrets):
  HF_TOKEN          — write-access token for pushing checkpoints
  MODEL_REPO        — e.g. bilalEthizo/BrainAge-SFCN
"""
from __future__ import annotations
import json, math, os, random, shutil, threading, time
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from huggingface_hub import HfApi

# ─────────────── config from env ──────────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", "")
MODEL_REPO = os.environ.get("MODEL_REPO", "bilalEthizo/BrainAge-SFCN")

MOUNT_DIR = Path("/dataset")
DATA_DIR = Path("/tmp/brainage")
CACHE_DIR = DATA_DIR / "cache"
SPLIT_CSV = DATA_DIR / "split.csv"
CKPT_DIR = DATA_DIR / "checkpoints"

# Live progress state (read by the web server thread)
PROGRESS: dict = {"phase": "init", "epoch": 0, "total_epochs": 0,
                  "train_loss": 0.0, "val_mae": 0.0, "best_mae": 999.0,
                  "best_epoch": 0, "message": "Initializing…", "log": []}

SITES = [
    "DataSet-1_BCP", "DataSet-1_BCP_patho",
    "DataSet-2_Calgary", "DataSet-3_ds002726",
    "DataSet-4_ds000248", "DataSet-5_PTBP", "DataSet-6_IXI",
    "DataSet-7_MPI-Leipzig", "DataSet-8_AOMIC-ID1000",
    "DataSet-9_NKI-Rockland",
    "DataSet-10_ABIDE-I", "DataSet-11_ABIDE-II", "DataSet-12_ADHD-200",
]
SEX_CATEGORIES = ["M", "F", "U"]
N_REGIONS = 70
TAB_DIM = N_REGIONS + len(SEX_CATEGORIES) + len(SITES)
VOLUME_SHAPE = (128, 144, 112)
AGE_BINS = [0, 2, 5, 12, 18, 25, 50, 80]


# ─────────────── model ────────────────────────────────────────────────────
class SFCN(nn.Module):
    def __init__(self, channels=(32, 64, 128, 256, 256, 64), dropout=0.3,
                 emb_dim=128):
        super().__init__()
        layers = []
        in_c = 1
        for i, c in enumerate(channels):
            layers.append(nn.Conv3d(in_c, c, 3, padding=1))
            layers.append(nn.BatchNorm3d(c))
            if i < len(channels) - 1:
                layers.append(nn.MaxPool3d(2, 2))
            layers.append(nn.ReLU(inplace=True))
            if dropout and i >= 3:
                layers.append(nn.Dropout3d(dropout))
            in_c = c
        self.features = nn.Sequential(*layers)
        self.pool = nn.AdaptiveAvgPool3d(1)
        self.fc = nn.Linear(channels[-1], emb_dim)

    def forward(self, x):
        x = self.features(x)
        x = self.pool(x).flatten(1)
        return self.fc(x)


class BrainAgeDual(nn.Module):
    def __init__(self, n_tabular: int = TAB_DIM, emb_dim: int = 128):
        super().__init__()
        self.img_branch = SFCN(emb_dim=emb_dim)
        self.tab_branch = nn.Sequential(
            nn.Linear(n_tabular, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, emb_dim))
        self.head = nn.Sequential(
            nn.Linear(emb_dim * 2, 64), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(64, 1))

    def forward(self, vol, tab):
        img_emb = self.img_branch(vol)
        tab_emb = self.tab_branch(tab)
        fused = torch.cat([img_emb, tab_emb], dim=1)
        return self.head(fused).squeeze(-1)


# ─────────────── dataset ──────────────────────────────────────────────────
class PreloadedDataset(Dataset):
    """Loads ALL tensors into CPU RAM (fp16) at init. Zero I/O at train time."""

    def __init__(self, subject_ids: list[str], cache_dir: Path,
                 augment: bool = False):
        self.augment = augment
        self._tx = self._build_tx() if augment else None
        n = len(subject_ids)

        # Pre-allocate as float16 to halve memory (~13 GB instead of ~27 GB)
        self.volumes = torch.empty(n, 1, *VOLUME_SHAPE, dtype=torch.float16)
        first = torch.load(cache_dir / f"{subject_ids[0]}.pt",
                           weights_only=False, map_location="cpu")
        tab_dim = first["tab"].numel()
        self.tabs = torch.empty(n, tab_dim, dtype=torch.float32)
        self.ages = torch.empty(n, dtype=torch.float32)

        t0 = time.time()
        for i, sid in enumerate(subject_ids):
            d = torch.load(cache_dir / f"{sid}.pt",
                           weights_only=False, map_location="cpu")
            self.volumes[i] = d["volume"].half().unsqueeze(0)
            self.tabs[i] = d["tab"].float()
            self.ages[i] = d["age"].float()
            if (i + 1) % 500 == 0:
                print(f"    preloaded {i+1}/{n}", flush=True)

        gb = self.volumes.nelement() * 2 / 1e9
        print(f"    preloaded {n} subjects ({gb:.1f} GB fp16) "
              f"in {time.time()-t0:.0f}s", flush=True)

    @staticmethod
    def _build_tx():
        from monai.transforms import (
            Compose, RandFlipd, RandAffined, RandBiasFieldd,
            RandGaussianNoised, RandAdjustContrastd, EnsureTyped,
        )
        return Compose([
            EnsureTyped(keys="volume", dtype=torch.float32),
            RandFlipd(keys="volume", prob=0.5, spatial_axis=2),
            RandAffined(
                keys="volume", prob=0.5,
                rotate_range=(math.radians(10), math.radians(10),
                              math.radians(10)),
                scale_range=(0.08, 0.08, 0.08),
                padding_mode="zeros"),
            RandBiasFieldd(keys="volume", prob=0.3, coeff_range=(0.0, 0.05)),
            RandAdjustContrastd(keys="volume", prob=0.3, gamma=(0.8, 1.2)),
            RandGaussianNoised(keys="volume", prob=0.3, std=0.01),
        ])

    def __len__(self):
        return len(self.ages)

    def __getitem__(self, idx):
        vol = self.volumes[idx].float()
        tab = self.tabs[idx]
        age = self.ages[idx]
        if self.augment and self._tx:
            vol = self._tx({"volume": vol})["volume"]
        return vol, tab, age


def age_bin_weights(ages, edges=(0, 2, 5, 12, 18, 25, 50, 80)):
    ages = np.asarray(ages, dtype=np.float32)
    bins = np.clip(np.digitize(ages, edges) - 1, 0, len(edges) - 2)
    cnt = np.bincount(bins, minlength=len(edges) - 1).astype(np.float32)
    cnt = np.where(cnt == 0, 1.0, cnt)
    w = (1.0 / cnt)[bins]
    return w / w.mean()


# ─────────────── web server (progress page) ──────────────────────────────
def start_web_server():
    """Tiny HTTP server on 7860 so HF Spaces knows the container is alive."""
    from http.server import HTTPServer, BaseHTTPRequestHandler
    import json as _json

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path == "/api/progress":
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(_json.dumps(PROGRESS).encode())
                return
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.end_headers()
            log_rows = PROGRESS.get("log", [])[-30:]
            log_html = ""
            for r in log_rows:
                log_html += (f"<tr><td>{r.get('epoch','')}</td>"
                             f"<td>{r.get('lr',''):.2e}</td>"
                             f"<td>{r.get('train_loss',''):.4f}</td>"
                             f"<td>{r.get('val_mae',''):.3f}</td>"
                             f"<td>{r.get('seconds',''):.0f}s</td></tr>")
            html = f"""<!DOCTYPE html><html><head>
<title>BrainAge Training</title>
<meta http-equiv="refresh" content="15">
<style>body{{font-family:monospace;background:#0f172a;color:#e2e8f0;padding:2rem}}
table{{border-collapse:collapse;width:100%}}th,td{{padding:4px 8px;text-align:right;border-bottom:1px solid #334155}}
th{{color:#94a3b8}}.badge{{display:inline-block;padding:2px 8px;border-radius:4px;background:#22c55e;color:#000;font-weight:bold}}
h1{{color:#38bdf8}}</style></head><body>
<h1>BrainAge Training Monitor</h1>
<p>Phase: <span class="badge">{PROGRESS['phase']}</span> &nbsp;
Epoch: <b>{PROGRESS['epoch']}/{PROGRESS['total_epochs']}</b> &nbsp;
Best MAE: <b>{PROGRESS['best_mae']:.3f}y</b> (ep {PROGRESS['best_epoch']})</p>
<p>{PROGRESS['message']}</p>
<table><thead><tr><th>Epoch</th><th>LR</th><th>Train Loss</th><th>Val MAE</th><th>Time</th></tr></thead>
<tbody>{log_html}</tbody></table>
<p style="color:#64748b;font-size:0.8em">Auto-refreshes every 15s &middot; JSON at /api/progress</p>
</body></html>"""
            self.wfile.write(html.encode())
        def log_message(self, format, *args):
            pass

    server = HTTPServer(("0.0.0.0", 7860), Handler)
    t = threading.Thread(target=server.serve_forever, daemon=True)
    t.start()
    print("[web] progress server on :7860")


# ─────────────── data from mounted volume ─────────────────────────────────
def setup_data():
    PROGRESS["message"] = "Setting up data from mounted volume…"
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    CKPT_DIR.mkdir(parents=True, exist_ok=True)

    # Locate cache dir in the mounted volume
    vol_cache = MOUNT_DIR / "cache"
    if vol_cache.is_dir():
        if CACHE_DIR.exists():
            if CACHE_DIR.is_symlink():
                CACHE_DIR.unlink()
            elif len(list(CACHE_DIR.glob("*.pt"))) > 100:
                count = len(list(CACHE_DIR.glob("*.pt")))
                print(f"[data] cache already present: {count} files")
                PROGRESS["message"] = f"Cache ready: {count} files"
                return
        CACHE_DIR.symlink_to(vol_cache)
        count = len(list(CACHE_DIR.glob("*.pt")))
        print(f"[data] symlinked {vol_cache} -> {CACHE_DIR}: {count} files")
    else:
        # Fallback: look for .pt files directly under mount
        pt_files = list(MOUNT_DIR.rglob("*.pt"))
        if pt_files:
            CACHE_DIR.mkdir(parents=True, exist_ok=True)
            parent = pt_files[0].parent
            if parent != CACHE_DIR:
                if CACHE_DIR.is_symlink():
                    CACHE_DIR.unlink()
                CACHE_DIR.symlink_to(parent)
            count = len(pt_files)
            print(f"[data] found {count} .pt files at {parent}")
        else:
            raise RuntimeError(f"No .pt files found under {MOUNT_DIR}")

    # Copy manifests
    for mf in MOUNT_DIR.glob("manifests/*.csv"):
        dst = DATA_DIR / mf.name
        if not dst.exists():
            shutil.copy2(str(mf), str(dst))

    count = len(list(CACHE_DIR.glob("*.pt")))
    PROGRESS["message"] = f"Data ready: {count} .pt files"
    print(f"[data] ready: {count} .pt files")


def generate_split():
    if SPLIT_CSV.exists():
        print(f"[split] already exists: {SPLIT_CSV}")
        return

    manifests = list(DATA_DIR.glob("*manifest*.csv"))
    if not manifests:
        raise RuntimeError("No manifest CSVs found")

    dfs = [pd.read_csv(p) for p in manifests]
    df = pd.concat(dfs, ignore_index=True)
    df = df[df["age_years"].notna()].copy()

    # Only healthy subjects
    if "healthy" in df.columns:
        df = df[df["healthy"] == True].copy()

    # Only subjects with cached .pt
    cached = {p.stem for p in CACHE_DIR.glob("*.pt")}
    df = df[df["subject_id"].astype(str).isin(cached)].copy()

    rng = np.random.default_rng(42)
    df["age_bin"] = pd.cut(df["age_years"], bins=AGE_BINS,
                           labels=[f"{AGE_BINS[i]}-{AGE_BINS[i+1]}"
                                   for i in range(len(AGE_BINS) - 1)],
                           include_lowest=True)
    df["split"] = ""
    for (site, bin_), grp in df.groupby(["dataset", "age_bin"], observed=True):
        n = len(grp)
        if n == 0:
            continue
        order = rng.permutation(n)
        n_tr = int(round(n * 0.75))
        n_va = int(round(n * 0.10))
        if n >= 3:
            n_tr = min(n_tr, n - 2)
            n_va = max(1, n_va)
        assign = np.array(["test"] * n, dtype=object)
        assign[order[:n_tr]] = "train"
        assign[order[n_tr:n_tr + n_va]] = "val"
        df.loc[grp.index, "split"] = assign

    keep = [c for c in ["subject_id", "dataset", "age_years", "age_bin",
                         "sex", "healthy", "split"] if c in df.columns]
    df[keep].to_csv(SPLIT_CSV, index=False)
    summary = df["split"].value_counts().to_dict()
    print(f"[split] created: {summary}")


# ─────────────── checkpoint push ──────────────────────────────────────────
def push_checkpoint(path: Path, msg: str):
    if not HF_TOKEN:
        print(f"[push] no HF_TOKEN, skipping push of {path.name}")
        return
    try:
        api = HfApi(token=HF_TOKEN)
        api.upload_file(
            path_or_fileobj=str(path),
            path_in_repo=path.name,
            repo_id=MODEL_REPO,
            repo_type="model",
            commit_message=msg,
        )
        print(f"[push] {path.name} -> {MODEL_REPO}")
    except Exception as e:
        print(f"[push] failed: {e}")


def push_log(path: Path):
    push_checkpoint(path, f"Update training log: {path.name}")


# ─────────────── training ─────────────────────────────────────────────────
def set_seed(seed=42):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True


def train_phase(
    phase_name: str,
    cache_dir: Path, split_csv: Path, out_ckpt: Path,
    epochs: int = 60, batch: int = 4, lr: float = 3e-4,
    warmup_epochs: int = 5, patience: int = 15,
    age_max: float | None = None,
    resume_ckpt: Path | None = None,
    num_workers: int = 4,
):
    print(f"\n{'='*60}")
    print(f"  PHASE: {phase_name}")
    print(f"{'='*60}\n")

    set_seed(42)
    df = pd.read_csv(split_csv)
    have = {p.stem for p in cache_dir.glob("*.pt")}
    df = df[df["subject_id"].astype(str).isin(have)].copy()
    if age_max is not None:
        df = df[df["age_years"] <= age_max].copy()

    tr_df = df[df["split"] == "train"].reset_index(drop=True)
    va_df = df[df["split"] == "val"].reset_index(drop=True)
    print(f"  subjects: {len(df)} (train={len(tr_df)}, val={len(va_df)})")

    PROGRESS["message"] = f"Preloading training data into RAM…"
    print("  preloading train set…")
    tr_ds = PreloadedDataset(tr_df["subject_id"].astype(str).tolist(),
                             cache_dir, augment=True)
    PROGRESS["message"] = f"Preloading validation data into RAM…"
    print("  preloading val set…")
    va_ds = PreloadedDataset(va_df["subject_id"].astype(str).tolist(),
                             cache_dir, augment=False)

    w = age_bin_weights(tr_df["age_years"].to_numpy())
    sampler = WeightedRandomSampler(w, num_samples=len(tr_ds), replacement=True)
    n_cpu = min(8, os.cpu_count() or 4)
    tr_ld = DataLoader(tr_ds, batch_size=batch, sampler=sampler,
                       num_workers=n_cpu, pin_memory=True,
                       persistent_workers=True, prefetch_factor=4)
    va_ld = DataLoader(va_ds, batch_size=batch, shuffle=False,
                       num_workers=n_cpu, pin_memory=True,
                       persistent_workers=True, prefetch_factor=4)

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"  device: {dev}")
    if dev == "cuda":
        print(f"  GPU: {torch.cuda.get_device_name(0)}")
        print(f"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

    model = BrainAgeDual(n_tabular=TAB_DIM).to(dev)
    if resume_ckpt and resume_ckpt.exists():
        sd = torch.load(resume_ckpt, map_location=dev, weights_only=False)
        state = sd.get("model", sd) if isinstance(sd, dict) else sd
        model.load_state_dict(state, strict=False)
        print(f"  resumed from: {resume_ckpt}")

    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)

    def lr_at(step, steps_per_epoch):
        ep = step / max(1, steps_per_epoch)
        if ep < warmup_epochs:
            return lr * (ep / warmup_epochs)
        prog = (ep - warmup_epochs) / max(1, epochs - warmup_epochs)
        return lr * 0.5 * (1 + math.cos(math.pi * min(1.0, prog)))

    loss_fn = nn.HuberLoss(delta=3.0)
    use_amp = dev == "cuda"
    scaler = torch.amp.GradScaler(enabled=use_amp)

    best_mae, best_epoch, no_improve = float("inf"), -1, 0
    log_rows = []
    t_start = time.time()
    steps_per_epoch = max(1, len(tr_ld))
    global_step = 0
    log_csv = out_ckpt.with_suffix(".log.csv")

    for ep in range(epochs):
        model.train()
        t_ep = time.time()
        tl_sum, tl_n = 0.0, 0
        for vol, tab, age in tr_ld:
            for g in opt.param_groups:
                g["lr"] = lr_at(global_step, steps_per_epoch)
            vol = vol.to(dev, non_blocking=True)
            tab = tab.to(dev, non_blocking=True)
            age = age.to(dev, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                pred = model(vol, tab)
                loss = loss_fn(pred, age)
            scaler.scale(loss).backward()
            scaler.unscale_(opt)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(opt); scaler.update()
            tl_sum += loss.detach().float().item() * len(age)
            tl_n += len(age)
            global_step += 1

        model.eval()
        errs = []
        with torch.inference_mode():
            for vol, tab, age in va_ld:
                vol, tab = vol.to(dev), tab.to(dev)
                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                    pred = model(vol, tab)
                errs += (pred.float().cpu() - age).abs().tolist()
        mae = float(np.mean(errs))
        tl = tl_sum / max(1, tl_n)
        took = time.time() - t_ep
        cur_lr = opt.param_groups[0]["lr"]

        print(f"  ep {ep+1:3d}/{epochs}  lr={cur_lr:.2e}  "
              f"train_loss={tl:.3f}  val_MAE={mae:.3f}y  ({took:.1f}s)",
              flush=True)
        row = {"epoch": ep + 1, "lr": cur_lr, "train_loss": tl,
               "val_mae": mae, "seconds": round(took, 1)}
        log_rows.append(row)
        PROGRESS.update({"phase": phase_name, "epoch": ep + 1,
                         "total_epochs": epochs, "train_loss": tl,
                         "val_mae": mae, "message": f"Epoch {ep+1}/{epochs}",
                         "log": log_rows[-50:]})

        # Save log every epoch
        pd.DataFrame(log_rows).to_csv(log_csv, index=False)

        if mae < best_mae:
            best_mae, best_epoch, no_improve = mae, ep + 1, 0
            PROGRESS.update({"best_mae": best_mae, "best_epoch": best_epoch})
            out_ckpt.parent.mkdir(parents=True, exist_ok=True)
            torch.save({
                "model": model.state_dict(),
                "n_tabular": TAB_DIM,
                "volume_shape": VOLUME_SHAPE,
                "best_val_mae": mae,
                "epoch": ep + 1,
                "phase": phase_name,
            }, out_ckpt)
            # Push best checkpoint every improvement
            push_checkpoint(out_ckpt,
                            f"{phase_name}: ep {ep+1} val_MAE={mae:.3f}y")
        else:
            no_improve += 1
            if no_improve >= patience:
                print(f"  early stop at ep {ep+1} "
                      f"(best {best_mae:.3f}y @ ep {best_epoch})")
                break

        # Push log every 5 epochs
        if (ep + 1) % 5 == 0:
            push_log(log_csv)

    # Final saves
    last_p = out_ckpt.with_suffix(".last.pt")
    torch.save({"model": model.state_dict(), "n_tabular": TAB_DIM,
                "volume_shape": VOLUME_SHAPE, "phase": phase_name}, last_p)
    push_checkpoint(last_p, f"{phase_name}: final (last epoch)")
    push_log(log_csv)

    total = time.time() - t_start
    print(f"\n  {phase_name} DONE: best_val_MAE={best_mae:.3f}y "
          f"@ ep {best_epoch} ({total:.0f}s total)")
    return out_ckpt


# ─────────────── main ─────────────────────────────────────────────────────
def main():
    print("="*60)
    print("  BrainAge Training Pipeline (Hugging Face)")
    print("="*60)
    print(f"  MODEL_REPO:   {MODEL_REPO}")
    print(f"  DEVICE:       {'cuda' if torch.cuda.is_available() else 'cpu'}")
    if torch.cuda.is_available():
        print(f"  GPU:          {torch.cuda.get_device_name(0)}")
    print()

    start_web_server()
    CKPT_DIR.mkdir(parents=True, exist_ok=True)

    # Step 1: Setup data from mounted volume
    setup_data()

    # Step 2: Generate split
    generate_split()

    # Step 3: Check which phases are already done by looking at HF model repo
    from huggingface_hub import hf_hub_download

    ckpt_a = CKPT_DIR / "brainage_sfcn.pt"
    ckpt_b = CKPT_DIR / "brainage_sfcn_finetune.pt"

    def file_exists_in_repo(filename):
        try:
            api = HfApi(token=HF_TOKEN)
            files = list(api.list_repo_files(MODEL_REPO, repo_type="model"))
            return filename in files
        except Exception as e:
            print(f"[state] error checking repo: {e}")
            return False

    phase_a_last_exists = file_exists_in_repo("brainage_sfcn.last.pt")
    phase_b_last_exists = file_exists_in_repo("brainage_sfcn_finetune.last.pt")
    print(f"[state] Phase A done={phase_a_last_exists}, Phase B done={phase_b_last_exists}")

    # ── Phase A ──────────────────────────────────────────────────
    if phase_a_last_exists:
        print("[skip] Phase A already completed — downloading best checkpoint")
        if not ckpt_a.exists():
            try:
                hf_hub_download(repo_id=MODEL_REPO, filename="brainage_sfcn.pt",
                                token=HF_TOKEN or None, local_dir=str(CKPT_DIR))
                print(f"[resume] downloaded Phase A best checkpoint")
            except Exception as e:
                print(f"[warn] could not download Phase A ckpt: {e}")
    else:
        prev_ckpt = CKPT_DIR / "brainage_sfcn_prev.pt"
        if not prev_ckpt.exists():
            try:
                p = hf_hub_download(repo_id=MODEL_REPO, filename="brainage_sfcn.pt",
                                    token=HF_TOKEN or None, local_dir=str(CKPT_DIR))
                Path(p).rename(prev_ckpt)
                print(f"[resume] downloaded previous best checkpoint")
            except Exception as e:
                print(f"[resume] no previous checkpoint: {e}")
                prev_ckpt = None

        train_phase(
            phase_name="PhaseA-Lifespan",
            cache_dir=CACHE_DIR, split_csv=SPLIT_CSV, out_ckpt=ckpt_a,
            epochs=60, batch=32, lr=1e-3, warmup_epochs=3, patience=15,
            num_workers=8,
            resume_ckpt=prev_ckpt if prev_ckpt and prev_ckpt.exists() else None,
        )

    # ── Phase B ──────────────────────────────────────────────────
    if phase_b_last_exists:
        print("[skip] Phase B already completed")
    else:
        train_phase(
            phase_name="PhaseB-Pediatric",
            cache_dir=CACHE_DIR, split_csv=SPLIT_CSV, out_ckpt=ckpt_b,
            epochs=30, batch=32, lr=3e-4, warmup_epochs=2, patience=10,
            age_max=25.0, resume_ckpt=ckpt_a, num_workers=8,
        )

    # Step 5: Push final model card
    if HF_TOKEN:
        card = f"""---
license: cc-by-nc-4.0
tags:
  - brain-age
  - neuroimaging
  - mri
  - 3d-cnn
  - pytorch
---

# BrainAge SFCN — Dual-Branch Brain Age Predictor

3D SFCN + tabular MLP trained on **6,050 healthy T1w brain MRIs** (0–86 years).

## Checkpoints

| File | Description |
|------|-------------|
| `brainage_sfcn.pt` | Phase A: lifespan pretraining (all ages) |
| `brainage_sfcn_finetune.pt` | Phase B: fine-tuned for pediatric (0–25 y) |

## Training data

[bilalahmad176176/BrainAge-Golden-Preprocessed](https://huggingface.co/datasets/bilalahmad176176/BrainAge-Golden-Preprocessed)

## Architecture

- Image: SFCN (Peng 2021), ~3M params, input 128×144×112
- Tabular: MLP over 86-dim vector (70 regions + sex + site)
- Fusion: concat → 64 → 1 (regression in years)
- Loss: HuberLoss(delta=3.0)
- Training: fp16 mixed precision, cosine LR, age-bin weighted sampler
"""
        try:
            api = HfApi(token=HF_TOKEN)
            api.upload_file(
                path_or_fileobj=card.encode(),
                path_in_repo="README.md",
                repo_id=MODEL_REPO,
                repo_type="model",
                commit_message="Add model card",
            )
            print("[push] model card uploaded")
        except Exception as e:
            print(f"[push] model card failed: {e}")

    print("\n" + "="*60)
    print("  ALL DONE")
    print("="*60)


if __name__ == "__main__":
    main()