rescore.py: resolve v5-S2 ckpts under checkpoints/v5s2/seed_*/best_phase2.pt (HF layout) with fallback to results/v5_training/10seed/... (dev layout)

code/scripts/rescore.py

@@ -51,9 +51,20 @@ from utils.pdb_utils import (
 # 3-seed ensemble (paper: v5-S2 TS-S2, seeds 1024 / 5555 / 789)
 # --------------------------------------------------------------------------
 SEEDS = [1024, 5555, 789]
-CKPT_TEMPLATE = str(BASE / "results/v5_training/10seed/seed_{seed}/best_phase2.pt")
+CKPT_LAYOUTS = [
+    str(BASE / "checkpoints/v5s2/seed_{seed}/best_phase2.pt"),         # HF release layout
+    str(BASE / "results/v5_training/10seed/seed_{seed}/best_phase2.pt"),  # in-repo dev layout
+]
 ESM_DIR = str(BASE / "data/esm2_embeddings")
 
+
+def _resolve_ckpt(seed):
+    for tmpl in CKPT_LAYOUTS:
+        p = tmpl.format(seed=seed)
+        if Path(p).exists():
+            return p
+    return None
+
 # Per-target holo/apo canonical PDB pair (extend as needed).
 TARGETS = {
     "cam":      {"holo": "data/pdbs/cam_holo/3CLN.pdb",   "apo": "data/pdbs/cam_apo/1CFD.pdb",   "chain": "A", "esm_target": "cam"},
@@ -197,9 +208,10 @@ def rescore(target, design_dir, gpu, out_path):
 
     results = {d["id"]: {} for d in extracted}
     for seed in SEEDS:
-        ckpt = CKPT_TEMPLATE.format(seed=seed)
-        if not Path(ckpt).exists():
-            logger.warning(f"Missing checkpoint {ckpt}; skip seed {seed}")
+        ckpt = _resolve_ckpt(seed)
+        if ckpt is None:
+            tried = [t.format(seed=seed) for t in CKPT_LAYOUTS]
+            logger.warning(f"Missing checkpoint for seed {seed} in any of: {tried}; skip")
             continue
         logger.info(f"Scoring seed {seed}...")
         dq = DifferentiableQTheta(checkpoint_path=ckpt, device=device, esm_dir=ESM_DIR)
@@ -237,7 +249,7 @@ def rescore(target, design_dir, gpu, out_path):
 
     out = {
         "target": target, "holo_pdb": holo_pdb, "apo_pdb": apo_pdb,
-        "seeds": SEEDS, "checkpoints": {str(s): CKPT_TEMPLATE.format(seed=s) for s in SEEDS},
+        "seeds": SEEDS, "checkpoints": {str(s): _resolve_ckpt(s) for s in SEEDS},
         "n_designs": len(extracted), "per_design": results,
     }
     Path(out_path).parent.mkdir(parents=True, exist_ok=True)