clean up

flare/models/__init__.py

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

flare/models/contrastive.py

@@ -30,12 +30,6 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         self.save_hyperparameters()
         self.external_test = external_test
 
-        if 'use_fp' not in self.hparams:
-            self.hparams.use_fp = False
-        if 'use_NL_spec' not in self.hparams:
-            self.hparams.use_NL_spec = False
-
-
         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)
   
@@ -59,13 +53,10 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         if g is None:
             mol_enc = None
             return spec_enc, mol_enc
-        
-        fp = batch['fp'] if self.hparams.use_fp else None
-
-
-        f = self.mol_enc_model.GNN(g, g.ndata['h'])
 
-        mol_enc = self.mol_enc_model(g, fp=fp)
+        # Match historical call pattern (some DGL/dgllife paths mutate graph state in-place).
+        _ = self.mol_enc_model.GNN(g, g.ndata["h"])
+        mol_enc = self.mol_enc_model(g)
 
         return spec_enc, mol_enc
 
@@ -91,16 +82,7 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         if stage == Stage.TEST:
             return dict(spec_enc=spec_enc, mol_enc=mol_enc)
         
-        # Aux tasks
         output = {}
-        if self.hparams.pred_fp:
-            output['fp'] = self.fp_pred_model(mol_enc)
-        
-        if self.hparams.use_cons_spec:
-            spec = batch[self.spec_view]
-            n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
-            output['ind_spec'] = self.spec_enc_model(spec, n_peaks)
-
         # Calculate loss
         losses = self.compute_loss(batch, spec_enc, mol_enc, output)
 
@@ -174,6 +156,9 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
 
         if not self.df_test_path:
             self.df_test_path = os.path.join(self.hparams['experiment_dir'], 'result.pkl')
+
+
+        os.makedirs(os.path.dirname(self.df_test_path), exist_ok=True)
         self.df_test.to_pickle(self.df_test_path)
 
     def get_checkpoint_monitors(self) -> T.List[dict]:

flare/models/mol_encoder.py

@@ -10,10 +10,25 @@ class MolEnc(nn.Module):
                  in_dim,):
         super().__init__()
 
-        self.return_emb = False
-
-        if args.model in ('filipContrastive', 'crossAttenContrastive', 'filipGlobalContrastive'):
-            self.return_emb = True
+        # Whether to return node-level embeddings (sum_nodes, D) instead of a
+        # pooled graph embedding (B, D).
+        #
+        # Backward compatible defaults:
+        # - Historically, this was inferred from `args.model` for FILIP-style models.
+        # - New: allow explicit override via `args.mol_return_emb` (preferred) or
+        #   `args.return_mol_emb` (legacy-friendly alias).
+        explicit = None
+        if hasattr(args, "mol_return_emb"):
+            explicit = getattr(args, "mol_return_emb")
+        elif hasattr(args, "return_mol_emb"):
+            explicit = getattr(args, "return_mol_emb")
+
+        if explicit is not None:
+            self.return_emb = bool(explicit)
+        else:
+            self.return_emb = False
+            if args.model in ("filipContrastive", "crossAttenContrastive", "filipGlobalContrastive"):
+                self.return_emb = True
 
         dropout = [args.gnn_dropout for _ in range(len(args.gnn_channels))]
         batchnorm = [True for _ in range(len(args.gnn_channels))]
@@ -30,7 +45,7 @@ class MolEnc(nn.Module):
             self.dropout = nn.Dropout(args.fc_dropout)
             self.relu = nn.ReLU()
 
-    def forward(self, g, fp=None) -> torch.Tensor:
+    def forward(self, g) -> torch.Tensor:
         g1 = g
         f1 = g.ndata['h']
 
@@ -38,8 +53,6 @@ class MolEnc(nn.Module):
         if self.return_emb:
             return f
         h = self.pool(g1, f)
-        if fp is not None:
-            h = torch.concat((h, fp), dim=-1)
         h1 = self.relu(self.fc1_graph(h))
         h1 = self.dropout(h1)
         h1 = self.fc2_graph(h1)