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

.agents/skills/chem-msms-predict/SKILL.md

@@ -59,6 +59,7 @@ python .agents/skills/chem-msms-predict/scripts/predict_msms.py \
     --inten_ckpt downloads/iceberg_dag_inten_msg_best.ckpt \
     --collision_energies 20 40 \
     --adduct "[M+H]+" \
+    --instrument "Orbitrap" \
     --output_dir results/msms_prediction
 ```
 
@@ -67,6 +68,7 @@ python .agents/skills/chem-msms-predict/scripts/predict_msms.py \
 - `--gen_ckpt` / `--inten_ckpt` — paths to ICEBERG checkpoints
 - `--collision_energies` — one or more collision energies in eV (e.g. `20 40 60`); model was trained on absolute eV values
 - `--adduct` — supported adducts: `[M+H]+`, `[M-H]-`, `[M+Na]+`, `[M+NH4]+`, and others from `ms_pred.common.ion2mass`
+- `--instrument` — instrument type for intensity prediction (e.g. `"Orbitrap"`, `"QTOF"`)
 - `--threshold` — confidence cutoff for DAG fragment generator (default `0.1`; lower = more fragments)
 - `--sparse_k` — maximum number of peaks returned (default `100`)
 - `--cuda_devices` — GPU device IDs (e.g. `"0"` or `"0,1"`); omit or set to `None` for CPU

.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/core_env.yaml

@@ -15,8 +15,11 @@ dependencies:
   - seaborn
   - h5py
   - pip:
-    - dgl -f https://data.dgl.ai/wheels/torch-2.4/repo.html
     - pytorch-lightning
+    - "ray[tune]==2.7.2"
+    - "setuptools<81"
+    - ipython
+    - pydantic
     - scikit-learn
     - pathos
     - CairoSVG

conda-envs/msms-agent/install.sh

@@ -22,18 +22,56 @@ fi
 source $(conda info --base)/etc/profile.d/conda.sh
 conda activate $ENV_NAME
 
-echo "Installing dgl from custom index..."
-uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.4/repo.html
+if [[ "$(uname)" == "Darwin" ]]; then
+    # macOS: DGL 2.2.0 graphbolt dylibs only go up to pytorch 2.3.0.
+    # torchdata>=0.8 dropped datapipes and pulls torch 2.12 as dep; pin 0.7.1 --no-deps.
+    # torch_scatter/sparse must match torch 2.3.0 exactly.
+    echo "macOS: installing torch 2.3.0 (CPU)..."
+    pip install "torch==2.3.0" --index-url https://download.pytorch.org/whl/cpu
 
-echo "Installing pip dependencies with uv..."
-uv pip install -r uv_requirements.txt
+    echo "macOS: installing DGL from torch-2.3 index..."
+    uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.3/repo.html
 
-# torch-scatter / torch-sparse: no generic arm64 wheels on PyPI.
-# Use the PyG find-links index pinned to torch 2.4.0+cpu.
-TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
-PYG_INDEX="https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html"
-echo "Installing torch-scatter and torch-sparse from PyG index (torch ${TORCH_VERSION})..."
-uv pip install torch-scatter torch-sparse --find-links "$PYG_INDEX"
+    echo "macOS: installing pip dependencies..."
+    uv pip install -r uv_requirements.txt
+
+    # DGL 2.2.0 graphbolt imports torchdata which is incompatible with torch 2.3.
+    # ICEBERG never uses graphbolt at runtime; stub all torchdata imports in graphbolt.
+    echo "macOS: patching DGL graphbolt to remove torchdata dependency..."
+    python - <<'DGLPATCH'
+import pathlib, sys, re
+site = pathlib.Path(sys.executable).parent.parent / "lib/python3.10/site-packages"
+graphbolt = site / "dgl/graphbolt"
+for f in graphbolt.glob("*.py"):
+    txt = f.read_text()
+    if "torchdata" not in txt:
+        continue
+    # Replace all torchdata imports with try/except stubs
+    patched = re.sub(
+        r'^((?:from|import) torchdata\S*.*)',
+        r'try:\n    \1\nexcept (ImportError, ModuleNotFoundError):\n    pass',
+        txt, flags=re.MULTILINE
+    )
+    f.write_text(patched)
+    print(f"Patched {f.name}")
+DGLPATCH
+
+    echo "macOS: installing torch-scatter and torch-sparse for torch 2.3.0..."
+    uv pip install --force-reinstall torch-scatter torch-sparse \
+        --find-links https://data.pyg.org/whl/torch-2.3.0+cpu.html
+else
+    echo "Installing dgl from custom index..."
+    uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.4/repo.html
+
+    echo "Installing pip dependencies with uv..."
+    uv pip install -r uv_requirements.txt
+
+    # torch-scatter / torch-sparse: use the PyG find-links index pinned to torch 2.4.0+cpu.
+    TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
+    PYG_INDEX="https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html"
+    echo "Installing torch-scatter and torch-sparse from PyG index (torch ${TORCH_VERSION})..."
+    uv pip install torch-scatter torch-sparse --find-links "$PYG_INDEX"
+fi
 
 # Install ms_pred from GitHub.
 # setup.py includes a Cython extension for massformer (not used by ICEBERG),
@@ -53,8 +91,60 @@ setup(
 )
 SETUP_EOF
 
+# Patch ms_pred source for upstream bugs before installing:
+#   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"
 
+# macOS only: create fake torchdata package so DGL graphbolt imports succeed.
+# DGL uses graphbolt only for distributed training; ICEBERG never triggers it.
+if [[ "$(uname)" == "Darwin" ]]; then
+    echo "macOS: creating torchdata compatibility shim for DGL..."
+    python - <<'TDPATCH'
+import pathlib, sys
+site = pathlib.Path(sys.executable).parent.parent / "lib/python3.10/site-packages"
+(site / "torchdata/datapipes/iter").mkdir(parents=True, exist_ok=True)
+(site / "torchdata/dataloader2").mkdir(parents=True, exist_ok=True)
+(site / "torchdata/__init__.py").write_text("")
+(site / "torchdata/datapipes/__init__.py").write_text("from . import iter\n")
+(site / "torchdata/datapipes/iter/__init__.py").write_text(
+    "import torch.utils.data\n"
+    "class IterDataPipe(torch.utils.data.IterableDataset): pass\n"
+    "class IterableWrapper(IterDataPipe):\n"
+    "    def __init__(self, iterable): self.iterable = iterable\n"
+    "    def __iter__(self): yield from self.iterable\n"
+    "class Mapper(IterDataPipe):\n"
+    "    def __init__(self, datapipe=None, fn=None): self.datapipe=datapipe; self.fn=fn\n"
+    "    def __iter__(self):\n"
+    "        for item in self.datapipe: yield self.fn(item)\n"
+)
+(site / "torchdata/dataloader2/__init__.py").write_text("")
+(site / "torchdata/dataloader2/graph.py").write_text(
+    "def traverse_dps(dp): return {}\n"
+    "def find_dps(g, t): return []\n"
+    "def replace_dp(g, old, new): return g\n"
+)
+print("Created torchdata shim")
+TDPATCH
+fi
+
 rm -f conda_only_env.yaml uv_requirements.txt
 echo "Environment $ENV_NAME created successfully!"