add support for augmented gnns

Author: sfluegel05

chebifier/ensemble/base_ensemble.py

@@ -83,7 +83,7 @@ def __init__(
 
     def gather_predictions(self, smiles_list):
         # get predictions from all models for the SMILES list
-        # order them by alphabetically by label class
+        # order them alphabetically by label class
         model_predictions = []
         predicted_classes = set()
         for model in self.models:

chebifier/prediction_models/gnn_predictor.py

@@ -1,19 +1,29 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
+
+import torch
 
 from .nn_predictor import NNPredictor
 
 if TYPE_CHECKING:
-    from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred
+    from chebai_graph.models.gat import GATGraphPred
+    from chebai_graph.models.resgated import ResGatedGraphPred
 
 
 class ResGatedPredictor(NNPredictor):
-    def __init__(self, model_name: str, ckpt_path: str, molecular_properties, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        ckpt_path: str,
+        molecular_properties,
+        dataset_cls: Optional[str] = None,
+        **kwargs,
+    ):
+        from chebai_graph.preprocessing.datasets.chebi import (
+            ChEBI50GraphProperties,
+            GraphPropertiesMixIn,
+        )
         from chebai_graph.preprocessing.properties import MolecularProperty
-        from chebai_graph.preprocessing.reader import GraphPropertyReader
 
-        super().__init__(
-            model_name, ckpt_path, reader_cls=GraphPropertyReader, **kwargs
-        )
         # molecular_properties is a list of class paths
         if molecular_properties is not None:
             properties = [self.load_class(prop)() for prop in molecular_properties]
@@ -22,22 +32,38 @@ def __init__(self, model_name: str, ckpt_path: str, molecular_properties, **kwar
             )
         else:
             properties = []
+        for property in properties:
+            property.encoder.eval = True
         self.molecular_properties = properties
         assert isinstance(self.molecular_properties, list) and all(
             isinstance(prop, MolecularProperty) for prop in self.molecular_properties
         )
+        # TODO it should not be necessary to refer to the whole dataset class, disentangle dataset and molecule reading
+        self.dataset_cls = (
+            self.load_class(dataset_cls)
+            if dataset_cls is not None
+            else ChEBI50GraphProperties
+        )
+        self.dataset: Optional[GraphPropertiesMixIn] = self.dataset_cls(
+            properties=molecular_properties
+        )
+
+        super().__init__(
+            model_name, ckpt_path, reader_cls=self.dataset.READER, **kwargs
+        )
+
         print(f"Initialised GNN model {self.model_name} (device: {self.device})")
 
     def load_class(self, class_path: str):
         module_path, class_name = class_path.rsplit(".", 1)
         module = __import__(module_path, fromlist=[class_name])
         return getattr(module, class_name)
 
-    def init_model(self, ckpt_path: str, **kwargs) -> "ResGatedGraphConvNetGraphPred":
+    def init_model(self, ckpt_path: str, **kwargs) -> "ResGatedGraphPred":
         import torch
-        from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred
+        from chebai_graph.models.resgated import ResGatedGraphPred
 
-        model = ResGatedGraphConvNetGraphPred.load_from_checkpoint(
+        model = ResGatedGraphPred.load_from_checkpoint(
             ckpt_path,
             map_location=torch.device(self.device),
             criterion=None,
@@ -49,67 +75,40 @@ def init_model(self, ckpt_path: str, **kwargs) -> "ResGatedGraphConvNetGraphPred
         return model
 
     def read_smiles(self, smiles):
-        import torch
-        from chebai_graph.preprocessing.properties import AtomProperty, BondProperty
-        from chebai_graph.preprocessing.property_encoder import (
-            IndexEncoder,
-            OneHotEncoder,
-        )
-        from torch_geometric.data.data import Data as GeomData
-
-        reader = self.reader_cls()
-        d = reader.to_data(dict(features=smiles, labels=None))
-        geom_data = d["features"]
-        edge_attr = geom_data.edge_attr
-        x = geom_data.x
-        molecule_attr = torch.empty((1, 0))
-        for prop in self.molecular_properties:
-            property_values = reader.read_property(smiles, prop)
-            encoded_values = []
-            for value in property_values:
-                # cant use standard encode for index encoder because model has been trained on a certain range of values
-                # use default value if we meet an unseen value
-                if isinstance(prop.encoder, IndexEncoder):
-                    if str(value) in prop.encoder.cache:
-                        index = prop.encoder.cache[str(value)] + prop.encoder.offset
-                    else:
-                        index = 0
-                        print(
-                            f"Unknown property value {value} for property {prop} at smiles {smiles}"
-                        )
-                    if isinstance(prop.encoder, OneHotEncoder):
-                        encoded_values.append(
-                            torch.nn.functional.one_hot(
-                                torch.tensor(index),
-                                num_classes=prop.encoder.get_encoding_length(),
-                            )
-                        )
-                    else:
-                        encoded_values.append(torch.tensor([index]))
-
-                else:
-                    encoded_values.append(prop.encoder.encode(value))
-            if len(encoded_values) > 0:
-                encoded_values = torch.stack(encoded_values)
-
-            if isinstance(encoded_values, torch.Tensor):
-                if len(encoded_values.size()) == 0:
-                    encoded_values = encoded_values.unsqueeze(0)
-                if len(encoded_values.size()) == 1:
-                    encoded_values = encoded_values.unsqueeze(1)
-            else:
-                encoded_values = torch.zeros((0, prop.encoder.get_encoding_length()))
-            if isinstance(prop, AtomProperty):
-                x = torch.cat([x, encoded_values], dim=1)
-            elif isinstance(prop, BondProperty):
-                edge_attr = torch.cat([edge_attr, encoded_values], dim=1)
+        d = self.dataset.READER().to_data(dict(features=smiles, labels=None))
+        property_data = d
+        # TODO merge props into base should not be a method of a dataset (or at least static)
+        for property in self.dataset.properties:
+            property.encoder.eval = True
+            property_value = self.reader.read_property(smiles, property)
+            if property_value is None or len(property_value) == 0:
+                encoded_value = None
             else:
-                molecule_attr = torch.cat([molecule_attr, encoded_values[0]], dim=1)
-
-        d["features"] = GeomData(
-            x=x,
-            edge_index=geom_data.edge_index,
-            edge_attr=edge_attr,
-            molecule_attr=molecule_attr,
+                encoded_value = torch.stack(
+                    [property.encoder.encode(v) for v in property_value]
+                )
+                if len(encoded_value.shape) == 3:
+                    encoded_value = encoded_value.squeeze(0)
+            property_data[property.name] = encoded_value
+        d["features"] = self.dataset._merge_props_into_base(
+            property_data, max_len_node_properties=self.model.gnn.in_channels
         )
         return d
+
+
+class GATPredictor(ResGatedPredictor):
+
+    def init_model(self, ckpt_path: str, **kwargs) -> "GATGraphPred":
+        import torch
+        from chebai_graph.models.gat import GATGraphPred
+
+        model = GATGraphPred.load_from_checkpoint(
+            ckpt_path,
+            map_location=torch.device(self.device),
+            criterion=None,
+            strict=False,
+            metrics=dict(train=dict(), test=dict(), validation=dict()),
+            pretrained_checkpoint=None,
+        )
+        model.eval()
+        return model

chebifier/prediction_models/nn_predictor.py

@@ -20,6 +20,7 @@ def __init__(
 
         super().__init__(model_name, **kwargs)
         self.reader_cls = reader_cls
+        self.reader = reader_cls()
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model = self.init_model(ckpt_path=ckpt_path)
@@ -49,8 +50,7 @@ def batchify(self, batch):
             yield cache
 
     def read_smiles(self, smiles):
-        reader = self.reader_cls()
-        d = reader.to_data(dict(features=smiles, labels=None))
+        d = self.reader.to_data(dict(features=smiles, labels=None))
         return d
 
     @modelwise_smiles_lru_cache.batch_decorator