1import torch
2import numpy as np
3import pandas as pd
4from instanovo.transformer.model import InstaNovo
5from instanovo.utils import SpectrumDataFrame
6from instanovo.transformer.dataset import SpectrumDataset, collate_batch
7from torch.utils.data import DataLoader
8from instanovo.inference import ScoredSequence
9from instanovo.inference import BeamSearchDecoder
10from instanovo.utils.metrics import Metrics
11from tqdm.notebook import tqdm
12
13# Load the model from the Hugging Face Hub
14model, config = InstaNovo.from_pretrained("InstaDeepAI/instanovo-v1.0.0")
15
16# Move the model to the GPU if available
17device = "cuda" if torch.cuda.is_available() else "cpu"
18model = model.to(device).eval()
19
20# Update the residue set with custom modifications
21model.residue_set.update_remapping(
22 {
23 "M(ox)": "M[UNIMOD:35]",
24 "M(+15.99)": "M[UNIMOD:35]",
25 "S(p)": "S[UNIMOD:21]", # Phosphorylation
26 "T(p)": "T[UNIMOD:21]",
27 "Y(p)": "Y[UNIMOD:21]",
28 "S(+79.97)": "S[UNIMOD:21]",
29 "T(+79.97)": "T[UNIMOD:21]",
30 "Y(+79.97)": "Y[UNIMOD:21]",
31 "Q(+0.98)": "Q[UNIMOD:7]", # Deamidation
32 "N(+0.98)": "N[UNIMOD:7]",
33 "Q(+.98)": "Q[UNIMOD:7]",
34 "N(+.98)": "N[UNIMOD:7]",
35 "C(+57.02)": "C[UNIMOD:4]", # Carboxyamidomethylation
36 "(+42.01)": "[UNIMOD:1]", # Acetylation
37 "(+43.01)": "[UNIMOD:5]", # Carbamylation
38 "(-17.03)": "[UNIMOD:385]",
39 }
40)
41
42# Load the test data
43sdf = SpectrumDataFrame.from_huggingface(
44 "InstaDeepAI/ms_ninespecies_benchmark",
45 is_annotated=True,
46 shuffle=False,
47 split="test[:10%]", # Let's only use a subset of the test data for faster inference
48)
49
50# Create the dataset
51ds = SpectrumDataset(
52 sdf,
53 model.residue_set,
54 config.get("n_peaks", 200),
55 return_str=True,
56 annotated=True,
57)
58
59# Create the data loader
60dl = DataLoader(ds, batch_size=64, shuffle=False, num_workers=0, collate_fn=collate_batch)
61
62# Create the decoder
63decoder = BeamSearchDecoder(model=model)
64
65# Initialize lists to store predictions and targets
66preds = []
67targs = []
68probs = []
69
70# Iterate over the data loader
71for _, batch in tqdm(enumerate(dl), total=len(dl)):
72 spectra, precursors, _, peptides, _ = batch
73 spectra = spectra.to(device)
74 precursors = precursors.to(device)
75
76 # Perform inference
77 with torch.no_grad():
78 p = decoder.decode(
79 spectra=spectra,
80 precursors=precursors,
81 beam_size=config["n_beams"],
82 max_length=config["max_length"],
83 )
84
85
86 preds += [x.sequence if isinstance(x, ScoredSequence) else [] for x in p]
87 probs += [
88 x.sequence_log_probability if isinstance(x, ScoredSequence) else -float("inf") for x in p
89 ]
90 targs += list(peptides)
91
92# Initialize metrics
93metrics = Metrics(model.residue_set, config["isotope_error_range"])
94
95
96# Compute precision and recall
97aa_precision, aa_recall, peptide_recall, peptide_precision = metrics.compute_precision_recall(
98 peptides, preds
99)
100
101# Compute amino acid error rate and AUC
102aa_error_rate = metrics.compute_aa_er(targs, preds)
103auc = metrics.calc_auc(targs, preds, np.exp(pd.Series(probs)))
104
105print(f"amino acid error rate: {aa_error_rate:.5f}")
106print(f"amino acid precision: {aa_precision:.5f}")
107print(f"amino acid recall: {aa_recall:.5f}")
108print(f"peptide precision: {peptide_precision:.5f}")
109print(f"peptide recall: {peptide_recall:.5f}")
110print(f"area under the PR curve: {auc:.5f}")