UltraMS

Pretrained models for tandem mass spectra. Turn MS/MS spectra into embeddings, compare spectra, or fine-tune the encoder for a measured property.

python -m pip install ultrams

Embed a spectrum

from ultrams import UltraMS

model = UltraMS.from_pretrained("unsupervised")
embedding = model.encode(
    mz=[100.1, 121.1, 150.0],
    intensity=[20, 100, 35],
    precursor_mz=301.2,
).embedding
print(embedding.shape)  # (1024,)

The first call downloads the selected checkpoint from Hugging Face. Use device="cuda" or device="mps" with from_pretrained to run on an available accelerator.

Embed a spectrum file

ultrams-embed spectra.mgf embeddings.npz --model unsupervised

The output has aligned ids and embeddings arrays. To use mzML, first run python -m pip install 'ultrams[io]'; the base installation reads MGF and .mgf.gz without an extra file-format package. The checked-in example includes five real MS/MS spectra and a complete Python batch workflow.

Fine-tune with PyTorch

UltraMS is a torch.nn.Module. Each dataset item contains peak mz, peak intensity, precursor_mz, and a measured numeric target. The two rows below are runnable example data; substitute your measured spectra and targets for research.

import torch
from torch.utils.data import DataLoader
from ultrams import UltraMS

dataset = [
    {"mz": [100.1, 121.1, 150.0], "intensity": [20, 100, 35], "precursor_mz": 301.2, "target": 0.5},
    {"mz": [102.1, 135.2, 167.3], "intensity": [40, 100, 25], "precursor_mz": 315.3, "target": 0.7},
]
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UltraMS.from_pretrained("unsupervised", device=device).train()
head = torch.nn.Linear(model.embedding_dim, 1).to(device)
loader = DataLoader(dataset, batch_size=2, collate_fn=model.batch_converter())
optimizer = torch.optim.AdamW([*model.parameters(), *head.parameters()], lr=1e-5)

for batch in loader:
    batch = {name: value.to(device) for name, value in batch.items()}
    prediction = head(model(batch["peaks"], batch["attention_mask"], batch["precursor_mz"]))
    loss = torch.nn.functional.mse_loss(prediction.squeeze(-1), batch["target"])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"loss: {loss.item():.4f}")

The complete fine-tuning example reads an MGF file and saves training history. For a shorter route, use model.finetune(labelled_spectra, task="regression") with a label field; it saves weights, configuration, and losses. See spectrum input for the input format.

Choose a model

Load any of them with UltraMS.from_pretrained("name"), or load a downloaded model.pt with UltraMS.from_checkpoint(path). Model selection explains the training purpose and output of each checkpoint. The Search checkpoint is the UltraAtlas application best.pt.

Pretraining

The UltraMSdata pretraining package contains the current model architecture, training entry points, configuration and history outputs. It installs separately from the inference package.