dynamic imports for nn models

Author: aditya0by0

chebifier/prediction_models/electra_predictor.py

@@ -1,6 +1,4 @@
 import numpy as np
-from chebai.models.electra import Electra
-from chebai.preprocessing.reader import EMBEDDING_OFFSET, ChemDataReader
 
 from .nn_predictor import NNPredictor
 
@@ -37,10 +35,14 @@ def build_graph_from_attention(att, node_labels, token_labels, threshold=0.0):
 
 class ElectraPredictor(NNPredictor):
     def __init__(self, model_name: str, ckpt_path: str, **kwargs):
+        from chebai.preprocessing.reader import ChemDataReader
+
         super().__init__(model_name, ckpt_path, reader_cls=ChemDataReader, **kwargs)
         print(f"Initialised Electra model {self.model_name} (device: {self.device})")
 
-    def init_model(self, ckpt_path: str, **kwargs) -> Electra:
+    def init_model(self, ckpt_path: str, **kwargs) -> "Electra":  # noqa: F821
+        from chebai.models.electra import Electra
+
         model = Electra.load_from_checkpoint(
             ckpt_path,
             map_location=self.device,
@@ -53,6 +55,8 @@ def init_model(self, ckpt_path: str, **kwargs) -> Electra:
         return model
 
     def explain_smiles(self, smiles) -> dict:
+        from chebai.preprocessing.reader import EMBEDDING_OFFSET
+
         reader = self.reader_cls()
         token_dict = reader.to_data(dict(features=smiles, labels=None))
         tokens = np.array(token_dict["features"]).astype(int).tolist()

chebifier/prediction_models/gnn_predictor.py

@@ -1,15 +1,11 @@
-import chebai_graph.preprocessing.properties as p
-import torch
-from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred
-from chebai_graph.preprocessing.property_encoder import IndexEncoder, OneHotEncoder
-from chebai_graph.preprocessing.reader import GraphPropertyReader
-from torch_geometric.data.data import Data as GeomData
-
 from .nn_predictor import NNPredictor
 
 
 class ResGatedPredictor(NNPredictor):
     def __init__(self, model_name: str, ckpt_path: str, molecular_properties, **kwargs):
+        from chebai_graph.preprocessing.properties import MolecularProperty
+        from chebai_graph.preprocessing.reader import GraphPropertyReader
+
         super().__init__(
             model_name, ckpt_path, reader_cls=GraphPropertyReader, **kwargs
         )
@@ -23,7 +19,7 @@ def __init__(self, model_name: str, ckpt_path: str, molecular_properties, **kwar
             properties = []
         self.molecular_properties = properties
         assert isinstance(self.molecular_properties, list) and all(
-            isinstance(prop, p.MolecularProperty) for prop in self.molecular_properties
+            isinstance(prop, MolecularProperty) for prop in self.molecular_properties
         )
         print(f"Initialised GNN model {self.model_name} (device: {self.device})")
 
@@ -32,7 +28,12 @@ def load_class(self, class_path: str):
         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
+    ) -> "ResGatedGraphConvNetGraphPred":  # noqa: F821
+        import torch
+        from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred
+
         model = ResGatedGraphConvNetGraphPred.load_from_checkpoint(
             ckpt_path,
             map_location=torch.device(self.device),
@@ -45,6 +46,14 @@ 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"]
@@ -87,9 +96,9 @@ def read_smiles(self, smiles):
                     encoded_values = encoded_values.unsqueeze(1)
             else:
                 encoded_values = torch.zeros((0, prop.encoder.get_encoding_length()))
-            if isinstance(prop, p.AtomProperty):
+            if isinstance(prop, AtomProperty):
                 x = torch.cat([x, encoded_values], dim=1)
-            elif isinstance(prop, p.BondProperty):
+            elif isinstance(prop, BondProperty):
                 edge_attr = torch.cat([edge_attr, encoded_values], dim=1)
             else:
                 molecule_attr = torch.cat([molecule_attr, encoded_values[0]], dim=1)

chebifier/prediction_models/nn_predictor.py

@@ -1,7 +1,6 @@
 from functools import lru_cache
 
 import numpy as np
-import torch
 import tqdm
 from rdkit import Chem
 
@@ -17,6 +16,8 @@ def __init__(
         target_labels_path: str,
         **kwargs,
     ):
+        import torch
+
         super().__init__(model_name, **kwargs)
         self.reader_cls = reader_cls
 
@@ -56,6 +57,8 @@ def read_smiles(self, smiles):
     def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
         """Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary
         Of classes and predicted values."""
+        import torch
+
         token_dicts = []
         could_not_parse = []
         index_map = dict()