cleaned up

flare/data/__init__.py

@@ -0,0 +1 @@
+

{mvp → flare}/data/data_module.py

@@ -1,6 +1,6 @@
 from torch.utils.data.dataloader import DataLoader
 from massspecgym.data.data_module import MassSpecDataModule
-from mvp.data.datasets import ContrastiveDataset
+from flare.data.datasets import ContrastiveDataset
 from functools import partial
 from massspecgym.models.base import Stage
 

{mvp → flare}/data/datasets.py

@@ -11,7 +11,7 @@ import dgl
 from collections import defaultdict
 from massspecgym.data.transforms import SpecTransform, MolTransform, MolToInChIKey
 from massspecgym.data.datasets import MassSpecDataset
-import mvp.utils.data as data_utils
+import flare.utils.data as data_utils
 from torch.nn.utils.rnn import pad_sequence
 from massspecgym.models.base import Stage
 import pickle
@@ -254,7 +254,7 @@ class ContrastiveDataset(Dataset):
         return item
 
     @staticmethod
-    def collate_fn(batch: T.Iterable[dict], spec_enc: str, spectra_view: str, stage=None, mask_peak_ratio: float = 0.0, aug_cands: bool = False) -> dict:
+    def collate_fn(batch: T.Iterable[dict], spec_enc: str, spectra_view: str, stage=None) -> dict:
         mol_key = 'cand' if stage == Stage.TEST else 'mol'
         non_standard_collate = ['mol', 'cand', 'aug_cands', 'cons_spec', 'aug_cands_fp', 'NL_spec']
         require_pad = False
@@ -314,25 +314,6 @@ class ContrastiveDataset(Dataset):
                     n_peaks.append(len(item['NL_spec']))
                 collated_batch['NL_spec'] = pad_sequence(peaks, batch_first=True, padding_value=padding_value)
                 collated_batch['NL_n_peaks'] = n_peaks
-
-
-        # mask peaks
-        if mask_peak_ratio > 0.0 and stage == Stage.TRAIN:
-            n_mask_peaks = [math.floor(n_peak* mask_peak_ratio) for n_peak in n_peaks]
-            mask_peak_idx = [np.random.choice(n_peak, n_mask, replace=False) for n_peak, n_mask in zip(n_peaks, n_mask_peaks)]
-            for i, peaks in enumerate(collated_batch[spectra_view]):
-                peaks[mask_peak_idx[i]] = -5.0
-        
-        # batch candidates
-        if aug_cands:
-            candidates = \
-                sum([item["aug_cands"] for item in batch], start=[])
-            collated_batch['aug_cands'] = dgl.batch(candidates)
-
-            if 'aug_cands_fp' in batch[0]:
-                cand_fp = [item['aug_cands_fp'] for item in batch]
-                collated_batch['aug_cands_fp'] = torch.flatten(torch.Tensor(cand_fp), end_dim=1)
-
         return collated_batch
     
  

{mvp → flare}/data/transforms.py

@@ -3,7 +3,7 @@ import torch
 import matchms
 from typing import Optional
 from rdkit.Chem import AllChem as Chem
-from mvp.definitions import CHEM_ELEMS_SMALL
+from flare.definitions import CHEM_ELEMS_SMALL
 from massspecgym.data.transforms import MolTransform, SpecTransform, default_matchms_transforms
 from massspecgym.data.transforms import SpecBinner
 import dgllife.utils as chemutils

{mvp → flare}/data_preprocess.py

@@ -1,5 +1,5 @@
 import argparse
-from mvp.utils.preprocessing import generate_cons_spec_formulas, generate_cons_spec
+from flare.utils.preprocessing import generate_cons_spec_formulas, generate_cons_spec
 import os
 import pickle
 import pandas as pd

{mvp → flare}/models/contrastive.py

@@ -10,11 +10,11 @@ from massspecgym.models.base import Stage
 from massspecgym import utils
 from torch.nn.utils.rnn import pad_sequence
 
-from mvp.utils.loss import contrastive_loss, cand_spec_sim_loss, fp_loss, cons_spec_loss, filip_loss_with_mask
-import mvp.utils.models as model_utils
-from mvp.utils.general import pad_graph_nodes, filip_similarity_batch
+from flare.utils.loss import contrastive_loss, cand_spec_sim_loss, fp_loss, cons_spec_loss, filip_loss_with_mask
+import flare.utils.models as model_utils
+from flare.utils.general import pad_graph_nodes, filip_similarity_batch
 
-from mvp.models.encoders import CrossAttention
+from flare.models.encoders import CrossAttention
 import torch.nn.functional as F
 
 from torch_geometric.nn import global_mean_pool
@@ -32,62 +32,21 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         if 'use_NL_spec' not in self.hparams:
             self.hparams.use_NL_spec = False
 
-        # if 'loss_strategy' not in self.hparams:
-        #     self.hparams.loss_strategy = 'static'
-        #     self.hparams.contr_wt = 1.0
-        #     self.hparams.use_contr = True
 
         self.spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
         self.mol_enc_model = model_utils.get_mol_encoder(self.hparams.mol_enc, self.hparams)
-        
-        # setup loss strategy
-        if self.hparams.model == 'contrastive':
-            self._loss_setup()
-        if self.hparams.pred_fp:
-            self.fp_loss = fp_loss(self.hparams.fp_loss_type)
-            self.fp_pred_model = model_utils.get_fp_pred_model(self.hparams)
-        if self.hparams.use_cons_spec:
-            self.cons_spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
-            self.cons_loss = cons_spec_loss(self.hparams.cons_loss_type)
-            
+  
         self.spec_view = self.hparams.spectra_view
         
         # result storage for testing results
         self.result_dct = defaultdict(lambda: defaultdict(list))
-
-
-    # def _loss_setup(self):
-    #     self.loss_wts = {}
-    #     self.loss_updates = {}
-
-
-    #     for p, loss in zip(['use_contr','pred_fp', 'use_cons_spec', 'aug_cands'], ['contr_wt','fp_wt','cons_spec_wt' ,'aug_cands_wt']):
-    #         if p not in self.hparams:
-    #             self.hparams[p] = False
-    #         if self.hparams[p]:
-    #             if self.hparams.loss_strategy == 'linear':
-    #                 start_wt = self.hparams[loss+'_update']['start']
-    #                 end_wt = self.hparams[loss+'_update']['end']
-    #                 change = (end_wt - start_wt)/self.hparams.max_epochs
-    #                 self.loss_updates[loss] = change
-    #                 self.loss_wts[loss] = start_wt
-    #             elif self.hparams.loss_strategy == 'manual':
-    #                 self.loss_updates[loss] = self.hparams[loss+'_update']
-    #                 self.loss_wts[loss] = self.hparams[loss]
-    #             else: 
-    #                 self.loss_wts[loss] = self.hparams[loss]
                     
     def forward(self, batch, stage):
         g = batch['cand'] if stage == Stage.TEST else batch['mol']
         
-        if self.hparams.use_cons_spec and stage != Stage.TEST:
-            spec = batch['cons_spec']
-            n_peaks = batch['cons_n_peaks'] if 'cons_n_peaks' in batch else None
-            spec_enc = self.cons_spec_enc_model(spec, n_peaks)
-        else:
-            spec = batch[self.spec_view]
-            n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
-            spec_enc = self.spec_enc_model(spec, n_peaks)
+        spec = batch[self.spec_view]
+        n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
+        spec_enc = self.spec_enc_model(spec, n_peaks)
 
         fp = batch['fp'] if self.hparams.use_fp else None
         mol_enc = self.mol_enc_model(g, fp=fp)
@@ -98,26 +57,24 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         loss = 0
         losses = {}
         contr_loss, _, _ = contrastive_loss(spec_enc, mol_enc, self.hparams.contr_temp)
-        # contr_loss = self.loss_wts['contr_wt'] *contr_loss
+
         losses['contr_loss'] = contr_loss.detach().item()
-        # losses['cong_loss'] = cong_loss.detach().item()
-        # losses['noncong_loss'] = noncong_loss.detach().item()
-        
+
         loss+=contr_loss
-        if self.hparams.pred_fp:
-            fp_loss_val = self.loss_wts['fp_wt'] *self.fp_loss(output['fp'], batch['fp'])
-            loss+= fp_loss_val
-            losses['fp_loss'] = fp_loss_val.detach().item()
-
-        if 'aug_cand_enc' in output:  
-            aug_cand_loss = self.loss_wts['aug_cand_wt'] * cand_spec_sim_loss(spec_enc, output['aug_cand_enc'])
-            loss+= aug_cand_loss
-            losses['aug_cand_loss'] = aug_cand_loss.detach().item()
+        # if self.hparams.pred_fp:
+        #     fp_loss_val = self.loss_wts['fp_wt'] *self.fp_loss(output['fp'], batch['fp'])
+        #     loss+= fp_loss_val
+        #     losses['fp_loss'] = fp_loss_val.detach().item()
+
+        # if 'aug_cand_enc' in output:  
+        #     aug_cand_loss = self.loss_wts['aug_cand_wt'] * cand_spec_sim_loss(spec_enc, output['aug_cand_enc'])
+        #     loss+= aug_cand_loss
+        #     losses['aug_cand_loss'] = aug_cand_loss.detach().item()
         
