use package namespace imports for prediction models

Author: aditya0by0

chebifier/ensemble/base_ensemble.py

@@ -4,10 +4,12 @@
 import torch
 import tqdm
 
-from chebifier.prediction_models.base_predictor import BasePredictor
-from chebifier.prediction_models.chemlog_predictor import ChemLogPredictor
-from chebifier.prediction_models.electra_predictor import ElectraPredictor
-from chebifier.prediction_models.gnn_predictor import ResGatedPredictor
+from chebifier.prediction_models import (
+    BasePredictor,
+    ChemLogPredictor,
+    ElectraPredictor,
+    ResGatedPredictor,
+)
 
 MODEL_TYPES = {
     "electra": ElectraPredictor,

chebifier/prediction_models/__init__.py

@@ -0,0 +1,6 @@
+from .base_predictor import BasePredictor
+from .chemlog_predictor import ChemLogPredictor
+from .electra_predictor import ElectraPredictor
+from .gnn_predictor import ResGatedPredictor
+
+__all__ = ["BasePredictor", "ChemLogPredictor", "ElectraPredictor", "ResGatedPredictor"]