Add MassSpecGym evaluation adapter and safetensors runtime loader

README.md

@@ -107,3 +107,8 @@ MS/MS structure inference can affect downstream scientific interpretation. Users
 ## Citation
 
 If you use this model, cite the NexaMass project release and the accompanying technical report when available. Relevant background work includes DreaMS for self-supervised MS/MS representation learning, MassSpecGym for benchmark framing, CSI:FingerID for fingerprint-mediated candidate search, and related spectra-structure retrieval and de novo generation systems such as MIST, MSNovelist, CMSSP, CSU-MS2, MSBERT, Spec2Mol, and MS2Mol.
+
+## MassSpecGym Adapter
+
+A safetensors-compatible MassSpecGym retrieval adapter is included under `evaluation/massspecgym/`. It loads `weights/NexaMass-V3-Struct-model_state.safetensors`, converts MassSpecGym tokenized spectra into the NexaMass batch contract, and reports Hit@k retrieval metrics through MassSpecGym's evaluator. The archived reference run reached test Hit@20 `0.3505` with the frozen projected-dot scorer. This should be read as evidence of transferable top-k signal, not solved molecular ranking or calibrated confidence.
+

config.json

@@ -16,6 +16,14 @@
   "architectures": [
     "NexaMassSpectralEncoder"
   ],
