Fix chem-msms-predict with public MassSpecGym checkpoints (instrument arg + ms_pred patches)

.agents/skills/chem-msms-predict/scripts/predict_msms.py

@@ -36,6 +36,7 @@ def run_iceberg(
     inten_ckpt: Path,
     collision_energies: list,
     adduct: str,
+    instrument: str,
     cuda_devices,
     batch_size: int,
     num_workers: int,
@@ -51,6 +52,7 @@ def run_iceberg(
         collision_energies=collision_energies,
         nce=False,
         adduct=adduct,
+        instrument=instrument,
         exp_name="skill_pred",
         python_path=sys.executable,
         gen_ckpt=str(gen_ckpt),
@@ -167,6 +169,10 @@ def parse_args() -> argparse.Namespace:
         help="Collision energies in eV (default: 20 40)",
     )
     p.add_argument("--adduct", default="[M+H]+", help="Adduct type (default: [M+H]+)")
+    p.add_argument(
+        "--instrument", default="Orbitrap", choices=["Orbitrap", "QTOF"],
+        help="Instrument the MSG checkpoints condition on (default: Orbitrap)",
+    )
     p.add_argument(
         "--output_dir", type=Path, default=Path("results/msms_prediction"),
         help="Output directory",
@@ -215,6 +221,7 @@ def main() -> None:
         inten_ckpt=args.inten_ckpt,
         collision_energies=args.collision_energies,
         adduct=args.adduct,
+        instrument=args.instrument,
         cuda_devices=args.cuda_devices,
         batch_size=args.batch_size,
         num_workers=args.num_workers,

conda-envs/msms-agent/install.sh

@@ -53,6 +53,27 @@ setup(
 )
 SETUP_EOF
 
+# Patch ms_pred for 3 upstream bugs that break ICEBERG forward prediction on a
+# pip install / CPU / pytorch_lightning 2.x (mirrors the runtime fixes):
+#   1. iceberg_elucidation.py calls predict_smis.py via a cwd-relative path
+#   2. predict_smis.py uses pl.utilities.seed.seed_everything (removed in PL 2.0)
+#   3. predict_smis.py calls torch.cuda.set_device(gpu_id) unconditionally (fails on CPU)
+python - "$TMP_DIR/ms-pred/src/ms_pred/dag_pred" <<'PYPATCH'
+import sys, pathlib
+d = pathlib.Path(sys.argv[1])
+el = d / "iceberg_elucidation.py"
+el.write_text(el.read_text().replace(
+    '{python_path} src/ms_pred/dag_pred/predict_smis.py',
+    '{python_path} {Path(__file__).resolve().parent / "predict_smis.py"}'))
+ps = d / "predict_smis.py"
+t = ps.read_text()
+t = t.replace('pl.utilities.seed.seed_everything', 'pl.seed_everything')
+t = t.replace('torch.cuda.set_device(gpu_id)',
+              '(torch.cuda.set_device(gpu_id) if (gpu and avail_gpu_num > 0) else None)')
+ps.write_text(t)
+print("Patched ms_pred dag_pred (pip-install path / PL2 seed / CPU cuda guard)")
+PYPATCH
+
 uv pip install "$TMP_DIR/ms-pred"
 rm -rf "$TMP_DIR"