dynamic imports for rest of predictors

Author: aditya0by0

chebifier/prediction_models/__init__.py

@@ -1,8 +1,9 @@
 from .base_predictor import BasePredictor
-from .chemlog_predictor import ChemlogPeptidesPredictor, ChemlogExtraPredictor
+from .c3p_predictor import C3PPredictor
+from .chebi_lookup import ChEBILookupPredictor
+from .chemlog_predictor import ChemlogExtraPredictor, ChemlogPeptidesPredictor
 from .electra_predictor import ElectraPredictor
 from .gnn_predictor import ResGatedPredictor
-from .chebi_lookup import ChEBILookupPredictor
 
 __all__ = [
     "BasePredictor",
@@ -11,4 +12,5 @@
     "ResGatedPredictor",
     "ChEBILookupPredictor",
     "ChemlogExtraPredictor",
+    "C3PPredictor",
 ]

chebifier/prediction_models/c3p_predictor.py

@@ -1,8 +1,6 @@
 from functools import lru_cache
-from typing import Optional, List
 from pathlib import Path
-
-from c3p import classifier as c3p_classifier
+from typing import List, Optional
 
 from chebifier.prediction_models import BasePredictor
 
@@ -26,6 +24,8 @@ def __init__(
 
     @lru_cache(maxsize=100)
     def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
+        from c3p import classifier as c3p_classifier
+
         result_list = c3p_classifier.classify(
             list(smiles_list),
             self.program_directory,
@@ -50,6 +50,8 @@ def explain_smiles(self, smiles):
         C3P provides natural language explanations for each prediction (positive or negative). Since there are more
         than 300 classes, only take the positive ones.
         """
+        from c3p import classifier as c3p_classifier
+
         highlights = []
         result_list = c3p_classifier.classify(
             [smiles], self.program_directory, self.chemical_classes, strict=False

chebifier/prediction_models/chebi_lookup.py

@@ -1,15 +1,14 @@
+import json
+import os
 from functools import lru_cache
 from typing import Optional
 
-from chebifier.prediction_models import BasePredictor
-import os
-import networkx as nx
 from rdkit import Chem
-import json
 
+from chebifier.prediction_models import BasePredictor
 
-class ChEBILookupPredictor(BasePredictor):
 
+class ChEBILookupPredictor(BasePredictor):
     def __init__(
         self,
         model_name: str,
@@ -49,6 +48,8 @@ def get_smiles_lookup(self):
         return smiles_lookup
 
     def build_smiles_lookup(self):
+        import networkx as nx
+
         smiles_lookup = dict()
         for chebi_id, smiles in nx.get_node_attributes(
             self.chebi_graph, "smiles"
@@ -152,7 +153,8 @@ def explain_smiles(self, smiles: str) -> dict:
     # Example usage
     smiles_list = [
         "CCO",
-        "C1=CC=CC=C1" "*C(=O)OC[C@H](COP(=O)([O-])OCC[N+](C)(C)C)OC(*)=O",
+        "C1=CC=CC=C1",
+        "*C(=O)OC[C@H](COP(=O)([O-])OCC[N+](C)(C)C)OC(*)=O",
     ]  # SMILES with 251 matches in ChEBI
     predictions = predictor.predict_smiles_list(smiles_list)
     print(predictions)

chebifier/prediction_models/chemlog_predictor.py

@@ -1,21 +1,7 @@
+from functools import lru_cache
 from typing import Optional
 
 import tqdm
-from chemlog.alg_classification.charge_classifier import get_charge_category
-from chemlog.alg_classification.peptide_size_classifier import get_n_amino_acid_residues
-from chemlog.alg_classification.proteinogenics_classifier import (
-    get_proteinogenic_amino_acids,
-)
-from chemlog.alg_classification.substructure_classifier import (
-    is_diketopiperazine,
-    is_emericellamide,
-)
-from chemlog.cli import CLASSIFIERS, _smiles_to_mol, strategy_call
-from chemlog_extra.alg_classification.by_element_classification import (
-    XMolecularEntityClassifier,
-    OrganoXCompoundClassifier,
-)
-from functools import lru_cache
 
 from .base_predictor import BasePredictor
 
@@ -47,15 +33,14 @@
 
 
 class ChemlogExtraPredictor(BasePredictor):
-
-    CHEMLOG_CLASSIFIER = None
-
     def __init__(self, model_name: str, **kwargs):
         super().__init__(model_name, **kwargs)
         self.chebi_graph = kwargs.get("chebi_graph", None)
-        self.classifier = self.CHEMLOG_CLASSIFIER()
+        self.classifier = None
 
     def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
+        from chemlog.cli import _smiles_to_mol
+
         mol_list = [_smiles_to_mol(smiles) for smiles in smiles_list]
         res = self.classifier.classify(mol_list)
         if self.chebi_graph is not None:
@@ -72,17 +57,29 @@ def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
 
 
 class ChemlogXMolecularEntityPredictor(ChemlogExtraPredictor):
+    def __init__(self, model_name: str, **kwargs):
+        from chemlog_extra.alg_classification.by_element_classification import (
+            XMolecularEntityClassifier,
+        )
 
-    CHEMLOG_CLASSIFIER = XMolecularEntityClassifier
+        super().__init__(model_name, **kwargs)
+        self.classifier = XMolecularEntityClassifier()
 
 
 class ChemlogOrganoXCompoundPredictor(ChemlogExtraPredictor):
+    def __init__(self, model_name: str, **kwargs):
+        from chemlog_extra.alg_classification.by_element_classification import (
+            OrganoXCompoundClassifier,
+        )
 
-    CHEMLOG_CLASSIFIER = OrganoXCompoundClassifier
+        super().__init__(model_name, **kwargs)
+        self.classifier = OrganoXCompoundClassifier()
 
 
 class ChemlogPeptidesPredictor(BasePredictor):
     def __init__(self, model_name: str, **kwargs):
+        from chemlog.cli import CLASSIFIERS
+
         super().__init__(model_name, **kwargs)
         self.strategy = "algo"
         self.chebi_graph = kwargs.get("chebi_graph", None)
@@ -99,6 +96,8 @@ def __init__(self, model_name: str, **kwargs):
 
     @lru_cache(maxsize=100)
     def predict_smiles(self, smiles: str) -> Optional[dict]:
+        from chemlog.cli import _smiles_to_mol, strategy_call
+
         mol = _smiles_to_mol(smiles)
         if mol is None:
             return None
@@ -134,6 +133,19 @@ def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
 
     def get_chemlog_result_info(self, smiles):
         """Get classification for single molecule with additional information."""
+        from chemlog.alg_classification.charge_classifier import get_charge_category
+        from chemlog.alg_classification.peptide_size_classifier import (
+            get_n_amino_acid_residues,
+        )
+        from chemlog.alg_classification.proteinogenics_classifier import (
+            get_proteinogenic_amino_acids,
+        )
+        from chemlog.alg_classification.substructure_classifier import (
+            is_diketopiperazine,
+            is_emericellamide,
+        )
+        from chemlog.cli import _smiles_to_mol
+
         mol = _smiles_to_mol(smiles)
         if mol is None or not smiles:
             return {"error": "Failed to parse SMILES"}

chebifier/prediction_models/electra_predictor.py

@@ -1,7 +1,12 @@
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from .nn_predictor import NNPredictor
 
+if TYPE_CHECKING:
+    from chebai.models.electra import Electra
+
 
 def build_graph_from_attention(att, node_labels, token_labels, threshold=0.0):
     n_nodes = len(node_labels)
@@ -40,7 +45,7 @@ def __init__(self, model_name: str, ckpt_path: str, **kwargs):
         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":  # noqa: F821
+    def init_model(self, ckpt_path: str, **kwargs) -> Electra:
         from chebai.models.electra import Electra
 
         model = Electra.load_from_checkpoint(

chebifier/prediction_models/gnn_predictor.py

@@ -1,5 +1,10 @@
+from typing import TYPE_CHECKING
+
 from .nn_predictor import NNPredictor
 
+if TYPE_CHECKING:
+    from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred
+
 
 class ResGatedPredictor(NNPredictor):
     def __init__(self, model_name: str, ckpt_path: str, molecular_properties, **kwargs):
@@ -28,9 +33,7 @@ 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":  # noqa: F821
+    def init_model(self, ckpt_path: str, **kwargs) -> ResGatedGraphConvNetGraphPred:
         import torch
         from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred