overlay/scripts/benchmark_checkpoint.py

#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path


def checkpoint_candidates(cache_dir: Path) -> list[Path]:
    return [
        cache_dir / "best_bpb.pt",
        cache_dir / "pretrain_final.pt",
        cache_dir / "latest.pt",
    ]


def choose_checkpoint_candidate(cache_dir: Path) -> Path | None:
    for path in checkpoint_candidates(cache_dir):
        if path.exists():
            return path
    return None