From 2e199d32ef3e7c13250079fc7cf2f817493defe4 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Mon, 14 Sep 2026 12:32:27 +0200 Subject: [PATCH 1/3] switch to more general fingerprint options --- .../DataGeneratorEmbeddingEvaluation.py | 56 +++++++++++++++---- 1 file changed, 45 insertions(+), 11 deletions(-) diff --git a/ms2deepscore/train_new_model/DataGeneratorEmbeddingEvaluation.py b/ms2deepscore/train_new_model/DataGeneratorEmbeddingEvaluation.py index 19ccf4a0..1af51caa 100644 --- a/ms2deepscore/train_new_model/DataGeneratorEmbeddingEvaluation.py +++ b/ms2deepscore/train_new_model/DataGeneratorEmbeddingEvaluation.py @@ -2,10 +2,12 @@ import numpy as np import pandas as pd -from torch import tensor +import torch from matchms import Spectrum -from matchms.similarity.vector_similarity_functions import jaccard_similarity_matrix - +from ms2deepscore.fingerprint_similarity_computations import ( + compute_fingerprint_similarity_matrix, + is_dense_fingerprint_type, +) from ms2deepscore.SettingsMS2Deepscore import SettingsEmbeddingEvaluator from ms2deepscore.models import SiameseSpectralModel from ms2deepscore.tensorize_spectra import tensorize_spectra @@ -54,12 +56,17 @@ def __init__( self.ms2ds_model.to(self.device) self.indexes = np.arange(len(self.spectrums)) self.batch_size = self.settings.evaluator_distribution_size - self.fingerprint_df = self.compute_fingerprint_dataframe( + self.fingerprints, fingerprint_inchikeys = compute_fingerprints_for_training( self.spectrums, fingerprint_type=self.ms2ds_model.model_settings.fingerprint_type, - fingerprint_nbits=self.ms2ds_model.model_settings.fingerprint_nbits, + nbits=self.ms2ds_model.model_settings.fingerprint_nbits, ) + self.fingerprint_index = { + inchikey: idx + for idx, inchikey in enumerate(fingerprint_inchikeys) + } + # Initialize random number generator self.rng = np.random.default_rng(self.settings.random_seed) @@ -87,17 +94,44 @@ def _compute_embeddings_and_scores(self, batch_index: int): spec_tensors, meta_tensors = tensorize_spectra( [self.spectrums[i] for i in indexes], self.ms2ds_model.model_settings ) - embeddings = self.ms2ds_model.encoder(spec_tensors.to(self.device), meta_tensors.to(self.device)) + with torch.no_grad(): + embeddings = self.ms2ds_model.encoder( + spec_tensors.to(self.device), + meta_tensors.to(self.device), + ) - ms2ds_scores = cosine_similarity_matrix(embeddings.cpu().detach().numpy(), embeddings.cpu().detach().numpy()) + embeddings_numpy = embeddings.cpu().numpy() + + ms2ds_scores = cosine_similarity_matrix(embeddings_numpy, embeddings_numpy) # Compute true scores - inchikeys = [self.inchikey14s[i] for i in indexes] - fingerprints = self.fingerprint_df.loc[inchikeys].to_numpy() + fingerprint_type = self.ms2ds_model.model_settings.fingerprint_type - tanimoto_scores = jaccard_similarity_matrix(fingerprints, fingerprints) + inchikeys = [self.inchikey14s[i] for i in indexes] + fingerprint_indexes = [ + self.fingerprint_index[inchikey] + for inchikey in inchikeys + ] + + if is_dense_fingerprint_type(fingerprint_type): + batch_fingerprints = self.fingerprints[fingerprint_indexes] + else: + batch_fingerprints = [ + self.fingerprints[i] + for i in fingerprint_indexes + ] + + tanimoto_scores = compute_fingerprint_similarity_matrix( + batch_fingerprints, + batch_fingerprints, + fingerprint_type=fingerprint_type, + ) - return tensor(tanimoto_scores), tensor(ms2ds_scores), embeddings.cpu().detach() + return ( + torch.as_tensor(tanimoto_scores), + torch.as_tensor(ms2ds_scores), + embeddings.cpu(), + ) def on_epoch_end(self): """Updates indexes after each epoch.""" From e382cb30bc948cb3f5a9361ba74062a8341fe3aa Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Mon, 14 Sep 2026 12:32:42 +0200 Subject: [PATCH 2/3] test different fingerprints --- tests/test_embedding_evaluator.py | 48 +++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/test_embedding_evaluator.py b/tests/test_embedding_evaluator.py index 0cced41d..802474aa 100644 --- a/tests/test_embedding_evaluator.py +++ b/tests/test_embedding_evaluator.py @@ -1,6 +1,8 @@ import os +from ms2deepscore.train_new_model.DataGeneratorEmbeddingEvaluation import DataGeneratorEmbeddingEvaluation import pytest import numpy as np +import torch from sklearn.datasets import make_regression from torch import randn from ms2deepscore.models import EmbeddingEvaluationModel, LinearModel @@ -131,3 +133,49 @@ def test_linear_model_save_load(tmp_path): assert np.array_equal(model.model.coef_, loaded_model.model.coef_), "Coefficients do not match." assert model.model.intercept_ == loaded_model.model.intercept_, "Intercepts do not match." assert model.degree == loaded_model.degree == 3, "Degree does not match." + + +@pytest.mark.parametrize( + "fingerprint_type", + [ + "rdkit_binary", + "rdkit_count", + "rdkit_logcount", + "rdkit_binary_unfolded", + "rdkit_count_unfolded", + "rdkit_logcount_unfolded", + ], +) +def test_embedding_evaluator_generator_supports_fingerprint_types( + fingerprint_type, +): + spectra = create_test_spectra( + num_of_unique_inchikeys=10, + num_of_spectra_per_inchikey=1, + ) + + model = MockMS2DSModel() + model.model_settings.fingerprint_type = fingerprint_type + model.model_settings.fingerprint_nbits = 256 + + generator = DataGeneratorEmbeddingEvaluation( + spectrums=spectra, + ms2ds_model=model, + settings=SettingsEmbeddingEvaluator( + evaluator_distribution_size=5, + ), + device="cpu", + ) + + tanimoto_scores, ms2ds_scores, embeddings = next(generator) + + assert tanimoto_scores.shape == (5, 5) + assert ms2ds_scores.shape == (5, 5) + assert embeddings.shape[0] == 5 + + assert torch.isfinite(tanimoto_scores).all() + + torch.testing.assert_close( + torch.diag(tanimoto_scores), + torch.ones(5), + ) From 49f15e4a935d38e61ee1662488975865a30e7272 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Mon, 14 Sep 2026 12:33:03 +0200 Subject: [PATCH 3/3] bump version --- ms2deepscore/__version__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ms2deepscore/__version__.py b/ms2deepscore/__version__.py index fb7a749a..0b288f83 100644 --- a/ms2deepscore/__version__.py +++ b/ms2deepscore/__version__.py @@ -1 +1 @@ -__version__ = '2.11.0' +__version__ = '2.12.0'