-        if 'ind_spec' in output:
-            spec_loss = self.loss_wts['cons_spec_wt'] * self.cons_loss(spec_enc, output['ind_spec'])
-            loss+=spec_loss
-            losses['cons_spec_loss'] = spec_loss.detach().item()
+        # if 'ind_spec' in output:
+        #     spec_loss = self.loss_wts['cons_spec_wt'] * self.cons_loss(spec_enc, output['ind_spec'])
+        #     loss+=spec_loss
+        #     losses['cons_spec_loss'] = spec_loss.detach().item()
 
         losses['loss'] = loss 
 
@@ -158,62 +115,6 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
             on_epoch=True,
             # on_step=True
         )
-
-        # contr loss
-        if self.hparams.use_contr:
-            self.log(
-                f'{stage.to_pref()}contr_loss',
-                outputs['contr_loss'],
-                batch_size=len(batch['identifier']),
-                sync_dist=True,
-                prog_bar=False,
-                on_epoch=True,
-                # on_step=True
-            )
-
-            # noncongruent pairs
-            self.log(
-                f'{stage.to_pref()}noncong_loss',
-                outputs['noncong_loss'],
-                batch_size=len(batch['identifier']),
-                sync_dist=True,
-                prog_bar=False,
-                on_epoch=True,
-                # on_step=True
-            )
-            
-            # congruent pairs
-            self.log(
-                f'{stage.to_pref()}cong_loss',
-                outputs['cong_loss'],
-                batch_size=len(batch['identifier']),
-                sync_dist=True,
-                prog_bar=False,
-                on_epoch=True,
-                # on_step=True
-            )
-
-
-        if self.hparams.pred_fp:
-
-            self.log(
-                f'{stage.to_pref()}_fp_loss',
-                outputs['fp_loss'],
-                batch_size=len(batch['identifier']),
-                sync_dist=True,
-                prog_bar=False,
-                on_epoch=True,
-            )
-        
-        if self.hparams.use_cons_spec:
-            self.log(
-                f'{stage.to_pref()}cons_loss',
-                outputs['cons_spec_loss'],
-                batch_size=len(batch['identifier']),
-                sync_dist=True,
-                prog_bar=False,
-                on_epoch=True,
-            )
     
     def test_step(self, batch, batch_idx):
         # Unpack inputs
@@ -275,172 +176,160 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
             {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": False}, # monitor val loss
         ]
         return monitors
-    
-    # def _update_loss_weights(self)-> None:
-    #     if self.hparams.loss_strategy == 'linear':
-    #         for loss in self.loss_wts:
-    #             self.loss_wts[loss] += self.loss_updates[loss]
-    #     elif self.hparams.loss_strategy == 'manual':
-    #         for loss in self.loss_wts:
-    #             if self.current_epoch in self.loss_updates[loss]:
-    #                 self.loss_wts[loss] = self.loss_updates[loss][self.current_epoch]
-
-    # def on_train_epoch_end(self) -> None:
-    #     self._update_loss_weights()      
-    
-class MultiViewContrastive(ContrastiveModel):
+       
+# class MultiViewContrastive(ContrastiveModel):
 
-    def __init__(self,
-                 **kwargs):
+#     def __init__(self,
+#                  **kwargs):
         
-        super().__init__(**kwargs)
+#         super().__init__(**kwargs)
     
-        # build fingerprint encoder model
-        if self.hparams.use_fp: 
-            self.fp_enc_model = model_utils.get_fp_enc_model(self.hparams)
+#         # build fingerprint encoder model
+#         if self.hparams.use_fp: 
+#             self.fp_enc_model = model_utils.get_fp_enc_model(self.hparams)
         
-        # build NL encoder model
-        if self.hparams.use_NL_spec:
-            self.NL_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
+#         # build NL encoder model
+#         if self.hparams.use_NL_spec:
+#             self.NL_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
             
-    def forward(self, batch, stage):
-        g = batch['cand'] if stage == Stage.TEST else batch['mol']
+#     def forward(self, batch, stage):
+#         g = batch['cand'] if stage == Stage.TEST else batch['mol']
         
-        spec = batch[self.spec_view]
-        n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
+#         spec = batch[self.spec_view]
+#         n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
 
-        spec_enc = self.spec_enc_model(spec, n_peaks)
-        mol_enc = self.mol_enc_model(g)
-        views = {'spec_enc': spec_enc, 'mol_enc': mol_enc}
+#         spec_enc = self.spec_enc_model(spec, n_peaks)
+#         mol_enc = self.mol_enc_model(g)
+#         views = {'spec_enc': spec_enc, 'mol_enc': mol_enc}
         
-        if self.hparams.use_fp:
-            fp_enc = self.fp_enc_model(batch['fp'])
-            views['fp_enc'] = fp_enc
-
-        if self.hparams.use_cons_spec:
-            spec = batch['cons_spec']
-            n_peaks = batch['cons_n_peaks'] if 'cons_n_peaks' in batch else None
-            spec_enc = self.cons_spec_enc_model(spec, n_peaks)
-            views['cons_spec_enc'] = spec_enc
-
-        if self.hparams.use_NL_spec:
-            spec = batch['NL_spec']
-            n_peaks = batch['NL_n_peaks'] if 'NL_n_peaks' in batch else None
-            spec_enc = self.NL_enc_model(spec, n_peaks) 
-            views['NL_spec_enc'] = spec_enc
-        return views
+#         if self.hparams.use_fp:
+#             fp_enc = self.fp_enc_model(batch['fp'])
+#             views['fp_enc'] = fp_enc
+
+#         if self.hparams.use_cons_spec:
+#             spec = batch['cons_spec']
+#             n_peaks = batch['cons_n_peaks'] if 'cons_n_peaks' in batch else None
+#             spec_enc = self.cons_spec_enc_model(spec, n_peaks)
+#             views['cons_spec_enc'] = spec_enc
+
+#         if self.hparams.use_NL_spec:
+#             spec = batch['NL_spec']
+#             n_peaks = batch['NL_n_peaks'] if 'NL_n_peaks' in batch else None
+#             spec_enc = self.NL_enc_model(spec, n_peaks) 
+#             views['NL_spec_enc'] = spec_enc
+#         return views
     
-    def step(
-        self, batch: dict, stage= Stage.NONE):
+#     def step(
+#         self, batch: dict, stage= Stage.NONE):
         
-        # Compute spectra and mol encoding
-        views = self.forward(batch, stage)
+#         # Compute spectra and mol encoding
+#         views = self.forward(batch, stage)
 
-        if stage == Stage.TEST:
-            return views
+#         if stage == Stage.TEST:
+#             return views
         
-        # Calculate loss
-        losses = self.compute_loss(batch, views)
+#         # Calculate loss
+#         losses = self.compute_loss(batch, views)
 
-        return losses
+#         return losses
     
