ultrams 0.1.2__tar.gz → 0.2.0__tar.gz

ultrams-0.2.0/MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include docs *.md
+recursive-include examples *.py *.md *.mgf

ultrams-0.2.0/PKG-INFO

@@ -0,0 +1,102 @@
+Metadata-Version: 2.4
+Name: ultrams
+Version: 0.2.0
+Summary: Pretrained UltraMS models for mass spectrometry
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/Dsadd4/UltraMS
+Project-URL: Source, https://github.com/Dsadd4/UltraMS
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<3,>=1.24
+Requires-Dist: torch>=2.1
+Requires-Dist: transformers<5,>=4.33
+Requires-Dist: huggingface-hub<1,>=0.23
+Provides-Extra: io
+Requires-Dist: pyteomics>=4.6; extra == "io"
+Dynamic: license-file
+
+# UltraMS
+
+**Pretrained models for tandem mass spectra.** Turn MS/MS spectra into embeddings, compare spectra, or fine-tune the encoder for a measured property.
+
+```bash
+python -m pip install ultrams
+```
+
+| I want to… | Start here |
+| --- | --- |
+| Embed one spectrum | [Python example](#embed-a-spectrum) |
+| Embed an MGF or mzML file | [File command](#embed-a-spectrum-file) |
+| Train on my labelled spectra | [PyTorch example](#fine-tune-with-pytorch) · [complete MGF example](https://github.com/Dsadd4/UltraMS/blob/main/examples/pytorch_finetune.py) |
+| Compare spectra | [Search example](https://github.com/Dsadd4/UltraMS/blob/main/examples/spectrum_search.py) |
+| Reproduce pretraining | [Training code](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) |
+
+## Embed a spectrum
+
+```python
+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
+
+```bash
+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](https://github.com/Dsadd4/UltraMS/tree/main/examples) 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.
+
+```python
+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](https://github.com/Dsadd4/UltraMS/blob/main/examples/pytorch_finetune.py) 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](https://github.com/Dsadd4/UltraMS/blob/main/docs/data-format.md) for the input format.
+
+## Choose a model
+
+| Model name | Representation | Dimension | Spectral peak limit | Weights |
+| --- | --- | ---: | ---: | --- |
+| `unsupervised` | Encoder spectrum embedding | 1024 | 150 | [Unsupervised](https://huggingface.co/dsadd4/UltraMS-Unsupervised) |
+| `mona` | MoNA contrastive projection | 1024 | 100 | [MoNA contrastive](https://huggingface.co/dsadd4/UltraMS-MoNA-Contrastive) |
+| `search` | UltraAtlas search projection | 512 | 150 | [Search](https://huggingface.co/dsadd4/UltraMS-Search) |
+
+Load any of them with `UltraMS.from_pretrained("name")`, or load a downloaded `model.pt` with `UltraMS.from_checkpoint(path)`. [Model selection](https://github.com/Dsadd4/UltraMS/blob/main/docs/model-selection.md) explains the training purpose and output of each checkpoint. The Search checkpoint is the **UltraAtlas application `best.pt`**.
+
+## Pretraining
+
+The [UltraMSdata pretraining package](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) contains the current model architecture, training entry points, configuration and history outputs. It installs separately from the inference package.

ultrams-0.2.0/README.md

@@ -0,0 +1,84 @@
+# UltraMS
+
+**Pretrained models for tandem mass spectra.** Turn MS/MS spectra into embeddings, compare spectra, or fine-tune the encoder for a measured property.
+
+```bash
+python -m pip install ultrams
+```
+
+| I want to… | Start here |
+| --- | --- |
+| Embed one spectrum | [Python example](#embed-a-spectrum) |
+| Embed an MGF or mzML file | [File command](#embed-a-spectrum-file) |
+| Train on my labelled spectra | [PyTorch example](#fine-tune-with-pytorch) · [complete MGF example](https://github.com/Dsadd4/UltraMS/blob/main/examples/pytorch_finetune.py) |
+| Compare spectra | [Search example](https://github.com/Dsadd4/UltraMS/blob/main/examples/spectrum_search.py) |
+| Reproduce pretraining | [Training code](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) |
+
+## Embed a spectrum
+
+```python
+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
+
+```bash
+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](https://github.com/Dsadd4/UltraMS/tree/main/examples) 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.
+
+```python
+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](https://github.com/Dsadd4/UltraMS/blob/main/examples/pytorch_finetune.py) 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](https://github.com/Dsadd4/UltraMS/blob/main/docs/data-format.md) for the input format.
+
+## Choose a model
+
+| Model name | Representation | Dimension | Spectral peak limit | Weights |
+| --- | --- | ---: | ---: | --- |
+| `unsupervised` | Encoder spectrum embedding | 1024 | 150 | [Unsupervised](https://huggingface.co/dsadd4/UltraMS-Unsupervised) |
+| `mona` | MoNA contrastive projection | 1024 | 100 | [MoNA contrastive](https://huggingface.co/dsadd4/UltraMS-MoNA-Contrastive) |
+| `search` | UltraAtlas search projection | 512 | 150 | [Search](https://huggingface.co/dsadd4/UltraMS-Search) |
+
+Load any of them with `UltraMS.from_pretrained("name")`, or load a downloaded `model.pt` with `UltraMS.from_checkpoint(path)`. [Model selection](https://github.com/Dsadd4/UltraMS/blob/main/docs/model-selection.md) explains the training purpose and output of each checkpoint. The Search checkpoint is the **UltraAtlas application `best.pt`**.
+
+## Pretraining
+
+The [UltraMSdata pretraining package](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) contains the current model architecture, training entry points, configuration and history outputs. It installs separately from the inference package.

ultrams-0.2.0/docs/data-format.md

@@ -0,0 +1,53 @@
+# Spectrum input
+
+UltraMS accepts an MS/MS spectrum as two equal-length, one-dimensional arrays and the precursor ion's measured `m/z`:
+
+```python
+from ultrams import UltraMS
+
+model = UltraMS.from_pretrained("unsupervised")
+result = model.encode(
+    mz=[100.1, 121.1, 150.0],
+    intensity=[20, 100, 35],
+    precursor_mz=301.2,
+)
+print(result.embedding.shape)  # (1024,)
+```
+
+`mz` and `intensity` describe the same spectral peaks in the same order. `precursor_mz` is a positive, finite number. Every spectral peak `m/z` and intensity must be finite. The input must contain at least three spectral peaks with positive `m/z` after preprocessing.
+
+## Preprocessing
+
+`encode()` and `batch_converter()` apply the same preparation:
+
+1. Discard spectral peaks with `m/z <= 0`.
+2. Divide intensities by the largest retained intensity when it is positive, then clip intensities to `[0, 1]`.
+3. If the spectrum exceeds the model's limit, keep the most intense spectral peaks.
+4. Sort the retained spectral peaks by `m/z`.
+
+The three released models accept up to **150** spectral peaks for `"unsupervised"`, **100** for `"mona"`, and **150** for `"search"`. You can inspect the loaded model's limit with `model.input_max_peaks`. A batch made with `model.batch_converter()` pads shorter spectra and supplies an `attention_mask`; padding does not count as a measured spectral peak.
+
+For a PyTorch `Dataset`, return one mapping per spectrum:
+
+```python
+{
+    "mz": [100.1, 121.1, 150.0],
+    "intensity": [20, 100, 35],
+    "precursor_mz": 301.2,
+    "target": 0.5,
+}
+```
+
+The `target` field is a numeric training value used by `model.batch_converter()` when present in every item of a batch. For `model.finetune(...)`, use `"label"` instead; see [PyTorch fine-tuning](../examples/pytorch_finetune.py) and the [README](../README.md).
+
+`model.encode(...).embedding` is a NumPy vector for inference. `model(peaks, attention_mask, precursor_mz)` returns a `[batch, embedding_dim]` PyTorch tensor with gradients for training. See [model selection](model-selection.md) for the representation returned by each checkpoint.
+
+## Spectrum files
+
+For MGF or `.mgf.gz`, UltraMS reads each `BEGIN IONS` block, uses `PEPMASS` for precursor-ion `m/z`, and takes the spectrum ID from `TITLE`, `NAME`, or `SCANS` when available. Extra columns after a spectral peak's `m/z` and intensity are ignored. The optional `ultrams[io]` installation adds mzML input; only MS2 spectra are read.
+
+```bash
+ultrams-embed examples/data/example_5_spectra.mgf embeddings.npz --model search
+```
+
+The `.npz` output contains aligned `ids` and `embeddings` arrays. See the [batch embedding example](../examples/batch_embeddings.py) to use the file reader and model directly in Python.

ultrams-0.2.0/docs/model-selection.md

@@ -0,0 +1,27 @@
+# Choose a model
+
+All three checkpoints use the UltraMS encoder. `UltraMS.from_pretrained(name)` downloads the selected weights from Hugging Face; `model.encode(...).embedding` returns the representation associated with that checkpoint.
+
+| Model | Load with | Training | `embedding` | Dimension | Maximum spectral peaks |
+| --- | --- | --- | --- | ---: | ---: |
+| [Unsupervised](https://huggingface.co/dsadd4/UltraMS-Unsupervised) | `"unsupervised"` | Self-supervised pretraining on UltraMSdata, followed by retention-time training | Encoder spectrum-level CLS embedding | 1024 | 150 |
+| [MoNA contrastive](https://huggingface.co/dsadd4/UltraMS-MoNA-Contrastive) | `"mona"` | Contrastive training on MoNA spectra | Projection of the encoder's CLS embedding | 1024 | 100 |
+| [Search](https://huggingface.co/dsadd4/UltraMS-Search) | `"search"` | Spectrum-to-spectrum contrastive training for the UltraAtlas search model | Projection of the encoder's CLS embedding | 512 | 150 |
+
+Use **Unsupervised** as the general starting point for embeddings or a new downstream task. Use **MoNA contrastive** for a representation trained on molecular-identity relationships in MoNA. Use **Search** when working in the spectrum similarity space used by the UltraAtlas application. A model alone does not include the UltraAtlas reference collection or its search index; the [spectrum search example](../examples/spectrum_search.py) shows how to rank a supplied spectral library.
+
+The Unsupervised release comes from the completed retention-time-only epoch 11 checkpoint (`rtonly11`). The Search release is the `best.pt` selected for the UltraAtlas application. These source names identify the released weights; the public API uses the names in the table.
+
+```python
+from ultrams import UltraMS
+
+model = UltraMS.from_pretrained("search")
+vector = model.encode(
+    mz=[100.1, 121.1, 150.0],
+    intensity=[20, 100, 35],
+    precursor_mz=301.2,
+).embedding
+print(vector.shape)  # (512,)
+```
+
+For every model, `result.cls` exposes the encoder's normalized CLS embedding. MoNA and Search also expose `result.projection`. The Unsupervised and Search default embeddings are unit-normalized; the MoNA default embedding is the learned projection before normalization. Use cosine similarity when comparing spectra. Input requirements and spectral peak preparation are in [Spectrum input](data-format.md).

ultrams-0.2.0/examples/README.md

@@ -0,0 +1,15 @@
+# Run UltraMS on MS/MS spectra
+
+These examples use five spectra from the [DreaMS project](data/NOTICE.md). Install UltraMS from the repository root first:
+
+```bash
+python -m pip install -e .
+```
+
+| Task | Run |
+| --- | --- |
+| Read an MGF file and save its embeddings | `python examples/batch_embeddings.py` |
+| Fine-tune with a native PyTorch loop | `python examples/pytorch_finetune.py` |
+| Rank spectra by cosine similarity | `python examples/spectrum_search.py` |
+
+Each command downloads its pretrained checkpoint from Hugging Face on first use. `batch_embeddings.py` saves an `.npz` containing spectrum IDs and one embedding per input spectrum. `spectrum_search.py` saves a CSV of ranked matches; its scores illustrate spectrum similarity and are not a benchmark. `pytorch_finetune.py` uses **placeholder regression targets based on file order** solely to show how to train the encoder and a task head. Replace these targets with measured labels for research. By default, output files go to a newly created temporary directory printed by each command; use `--output` or `--output-dir` to choose a persistent destination.

ultrams-0.2.0/examples/batch_embeddings.py

@@ -0,0 +1,36 @@
+"""Read an MGF file and save one UltraMS embedding per spectrum."""
+
+import argparse
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from ultrams import UltraMS, read_spectra
+
+
+EXAMPLE_MGF = Path(__file__).parent / "data" / "example_5_spectra.mgf"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input", type=Path, default=EXAMPLE_MGF, help="MGF input file")
+    parser.add_argument("--model", choices=("unsupervised", "mona", "search"), default="unsupervised")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--output", type=Path, help="Output .npz path")
+    args = parser.parse_args()
+
+    spectra = list(read_spectra(args.input))
+    model = UltraMS.from_pretrained(args.model, device=args.device)
+    embeddings = model.encode_batch(spectra, batch_size=args.batch_size)
+    output = args.output or Path(tempfile.mkdtemp(prefix="ultrams-embeddings-")) / "embeddings.npz"
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, ids=np.asarray([spectrum["id"] for spectrum in spectra]), embeddings=embeddings)
+    print(f"Encoded {len(spectra)} spectra into {embeddings.shape} embeddings")
+    print(output)
+
+
+if __name__ == "__main__":
+    main()

ultrams-0.2.0/examples/data/NOTICE.md

@@ -0,0 +1,29 @@
+# Example spectra
+
+`example_5_spectra.mgf` is copied without modification from the [DreaMS repository](https://github.com/pluskal-lab/DreaMS/blob/main/data/examples/example_5_spectra.mgf). It contains five example MS/MS spectra. SHA-256 of this copy: `a68132b1528dd818a4297a85463c21384bc80b8cad78231272e43528eae44d30`.
+
+The upstream repository publishes this file under the following MIT license:
+
+```text
+MIT License
+
+Copyright (c) 2024 Roman Bushuiev
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+```