Fix MoleRec fallback predictor never training

pyhealth/models/molerec.py

@@ -628,7 +628,14 @@ def __init__(
 
         # check if we have valid molecular data (not just dummy data)
         self.has_valid_smiles = len(self.substructure_smiles) > 1
-
+
+        # Fallback predictor used when there is no valid molecular data. It must
+        # be created here (not lazily in forward) so its parameters are
+        # registered before the optimizer is built; otherwise it never trains.
+        # patient_emb is [condition_emb | procedure_emb], i.e. hidden_dim * 2.
+        if not self.has_valid_smiles:
+            self.simple_predictor = torch.nn.Linear(hidden_dim * 2, self.label_size)
+
         self.substructure_graphs = StaticParaDict(
             **graph_batch_from_smiles(self.substructure_smiles)
         )
@@ -810,12 +817,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
             batch_indices = torch.arange(patient_emb.size(0)).to(self.device)
             last_patient_emb = patient_emb[batch_indices, last_visit_mask]
 
-            # simple linear projection from patient embedding to drug predictions
-            if not hasattr(self, 'simple_predictor'):
-                self.simple_predictor = torch.nn.Linear(
-                    patient_emb.size(-1), self.label_size
-                ).to(self.device)
-
+            # simple linear projection from patient embedding to drug
+            # predictions (self.simple_predictor is created in __init__ so its
+            # parameters are registered with the optimizer).
             logits = self.simple_predictor(last_patient_emb)
             y_prob = torch.sigmoid(logits)
             loss = binary_cross_entropy_with_logits(logits, y_true)

tests/core/test_molerec.py

@@ -133,6 +133,35 @@ def test_output_shapes(self):
         self.assertEqual(ret["y_prob"].shape, (batch_size, num_labels))
         self.assertEqual(ret["y_true"].shape, (batch_size, num_labels))
 
+    def test_fallback_predictor_registered_and_trains(self):
+        """The fallback predictor must be registered before the optimizer.
+
+        Regression test: on the no-valid-SMILES fallback path, simple_predictor
+        was created lazily inside forward(). An optimizer built from
+        model.parameters() beforehand (the standard training pattern) therefore
+        never saw its parameters and never trained it.
+        """
+        self.assertFalse(self.model.has_valid_smiles)
+
+        # simple_predictor exists (and is in named_parameters) before any forward
+        param_names = [n for n, _ in self.model.named_parameters()]
+        self.assertTrue(any("simple_predictor" in n for n in param_names))
+
+        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.1)
+        before = self.model.simple_predictor.weight.detach().clone()
+
+        train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False)
+        data_batch = next(iter(train_loader))
+        ret = self.model(**data_batch)
+        optimizer.zero_grad()
+        ret["loss"].backward()
+        optimizer.step()
+
+        self.assertFalse(
+            torch.equal(before, self.model.simple_predictor.weight.detach()),
+            "fallback predictor was not updated by the optimizer step",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()