+  "evaluation_adapters": {
+    "massspecgym": {
+      "benchmark": "MassSpecGym molecule retrieval",
+      "claim_boundary": "top-k transfer signal; ranking and confidence remain open decision-layer problems",
+      "path": "evaluation/massspecgym/run_massspecgym_retrieval_hf.py",
+      "reference_result": "test Hit@20 0.3505 with frozen V3 projected-dot scorer under Hit@k-only evaluation"
+    }
+  },
   "foundation_checkpoint": "weights/Final_V3-model_state.safetensors",
   "foundation_checkpoint_format": "safetensors",
   "full_training_checkpoints": {

evaluation/massspecgym/README.md

@@ -0,0 +1,35 @@
+# MassSpecGym Evaluation Adapter
+
+This directory contains the public Hugging Face adapter used to position `NexaMass-V3-Struct` on the MassSpecGym molecule-retrieval task.
+
+The adapter loads the safetensors-only public checkpoint and wraps MassSpecGym's own `RetrievalDataset`, `MassSpecDataModule`, and retrieval evaluator. It is meant for external benchmark positioning, not for claiming that ranking or confidence are solved.
+
+## Install
+
+Use an isolated environment because MassSpecGym has its own dependency surface:
+
+```bash
+python -m pip install torch safetensors huggingface_hub massspecgym==1.3.1 pytorch-lightning
+```
+
+## Run From A Clone Of This HF Repo
+
+```bash
+python evaluation/massspecgym/run_massspecgym_retrieval_hf.py \
+  --checkpoint weights/NexaMass-V3-Struct-model_state.safetensors \
+  --config config.json \
+  --split test \
+  --scorer projected_dot \
+  --hit-only \
+  --batch-size 32 \
+  --num-workers 25 \
+  --output-json evaluation/massspecgym/results/local_massspecgym_test.json
+```
+
+If the checkpoint is not present locally, the script can download it from this repo through `huggingface_hub`.
+
+## Reported Reference Result
+
+The archived adapter run reached MassSpecGym test Hit@20 `0.3505` under Hit@k-only evaluation using the frozen V3 projected-dot scorer. This put the model above lower baselines such as Random, DeepSets, Fingerprint FFN, and DeepSets+Fourier, while remaining below specialized retrieval systems such as MIST.
+
+Interpretation: the encoder transfers real top-k structure signal to retrieval, but exact local ranking and calibrated confidence remain separate downstream problems.

evaluation/massspecgym/results/massspecgym_hitk_summary.json

@@ -0,0 +1,16 @@
+{
+  "benchmark": "MassSpecGym molecule retrieval",
+  "adapter": "evaluation/massspecgym/run_massspecgym_retrieval_hf.py",
+  "checkpoint": "weights/NexaMass-V3-Struct-model_state.safetensors",
+  "scorer": "projected_dot",
+  "evaluation_mode": "test dataloader through validation loop, Hit@k-only",
+  "metrics": {
+    "test_hit_at_1": 0.0627,
+    "test_hit_at_5": 0.1753,
+    "test_hit_at_20": 0.3505,
+    "val_hit_at_1": 0.1162,
+    "val_hit_at_5": 0.1915,
+    "val_hit_at_20": 0.3328
+  },
+  "claim_boundary": "External positioning sanity check; demonstrates top-k transfer signal, not solved ranking or confidence."
+}

evaluation/massspecgym/run_massspecgym_retrieval_hf.py

@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+"""Evaluate NexaMass-V3-Struct on MassSpecGym retrieval.
+
+This is the Hugging Face release adapter. It loads the public safetensors
+checkpoint from this repository and wraps MassSpecGym's official retrieval data
+module/evaluator. The adapter is for external benchmark positioning, not for
+claiming that ranking or confidence are solved.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from runtime.nexamass_encoder import ModelConfig, NexaMassSpectralEncoder, load_nexamass_model_state  # noqa: E402
+
+
+def _require_massspecgym() -> tuple[Any, Any, Any, Any, Any]:
+    try:
+        from massspecgym.data import MassSpecDataModule, RetrievalDataset
+        from massspecgym.data.transforms import MolFingerprinter, SpecTokenizer
+        from massspecgym.models.retrieval.base import RetrievalMassSpecGymModel
+        from pytorch_lightning import Trainer
+    except ImportError as exc:
+        raise SystemExit(
+            "MassSpecGym dependencies are missing. Install in an isolated env with: "
+            "python -m pip install massspecgym==1.3.1 pytorch-lightning safetensors huggingface_hub"
+        ) from exc
+    return Trainer, MassSpecDataModule, RetrievalDataset, MolFingerprinter, SpecTokenizer, RetrievalMassSpecGymModel
+
+
+def _cfg_from_json(path: Path) -> ModelConfig:
+    if not path.exists():
+        return ModelConfig()
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    arch = payload.get("architecture_config", payload)
+    allowed = ModelConfig.__dataclass_fields__.keys()
+    return ModelConfig(**{key: arch[key] for key in allowed if key in arch})
+
+
+def _resolve_checkpoint(path: Path, repo_id: str, filename: str) -> Path:
+    if path.exists():
+        return path
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise SystemExit("Checkpoint was not found locally and huggingface_hub is not installed.") from exc
+    return Path(hf_hub_download(repo_id=repo_id, repo_type="model", filename=filename))
+
+
+def _parse_limit_batches(raw: str) -> int | float:
+    value = raw.strip()
+    if value.isdigit():
+        return int(value)
+    return float(value)
+
+
+def _batch_from_massspecgym_spec(
+    spec: torch.Tensor,
+    cfg: ModelConfig,
+    device: torch.device,
+    *,
+    precursor_mz: torch.Tensor | None = None,
+) -> dict[str, torch.Tensor]:
+    """Convert MassSpecGym tokenized spectra into NexaMass' encoder batch contract."""
+
+    if spec.ndim != 3 or spec.shape[-1] < 2:
+        raise ValueError(f"Expected MassSpecGym spec shape [batch, peaks, >=2], got {tuple(spec.shape)}")
+
+    spec = spec.to(device=device, dtype=torch.float32)
+    mzs_raw = spec[..., 0].clamp(min=0.0)
+    ints_raw = spec[..., 1].clamp(min=0.0)
+    batch_size, peak_count = mzs_raw.shape
+    if peak_count > cfg.max_peaks:
+        mzs_raw = mzs_raw[:, : cfg.max_peaks]
+        ints_raw = ints_raw[:, : cfg.max_peaks]
+        peak_count = cfg.max_peaks
+
+    mask = (mzs_raw > 0) & torch.isfinite(mzs_raw) & torch.isfinite(ints_raw)
+    max_intensity = ints_raw.masked_fill(~mask, 0.0).amax(dim=1, keepdim=True).clamp(min=1e-6)
+    mzs_norm = (mzs_raw / cfg.mz_max).clamp(0.0, 1.5)
+    ints_norm = (ints_raw / max_intensity).masked_fill(~mask, 0.0)
+    if precursor_mz is not None:
+        precursor_raw = precursor_mz.to(device=device, dtype=torch.float32).view(-1).clamp(min=1e-6)
+        if precursor_raw.numel() != batch_size:
+            raise ValueError(f"Expected {batch_size} precursor_mz values, got {precursor_raw.numel()}")
+    else:
+        precursor_raw = mzs_raw.masked_fill(~mask, 0.0).amax(dim=1).clamp(min=1e-6)
+    mz_to_precursor = (mzs_raw / precursor_raw[:, None]).clamp(0.0, 2.0).masked_fill(~mask, 0.0)
+    ranks = torch.linspace(0.0, 1.0, peak_count, device=device, dtype=torch.float32)[None, :].expand(batch_size, -1)
+
+    if peak_count < cfg.max_peaks:
+        pad_width = cfg.max_peaks - peak_count
+
+        def pad(values: torch.Tensor, value: float = 0.0) -> torch.Tensor:
+            return F.pad(values, (0, pad_width), value=value)
+
+        mzs_norm = pad(mzs_norm)
+        ints_norm = pad(ints_norm)
+        mz_to_precursor = pad(mz_to_precursor)
+        ranks = pad(ranks)
+        mask = F.pad(mask, (0, pad_width), value=False)
+
+    observed_peak_count = mask.sum(dim=1).to(dtype=torch.float32).clamp(min=1.0)
+    return {
+        "mzs": mzs_norm,
+        "ints": ints_norm,
+        "mz_to_precursor": mz_to_precursor,
+        "peak_rank": ranks,
+        "mask": mask.to(dtype=torch.bool),
+        "precursor_mz": (precursor_raw / cfg.mz_max).clamp(max=2.0),
+        "charge": torch.zeros(batch_size, device=device, dtype=torch.float32),
+        "collision_energy": torch.zeros(batch_size, device=device, dtype=torch.float32),
+        "adduct_id": torch.zeros(batch_size, device=device, dtype=torch.long),
+        "instrument_id": torch.zeros(batch_size, device=device, dtype=torch.long),
+        "peak_count": observed_peak_count / float(cfg.max_peaks),
+    }
+
+
+def _scores_for_batch(
+    *,
+    scorer: str,
+    model: NexaMassSpectralEncoder,
+    cfg: ModelConfig,
+    spec: torch.Tensor,
+    candidates: torch.Tensor,
+    batch_ptr: torch.Tensor,
+    precursor_mz: torch.Tensor | None,
+    device: torch.device,
+) -> torch.Tensor:
+    batch = _batch_from_massspecgym_spec(spec, cfg, device, precursor_mz=precursor_mz)
+    candidates = candidates.to(device=device, dtype=torch.float32)
+    with torch.no_grad():
+        _embedding, _raw_projected, logits, query_raw = model.forward_with_heads(batch)
+        pred_probs = torch.sigmoid(logits)
+        if scorer == "predicted_fingerprint":
+            query_repeated = F.normalize(pred_probs, dim=-1).repeat_interleave(batch_ptr.to(device), dim=0)
+            return F.cosine_similarity(query_repeated, F.normalize(candidates, dim=-1), dim=-1).detach()
+        if scorer == "projected_dot":
+            query_repeated = F.normalize(query_raw, dim=-1).repeat_interleave(batch_ptr.to(device), dim=0)
+            target_projection = F.normalize(model.project_structure_targets(candidates), dim=-1)
+            return (query_repeated * target_projection).sum(dim=-1).detach()
+    raise ValueError(f"Unsupported scorer: {scorer}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--repo-id", default="AethronPhantom/NexaMass-V3-Struct")
+    parser.add_argument("--checkpoint", type=Path, default=REPO_ROOT / "weights/NexaMass-V3-Struct-model_state.safetensors")
+    parser.add_argument("--checkpoint-filename", default="weights/NexaMass-V3-Struct-model_state.safetensors")
+    parser.add_argument("--config", type=Path, default=REPO_ROOT / "config.json")
+    parser.add_argument("--scorer", choices=["projected_dot", "predicted_fingerprint"], default="projected_dot")
+    parser.add_argument("--split", choices=["val", "test"], default="test")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--n-peaks", type=int, default=256)
+    parser.add_argument("--accelerator", default="gpu")
+    parser.add_argument("--devices", default="1")
+    parser.add_argument("--limit-batches", default="1.0")
+    parser.add_argument("--hit-only", action="store_true", help="Use validation loop over test dataloader for Hit@k-only scoring.")
+    parser.add_argument("--inspect-batch-only", action="store_true")
+    parser.add_argument("--output-json", type=Path)
+    args = parser.parse_args()
+
+    Trainer, MassSpecDataModule, RetrievalDataset, MolFingerprinter, SpecTokenizer, RetrievalMassSpecGymModel = (
+        _require_massspecgym()
+    )
+
+    torch.set_float32_matmul_precision("high")
+    limit_batches = _parse_limit_batches(args.limit_batches)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    cfg = _cfg_from_json(args.config)
+    checkpoint = _resolve_checkpoint(args.checkpoint.expanduser(), args.repo_id, args.checkpoint_filename)
+    v3_model = load_nexamass_model_state(str(checkpoint), cfg=cfg, map_location="cpu")
+    v3_model.to(device)
+    v3_model.eval()
+
+    class NexaMassRetrievalModel(RetrievalMassSpecGymModel):  # type: ignore[misc, valid-type]
+        def __init__(self) -> None:
+            super().__init__()
+            self._inspected = False
+
+        def forward(self, spec: torch.Tensor) -> torch.Tensor:
+            batch = _batch_from_massspecgym_spec(spec, cfg, device)
+            with torch.no_grad():
+                _embedding, _raw, logits, query_raw = v3_model.forward_with_heads(batch)
+            return query_raw if args.scorer == "projected_dot" else torch.sigmoid(logits)
+
+        def step(self, batch: dict[str, Any], stage: Any) -> dict[str, torch.Tensor]:
+            if args.inspect_batch_only and not self._inspected:
+                print(
+                    json.dumps(
+                        {
+                            "batch_keys": sorted(batch.keys()),
+                            "spec_shape": list(batch["spec"].shape),
+                            "candidates_mol_shape": list(batch["candidates_mol"].shape),
+                            "batch_ptr_head": batch["batch_ptr"].detach().cpu().tolist()[:8],
+                        },
+                        indent=2,
+                    ),
+                    flush=True,
+                )
+                self._inspected = True
+            scores = _scores_for_batch(
+                scorer=args.scorer,
+                model=v3_model,
+                cfg=cfg,
+                spec=batch["spec"],
+                candidates=batch["candidates_mol"],
+                batch_ptr=batch["batch_ptr"],
+                precursor_mz=batch.get("precursor_mz"),
+                device=device,
+            )
+            return {"loss": torch.zeros((), device=scores.device), "scores": scores}
+
+    dataset = RetrievalDataset(
+        spec_transform=SpecTokenizer(n_peaks=args.n_peaks),
+        mol_transform=MolFingerprinter(fp_size=cfg.fingerprint_dim),
+    )
+    data_module = MassSpecDataModule(dataset=dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+    data_module.prepare_data()
+    data_module.setup(None if args.split == "val" else "test")
+    model = NexaMassRetrievalModel()
+    trainer = Trainer(
+        accelerator=args.accelerator,
+        devices=args.devices,
+        logger=False,
+        enable_checkpointing=False,
+        limit_val_batches=limit_batches if args.split == "val" or args.hit_only else 1.0,
+        limit_test_batches=limit_batches if args.split == "test" else 1.0,
+    )
+    if args.split == "val":
+        metrics = trainer.validate(model, datamodule=data_module)
+    elif args.hit_only:
+        metrics = trainer.validate(model, dataloaders=data_module.test_dataloader())
+    else:
+        metrics = trainer.test(model, datamodule=data_module)
+
+    payload = {
+        "checkpoint": str(checkpoint),
+        "scorer": args.scorer,
+        "split": args.split,
+        "metrics": metrics,
+        "massspecgym_adapter": {
+            "repo_id": args.repo_id,
+            "n_peaks": args.n_peaks,
+            "fingerprint_dim": cfg.fingerprint_dim,
+            "limit_batches": limit_batches,
+            "hit_only": args.hit_only,
+            "metadata_defaults": "charge/collision/adduct/instrument set to zero when absent from MassSpecGym batch",
+        },
+    }
+    print(json.dumps(payload, indent=2), flush=True)
+    if args.output_json:
+        args.output_json.parent.mkdir(parents=True, exist_ok=True)
+        args.output_json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

runtime/nexamass_encoder.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from pathlib import Path
 
 import torch
 import torch.nn as nn
@@ -126,14 +127,49 @@ class NexaMassSpectralEncoder(nn.Module):
         return F.normalize(self.target_projection(targets), dim=-1)
 
 
+
+
+def load_nexamass_state_dict(
+    checkpoint_path: str,
+    map_location: str | torch.device = "cpu",
+) -> dict[str, torch.Tensor]:
+    """Load public NexaMass model-state weights from Safetensors or PyTorch.
+
+    Hugging Face public release weights are Safetensors-only. The PyTorch branch is
+    kept for internal/object-storage compatibility with full training checkpoints
+    and model-state fallbacks.
+    """
+
+    path = Path(checkpoint_path)
+    if path.suffix == ".safetensors":
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:  # pragma: no cover - dependency message path
+            raise RuntimeError("Install safetensors to load NexaMass public weights: pip install safetensors") from exc
+        device = str(map_location) if isinstance(map_location, str) else "cpu"
+        if device not in {"cpu", "cuda"} and not device.startswith("cuda:"):
+            device = "cpu"
+        return load_file(str(path), device=device)
+
+    try:
+        payload = torch.load(path, map_location=map_location, weights_only=True)
+    except TypeError:  # older PyTorch
+        payload = torch.load(path, map_location=map_location)
+    if isinstance(payload, dict) and "model_state" in payload:
+        return payload["model_state"]
+    if isinstance(payload, dict):
+        return payload
+    raise TypeError(f"Unsupported NexaMass checkpoint payload type: {type(payload)!r}")
+
+
 def load_nexamass_model_state(
     checkpoint_path: str,
     cfg: ModelConfig | None = None,
     map_location: str | torch.device = "cpu",
 ) -> NexaMassSpectralEncoder:
-    payload = torch.load(checkpoint_path, map_location=map_location)
+    state_dict = load_nexamass_state_dict(checkpoint_path, map_location=map_location)
     cfg = cfg or ModelConfig()
     model = NexaMassSpectralEncoder(cfg)
-    model.load_state_dict(payload["model_state"], strict=True)
+    model.load_state_dict(state_dict, strict=True)
     model.eval()
     return model