-    def compute_loss(self, batch: dict, views: dict):
-        loss = 0
-        losses = {}
-        for v1, v2 in self.hparams.contr_views:
-            contr_loss, cong_loss, noncong_loss = contrastive_loss(views[v1], views[v2], self.hparams.contr_temp)
-            loss+=contr_loss
-
-            losses[f'{v1[:-4]}-{v2[:-4]}_contr_loss'] = contr_loss.detach().item()
-            losses[f'{v1[:-4]}-{v2[:-4]}_cong_loss'] = cong_loss.detach().item()
-            losses[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'] = noncong_loss.detach().item()
+#     def compute_loss(self, batch: dict, views: dict):
+#         loss = 0
+#         losses = {}
+#         for v1, v2 in self.hparams.contr_views:
+#             contr_loss, cong_loss, noncong_loss = contrastive_loss(views[v1], views[v2], self.hparams.contr_temp)
+#             loss+=contr_loss
+
+#             losses[f'{v1[:-4]}-{v2[:-4]}_contr_loss'] = contr_loss.detach().item()
+#             losses[f'{v1[:-4]}-{v2[:-4]}_cong_loss'] = cong_loss.detach().item()
+#             losses[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'] = noncong_loss.detach().item()
         
-        losses['loss'] = loss 
+#         losses['loss'] = loss 
 
-        return losses
+#         return losses
     
-    def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
-        # total loss
-        self.log(
-            f'{stage.to_pref()}loss',
-            outputs['loss'],
-            batch_size=len(batch['identifier']),
-            sync_dist=True,
-            prog_bar=True,
-            on_epoch=True,
-            # on_step=True
-        )
-
-        for v1, v2 in self.hparams.contr_views:
-            self.log(
-            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_contr_loss',
-            outputs[f'{v1[:-4]}-{v2[:-4]}_contr_loss'],
-            batch_size=len(batch['identifier']),
-            sync_dist=True,
-            on_epoch=True,
-        )
-            self.log(
-            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_cong_loss',
-            outputs[f'{v1[:-4]}-{v2[:-4]}_cong_loss'],
-            batch_size=len(batch['identifier']),
-            sync_dist=True,
-            on_epoch=True,
-        )
-            self.log(
-            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_noncong_loss',
-            outputs[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'],
-            batch_size=len(batch['identifier']),
-            sync_dist=True,
-            on_epoch=True,
-        )
+#     def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
+#         # total loss
+#         self.log(
+#             f'{stage.to_pref()}loss',
+#             outputs['loss'],
+#             batch_size=len(batch['identifier']),
+#             sync_dist=True,
+#             prog_bar=True,
+#             on_epoch=True,
+#             # on_step=True
+#         )
+
+#         for v1, v2 in self.hparams.contr_views:
+#             self.log(
+#             f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_contr_loss',
+#             outputs[f'{v1[:-4]}-{v2[:-4]}_contr_loss'],
+#             batch_size=len(batch['identifier']),
+#             sync_dist=True,
+#             on_epoch=True,
+#         )
+#             self.log(
+#             f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_cong_loss',
+#             outputs[f'{v1[:-4]}-{v2[:-4]}_cong_loss'],
+#             batch_size=len(batch['identifier']),
+#             sync_dist=True,
+#             on_epoch=True,
+#         )
+#             self.log(
+#             f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_noncong_loss',
+#             outputs[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'],
+#             batch_size=len(batch['identifier']),
+#             sync_dist=True,
+#             on_epoch=True,
+#         )
             
-    def test_step(self, batch):
-        # Unpack inputs
-        identifiers = batch['identifier']
-        cand_smiles = batch['cand_smiles']
-        id_to_ct = defaultdict(int)
-        for i in identifiers: id_to_ct[i]+=1
-        batch_ptr = torch.tensor(list(id_to_ct.values()))
-
-        outputs = self.step(batch, stage=Stage.TEST)
-        scores = {}
-        for v1, v2 in self.hparams.contr_views:
-            # if 'cons_spec_enc' in (v1, v2):
-            #     continue
-            v1_enc = outputs[v1]
-            v2_enc = outputs[v2]
+#     def test_step(self, batch):
+#         # Unpack inputs
+#         identifiers = batch['identifier']
+#         cand_smiles = batch['cand_smiles']
+#         id_to_ct = defaultdict(int)
+#         for i in identifiers: id_to_ct[i]+=1
+#         batch_ptr = torch.tensor(list(id_to_ct.values()))
+
+#         outputs = self.step(batch, stage=Stage.TEST)
+#         scores = {}
+#         for v1, v2 in self.hparams.contr_views:
+#             # if 'cons_spec_enc' in (v1, v2):
+#             #     continue
+#             v1_enc = outputs[v1]
+#             v2_enc = outputs[v2]
             
-            s = nn.functional.cosine_similarity(v1_enc, v2_enc)
-            scores[f'{v1[:-4]}-{v2[:-4]}_scores'] = torch.split(s, list(id_to_ct.values()))
+#             s = nn.functional.cosine_similarity(v1_enc, v2_enc)
+#             scores[f'{v1[:-4]}-{v2[:-4]}_scores'] = torch.split(s, list(id_to_ct.values()))
 
-        indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
+#         indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
 
-        cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
-        labels = utils.unbatch_list(batch['label'], indexes)
+#         cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
+#         labels = utils.unbatch_list(batch['label'], indexes)
         
-        return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
+#         return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
     
-    def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
+#     def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
 
-        # save scores
-        for i, cands, l in zip(outputs['identifiers'], outputs['cand_smiles'], outputs['labels']):
-            self.result_dct[i]['candidates'].extend(cands)
-            self.result_dct[i]['labels'].extend([x.cpu().item() for x in l])
+#         # save scores
+#         for i, cands, l in zip(outputs['identifiers'], outputs['cand_smiles'], outputs['labels']):
+#             self.result_dct[i]['candidates'].extend(cands)
+#             self.result_dct[i]['labels'].extend([x.cpu().item() for x in l])
 
-        for v1, v2 in self.hparams.contr_views:
-            for i, scores in zip(outputs['identifiers'], outputs['scores'][f'{v1[:-4]}-{v2[:-4]}_scores']):
-                self.result_dct[i][f'{v1[:-4]}-{v2[:-4]}_scores'].extend(scores.cpu().tolist())
+#         for v1, v2 in self.hparams.contr_views:
+#             for i, scores in zip(outputs['identifiers'], outputs['scores'][f'{v1[:-4]}-{v2[:-4]}_scores']):
+#                 self.result_dct[i][f'{v1[:-4]}-{v2[:-4]}_scores'].extend(scores.cpu().tolist())
 
 
-    def on_test_epoch_end(self) -> None:
+#     def on_test_epoch_end(self) -> None:
 
-        self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
+#         self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
 
-        # Compute rank
-        for v1, v2 in self.hparams.contr_views:
-            self.df_test[f'{v1[:-4]}-{v2[:-4]}_rank'] = self.df_test.apply(lambda row: self._compute_rank(row[f'{v1[:-4]}-{v2[:-4]}_scores'], row['labels']), axis=1)
+#         # Compute rank
+#         for v1, v2 in self.hparams.contr_views:
+#             self.df_test[f'{v1[:-4]}-{v2[:-4]}_rank'] = self.df_test.apply(lambda row: self._compute_rank(row[f'{v1[:-4]}-{v2[:-4]}_scores'], row['labels']), axis=1)
 
-        self.df_test.to_pickle(self.df_test_path)
+#         self.df_test.to_pickle(self.df_test_path)
 
 class FilipContrastive(ContrastiveModel):
     def __init__(self,
@@ -492,7 +381,7 @@ class FilipContrastive(ContrastiveModel):
         # Calculate scores
         indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
         
-        scores = filip_similarity_batch(spec_enc, mol_enc, spec_mask, mol_mask)
+        scores = filip_similarity_batch(spec_enc, mol_enc, spec_mask, mol_mask, reduction='geom', temperature=0.05)
         scores = torch.split(scores, list(id_to_ct.values()))
 
         cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
@@ -500,248 +389,248 @@ class FilipContrastive(ContrastiveModel):
         
         return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
     
-class MultiViewFineTuning(MultiViewContrastive):
-    def __init__(self,
-                 **kwargs):
-        super().__init__(**kwargs)
-
-        # load preptrained spec, mol, fp encoders
-        checkpoint = torch.load(self.hparams.partial_checkpoint)
-        state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
-        self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
-
-        state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
-        self.mol_enc_model.load_state_dict(state_dict)
-
-        state_dict = state_dict = {k[len("fp_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("fp_enc_model")}
-        self.fp_enc_model.load_state_dict(state_dict)
-
-        self.encoding_views = ['spec_enc', 'mol_enc', 'fp_enc']
-        self.loss_fn = nn.BCELoss()
-
-        # freeze encoders
-        for param in self.mol_enc_model.parameters():
-            param.requires_grad = False
-        for param in self.spec_enc_model.parameters():
-            param.requires_grad = False
-        for param in self.fp_enc_model.parameters():
-            param.requires_grad = False
-        for param in self.cons_spec_enc_model.parameters():
-            param.requires_grad = False
-
-        # n_views = 2
-        # if self.hparams.use_fp:
-        #     n_views+=1
-
-        # in_dim = self.hparams.final_embedding_dim*n_views
-        in_dim = self.hparams.final_embedding_dim *2 + 2
-
-        self.classifier_model = nn.Sequential(
-            nn.Linear(in_dim, 512),
-            nn.ReLU(),
-            nn.BatchNorm1d(512),
-            nn.Dropout(0.3),
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.BatchNorm1d(256),
-            nn.Dropout(0.3),
-            nn.Linear(256, 1),
-            nn.Sigmoid()
-        )
-        self.noise_std = 0.01
-
-    def _add_noise(self, x):
-        noise = torch.randn_like(x) * self.noise_std
-        return x + noise
-
-    def forward(self, batch, stage):
-
-        matching_views = super().forward(batch, stage)
-        # matching_enc = torch.concat((matching_views['spec_enc'], matching_views['mol_enc'], matching_views['fp_enc']), dim=-1)
-        # enc1 = matching_views['spec_enc'] - matching_views['mol_enc']
-        # enc2 = matching_views['spec_enc'] - matching_views['fp_enc']
-        # matching_enc = torch.concat((enc1, enc2), dim=-1)
-        view1 = matching_views['spec_enc']
-        view2 = matching_views['mol_enc']
-        view3 = matching_views['fp_enc']
-
-        if stage == Stage.TRAIN:
-            view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
-
-        pairwise_diffs = torch.cat([
-            torch.abs(view1 - view2),
-            torch.abs(view1 - view3),
-        ], dim=-1)
-
-        pairwise_sims = torch.cat([
-            (view1 * view2).sum(dim=-1, keepdim=True),
-            (view1 * view3).sum(dim=-1, keepdim=True),
-        ], dim=-1)
-
-        matching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
-        matching_scores = self.classifier_model(matching_enc)
-
-        if stage == Stage.TEST:
-            return dict(matching_scores = matching_scores)
+# class MultiViewFineTuning(MultiViewContrastive):
+#     def __init__(self,
+#                  **kwargs):
+#         super().__init__(**kwargs)
+
+#         # load preptrained spec, mol, fp encoders
+#         checkpoint = torch.load(self.hparams.partial_checkpoint)
+#         state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
+#         self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
+
+#         state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
+#         self.mol_enc_model.load_state_dict(state_dict)
+
+#         state_dict = state_dict = {k[len("fp_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("fp_enc_model")}
+#         self.fp_enc_model.load_state_dict(state_dict)
+
+#         self.encoding_views = ['spec_enc', 'mol_enc', 'fp_enc']
+#         self.loss_fn = nn.BCELoss()
+
+#         # freeze encoders
+#         for param in self.mol_enc_model.parameters():
+#             param.requires_grad = False
+#         for param in self.spec_enc_model.parameters():
+#             param.requires_grad = False
+#         for param in self.fp_enc_model.parameters():
+#             param.requires_grad = False
+#         for param in self.cons_spec_enc_model.parameters():
+#             param.requires_grad = False
+
+#         # n_views = 2
+#         # if self.hparams.use_fp:
+#         #     n_views+=1
+
+#         # in_dim = self.hparams.final_embedding_dim*n_views
+#         in_dim = self.hparams.final_embedding_dim *2 + 2
+
+#         self.classifier_model = nn.Sequential(
+#             nn.Linear(in_dim, 512),
+#             nn.ReLU(),
+#             nn.BatchNorm1d(512),
+#             nn.Dropout(0.3),
+#             nn.Linear(512, 256),
+#             nn.ReLU(),
+#             nn.BatchNorm1d(256),
+#             nn.Dropout(0.3),
+#             nn.Linear(256, 1),
+#             nn.Sigmoid()
+#         )
+#         self.noise_std = 0.01
+
+#     def _add_noise(self, x):
+#         noise = torch.randn_like(x) * self.noise_std
+#         return x + noise
+
+#     def forward(self, batch, stage):
+
+#         matching_views = super().forward(batch, stage)
+#         # matching_enc = torch.concat((matching_views['spec_enc'], matching_views['mol_enc'], matching_views['fp_enc']), dim=-1)
+#         # enc1 = matching_views['spec_enc'] - matching_views['mol_enc']
+#         # enc2 = matching_views['spec_enc'] - matching_views['fp_enc']
+#         # matching_enc = torch.concat((enc1, enc2), dim=-1)
+#         view1 = matching_views['spec_enc']
+#         view2 = matching_views['mol_enc']
+#         view3 = matching_views['fp_enc']
+
+#         if stage == Stage.TRAIN:
+#             view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
+
+#         pairwise_diffs = torch.cat([
+#             torch.abs(view1 - view2),
+#             torch.abs(view1 - view3),
+#         ], dim=-1)
+
+#         pairwise_sims = torch.cat([
+#             (view1 * view2).sum(dim=-1, keepdim=True),
+#             (view1 * view3).sum(dim=-1, keepdim=True),
+#         ], dim=-1)
+
+#         matching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
+#         matching_scores = self.classifier_model(matching_enc)
+
+#         if stage == Stage.TEST:
+#             return dict(matching_scores = matching_scores)
         
-        view1 = view1.repeat_interleave(self.hparams.aug_cands_size, dim=0)
-        view2 = self.mol_enc_model(batch['aug_cands'])
-        view3= self.fp_enc_model(batch['aug_cands_fp'])
-        if stage == Stage.TRAIN:
-            view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
+#         view1 = view1.repeat_interleave(self.hparams.aug_cands_size, dim=0)
+#         view2 = self.mol_enc_model(batch['aug_cands'])
+#         view3= self.fp_enc_model(batch['aug_cands_fp'])
+#         if stage == Stage.TRAIN:
+#             view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
 
-        pairwise_diffs = torch.cat([
-            torch.abs(view1 - view2),
-            torch.abs(view1 - view3),
-        ], dim=-1)
+#         pairwise_diffs = torch.cat([
+#             torch.abs(view1 - view2),
+#             torch.abs(view1 - view3),
+#         ], dim=-1)
 
-        pairwise_sims = torch.cat([
-            (view1 * view2).sum(dim=-1, keepdim=True),
-            (view1 * view3).sum(dim=-1, keepdim=True),
-        ], dim=-1)
+#         pairwise_sims = torch.cat([
+#             (view1 * view2).sum(dim=-1, keepdim=True),
+#             (view1 * view3).sum(dim=-1, keepdim=True),
+#         ], dim=-1)
 
-        nonmatching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
+#         nonmatching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
 
-        nonmatching_scores = self.classifier_model(nonmatching_enc)
+#         nonmatching_scores = self.classifier_model(nonmatching_enc)
 
-        return dict(matching_scores=matching_scores, nonmatching_scores=nonmatching_scores)
+#         return dict(matching_scores=matching_scores, nonmatching_scores=nonmatching_scores)
     
-    def compute_loss(self, matching_scores, nonmatching_scores):
+#     def compute_loss(self, matching_scores, nonmatching_scores):
         
-        matching_loss = self.loss_fn(matching_scores, torch.ones_like(matching_scores).to(matching_scores.device))
-        nonmatching_loss = self.loss_fn(nonmatching_scores, torch.zeros_like(nonmatching_scores).to(nonmatching_scores.device))
+#         matching_loss = self.loss_fn(matching_scores, torch.ones_like(matching_scores).to(matching_scores.device))
+#         nonmatching_loss = self.loss_fn(nonmatching_scores, torch.zeros_like(nonmatching_scores).to(nonmatching_scores.device))
 
-        loss = matching_loss + (1/self.hparams.aug_cands_size)*nonmatching_loss
+#         loss = matching_loss + (1/self.hparams.aug_cands_size)*nonmatching_loss
 
-        return dict(loss=loss)
+#         return dict(loss=loss)
     
-    def step(
-        self, batch: dict, stage= Stage.NONE):
+#     def step(
+#         self, batch: dict, stage= Stage.NONE):
         
-        output = self.forward(batch, stage)
+#         output = self.forward(batch, stage)
 
-        if stage == Stage.TEST:
-            return output
+#         if stage == Stage.TEST:
+#             return output
         
-        # Calculate loss
-        losses = self.compute_loss(output['matching_scores'], output['nonmatching_scores'])
+#         # Calculate loss
+#         losses = self.compute_loss(output['matching_scores'], output['nonmatching_scores'])
 
-        return losses
+#         return losses
     
-    def test_step(self, batch):
-        # Unpack inputs
-        identifiers = batch['identifier']
-        cand_smiles = batch['cand_smiles']
-        id_to_ct = defaultdict(int)
-        for i in identifiers: id_to_ct[i]+=1
-        batch_ptr = torch.tensor(list(id_to_ct.values()))
+#     def test_step(self, batch):
+#         # Unpack inputs
+#         identifiers = batch['identifier']
+#         cand_smiles = batch['cand_smiles']
+#         id_to_ct = defaultdict(int)
+#         for i in identifiers: id_to_ct[i]+=1
+#         batch_ptr = torch.tensor(list(id_to_ct.values()))
 
-        outputs = self.step(batch, stage=Stage.TEST)
-        scores = outputs['matching_scores']
+#         outputs = self.step(batch, stage=Stage.TEST)
+#         scores = outputs['matching_scores']
 
-        indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
+#         indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
 
-        cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
-        labels = utils.unbatch_list(batch['label'], indexes)
+#         cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
+#         labels = utils.unbatch_list(batch['label'], indexes)
         
-        return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
+#         return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
     
-    def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
-        # total loss
-        self.log(
-            f'{stage.to_pref()}loss',
-            outputs['loss'],
-            batch_size=len(batch['identifier']),
-            sync_dist=True,
-            prog_bar=True,
-            on_epoch=True,
-            # on_step=True
-        )
-
-    def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
-        ContrastiveModel.on_test_batch_end(self, outputs, batch, batch_idx, stage)
-
-    def on_test_epoch_end(self):
-        self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
-        # self.df_test.to_csv(self.hparams.resutl)
-        print(self.df_test_path)
-        self.df_test.to_pickle(self.df_test_path)
-        # ContrastiveModel.on_test_epoch_end(self)
-
-    def get_checkpoint_monitors(self) -> T.List[dict]:
-        monitors = [
-            {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
-        ]
-        return monitors
-    def configure_optimizers(self):
-        return torch.optim.Adam(
-            self.classifier_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
-        )
-
-class IndSpecEncoder(ContrastiveModel):
-    """ Trains a spectra encoder that maps to a pretrained spec encoder"""
-    def __init__(
-            self,
-            **kwargs
-    ):
-        super().__init__(**kwargs)
-
-        # initialize ind_spec_encoder and loss
-        self.ind_spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
-        self.cons_loss = cons_spec_loss(self.hparams.cons_loss_type)
-
-        # load preptrained spec and mol encoders
-        checkpoint = torch.load(self.hparams.partial_checkpoint)
-        state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
-        self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
-
-        state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
-        self.mol_enc_model.load_state_dict(state_dict)
-
-        # freeze cons spec and mol encoders
-        for param in self.mol_enc_model.parameters():
-            param.requires_grad = False
-        for param in self.spec_enc_model.parameters():
-            param.requires_grad = False
-
-    def forward(self, batch, stage):
-
-        spec = batch[self.spec_view]
-        n_peaks = batch['n_peaks']
-        spec_enc = self.ind_spec_enc_model(spec, n_peaks)
-
-        return spec_enc
+#     def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
+#         # total loss
+#         self.log(
+#             f'{stage.to_pref()}loss',
+#             outputs['loss'],
+#             batch_size=len(batch['identifier']),
+#             sync_dist=True,
+#             prog_bar=True,
+#             on_epoch=True,
+#             # on_step=True
+#         )
+
+#     def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
+#         ContrastiveModel.on_test_batch_end(self, outputs, batch, batch_idx, stage)
+
+#     def on_test_epoch_end(self):
+#         self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
+#         # self.df_test.to_csv(self.hparams.resutl)
+#         print(self.df_test_path)
+#         self.df_test.to_pickle(self.df_test_path)
+#         # ContrastiveModel.on_test_epoch_end(self)
+
+#     def get_checkpoint_monitors(self) -> T.List[dict]:
+#         monitors = [
+#             {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
+#         ]
+#         return monitors
+#     def configure_optimizers(self):
+#         return torch.optim.Adam(
+#             self.classifier_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
+#         )
+
+# class IndSpecEncoder(ContrastiveModel):
+#     """ Trains a spectra encoder that maps to a pretrained spec encoder"""
+#     def __init__(
+#             self,
+#             **kwargs
+#     ):
+#         super().__init__(**kwargs)
+
+#         # initialize ind_spec_encoder and loss
+#         self.ind_spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
+#         self.cons_loss = cons_spec_loss(self.hparams.cons_loss_type)
+
+#         # load preptrained spec and mol encoders
+#         checkpoint = torch.load(self.hparams.partial_checkpoint)
+#         state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
+#         self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
+
+#         state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
+#         self.mol_enc_model.load_state_dict(state_dict)
+
+#         # freeze cons spec and mol encoders
+#         for param in self.mol_enc_model.parameters():
+#             param.requires_grad = False
+#         for param in self.spec_enc_model.parameters():
+#             param.requires_grad = False
+
+#     def forward(self, batch, stage):
+
+#         spec = batch[self.spec_view]
+#         n_peaks = batch['n_peaks']
+#         spec_enc = self.ind_spec_enc_model(spec, n_peaks)
+
+#         return spec_enc
     
-    def compute_loss(self, spec_enc, cons_spec_enc):
-        loss = self.cons_loss(spec_enc, cons_spec_enc)
-        return dict(loss=loss)
+#     def compute_loss(self, spec_enc, cons_spec_enc):
+#         loss = self.cons_loss(spec_enc, cons_spec_enc)
+#         return dict(loss=loss)
 
-    def step(self, batch: dict, stage=Stage.NONE):
-        self.spec_enc_model.eval()
-        self.mol_enc_model.eval()
+#     def step(self, batch: dict, stage=Stage.NONE):
+#         self.spec_enc_model.eval()
+#         self.mol_enc_model.eval()
 
-        spec_enc = self.forward(batch, stage)
+#         spec_enc = self.forward(batch, stage)
 
-        if stage == Stage.TEST:
-            mol_enc = self.mol_enc_model(batch['cand'])
-            return dict(spec_enc=spec_enc, mol_enc=mol_enc)
+#         if stage == Stage.TEST:
+#             mol_enc = self.mol_enc_model(batch['cand'])
+#             return dict(spec_enc=spec_enc, mol_enc=mol_enc)
         
-        cons_spec_enc = self.spec_enc_model(batch['cons_spec'], batch['cons_n_peaks'])
+#         cons_spec_enc = self.spec_enc_model(batch['cons_spec'], batch['cons_n_peaks'])
 
-        losses = self.compute_loss(spec_enc, cons_spec_enc)
+#         losses = self.compute_loss(spec_enc, cons_spec_enc)
 
-        return losses
+#         return losses
     
     
-    def configure_optimizers(self):
-        return torch.optim.Adam(
-            self.ind_spec_enc_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
-        )
-    def get_checkpoint_monitors(self) -> T.List[dict]:
-        monitors = [
-            {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
-        ]
-        return monitors
+#     def configure_optimizers(self):
+#         return torch.optim.Adam(
+#             self.ind_spec_enc_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
+#         )
+#     def get_checkpoint_monitors(self) -> T.List[dict]:
+#         monitors = [
+#             {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
+#         ]
+#         return monitors
     
 class CrossAttenContrastive(ContrastiveModel):
     def __init__(

{mvp → flare}/models/spec_encoder.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 import torch
-from mvp.models.encoders import MLP
+from flare.models.encoders import MLP
 from torch_geometric.nn import global_mean_pool
 
 
@@ -41,14 +41,14 @@ class SpecMzIntTokenTransformer(nn.Module):
         if args.model in ('crossAttenContrastive', 'filipContrastive'):
             self.returnEmb = True
             assert(args.use_cls == False)
+        else:
+            self.specEncoder = nn.Sequential(nn.Linear(args.hidden_dims[-1], args.final_embedding_dim), nn.Dropout(args.fc_dropout))
 
         self.use_cls = args.use_cls
         if self.use_cls:
             self.cls_embed = torch.nn.Embedding(1,args.hidden_dims[-1])
         encoder_layer = nn.TransformerEncoderLayer(d_model=args.hidden_dims[-1], nhead=2, batch_first=True)
         self.tokenTransformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
-
-        self.specEncoder = nn.Sequential(nn.Linear(args.hidden_dims[-1], args.final_embedding_dim), nn.Dropout(args.fc_dropout))
     
     def forward(self, spec, n_peaks=None):
         h = self.tokenEnc(spec)
@@ -61,11 +61,10 @@ class SpecMzIntTokenTransformer(nn.Module):
             pad = torch.concat((torch.tensor(False).repeat(pad.shape[0],1).to(spec.device), pad), dim=1)
             h = self.tokenTransformer(h, src_key_padding_mask=pad)
             h = h[:,0,:]
-        else:
             
+        else:
             # mean
             h = self.tokenTransformer(h, src_key_padding_mask=pad)
-
             if self.returnEmb:
                 # repad h
                 h[pad] = -5
@@ -123,11 +122,10 @@ class SpecFormulaTransformer(nn.Module):
             self.cls_embed = torch.nn.Embedding(1,args.formula_dims[-1])
         encoder_layer = nn.TransformerEncoderLayer(d_model=args.formula_dims[-1], nhead=args.formula_attn_heads, batch_first=True)
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=args.formula_transformer_layers)
-
-        if not out_dim:
-            out_dim = args.final_embedding_dim
-
+        
         if not self.returnEmb:
+            if not out_dim:
+                out_dim = args.final_embedding_dim
             self.fc = nn.Linear(args.formula_dims[-1], out_dim)
         
     def forward(self, spec, n_peaks):

{mvp → flare}/params_formSpec.yaml

@@ -1,11 +1,11 @@
 # Experiment setup 
 job_key: ''
-run_name: 'optimized_flare'
+run_name: 'flare_sirius_labels_42'
 run_details: ""
 project_name: ''
 wandb_entity_name: 'mass-spec-ml'
 no_wandb: True
-seed: 0
+seed: 42
 debug: False
 checkpoint_pth:
 
@@ -19,27 +19,23 @@ val_check_interval: 1.0
 # Data paths
 candidates_pth: /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_mass.json # "../data/MassSpecGym/data/molecules/MassSpecGym_retrieval_candidates_formula.json"
 dataset_pth: /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv # /data/yzhouc01/MVP/data/sample/data.tsv #/r/hassounlab/spectra_data/msgym/MassSpecGym.tsv #/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv # /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv # "../data/MassSpecGym/data/sample_data.tsv" 
-subformula_dir_pth: /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default # /data/yzhouc01/FILIP-MS/data/magma # /r/hassounlab/msgym_sirius # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
+subformula_dir_pth: /r/hassounlab/msgym_sirius # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default # /data/yzhouc01/FILIP-MS/data/magma # /r/hassounlab/msgym_sirius # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
 split_pth: 
 fp_dir_pth:
-cons_spec_dir_pth: 
-NL_spec_dir_pth: ""
 partial_checkpoint: ""
 
 # General hyperparameters
-batch_size: 32 #64
-lr: 7.092216555765765e-05 #2.881339661302105e-05 # 5.0e-05 
+batch_size: 64 #64
+lr: 2.881339661302105e-05 # 5.0e-05 
 weight_decay: 1.8376229667330708e-05
-contr_temp: 0.043339030104611806 # 0.022772534845886608 # 0.05
-early_stopping_patience: 300
-loss_strategy: 'static'
+contr_temp: 0.022772534845886608 # 0.022772534845886608 # 0.05
 num_workers: 50
 
 
 ############################## Data transforms ##############################
 # - Spectra 
 spectra_view: SpecFormula #SpecMzIntTokens #SpecFormula 
-formula_source: 'default' # magma_1, magma_all, sirius, default
+formula_source: 'sirius' # magma_1, magma_all, sirius, default
 #  1. Binner
 max_mz: 1000
 bin_width: 1
@@ -48,7 +44,6 @@ mask_peak_ratio: 0.00
 # 2. SpecFormula
 element_list: ['H', 'C',  'O', 'N', 'P', 'S', 'Cl', 'F', 'Br', 'I', 'B', 'As', 'Si', 'Se']
 add_intensities: True
-mask_precursor: False
 
 # - Molecule
 molecule_view: "MolGraph"
@@ -58,34 +53,34 @@ bond_feature: 'full'
 
 ############################## Views ##############################
 # contrastive
-use_contr: False
-contr_wt: 1
-contr_wt_update: {} 
+# use_contr: False
+# contr_wt: 1
+# contr_wt_update: {} 
 
 # consensus spectra
-use_cons_spec: False
-cons_spec_wt: 3
-cons_spec_wt_update: {} 
-cons_loss_type: 'l2' # cosine, l2
+# use_cons_spec: False
+# cons_spec_wt: 3
+# cons_spec_wt_update: {} 
+# cons_loss_type: 'l2' # cosine, l2
 
 # fp prediction/usage
-pred_fp: False
-use_fp: False
-fp_loss_type: 'cosine' #cosine, bce
-fp_wt: 3
-fp_wt_update: {} 
-fp_size: 1024
-fp_radius: 5
-fp_dropout: 0.4
+# pred_fp: False
+# use_fp: False
+# fp_loss_type: 'cosine' #cosine, bce
+# fp_wt: 3
+# fp_wt_update: {} 
+# fp_size: 1024
+# fp_radius: 5
+# fp_dropout: 0.4
 
 # candidates
-aug_cands: False
-aug_cands_wt: 0.1
-aug_cands_update: {} 
-aug_cands_size: 3
+# aug_cands: False
+# aug_cands_wt: 0.1
+# aug_cands_update: {} 
+# aug_cands_size: 3
 
 # neutral loss
-use_NL: False
+# use_NL: False
 
 
 ############################## Task and model ##############################
@@ -93,33 +88,34 @@ task: 'retrieval'
 spec_enc: Transformer_Formula # Transformer_MzInt #Transformer_Formula 
 mol_enc: "GNN"
 model:  filipContrastive # "MultiviewContrastive"
-contr_views: [['spec_enc', 'mol_enc']] #[['spec_enc', 'mol_enc'], ['spec_enc', 'NL_spec_enc'], ['mol_enc', 'NL_spec_enc']] #[['spec_enc', 'mol_enc'], ['mol_enc', 'cons_spec_enc'], ['cons_spec_enc', 'spec_enc'], ['fp_enc', 'mol_enc'], ['fp_enc', 'spec_enc'], ['fp_enc', 'cons_spec_enc']]
+contr_views: [['spec_enc', 'mol_enc']]
 log_only_loss_at_stages: []
 df_test_path: ""
 
-# - Spectra encoder 
-final_embedding_dim: 512
-fc_dropout: 0.4
-
-# - Spectra Token encoder
-hidden_dims: [64, 128]
-
 
 # - Formula-based spec encoders
 formula_dropout:  0.2
-formula_dims: [256,512,256] #[512, 256, 512] #[64, 128, 256]
+formula_dims: [512,256,512] #[512, 256, 512] #[64, 128, 256]
 cross_attn_heads: 2
 use_cls: False
 peak_dropout: 0.2
 formula_attn_heads: 4 # 2
-formula_transformer_layers: 1 #2
+formula_transformer_layers: 2 #2
 
 # -- GAT params
 attn_heads: [12,12,12]
 
 # - Molecule encoder (GNN)
-gnn_channels: [128, 256, 256] #[64,128,512] 
+gnn_channels: [128, 256, 512] #[64,128,512] 
 gnn_type: "gcn"
-num_gnn_layers: 3
-gnn_hidden_dim: 512
-gnn_dropout: 0.157104273477570 #0.23234950970370824 #0.3
+# num_gnn_layers: 3
+# gnn_hidden_dim: 512
+gnn_dropout: 0.23234950970370824 #0.3
+
+
+# - Spectra encoder (cross attention model)
+# final_embedding_dim: 512
+# fc_dropout: 0.4
+
+# - Spectra Token encoder (mz-int token model)
+# hidden_dims: [64, 256]

flare/run.sh

@@ -0,0 +1,3 @@
+# python train.py
+python test.py --param_pth ../hparams.yaml
+# python test.py --candidates_pth /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_formula.json

{mvp → flare}/test.py

@@ -10,12 +10,12 @@ from pytorch_lightning import Trainer
 from massspecgym.models.base import Stage
 import os
 
-from mvp.data.data_module import TestDataModule
-from mvp.data.datasets import ContrastiveDataset
-from mvp.utils.data import get_spec_featurizer, get_mol_featurizer, get_test_ms_dataset
-from mvp.utils.models import get_model
+from flare.data.data_module import TestDataModule
+from flare.data.datasets import ContrastiveDataset
+from flare.utils.data import get_spec_featurizer, get_mol_featurizer, get_test_ms_dataset
+from flare.utils.models import get_model
 
-from mvp.definitions import TEST_RESULTS_DIR
+from flare.definitions import TEST_RESULTS_DIR
 import yaml
 from functools import partial
 # Suppress RDKit warnings and errors

{mvp → flare}/train.py

@@ -11,15 +11,15 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 
 
-from mvp.data.data_module import ContrastiveDataModule
+from flare.data.data_module import ContrastiveDataModule
 
-from mvp.definitions import TEST_RESULTS_DIR
+from flare.definitions import TEST_RESULTS_DIR
 import yaml
-from mvp.data.datasets import ContrastiveDataset
+from flare.data.datasets import ContrastiveDataset
 from functools import partial
 
-from mvp.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
-from mvp.utils.models import get_model
+from flare.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
+from flare.utils.models import get_model
 # Suppress RDKit warnings and errors
 lg = RDLogger.logger()
 lg.setLevel(RDLogger.CRITICAL)
@@ -43,7 +43,7 @@ def main(params):
     dataset = get_ms_dataset(params['spectra_view'], params['molecule_view'], spec_featurizer, mol_featurizer, params)
     
     # Init data module
-    collate_fn = partial(ContrastiveDataset.collate_fn, spec_enc=params['spec_enc'], spectra_view=params['spectra_view'], mask_peak_ratio=params['mask_peak_ratio'], aug_cands=params['aug_cands'])
+    collate_fn = partial(ContrastiveDataset.collate_fn, spec_enc=params['spec_enc'], spectra_view=params['spectra_view'])
     data_module = ContrastiveDataModule(
         dataset=dataset,
         collate_fn=collate_fn,

{mvp → flare}/tune.py

@@ -15,11 +15,11 @@ from pytorch_lightning import Trainer
 from optuna.integration import PyTorchLightningPruningCallback
 from pytorch_lightning.callbacks import Callback
 
-from mvp.data.data_module import ContrastiveDataModule
-from mvp.data.datasets import ContrastiveDataset
-from mvp.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
-from mvp.utils.models import get_model
-from mvp.definitions import TEST_RESULTS_DIR
+from flare.data.data_module import ContrastiveDataModule
+from flare.data.datasets import ContrastiveDataset
+from flare.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
+from flare.utils.models import get_model
+from flare.definitions import TEST_RESULTS_DIR
 from functools import partial
 from rdkit import RDLogger
 from massspecgym.models.base import Stage

{mvp → flare}/utils/data.py

@@ -2,12 +2,12 @@ import os
 import json
 import numpy as np
 
-from mvp.data.transforms import SpecBinner, SpecBinnerLog, SpecFormulaFeaturizer, SpecFormulaMzFeaturizer, SpecMzIntTokenizer
+from flare.data.transforms import SpecBinner, SpecBinnerLog, SpecFormulaFeaturizer, SpecFormulaMzFeaturizer, SpecMzIntTokenizer
 from massspecgym.data.transforms import SpecTransform, MolTransform
-from mvp.data.transforms import MolToGraph
-import mvp.data.datasets as jestr_datasets
+from flare.data.transforms import MolToGraph
+import flare.data.datasets as jestr_datasets
 import typing as T
-from mvp.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH, PRECURSOR_INTENSITY
+from flare.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH, PRECURSOR_INTENSITY
 import matchms
 import tqdm
 
@@ -42,9 +42,6 @@ class Subformula_Loader:
         '''MIST subformula format:https://github.com/samgoldman97/mist/blob/main_v2/src/mist/utils/spectra_utils.py
         '''
         try:
-            # file = os.path.join(self.dir_path, spec_id+".json")
-            # with open(file) as f:
-            #     data = json.load(f)
             mzs = np.array(data['output_tbl']['mz'])
             formulas = np.array(data['output_tbl']['formula'])
             intensities = np.array(data['output_tbl']['ms2_inten'])
@@ -271,12 +268,12 @@ def get_test_ms_dataset(spectra_view: T.Union[str, T.List[str]],
         dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth'], 'use_magma': params['formula_source'].startswith('magma'), 'formula_source':params['formula_source']})
         use_formulas = True 
     
-    if params['use_cons_spec']:
-        dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
-    if 'use_NL_spec' in params and params['use_NL_spec']:
-        dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
-    if params['pred_fp'] or params['use_fp']:
-        dataset_params.update({'fp_dir_pth': '', 'fp_size': params['fp_size'], 'fp_radius': params['fp_radius']})
+    # if params['use_cons_spec']:
+    #     dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
+    # if 'use_NL_spec' in params and params['use_NL_spec']:
+    #     dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
+    # if params['pred_fp'] or params['use_fp']:
+    #     dataset_params.update({'fp_dir_pth': '', 'fp_size': params['fp_size'], 'fp_radius': params['fp_radius']})
 
     return jestr_datasets.ExpandedRetrievalDataset(use_formulas=use_formulas, **dataset_params)
     
@@ -294,24 +291,8 @@ def get_ms_dataset(spectra_view: str,
         dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth'], 'formula_source': params['formula_source']})
         use_formulas = True
 
-    if params['pred_fp'] or params['use_fp']:
-        dataset_params.update({'fp_dir_pth': params['fp_dir_pth']})
-
-    if params['aug_cands']:
-        dataset_params.update({'aug_cands_dir_pth': params['aug_cands_dir_pth'],
-                               'use_formulas':use_formulas,
-                               "aug_cands_size": params['aug_cands_size']})
-        
-    if params['use_cons_spec']:
-        dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
-
-    if 'use_NL_spec' in params and params['use_NL_spec']:
-        dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
-
     # select dataset
-    if params['aug_cands']:
-        return jestr_datasets.MassSpecDataset_Candidates(**dataset_params)
-    elif use_formulas:
+    if use_formulas:
         return jestr_datasets.MassSpecDataset_PeakFormulas(**dataset_params)
     
     return jestr_datasets.JESTR1_MassSpecDataset(**dataset_params)

{mvp → flare}/utils/eval.py

@@ -1,8 +1,8 @@
-from MassSpecGym.massspecgym.utils import MyopicMCES
+from massspecgym.utils import MyopicMCES
 import numpy as np
 import tqdm
 from multiprocessing import Pool
-
+from scipy.stats import bootstrap
 import os
 import pandas as pd
 
@@ -51,29 +51,6 @@ class Compute_Myopic_MCES_timeout:
 
         return results
 
-    
-def get_result_files(exp_dir, spec_type, views_type):
-    files = os.listdir(exp_dir)
-    mass_result = ''
-    form_result = ''
-
-    for f in files:
-        try:
-            _, s, views = f.split('_')
-        except:
-            continue
-        
-        if s == spec_type and views == views_type:
-            print(exp_dir / f)
-
-            files = os.listdir(exp_dir / f)
-            for fr in files:
-                if 'mass_result' in fr:
-                    mass_result = exp_dir / f / fr
-                elif 'result' in fr:
-                    form_result = exp_dir / f/ fr
-            
-    return mass_result, form_result
 
 # get target
 def get_target(candidates, labels):
@@ -85,73 +62,17 @@ def get_top_cand(candidates, scores):
 
 # split into hit rates
 def convert_rank_to_hit_rates(row, rank_col ,top_k=[1,5,20]):
-    top_k_hits ={}
+    top_k_hits = []
     rank = row[rank_col]
     for k in top_k:
         if rank <= k:
-            top_k_hits[f'{rank_col}-hit_rate@{k}'] = 1
+            top_k_hits.append(1)
         else:
-            top_k_hits[f'{rank_col}-hit_rate@{k}'] = 0
-    return pd.Series(top_k_hits)
-
-#################### Rank aggregation #######################
-from collections import defaultdict
-import numpy as np
-from scipy.stats import rankdata
-
-def borda_count(candidates, score_lists, target):
-    scores = defaultdict(int)
-    N = len(candidates)
-    for score_list in score_lists:
-        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
-        for rank, (mol, _) in enumerate(ranked_list, start=1):
-            scores[mol] += N - rank + 1
-    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
-
-def average_rank(candidates, score_lists, target):
-    rank_sums = defaultdict(list)
-    for score_list in score_lists:
-        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
-        for rank, (mol, _) in enumerate(ranked_list, start=1):
-            rank_sums[mol].append(rank)
-    avg_ranks = {mol: np.mean(ranks) for mol, ranks in rank_sums.items()}
-    ranked_candidates = [mol for mol, _ in sorted(avg_ranks.items(), key=lambda x: x[1])]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
-
-def reciprocal_rank_aggregation(candidates, score_lists, target):
-    scores = defaultdict(float)
-    for score_list in score_lists:
-        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
-        for rank, (mol, _) in enumerate(ranked_list, start=1):
-            scores[mol] += 1 / rank
-    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
-
-def weighted_voting(candidates, score_lists, weights, target):
-    scores = defaultdict(float)
-    for weight, score_list in zip(weights, score_lists):
-        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
-        for rank, (mol, _) in enumerate(ranked_list, start=1):
-            scores[mol] += weight / rank
-    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+            top_k_hits.append(0)
+    return top_k_hits
 
-def median_rank(candidates, score_lists, target):
-    rank_sums = defaultdict(list)
-    for score_list in score_lists:
-        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
-        for rank, (mol, _) in enumerate(ranked_list, start=1):
-            rank_sums[mol].append(rank)
-    median_ranks = {mol: np.median(ranks) for mol, ranks in rank_sums.items()}
-    ranked_candidates = [mol for mol, _ in sorted(median_ranks.items(), key=lambda x: x[1])]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
 
-def score_based_aggregation(candidates, score_lists, target):
-    scores = defaultdict(list)
-    for score_list in score_lists:
-        for mol, score in zip(candidates, score_list):
-            scores[mol].append(score)
-    avg_scores = {mol: np.mean(vals) for mol, vals in scores.items()}
-    ranked_candidates = [mol for mol, _ in sorted(avg_scores.items(), key=lambda x: x[1], reverse=True)]
-    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+def get_ci(col_vals, confidence_level=0.999, n_resamples=20_000, seed=0):
+    res = bootstrap((col_vals,), np.mean, confidence_level=confidence_level, n_resamples=n_resamples, random_state=seed)
+    ci = res.confidence_interval
+    return f'{ci.low:.2f}-{ci.high:.2f}'

flare/utils/general.py

@@ -0,0 +1,186 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+def pad_graph_nodes(mol_enc, g_n_nodes):
+    """
+    Args:
+        mol_enc: 2D tensor of shape (sum_nodes, D)
+                 Node embeddings for each molecule.
+        g_n_nodes: list[int]  Number of nodes per graph (len = B)
+
+    Returns:
+        padded: (B, max_nodes, D) tensor
+        mask:   (B, max_nodes) bool tensor, True for valid nodes
+    """
+
+    # Already concatenated: shape (sum_nodes, D)
+    B = len(g_n_nodes)
+    D = mol_enc.shape[1]
+    max_nodes = max(g_n_nodes)
+    padded = mol_enc.new_zeros((B, max_nodes, D))
+    mask = torch.zeros((B, max_nodes), dtype=torch.bool, device=mol_enc.device)
+
+    idx = 0
+    for i, n in enumerate(g_n_nodes):
+        padded[i, :n] = mol_enc[idx:idx+n]
+        mask[i, :n] = True
+        idx += n
+    return padded, mask
+
+import torch
+import torch.nn.functional as F
+
+import torch
+import torch.nn.functional as F
+
+def filip_similarity_batch(
+    image_tokens,
+    text_tokens,
+    mask_image,
+    mask_text,
+    reduction="mean",  # "mean", "topk", "softmax", or "geom"
+    k=5,
+    temperature=0.05,
+    eps=1e-6
+):
+    """
+    Compute FILIP similarity for batches of image and text token embeddings.
+
+    Args:
+        image_tokens: (B, N_img, D) float tensor
+        text_tokens:  (B, N_text, D) float tensor
+        mask_image:   (B, N_img) bool tensor
+        mask_text:    (B, N_text) bool tensor
+        reduction:    str, aggregation strategy: "mean", "topk", "softmax", or "geom"
+        k:            int, used if reduction == "topk"
+        temperature:  float, used if reduction == "softmax"
+        eps:          float, small constant for numerical stability
+
+    Returns:
+        similarities: (B,) float tensor of similarity scores
+    """
+    B, N_img, D = image_tokens.shape
+    N_text = text_tokens.shape[1]
+
+    # Normalize tokens
+    image_norm = F.normalize(image_tokens, p=2, dim=-1)
+    text_norm = F.normalize(text_tokens, p=2, dim=-1)
+
+    # Compute cosine similarity matrices
+    sim_matrix = torch.bmm(image_norm, text_norm.transpose(1, 2))
+
+    # Expand masks
+    mask_image_exp = mask_image.unsqueeze(2)
+    mask_text_exp = mask_text.unsqueeze(1)
+    valid_mask = mask_image_exp & mask_text_exp
+
+    # Mask invalid positions
+    sim_matrix_masked = sim_matrix.masked_fill(~valid_mask, float('-inf'))
+
+    # Max per image/text token
+    max_sim_img, _ = sim_matrix_masked.max(dim=2)
+    max_sim_text, _ = sim_matrix_masked.max(dim=1)
+
+    # Replace -inf with zeros
+    max_sim_img[max_sim_img == float('-inf')] = 0
+    max_sim_text[max_sim_text == float('-inf')] = 0
+
+    # Helper: aggregate with chosen strategy
+    def aggregate(max_sim, mask):
+        count = mask.sum(dim=1).clamp(min=1).float()
+
+        if reduction == "mean":
+            return (max_sim * mask).sum(dim=1) / count
+
+        elif reduction == "topk":
+            k_eff = min(k, max_sim.size(1))
+            # Mask invalid tokens to large negative before topk
+            masked_vals = max_sim.masked_fill(~mask, float('-inf'))
+            topk_vals, _ = torch.topk(masked_vals, k_eff, dim=1)
+            topk_vals[topk_vals == float('-inf')] = 0
+            return topk_vals.sum(dim=1) / k_eff
+
+        elif reduction == "softmax":
+            masked_vals = max_sim.masked_fill(~mask, float('-inf'))
+            weights = torch.softmax(masked_vals / temperature, dim=1)
+            weights = weights * mask
+            weights = weights / weights.sum(dim=1, keepdim=True).clamp(min=eps)
+            return (weights * max_sim).sum(dim=1)
+
+        elif reduction == "geom":
+            # Use log-sum-exp trick for geometric mean stability
+            masked_vals = (max_sim * mask).clamp(min=eps)
+            log_vals = torch.log(masked_vals)
+            geom_mean = torch.exp((log_vals.sum(dim=1)) / count)
+            return geom_mean
+
+        else:
+            raise ValueError(f"Unknown reduction type: {reduction}")
+
+    # Aggregate both sides
+    avg_img = aggregate(max_sim_img, mask_image)
+    avg_text = aggregate(max_sim_text, mask_text)
+
+    # Final similarity
+    similarity = (avg_img + avg_text) / 2
+    return similarity
+
+
+
+# def filip_similarity_batch(image_tokens, text_tokens, mask_image, mask_text):
+#     """
+#     Compute FILIP similarity for batches of image and text token embeddings.
+
+#     Args:
+#         image_tokens: (B, N_img, D) float tensor
+#         text_tokens:  (B, N_text, D) float tensor
+#         mask_image:   (B, N_img) bool tensor
+#         mask_text:    (B, N_text) bool tensor
+
+#     Returns:
+#         similarities: (B,) float tensor of similarity scores
+#     """
+#     B, N_img, D = image_tokens.shape
+#     N_text = text_tokens.shape[1]
+
+#     # Normalize tokens
+#     image_norm = F.normalize(image_tokens, p=2, dim=-1)  # (B, N_img, D)
+#     text_norm = F.normalize(text_tokens, p=2, dim=-1)    # (B, N_text, D)
+
+#     # Compute batched cosine similarity matrices
+#     # Result shape: (B, N_img, N_text)
+#     sim_matrix = torch.bmm(image_norm, text_norm.transpose(1, 2))
+
+#     # Expand masks for broadcasting
+#     mask_image_exp = mask_image.unsqueeze(2)  # (B, N_img, 1)
+#     mask_text_exp = mask_text.unsqueeze(1)    # (B, 1, N_text)
+#     valid_mask = mask_image_exp & mask_text_exp  # (B, N_img, N_text)
+
+#     # Mask invalid positions by setting them to -inf
+#     sim_matrix_masked = sim_matrix.masked_fill(~valid_mask, float('-inf'))
+
+#     # Max over text tokens per image token: (B, N_img)
+#     max_sim_img, _ = sim_matrix_masked.max(dim=2)
+
+#     # Max over image tokens per text token: (B, N_text)
+#     max_sim_text, _ = sim_matrix_masked.max(dim=1)
+
+#     # Replace -inf (no valid tokens) with zeros to avoid NaNs
+#     max_sim_img[max_sim_img == float('-inf')] = 0
+#     max_sim_text[max_sim_text == float('-inf')] = 0
+
+#     # Sum over valid tokens and divide by number of valid tokens (avoid division by zero)
+#     sum_img = (max_sim_img * mask_image).sum(dim=1)
+#     count_img = mask_image.sum(dim=1).clamp(min=1).float()
+
+#     sum_text = (max_sim_text * mask_text).sum(dim=1)
+#     count_text = mask_text.sum(dim=1).clamp(min=1).float()
+
+#     avg_img = sum_img / count_img
+#     avg_text = sum_text / count_text
+
+#     # Final similarity per batch element
+#     similarity = (avg_img + avg_text) / 2
+
+#     return similarity

{mvp → flare}/utils/models.py

@@ -1,7 +1,7 @@
-from mvp.models.spec_encoder import SpecEncMLP_BIN, SpecFormulaEncMLP, SpecFormulaTransformer,SpecFormula_mz_Encoder, SpecMzIntTokenTransformer
-from mvp.models.mol_encoder import MolEnc
-from mvp.models.encoders import MLP
-from mvp.models.contrastive import ContrastiveModel, CrossAttenContrastive, IndSpecEncoder, MultiViewContrastive, MultiViewFineTuning, FilipContrastive
+from flare.models.spec_encoder import SpecEncMLP_BIN, SpecFormulaEncMLP, SpecFormulaTransformer,SpecFormula_mz_Encoder, SpecMzIntTokenTransformer
+from flare.models.mol_encoder import MolEnc
+from flare.models.encoders import MLP
+from flare.models.contrastive import ContrastiveModel, CrossAttenContrastive, FilipContrastive
 
 def get_spec_encoder(spec_enc:str, args):
     return {"MLP_BIN": SpecEncMLP_BIN,
@@ -26,14 +26,6 @@ def get_model(model:str,
         model= ContrastiveModel(**params)
     elif model =='crossAttenContrastive':
         model = CrossAttenContrastive(**params)
-    elif model == 'IndSpecEncoder':
-        params['pred_fp'] = False
-        params['use_cons_spec'] = False
-        model = IndSpecEncoder(**params)
-    elif model == "MultiviewContrastive":
-        model = MultiViewContrastive(**params)
-    elif model == "MultiViewFineTuning":
-        model = MultiViewFineTuning(**params)
     elif model == "filipContrastive":
         model = FilipContrastive(**params)
     else:

{mvp → flare}/utils/preprocessing.py

@@ -1,7 +1,7 @@
 import pandas as pd
 import pickle
 import numpy as np
-import mvp.utils.data as data_utils
+import flare.utils.data as data_utils
 import collections
 import os
 import requests

mvp/data/__init__.py

@@ -1,3 +0,0 @@
-import sys
-sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
-from massspecgym.data import *

mvp/run.sh

@@ -1,3 +0,0 @@
-python train.py 
-python test.py 
-python test.py --candidates_pth /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_formula.json

mvp/utils/general.py

@@ -1,87 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-def pad_graph_nodes(mol_enc, g_n_nodes):
-    """
-    Args:
-        mol_enc: 2D tensor of shape (sum_nodes, D)
-                 Node embeddings for each molecule.
-        g_n_nodes: list[int]  Number of nodes per graph (len = B)
-
-    Returns:
-        padded: (B, max_nodes, D) tensor
-        mask:   (B, max_nodes) bool tensor, True for valid nodes
-    """
-
-    # Already concatenated: shape (sum_nodes, D)
-    B = len(g_n_nodes)
-    D = mol_enc.shape[1]
-    max_nodes = max(g_n_nodes)
-    padded = mol_enc.new_zeros((B, max_nodes, D))
-    mask = torch.zeros((B, max_nodes), dtype=torch.bool, device=mol_enc.device)
-
-    idx = 0
-    for i, n in enumerate(g_n_nodes):
-        padded[i, :n] = mol_enc[idx:idx+n]
-        mask[i, :n] = True
-        idx += n
-    return padded, mask
-
-
-def filip_similarity_batch(image_tokens, text_tokens, mask_image, mask_text):
-    """
-    Compute FILIP similarity for batches of image and text token embeddings.
-
-    Args:
-        image_tokens: (B, N_img, D) float tensor
-        text_tokens:  (B, N_text, D) float tensor
-        mask_image:   (B, N_img) bool tensor
-        mask_text:    (B, N_text) bool tensor
-
-    Returns:
-        similarities: (B,) float tensor of similarity scores
-    """
-    B, N_img, D = image_tokens.shape
-    N_text = text_tokens.shape[1]
-
-    # Normalize tokens
-    image_norm = F.normalize(image_tokens, p=2, dim=-1)  # (B, N_img, D)
-    text_norm = F.normalize(text_tokens, p=2, dim=-1)    # (B, N_text, D)
-
-    # Compute batched cosine similarity matrices
-    # Result shape: (B, N_img, N_text)
-    sim_matrix = torch.bmm(image_norm, text_norm.transpose(1, 2))
-
-    # Expand masks for broadcasting
-    mask_image_exp = mask_image.unsqueeze(2)  # (B, N_img, 1)
-    mask_text_exp = mask_text.unsqueeze(1)    # (B, 1, N_text)
-    valid_mask = mask_image_exp & mask_text_exp  # (B, N_img, N_text)
-
-    # Mask invalid positions by setting them to -inf
-    sim_matrix_masked = sim_matrix.masked_fill(~valid_mask, float('-inf'))
-
-    # Max over text tokens per image token: (B, N_img)
-    max_sim_img, _ = sim_matrix_masked.max(dim=2)
-
-    # Max over image tokens per text token: (B, N_text)
-    max_sim_text, _ = sim_matrix_masked.max(dim=1)
-
-    # Replace -inf (no valid tokens) with zeros to avoid NaNs
-    max_sim_img[max_sim_img == float('-inf')] = 0
-    max_sim_text[max_sim_text == float('-inf')] = 0
-
-    # Sum over valid tokens and divide by number of valid tokens (avoid division by zero)
-    sum_img = (max_sim_img * mask_image).sum(dim=1)
-    count_img = mask_image.sum(dim=1).clamp(min=1).float()
-
-    sum_text = (max_sim_text * mask_text).sum(dim=1)
-    count_text = mask_text.sum(dim=1).clamp(min=1).float()
-
-    avg_img = sum_img / count_img
-    avg_text = sum_text / count_text
-
-    # Final similarity per batch element
-    similarity = (avg_img + avg_text) / 2
-
-    return similarity