ultrams 0.2.0__tar.gz → 0.3.0__tar.gz

ultrams-0.3.0/CITATION.cff

@@ -0,0 +1,15 @@
+cff-version: 1.2.0
+message: "If you use UltraMS in your research, please cite this software."
+type: software
+title: UltraMS
+authors:
+  - family-names: Li
+    given-names: Yifan
+  - family-names: Zhong
+    given-names: Yunhua
+  - family-names: Liu
+    given-names: Pan
+  - family-names: Xia
+    given-names: Jun
+repository-code: https://github.com/Dsadd4/UltraMS
+license: Apache-2.0

ultrams-0.3.0/MANIFEST.in

@@ -0,0 +1,4 @@
+recursive-include docs *.md
+recursive-include examples *.py *.md *.mgf
+recursive-include cookbook/tutorials *.ipynb *.md *.py
+include CITATION.cff

{ultrams-0.2.0/src/ultrams.egg-info → ultrams-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ultrams
-Version: 0.2.0
+Version: 0.3.0
 Summary: Pretrained UltraMS models for mass spectrometry
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/Dsadd4/UltraMS
@@ -30,7 +30,8 @@ python -m pip install ultrams
 | 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) |
+| Learn in a notebook | [UltraMS tutorials](https://github.com/Dsadd4/UltraMS/tree/main/cookbook/tutorials) |
+| Inspect or run pretraining | [Training code](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) |
 
 ## Embed a spectrum
 
@@ -48,6 +49,8 @@ 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.
 
+The three checkpoints are collected on [Hugging Face](https://huggingface.co/collections/dsadd4/ultrams-6ab4cfaa860cbbd1f9a6b9b6).
+
 ## Embed a spectrum file
 
 ```bash
@@ -87,6 +90,8 @@ for batch in loader:
 
 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.
 
+For measured molecular labels, the [MassSpecGym contrastive tutorial](https://github.com/Dsadd4/UltraMS/blob/main/cookbook/tutorials/MASSSPECGYM_CONTRASTIVE.md) downloads the dataset, keeps its train/validation/test folds, fine-tunes the encoder, and reloads the selected checkpoint.
+
 ## Choose a model
 
 | Model name | Representation | Dimension | Spectral peak limit | Weights |
@@ -100,3 +105,7 @@ Load any of them with `UltraMS.from_pretrained("name")`, or load a downloaded `m
 ## 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.
+
+## Cite
+
+Citation metadata for this software is in [CITATION.cff](https://github.com/Dsadd4/UltraMS/blob/main/CITATION.cff).

{ultrams-0.2.0 → ultrams-0.3.0}/README.md

@@ -12,7 +12,8 @@ python -m pip install ultrams
 | 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) |
+| Learn in a notebook | [UltraMS tutorials](https://github.com/Dsadd4/UltraMS/tree/main/cookbook/tutorials) |
+| Inspect or run pretraining | [Training code](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) |
 
 ## Embed a spectrum
 
@@ -30,6 +31,8 @@ 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.
 
+The three checkpoints are collected on [Hugging Face](https://huggingface.co/collections/dsadd4/ultrams-6ab4cfaa860cbbd1f9a6b9b6).
+
 ## Embed a spectrum file
 
 ```bash
@@ -69,6 +72,8 @@ for batch in loader:
 
 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.
 
+For measured molecular labels, the [MassSpecGym contrastive tutorial](https://github.com/Dsadd4/UltraMS/blob/main/cookbook/tutorials/MASSSPECGYM_CONTRASTIVE.md) downloads the dataset, keeps its train/validation/test folds, fine-tunes the encoder, and reloads the selected checkpoint.
+
 ## Choose a model
 
 | Model name | Representation | Dimension | Spectral peak limit | Weights |
@@ -82,3 +87,7 @@ Load any of them with `UltraMS.from_pretrained("name")`, or load a downloaded `m
 ## 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.
+
+## Cite
+
+Citation metadata for this software is in [CITATION.cff](https://github.com/Dsadd4/UltraMS/blob/main/CITATION.cff).

ultrams-0.3.0/cookbook/tutorials/MASSSPECGYM_CONTRASTIVE.md

@@ -0,0 +1,22 @@
+# Contrastive fine-tuning with MassSpecGym
+
+This tutorial fine-tunes the **UltraMS Unsupervised** encoder on MassSpecGym. Two distinct MS/MS spectra with the same `inchikey` are a positive pair; spectra from different molecules in a batch are negatives. The dataset's `train`, `val`, and `test` folds are retained, and the script checks that no molecular key appears in more than one fold. It keeps spectra with at least three positive-intensity, positive-m/z peaks and records the source and selected row and molecule counts in `data_audit.json`.
+
+Install UltraMS, then run the small teaching subset from the repository root:
+
+```bash
+python -m pip install -e .
+python cookbook/tutorials/massspecgym_contrastive.py --output-dir /path/to/my_run
+```
+
+The default selects up to 64 train molecules and 24 molecules from each evaluation fold, with two distinct spectra per molecule. Selection is deterministic. To use every training molecule with at least two distinct spectra and all valid evaluation spectra, after removing exact spectral duplicates:
+
+```bash
+python cookbook/tutorials/massspecgym_contrastive.py --full --output-dir /path/to/full_run
+```
+
+The full run can take substantially longer. `--data-file /path/to/MassSpecGym.tsv` reuses a local copy; otherwise the script downloads the pinned file through Hugging Face Hub. In both cases it checks SHA-256 before reading. The source is [`roman-bushuiev/MassSpecGym`](https://huggingface.co/datasets/roman-bushuiev/MassSpecGym), revision `d2e86d0c3bd905a6d578c0dd6053ed2bd41f9c2a`, file `data/MassSpecGym.tsv`, SHA-256 `0c9cc50450def3f0d4fe2dc09dea1105fc15e635db8c6656bc3e3be37a3bcd95`.
+
+The script writes `best.pt`, `config.json`, `data_audit.json`, `history.json`, and `test.json`. Validation selects the checkpoint; the held-out test fold is evaluated once afterward. Top1 means nearest *other spectrum* within the same fold has the same `inchikey`; singleton molecules are candidates but cannot be queries. This is a **tutorial retrieval task**, not an official MassSpecGym benchmark or a reported UltraMS result. The label is molecular identity from the dataset, not measured retention time or logP.
+
+For the original dataset and citation, see [MassSpecGym](https://github.com/pluskal-lab/MassSpecGym).

ultrams-0.3.0/cookbook/tutorials/README.md

@@ -0,0 +1,13 @@
+# UltraMS tutorials
+
+Each notebook runs from top to bottom on the [five-spectrum MGF example](../../examples/data/NOTICE.md). It installs UltraMS from PyPI and downloads the selected model from Hugging Face on first use.
+
+| Tutorial | Notebook | Colab |
+| --- | --- | --- |
+| Read MGF and obtain embeddings | [Get embeddings](embed.ipynb) | [Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/embed.ipynb) |
+| Fine-tune in a PyTorch loop | [PyTorch fine-tuning](pytorch_finetune.ipynb) | [Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/pytorch_finetune.ipynb) |
+| Rank spectra with the Search model | [Spectrum search](spectrum_search.ipynb) | [Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/spectrum_search.ipynb) |
+
+The fine-tuning labels are placeholders for learning the API. They do not represent measured properties or a model performance result.
+
+For a complete fine-tuning example with real molecular labels and the original dataset folds, follow the [MassSpecGym contrastive tutorial](MASSSPECGYM_CONTRASTIVE.md).

ultrams-0.3.0/cookbook/tutorials/embed.ipynb

@@ -0,0 +1,83 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "introduction",
+   "metadata": {},
+   "source": "# Get MS/MS embeddings\n\n[Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/embed.ipynb)\n\nRead a real MGF file and obtain one UltraMS embedding per spectrum. The input file is the [five-spectrum DreaMS example](../../examples/data/NOTICE.md). Run cells from top to bottom."
+  },
+  {
+   "cell_type": "markdown",
+   "id": "installation",
+   "metadata": {},
+   "source": "## Install"
+  },
+  {
+   "cell_type": "code",
+   "id": "install-ultrams",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%pip install -q ultrams\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "read-spectra",
+   "metadata": {},
+   "source": "## Read spectra"
+  },
+  {
+   "cell_type": "code",
+   "id": "load-mgf",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "from pathlib import Path\nfrom urllib.request import urlretrieve\nimport torch\nfrom ultrams import UltraMS, read_spectra\n\nsample = Path(\"examples/data/example_5_spectra.mgf\")\nif not sample.exists():\n    sample = Path(\"example_5_spectra.mgf\")\n    if not sample.exists():\n        urlretrieve(\n            \"https://raw.githubusercontent.com/Dsadd4/UltraMS/main/examples/data/example_5_spectra.mgf\",\n            sample,\n        )\nspectra = list(read_spectra(sample))\nassert len(spectra) == 5\nprint(f\"Loaded {len(spectra)} MS/MS spectra\")"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "encode-heading",
+   "metadata": {},
+   "source": "## Encode in batches"
+  },
+  {
+   "cell_type": "code",
+   "id": "encode-batch",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = UltraMS.from_pretrained(\"unsupervised\", device=device)\nembeddings = model.encode_batch(spectra, batch_size=5)\nassert embeddings.shape == (len(spectra), model.embedding_dim)\nprint(\"Embedding matrix:\", embeddings.shape)\nprint(\"First spectrum:\", spectra[0][\"id\"])"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "save-heading",
+   "metadata": {},
+   "source": "## Save embeddings"
+  },
+  {
+   "cell_type": "code",
+   "id": "save-embeddings",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "import numpy as np\n\noutput = Path(\"ultrams_embeddings.npz\")\nnp.savez_compressed(output, ids=[spectrum[\"id\"] for spectrum in spectra], embeddings=embeddings)\nprint(output.resolve())"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "colab": {
+   "provenance": []
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

ultrams-0.3.0/cookbook/tutorials/massspecgym_contrastive.py

@@ -0,0 +1,301 @@
+"""Fine-tune UltraMS on MassSpecGym spectra sharing a molecular InChIKey."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import random
+import tempfile
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+from ultrams import UltraMS
+
+
+DATASET = "roman-bushuiev/MassSpecGym"
+REVISION = "d2e86d0c3bd905a6d578c0dd6053ed2bd41f9c2a"
+FILENAME = "data/MassSpecGym.tsv"
+SHA256 = "0c9cc50450def3f0d4fe2dc09dea1105fc15e635db8c6656bc3e3be37a3bcd95"
+FOLDS = ("train", "val", "test")
+
+
+def checked_data_file(path: Path | None) -> Path:
+    if path is None:
+        path = Path(hf_hub_download(repo_id=DATASET, repo_type="dataset", filename=FILENAME, revision=REVISION))
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
+            digest.update(chunk)
+    if digest.hexdigest() != SHA256:
+        raise ValueError(f"MassSpecGym SHA-256 mismatch for {path}")
+    return path
+
+
+def rows_in(path: Path):
+    with path.open(newline="", encoding="utf-8") as handle:
+        reader = csv.DictReader(handle, delimiter="\t")
+        needed = {"identifier", "mzs", "intensities", "inchikey", "precursor_mz", "fold"}
+        missing = needed - set(reader.fieldnames or [])
+        if missing:
+            raise ValueError(f"MassSpecGym is missing columns: {sorted(missing)}")
+        yield from reader
+
+
+def parse_spectrum(row: dict[str, str]):
+    key = row["inchikey"].strip()
+    if not key:
+        return None
+    try:
+        mz = np.fromstring(row["mzs"], sep=",", dtype=np.float32)
+        intensity = np.fromstring(row["intensities"], sep=",", dtype=np.float32)
+        precursor = float(row["precursor_mz"])
+    except ValueError:
+        return None
+    if mz.size != intensity.size or not np.isfinite(precursor) or precursor <= 0:
+        return None
+    valid = np.isfinite(mz) & np.isfinite(intensity) & (mz > 0)
+    effective = valid & (intensity > 0)
+    if int(effective.sum()) < 3:
+        return None
+    # Keep zero-intensity peaks, as UltraMS inference does.
+    mz, intensity = mz[valid], intensity[valid]
+    fingerprint = hashlib.blake2b(
+        mz.tobytes() + intensity.tobytes() + np.float32(precursor).tobytes(), digest_size=16
+    ).hexdigest()
+    return {
+        "id": row["identifier"],
+        "inchikey": key,
+        "mz": mz,
+        "intensity": intensity,
+        "precursor_mz": precursor,
+        "fingerprint": fingerprint,
+    }
+
+
+def selected_keys(groups: dict[str, set[str]], fold: str, seed: int, limit: int | None):
+    keys = sorted(groups, key=lambda key: hashlib.sha256(f"{seed}:{fold}:{key}".encode()).digest())
+    return set(keys if limit is None else keys[:limit])
+
+
+def prepare(path: Path, *, seed: int, full: bool, limits: dict[str, int]):
+    raw = Counter()
+    valid = Counter()
+    fingerprints = {fold: defaultdict(set) for fold in FOLDS}
+    for row in rows_in(path):
+        fold = row["fold"]
+        if fold not in FOLDS:
+            raise ValueError(f"unexpected MassSpecGym fold: {fold!r}")
+        raw[fold] += 1
+        spectrum = parse_spectrum(row)
+        if spectrum is None:
+            continue
+        valid[fold] += 1
+        fingerprints[fold][spectrum["inchikey"]].add(spectrum["fingerprint"])
+
+    keys = {fold: set(fingerprints[fold]) for fold in FOLDS}
+    for i, left in enumerate(FOLDS):
+        for right in FOLDS[i + 1 :]:
+            if keys[left] & keys[right]:
+                raise ValueError(f"molecular overlap between {left} and {right}")
+
+    eligible = {
+        fold: {key: values for key, values in fingerprints[fold].items() if len(values) >= 2}
+        for fold in FOLDS
+    }
+    chosen = {
+        "train": selected_keys(eligible["train"], "train", seed, None if full else limits["train"]),
+        "val": selected_keys(fingerprints["val"] if full else eligible["val"], "val", seed, None if full else limits["val"]),
+        "test": selected_keys(fingerprints["test"] if full else eligible["test"], "test", seed, None if full else limits["test"]),
+    }
+    if any(not chosen[fold] for fold in FOLDS):
+        raise ValueError("at least one fold has no selected molecules with valid spectra")
+
+    spectra = {fold: defaultdict(list) for fold in FOLDS}
+    seen = {fold: defaultdict(set) for fold in FOLDS}
+    duplicates = Counter()
+    for row in rows_in(path):
+        fold, key = row["fold"], row["inchikey"].strip()
+        if key not in chosen[fold]:
+            continue
+        spectrum = parse_spectrum(row)
+        if spectrum is None:
+            continue
+        fingerprint = spectrum.pop("fingerprint")
+        if fingerprint in seen[fold][key]:
+            duplicates[fold] += 1
+            continue
+        seen[fold][key].add(fingerprint)
+        if not full and len(spectra[fold][key]) >= 2:
+            continue
+        spectra[fold][key].append(spectrum)
+
+    for fold in FOLDS:
+        if fold == "train" or not full:
+            if any(len(group) < 2 for group in spectra[fold].values()):
+                raise ValueError(f"{fold} contains a selected molecule without two distinct spectra")
+
+    audit = {
+        "source": {"dataset": DATASET, "revision": REVISION, "file": FILENAME, "sha256": SHA256},
+        "mode": "full" if full else "teaching subset",
+        "selection_seed": seed,
+        "folds": {
+            fold: {
+                "source_rows": raw[fold],
+                "valid_rows_at_least_three_peaks": valid[fold],
+                "source_molecules_after_filter": len(fingerprints[fold]),
+                "molecules_with_two_distinct_spectra": len(eligible[fold]),
+                "selected_rows": sum(map(len, spectra[fold].values())),
+                "selected_molecules": len(spectra[fold]),
+                "selected_molecules_with_two_spectra": sum(len(group) >= 2 for group in spectra[fold].values()),
+                "retrieval_queries_with_positive": sum(
+                    len(group) for group in spectra[fold].values() if len(group) >= 2
+                ),
+                "exact_duplicate_spectra_removed": duplicates[fold],
+            }
+            for fold in FOLDS
+        },
+        "molecule_overlap_between_folds": 0,
+    }
+    return spectra, audit
+
+
+def pairs_for_epoch(groups, seed: int, epoch: int):
+    rng = random.Random(seed + epoch)
+    keys = list(groups)
+    rng.shuffle(keys)
+    return [rng.sample(groups[key], 2) for key in keys]
+
+
+def pair_batches(pairs, batch_size: int):
+    batches = [pairs[start : start + batch_size] for start in range(0, len(pairs), batch_size)]
+    if len(batches) > 1 and len(batches[-1]) == 1:
+        batches[-2].extend(batches.pop())
+    return batches
+
+
+def top1(model: UltraMS, groups, *, batch_size: int, device: torch.device):
+    records = [record for molecule in groups.values() for record in molecule]
+    labels = [record["inchikey"] for record in records]
+    counts = Counter(labels)
+    embeddings = model.encode_batch(records, batch_size=batch_size)
+    vectors = F.normalize(torch.from_numpy(embeddings).to(device), dim=-1)
+    correct = queries = 0
+    for start in range(0, len(records), 256):
+        end = min(start + 256, len(records))
+        scores = vectors[start:end] @ vectors.T
+        scores[torch.arange(end - start, device=device), torch.arange(start, end, device=device)] = -float("inf")
+        nearest = scores.argmax(dim=1).cpu().tolist()
+        for offset, neighbour in enumerate(nearest):
+            index = start + offset
+            if counts[labels[index]] >= 2:
+                queries += 1
+                correct += labels[index] == labels[neighbour]
+    return {"top1": correct / queries, "correct": correct, "queries": queries, "candidates": len(records)}
+
+
+def write_json(path: Path, content) -> None:
+    path.write_text(json.dumps(content, indent=2) + "\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data-file", type=Path, help="Use an already downloaded copy of the pinned TSV")
+    parser.add_argument("--output-dir", type=Path, help="Keep run files here; otherwise use a temporary directory")
+    parser.add_argument("--full", action="store_true", help="Use all valid spectra, not the teaching subset")
+    parser.add_argument("--train-molecules", type=int, default=64)
+    parser.add_argument("--val-molecules", type=int, default=24)
+    parser.add_argument("--test-molecules", type=int, default=24)
+    parser.add_argument("--epochs", type=int, default=2)
+    parser.add_argument("--batch-size", type=int, default=8, help="Number of positive pairs per training batch")
+    parser.add_argument("--lr", type=float, default=1e-5)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--prepare-only", action="store_true", help="Audit and select data without loading the model")
+    args = parser.parse_args()
+    if any(value < 1 for value in (args.train_molecules, args.val_molecules, args.test_molecules, args.epochs)) or args.batch_size < 2:
+        parser.error("molecule limits and epochs must be positive; batch size needs at least two positive pairs")
+    if args.lr <= 0 or args.temperature <= 0:
+        parser.error("learning rate and temperature must be positive")
+
+    output = args.output_dir or Path(tempfile.mkdtemp(prefix="ultrams-massspecgym-"))
+    output.mkdir(parents=True, exist_ok=True)
+    path = checked_data_file(args.data_file)
+    groups, audit = prepare(
+        path,
+        seed=args.seed,
+        full=args.full,
+        limits={"train": args.train_molecules, "val": args.val_molecules, "test": args.test_molecules},
+    )
+    if len(groups["train"]) < 2:
+        raise ValueError("contrastive training needs at least two distinct training molecules")
+    write_json(output / "data_audit.json", audit)
+    config = {
+        "pretrained_model": "unsupervised",
+        "objective": "in-batch contrastive learning; positives are two distinct spectra with the same MassSpecGym inchikey",
+        "evaluation": "same-fold spectrum retrieval Top1, self-match excluded; singleton queries excluded",
+        "dataset": audit["source"],
+        "mode": audit["mode"],
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size_in_positive_pairs": args.batch_size,
+        "learning_rate": args.lr,
+        "temperature": args.temperature,
+        "device": args.device,
+        "output_dir": str(output.resolve()),
+    }
+    write_json(output / "config.json", config)
+    print(json.dumps(audit["folds"], indent=2), flush=True)
+    print(f"Run directory: {output.resolve()}", flush=True)
+    if args.prepare_only:
+        return
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    device = torch.device(args.device)
+    model = UltraMS.from_pretrained("unsupervised", device=device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+    history = []
+    best_top1 = -1.0
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        pairs = pairs_for_epoch(groups["train"], args.seed, epoch)
+        sum_loss = 0.0
+        for mini in pair_batches(pairs, args.batch_size):
+            batch = model.collate([spectrum for pair in mini for spectrum in pair])
+            embeddings = model(
+                batch["peaks"].to(device), batch["attention_mask"].to(device), batch["precursor_mz"].to(device)
+            )
+            logits = embeddings @ embeddings.T / args.temperature
+            logits = logits.masked_fill(torch.eye(len(embeddings), device=device, dtype=torch.bool), -float("inf"))
+            target = torch.arange(len(embeddings), device=device) ^ 1
+            loss = F.cross_entropy(logits, target)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            sum_loss += float(loss.detach()) * len(mini)
+        validation = top1(model, groups["val"], batch_size=args.batch_size * 2, device=device)
+        result = {"epoch": epoch, "train_loss": sum_loss / len(pairs), "validation": validation}
+        history.append(result)
+        write_json(output / "history.json", history)
+        print(json.dumps(result), flush=True)
+        if validation["top1"] > best_top1:
+            best_top1 = validation["top1"]
+            torch.save({"model_state_dict": model.model.state_dict(), "config": model.config}, output / "best.pt")
+
+    best = UltraMS.from_checkpoint(output / "best.pt", device=device)
+    test = top1(best, groups["test"], batch_size=args.batch_size * 2, device=device)
+    write_json(output / "test.json", test)
+    print("Test on the best validation checkpoint:", json.dumps(test), flush=True)
+
+
+if __name__ == "__main__":
+    main()

ultrams-0.3.0/cookbook/tutorials/pytorch_finetune.ipynb

@@ -0,0 +1,97 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "introduction",
+   "metadata": {},
+   "source": "# Fine-tune with PyTorch\n\n[Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/pytorch_finetune.ipynb)\n\nUse UltraMS as a standard `torch.nn.Module` in a training loop. The input file is the [five-spectrum DreaMS example](../../examples/data/NOTICE.md). Run cells from top to bottom."
+  },
+  {
+   "cell_type": "markdown",
+   "id": "installation",
+   "metadata": {},
+   "source": "## Install"
+  },
+  {
+   "cell_type": "code",
+   "id": "install-ultrams",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%pip install -q ultrams\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "read-spectra",
+   "metadata": {},
+   "source": "## Read spectra"
+  },
+  {
+   "cell_type": "code",
+   "id": "load-mgf",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "from pathlib import Path\nfrom urllib.request import urlretrieve\nimport torch\nfrom ultrams import UltraMS, read_spectra\n\nsample = Path(\"examples/data/example_5_spectra.mgf\")\nif not sample.exists():\n    sample = Path(\"example_5_spectra.mgf\")\n    if not sample.exists():\n        urlretrieve(\n            \"https://raw.githubusercontent.com/Dsadd4/UltraMS/main/examples/data/example_5_spectra.mgf\",\n            sample,\n        )\nspectra = list(read_spectra(sample))\nassert len(spectra) == 5\nprint(f\"Loaded {len(spectra)} MS/MS spectra\")"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "labels-heading",
+   "metadata": {},
+   "source": "## Add demonstration labels\n\nThe five targets below are **placeholders made from file order**. They are not measured properties or research labels; replace them with your own labels for a scientific task."
+  },
+  {
+   "cell_type": "code",
+   "id": "demo-labels",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "from torch.utils.data import DataLoader\n\nrecords = [\n    {**spectrum, \"target\": index / (len(spectra) - 1)}\n    for index, spectrum in enumerate(spectra)\n]\nprint([record[\"target\"] for record in records])"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "train-heading",
+   "metadata": {},
+   "source": "## Train encoder and task head"
+  },
+  {
+   "cell_type": "code",
+   "id": "native-training",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = UltraMS.from_pretrained(\"unsupervised\", device=device).train()\nhead = torch.nn.Linear(model.embedding_dim, 1).to(device)\nloader = DataLoader(records, batch_size=2, collate_fn=model.batch_converter())\noptimizer = torch.optim.AdamW([*model.parameters(), *head.parameters()], lr=1e-5)\n\nlosses = []\nfor batch in loader:\n    batch = {name: value.to(device) for name, value in batch.items()}\n    prediction = head(model(batch[\"peaks\"], batch[\"attention_mask\"], batch[\"precursor_mz\"]))\n    loss = torch.nn.functional.mse_loss(prediction.squeeze(-1), batch[\"target\"])\n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n    losses.append(float(loss.detach()))\nprint(\"Batch losses (demonstration only):\", losses)"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "save-heading",
+   "metadata": {},
+   "source": "## Save the run configuration and history\n\nThe weights trained on placeholder labels are only a code demonstration."
+  },
+  {
+   "cell_type": "code",
+   "id": "save-history",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "import json\n\noutput = Path(\"ultrams_finetune_demo.json\")\noutput.write_text(json.dumps({\n    \"model\": \"unsupervised\",\n    \"target\": \"placeholder labels from file order\",\n    \"batch_size\": 2,\n    \"learning_rate\": 1e-5,\n    \"loss_per_batch\": losses,\n}, indent=2) + \"\\n\")\nprint(output.resolve())"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "colab": {
+   "provenance": []
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

ultrams-0.3.0/cookbook/tutorials/spectrum_search.ipynb

@@ -0,0 +1,69 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "introduction",
+   "metadata": {},
+   "source": "# Search similar MS/MS spectra\n\n[Open in Colab](https://colab.research.google.com/github/Dsadd4/UltraMS/blob/main/cookbook/tutorials/spectrum_search.ipynb)\n\nRank the other spectra in the same example file by cosine similarity using the UltraMS Search model. The input file is the [five-spectrum DreaMS example](../../examples/data/NOTICE.md). Run cells from top to bottom."
+  },
+  {
+   "cell_type": "markdown",
+   "id": "installation",
+   "metadata": {},
+   "source": "## Install"
+  },
+  {
+   "cell_type": "code",
+   "id": "install-ultrams",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%pip install -q ultrams\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "read-spectra",
+   "metadata": {},
+   "source": "## Read spectra"
+  },
+  {
+   "cell_type": "code",
+   "id": "load-mgf",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "from pathlib import Path\nfrom urllib.request import urlretrieve\nimport torch\nfrom ultrams import UltraMS, read_spectra\n\nsample = Path(\"examples/data/example_5_spectra.mgf\")\nif not sample.exists():\n    sample = Path(\"example_5_spectra.mgf\")\n    if not sample.exists():\n        urlretrieve(\n            \"https://raw.githubusercontent.com/Dsadd4/UltraMS/main/examples/data/example_5_spectra.mgf\",\n            sample,\n        )\nspectra = list(read_spectra(sample))\nassert len(spectra) == 5\nprint(f\"Loaded {len(spectra)} MS/MS spectra\")"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "rank-heading",
+   "metadata": {},
+   "source": "## Encode and rank\n\nThis small in-file ranking illustrates the Search representation. It is not an UltraAtlas search or a retrieval benchmark."
+  },
+  {
+   "cell_type": "code",
+   "id": "cosine-ranking",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": "import numpy as np\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = UltraMS.from_pretrained(\"search\", device=device)\nembeddings = model.encode_batch(spectra, batch_size=5)\nunit = embeddings / np.maximum(np.linalg.norm(embeddings, axis=1, keepdims=True), 1e-12)\nquery_index = 0\nscores = unit @ unit[query_index]\nranked = [index for index in np.argsort(-scores) if index != query_index]\nprint(\"Query:\", spectra[query_index][\"id\"])\nfor rank, index in enumerate(ranked, start=1):\n    print(rank, spectra[index][\"id\"], float(scores[index]))"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "colab": {
+   "provenance": []
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

{ultrams-0.2.0 → ultrams-0.3.0}/docs/data-format.md

@@ -42,6 +42,21 @@ The `target` field is a numeric training value used by `model.batch_converter()`
 
 `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.
 
+To inspect representations for individual spectral peaks, pass `return_peaks=True`:
+
+```python
+result = model.encode(
+    mz=[150.0, 100.1, 121.1],
+    intensity=[35, 20, 100],
+    precursor_mz=301.2,
+    return_peaks=True,
+)
+print(result.peak_mz)                 # processed peak m/z, sorted ascending
+print(result.peak_embeddings.shape)  # (3, 1024) for the Unsupervised model
+```
+
+`peak_embeddings[i]` is the encoder's final contextual representation for `peak_mz[i]` and `peak_intensity[i]`. These arrays reflect the peak selection, intensity normalization and sorting described above. They are returned only when requested.
+
 ## 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.

{ultrams-0.2.0 → ultrams-0.3.0}/examples/README.md

@@ -13,3 +13,16 @@ python -m pip install -e .
 | 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.
+
+The PyTorch example saves both the encoder and task head in `model.pt`. Load them for prediction with:
+
+```python
+import torch
+from ultrams import UltraMS
+
+weights = torch.load("/path/to/model.pt", map_location="cpu", weights_only=True)
+model = UltraMS.from_pretrained("unsupervised").eval()
+head = torch.nn.Linear(model.embedding_dim, 1).eval()
+model.load_state_dict(weights["encoder"])
+head.load_state_dict(weights["head"])
+```

{ultrams-0.2.0 → ultrams-0.3.0}/examples/pytorch_finetune.py

@@ -47,6 +47,8 @@ def main() -> None:
 
     output_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="ultrams-finetune-"))
     output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint = output_dir / "model.pt"
+    torch.save({"encoder": model.state_dict(), "head": head.state_dict()}, checkpoint)
     record = {
         "model": "unsupervised",
         "target": "demonstration labels from input order, scaled to [0, 1]",
@@ -56,6 +58,7 @@ def main() -> None:
     }
     (output_dir / "training.json").write_text(json.dumps(record, indent=2) + "\n")
     print(f"Trained on {len(spectra)} spectra with demonstration targets; final loss: {history[-1]:.4f}")
+    print(checkpoint)
     print(output_dir / "training.json")
 
 

{ultrams-0.2.0 → ultrams-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ultrams"
-version = "0.2.0"
+version = "0.3.0"
 description = "Pretrained UltraMS models for mass spectrometry"
 readme = "README.md"
 requires-python = ">=3.10"

{ultrams-0.2.0 → ultrams-0.3.0}/src/ultrams/inference.py

@@ -29,6 +29,9 @@ class SpectrumEmbedding:
     fusion: np.ndarray
     projection: np.ndarray | None = None
     rt_seconds: float | None = None
+    peak_embeddings: np.ndarray | None = None
+    peak_mz: np.ndarray | None = None
+    peak_intensity: np.ndarray | None = None
 
     @property
     def embedding(self) -> np.ndarray:
@@ -323,7 +326,18 @@ class UltraMS(torch.nn.Module):
 
     @torch.inference_mode()
     def encode(
-        self, mz: Any, intensity: Any, *, precursor_mz: float
+        self, mz: Any, intensity: Any, *, precursor_mz: float, return_peaks: bool = False
+    ) -> SpectrumEmbedding:
+        """Encode one MS/MS spectrum; optionally return aligned peak embeddings."""
+        was_training = self.training
+        self.eval()
+        try:
+            return self._encode_one(mz, intensity, precursor_mz, return_peaks)
+        finally:
+            self.train(was_training)
+
+    def _encode_one(
+        self, mz: Any, intensity: Any, precursor_mz: float, return_peaks: bool
     ) -> SpectrumEmbedding:
         peaks, attention = _prepare_spectrum(mz, intensity, self.input_max_peaks)
         if not np.isfinite(precursor_mz) or precursor_mz <= 0:
@@ -356,4 +370,7 @@ class UltraMS(torch.nn.Module):
             fusion=fusion[0].cpu().numpy(),
             projection=projected[0].cpu().numpy() if projected is not None else None,
             rt_seconds=rt_seconds,
+            peak_embeddings=peak_hidden[0].float().cpu().numpy() if return_peaks else None,
+            peak_mz=peaks[:, 0].copy() if return_peaks else None,
+            peak_intensity=peaks[:, 1].copy() if return_peaks else None,
         )

{ultrams-0.2.0 → ultrams-0.3.0/src/ultrams.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ultrams
-Version: 0.2.0
+Version: 0.3.0
 Summary: Pretrained UltraMS models for mass spectrometry
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/Dsadd4/UltraMS
@@ -30,7 +30,8 @@ python -m pip install ultrams
 | 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) |
+| Learn in a notebook | [UltraMS tutorials](https://github.com/Dsadd4/UltraMS/tree/main/cookbook/tutorials) |
+| Inspect or run pretraining | [Training code](https://github.com/Dsadd4/UltraMS/blob/main/training/README.md) |
 
 ## Embed a spectrum
 
@@ -48,6 +49,8 @@ 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.
 
+The three checkpoints are collected on [Hugging Face](https://huggingface.co/collections/dsadd4/ultrams-6ab4cfaa860cbbd1f9a6b9b6).
+
 ## Embed a spectrum file
 
 ```bash
@@ -87,6 +90,8 @@ for batch in loader:
 
 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.
 
+For measured molecular labels, the [MassSpecGym contrastive tutorial](https://github.com/Dsadd4/UltraMS/blob/main/cookbook/tutorials/MASSSPECGYM_CONTRASTIVE.md) downloads the dataset, keeps its train/validation/test folds, fine-tunes the encoder, and reloads the selected checkpoint.
+
 ## Choose a model
 
 | Model name | Representation | Dimension | Spectral peak limit | Weights |
@@ -100,3 +105,7 @@ Load any of them with `UltraMS.from_pretrained("name")`, or load a downloaded `m
 ## 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.
+
+## Cite
+
+Citation metadata for this software is in [CITATION.cff](https://github.com/Dsadd4/UltraMS/blob/main/CITATION.cff).

{ultrams-0.2.0 → ultrams-0.3.0}/src/ultrams.egg-info/SOURCES.txt

@@ -1,7 +1,14 @@
+CITATION.cff
 LICENSE
 MANIFEST.in
 README.md
 pyproject.toml
+cookbook/tutorials/MASSSPECGYM_CONTRASTIVE.md
+cookbook/tutorials/README.md
+cookbook/tutorials/embed.ipynb
+cookbook/tutorials/massspecgym_contrastive.py
+cookbook/tutorials/pytorch_finetune.ipynb
+cookbook/tutorials/spectrum_search.ipynb
 docs/data-format.md
 docs/model-selection.md
 examples/README.md

{ultrams-0.2.0 → ultrams-0.3.0}/tests/test_model_api.py

@@ -53,6 +53,21 @@ class ModelAPITests(unittest.TestCase):
         actual = self.model.encode_batch([], batch_size=2)
         self.assertEqual(actual.shape, (0, self.model.embedding_dim))
 
+    def test_peak_embeddings_align_with_selected_sorted_peaks(self) -> None:
+        mz = [180, 100, 150, 110, 140, 120, 170, 130, 190, 160]
+        intensity = [8, 1, 5, 2, 4, 3, 7, 6, 10, 9]
+        result = self.model.encode(mz, intensity, precursor_mz=301.2, return_peaks=True)
+        self.assertEqual(result.peak_embeddings.shape, (8, 16))
+        np.testing.assert_array_equal(
+            result.peak_mz,
+            np.asarray([120, 130, 140, 150, 160, 170, 180, 190], dtype=np.float32),
+        )
+        np.testing.assert_array_equal(
+            result.peak_intensity,
+            np.asarray([3, 6, 4, 5, 9, 7, 8, 10], dtype=np.float32) / 10,
+        )
+        self.assertIsNone(self.model.encode(mz, intensity, precursor_mz=301.2).peak_embeddings)
+
     def test_real_mgf_to_npz_with_local_checkpoint(self) -> None:
         source = Path(__file__).resolve().parents[1] / "examples/data/example_5_spectra.mgf"
         with tempfile.TemporaryDirectory() as directory:

ultrams-0.2.0/MANIFEST.in

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