Repository: Oloren-AI/olorenchemengine (language: python)

Issue:
Provide an option to display structures as Kekule
A picky point from an old organic chemist.  No one I know wants to look at aromatic rings with circles in the center.   IMO, structure display should default to Kekule.  At the very least, an option should be provided to display structures as Kekule.  :-)


Discussion / hints:
Will do, looking into the best place to do so.

Fix patch:
diff --git a/.gitignore b/.gitignore
index d4a8b49f..0cd65835 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,7 @@
 # Temporary and binary files
 
+demo.py
+
 gcloud-auth.json
 
 experiments
diff --git a/docker/Dockerfile.gpu b/docker/Dockerfile.gpu
index de5b756b..ce7db7d0 100644
--- a/docker/Dockerfile.gpu
+++ b/docker/Dockerfile.gpu
@@ -32,6 +32,6 @@ RUN conda install python=3.8
 
 # Copy the library over
 COPY install.sh install.sh
-RUN bash install.sh --docker
+RUN bash install.sh --docker --gpu
 COPY . .
 RUN pip install .
\ No newline at end of file
diff --git a/docs/conf.py b/docs/conf.py
index dcdb05a9..e2cb251a 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -8,8 +8,8 @@
 # serve to show the default.
 
 import os
-import sys
 import shutil
+import sys
 
 # -- Path setup --------------------------------------------------------------
 
@@ -94,7 +94,7 @@
 print(autodoc_mock_imports)
 
 gtagjs_ids = [
-    'G-R5E6ME01BB',
+    "G-R5E6ME01BB",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
diff --git a/install.sh b/install.sh
index 00435487..2e9c4786 100755
--- a/install.sh
+++ b/install.sh
@@ -1,5 +1,13 @@
 #!/bin/bash
 
+GPU=0
+if  nvidia-smi 2>/dev/null | grep -q "Driver Version"; then
+    GPU=1
+fi
+if [[ $* == *--gpu* ]]; then
+    GPU=1
+fi
+
 if [[ $(python -c "import sys, os; print(os.path.exists(os.path.join(sys.prefix, 'conda-meta')))") = *True* ]]; then
     echo "Detected conda environment, using conda commands going forward..."
     CONDA=1
@@ -16,11 +24,12 @@ else
     fi
 fi
 
+
 # Check if pytorch is installed
 if ! python -c "import torch; print(torch.__version__)" 2>/dev/null | grep -q "1.11.0" ; then
     echo "Pytorch 1.11.0 (required for torch geometric) not found, installing..."
 
-    if  nvidia-smi 2>/dev/null | grep -q "Driver Version"; then
+    if [ $GPU == 1 ]; then
         echo "Detected NVIDIA GPU and Drivers, installing CUDA enabled PyTorch with:"
         if [ $CONDA == 1 ]; then
             echo ">> conda install pytorch==1.11.0 cudatoolkit=11.3 -c pytorch"
@@ -72,13 +81,16 @@ fi
 
 
 if [[ $* == *--docker* ]]; then
+        python3.8 -m pip install olorenchemengine
     echo "Docker argument passed, skipping installing chemengine package."
 else
     if [[ $* == *--dev* ]]; then
         SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
-        python3.8 -m pip install -e $SCRIPT_DIR # install editable copy of the package for dev
+        python3.8 -m pip install -e $SCRIPT_DIR
+        python3.8 -m pip install -e "${SCRIPT_DIR}[full]"
     else
         python3.8 -m pip install olorenchemengine
+        python3.8 -m pip install olorenchemengine[full]
     fi
 fi
 
diff --git a/setup.cfg b/setup.cfg
index 62345ad3..858e7524 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -53,6 +53,24 @@ package_dir =
 # For more information, check out https://semver.org/.
 install_requires =
     importlib-metadata; python_version<"3.8"
+    ipython
+    typing_extensions
+    pandas
+    google-cloud-firestore
+    pyrebase4
+
+[options.packages.find]
+where = src
+exclude =
+    tests
+
+[options.extras_require]
+full =
+    ipython
+    typing_extensions
+    pandas
+    google-cloud-firestore
+    pyrebase4
     google-auth-oauthlib<0.5
     pytorch-lightning<=1.6.2
     Cython
@@ -72,25 +90,17 @@ install_requires =
     protobuf==3.20.1
     pytorch_lightning==1.4.1
     torchmetrics==0.6.0
-
-[options.packages.find]
-where = src
-exclude =
-    tests
-
-[options.extras_require]
-testing =
-    setuptools
-    pytest
-    pytest-cov
-tensorflow = 
     tensorflow
     tensorboard
-representations = 
     pubchempy
     gensim
     mordred
 
+testing =
+    setuptools
+    pytest
+    pytest-cov
+
 [options.entry_points]
 # Add here console scripts like:
 # console_scripts =
@@ -148,4 +158,4 @@ version = 4.1.1
 package = olorenchemengine
 
 [pytest]
-timeout = 60
+timeout = 180
diff --git a/setup.py b/setup.py
index 091e9751..a5d100df 100644
--- a/setup.py
+++ b/setup.py
@@ -6,19 +6,21 @@
     PyScaffold helps you to put up the scaffold of your new Python project.
     Learn more under: https://pyscaffold.org/
 """
-from setuptools import setup
-import os
 import codecs
+import os
+
+from setuptools import setup
 
 
 def read(rel_path):
     here = os.path.abspath(os.path.dirname(__file__))
-    with codecs.open(os.path.join(here, rel_path), 'r') as fp:
+    with codecs.open(os.path.join(here, rel_path), "r") as fp:
         return fp.read()
 
+
 def get_version(rel_path):
     for line in read(rel_path).splitlines():
-        if line.startswith('__version__'):
+        if line.startswith("__version__"):
             delim = '"' if '"' in line else "'"
             return line.split(delim)[1]
         else:
@@ -27,8 +29,10 @@ def get_version(rel_path):
 
 if __name__ == "__main__":
     try:
-        setup(version=get_version("src/olorenchemengine/__version__.py"),
-              use_scm_version={"version_scheme": "no-guess-dev"})
+        setup(
+            version=get_version("src/olorenchemengine/__version__.py"),
+            use_scm_version={"version_scheme": "no-guess-dev"},
+        )
     except:  # noqa
         print(
             "\n\nAn error occurred while building the project, "
@@ -36,4 +40,4 @@ def get_version(rel_path):
             "setuptools_scm and wheel with:\n"
             "   pip install -U setuptools setuptools_scm wheel\n\n"
         )
-        raise
\ No newline at end of file
+        raise
diff --git a/src/olorenchemengine/__init__.py b/src/olorenchemengine/__init__.py
index 7a590578..ef4e00b9 100644
--- a/src/olorenchemengine/__init__.py
+++ b/src/olorenchemengine/__init__.py
@@ -1,7 +1,6 @@
 import contextlib
-from locale import D_FMT
 import sys
-import contextlib
+from locale import D_FMT
 
 if sys.version_info[:2] >= (3, 8):
     # TODO: Import directly (no need for conditional) when `python_requires = >= 3.8`
@@ -18,6 +17,63 @@
 finally:
     del version, PackageNotFoundError
 
+
+_OPTIONAL_IMPORTS_FOR_OCE_ONLINE = [
+    "sklearn",
+    "torch",
+    "torch_geometric",
+    "torch_scatter",
+    "torch_sparse",
+    "torch_cluster",
+    "torch_spline_conv",
+    "rdkit",
+    "sklearn.metrics",
+    "scipy",
+    "scipy.stats",
+    "joblib",
+    "sklearn.model_selection",
+    "tqdm",
+    "sklearn.preprocessing",
+    "rdkit.Chem",
+    "torch.nn",
+    "torch_geometric.nn",
+    "torch_geometric.data",
+    "torch.optim",
+    "torch_geometric.utils",
+    "torch.nn.functional",
+    "torch_geometric.nn.inits",
+    "rdkit.DataStructs",
+    "rdkit.DataStructs.cDataStructs",
+    "rdkit.Chem.AtomPairs",
+    "rdkit.Chem.AtomPairs.Sheridan",
+    "rdkit.Chem.Pharm2D",
+    "torch_geometric.loader",
+    "torch_geometric.data.data",
+    "torch.utils",
+    "torch.utils.data",
+    "torch.autograd",
+    "rdkit.Chem.rdchem",
+    "pytorch_lightning",
+    "rdkit.Chem.Scaffolds",
+    "PIL",
+    "selfies",
+    "sklearn.cross_decomposition",
+    "sklearn.decomposition",
+    "sklearn.neighbors",
+    "sklearn.manifold",
+    "hyperopt",
+    "hyperopt.pyll",
+]
+
+for imp in _OPTIONAL_IMPORTS_FOR_OCE_ONLINE:
+    try:
+        __import__(imp)
+    except ImportError:
+        import sys
+        from unittest.mock import MagicMock
+
+        sys.modules[imp] = MagicMock()
+
 from rdkit import RDLogger
 
 RDLogger.DisableLog("rdApp.*")
@@ -38,22 +94,28 @@
 if not path.exists(path.join(path.expanduser("~"), f".oce/cache/vecrep/")):
     os.mkdir(path.join(path.expanduser("~"), f".oce/cache/vecrep/"))
 
-import pandas as pd
-
 import json
 
+import pandas as pd
 
 CONFIG_PATH = os.path.join(os.path.expanduser("~"), ".oce/CONFIG.json")
 if os.path.exists(CONFIG_PATH):
     with open(CONFIG_PATH) as f:
         CONFIG_ = json.load(f)
 else:
-    CONFIG_ = {"MAP_LOCATION": "cpu", "USE_CUDA": False, "CDD_TOKEN": None, "VAULT_ID": None, "NUM_WORKERS": 4}
+    CONFIG_ = {
+        "MAP_LOCATION": "cpu",
+        "USE_CUDA": False,
+        "CDD_TOKEN": None,
+        "VAULT_ID": None,
+        "NUM_WORKERS": 0,
+    }
     with open(CONFIG_PATH, "w+") as f:
         json.dump(CONFIG_, f)
 
 CONFIG = CONFIG_.copy()
 
+
 def update_config():
     """Update the configuration file.
 
@@ -65,6 +127,7 @@ def update_config():
 
     with contextlib.suppress(ImportError):
         import torch
+
         CONFIG["DEVICE"] = torch.device(CONFIG["MAP_LOCATION"])
         CONFIG["USE_CUDA"] = "cuda" in CONFIG["MAP_LOCATION"]
 
@@ -95,37 +158,51 @@ def remove_config_param(param):
 
 update_config()
 
+
 def ExampleDataFrame():
-    return pd.read_csv("gs://oloren-public-data/sample-csvs/sample_data3.csv")
+    return pd.read_csv(
+        "https://storage.googleapis.com/oloren-public-data/sample-csvs/sample_data3.csv"
+    )
+
 
 from .base_class import *
 from .basics import *
+from .benchmarks import *
+from .dataset import *
 from .ensemble import *
-from .gnn import *
-from .representations import *
 from .external import *
-from .splitters import *
-from .dataset import *
+from .gnn import *
+from .hyperparameters import *
+from .internal import *
 from .interpret import *
 from .manager import *
-from .visualizations import *
+from .representations import *
+from .splitters import *
 from .uncertainty import *
-from .benchmarks import *
-from .hyperparameters import *
+from .visualizations import *
+
 
 def ExampleDataset():
-    if os.path.exists(path.join(path.expanduser("~"), f".oce/exampledataset.oce")):
-        return load(path.join(path.expanduser("~"), f".oce/exampledataset.oce"))
-    else:
-        dataset =  BaseDataset(data = ExampleDataFrame().to_csv(), structure_col = "Smiles",
-            property_col = "pChEMBL Value") + RandomSplit()
-        save(dataset, path.join(path.expanduser("~"), f".oce/exampledataset.oce"))
-        return dataset
+    dataset = (
+        BaseDataset(
+            data=ExampleDataFrame().to_csv(),
+            structure_col="Smiles",
+            property_col="pChEMBL Value",
+        )
+        + RandomSplit()
+    )
+    return dataset
+
 
 def BACEDataset():
     df = pd.read_csv(download_public_file("MoleculeNet/load_bace_regression.csv"))
-    df["split"] = df["split"].replace({"Train": "train", "Valid": "valid", "Test": "test"})
-    return oce.BaseDataset(data = df.to_csv(), structure_col = "smiles", property_col = "pIC50")
+    df["split"] = df["split"].replace(
+        {"Train": "train", "Valid": "valid", "Test": "test"}
+    )
+    return oce.BaseDataset(
+        data=df.to_csv(), structure_col="smiles", property_col="pIC50"
+    )
+
 
 def test_oce():
     """Convenience function to test all functions of the oce package."""
@@ -134,8 +211,12 @@ def test_oce():
 
     model = BaseBoosting(
         [
-            RandomForestModel(DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=1),
-            BaseTorchGeometricModel(TLFromCheckpoint("default"), batch_size=8, epochs=1, preinitialized=True),
+            RandomForestModel(
+                DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=1
+            ),
+            BaseTorchGeometricModel(
+                TLFromCheckpoint("default"), batch_size=8, epochs=1, preinitialized=True
+            ),
             RandomForestModel(OlorenCheckpoint("default"), n_estimators=1),
         ]
     )
@@ -145,4 +226,3 @@ def test_oce():
     save(model, "model.oce")
     _ = load("model.oce")
     os.remove("model.oce")
-
diff --git a/src/olorenchemengine/__version__.py b/src/olorenchemengine/__version__.py
index 92192eed..68cdeee4 100644
--- a/src/olorenchemengine/__version__.py
+++ b/src/olorenchemengine/__version__.py
@@ -1,1 +1,1 @@
-__version__ = "1.0.4"
+__version__ = "1.0.5"
diff --git a/src/olorenchemengine/base_class.py b/src/olorenchemengine/base_class.py
index 6af6f873..9ad617f3 100644
--- a/src/olorenchemengine/base_class.py
+++ b/src/olorenchemengine/base_class.py
@@ -2,30 +2,29 @@
 """
 from __future__ import annotations
 
-import os
-from abc import abstractmethod
-from typing import Any, Callable, Tuple, Union
-import pandas as pd
-import numpy as np
 import inspect
 import io
-import sys
-
-import olorenchemengine as oce
-from olorenchemengine.internal import *
-
 import logging
+import os
+import sys
+from abc import abstractmethod
+from typing import Any, Callable, Tuple, Union
 
+import numpy as np
+import pandas as pd
 from sklearn.metrics import (
     average_precision_score,
-    roc_auc_score,
-    r2_score,
     explained_variance_score,
     max_error,
     mean_absolute_error,
     mean_squared_error,
+    r2_score,
+    roc_auc_score,
 )
 
+import olorenchemengine as oce
+from olorenchemengine.internal import *
+
 # List of metrics that can be used for evaluating classification models
 classification_metrics = {
     "Average Precision": average_precision_score,
@@ -34,6 +33,7 @@
 
 # List of metrics that can be used for evaluating regression models
 from scipy.stats import spearmanr
+
 regression_metrics = {
     "r2": r2_score,
     "Spearman": lambda y, y_pred: spearmanr(y, y_pred)[0],
@@ -60,6 +60,7 @@
     "Spearman": "higher",
 }
 
+
 class BaseObject(BaseClass):
     """BaseObject is the parent class for all classes which directly wrap some object to be saved via joblib.
 
@@ -90,25 +91,25 @@ class BaseEstimator(BaseObject):
     """Utility class used to wrap any object with a fit and predict method"""
 
     def fit(self, X, y):
-        """ Fit the estimator to the data
+        """Fit the estimator to the data
 
         Parameters:
             X (np.array): The data to fit the estimator to
             y (np.array): The target data to fit the estimator to
         Returns:
             self (object): The estimator object fit to the data
-            """
+        """
 
         return self.obj.fit(X, y)
 
     def predict(self, X):
-        """ Predict the output of the estimator
+        """Predict the output of the estimator
 
-            Parameters:
-                X (np.array): The data to predict the output of the estimator on
-            Returns:
-                y (np.array): The predicted output of the estimator
-            """
+        Parameters:
+            X (np.array): The data to predict the output of the estimator on
+        Returns:
+            y (np.array): The predicted output of the estimator
+        """
         pred = self.obj.predict(X)
         pred = np.array(pred)
         if pred.ndim == 1:
@@ -169,7 +170,7 @@ def fit_transform(self, X):
         return X
 
     def transform(self, X):
-        """ Returns the transformed values of the dataset as a numpy array.
+        """Returns the transformed values of the dataset as a numpy array.
 
         Parameters:
             X (np.ndarray): the dataset
@@ -184,7 +185,7 @@ def transform(self, X):
         return X
 
     def inverse_transform(self, X):
-        """ Returns the original values from the transformed values.
+        """Returns the original values from the transformed values.
 
         Parameters:
             X (np.ndarray): the transformed values
@@ -207,7 +208,11 @@ class QuantileTransformer(BasePreprocessor):
 
     @log_arguments
     def __init__(
-        self, n_quantiles=1000, output_distribution="normal", subsample=1e5, random_state=None,
+        self,
+        n_quantiles=1000,
+        output_distribution="normal",
+        subsample=1e5,
+        random_state=None,
     ):
 
         from sklearn.preprocessing import QuantileTransformer
@@ -236,8 +241,7 @@ def __init__(self, with_mean=True, with_std=True):
 
 
 class LogScaler(BasePreprocessor):
-    """LogScaler is a BasePreprocessor which standardizes the data by taking the log and then removing the mean and scaling to unit variance.
-    """
+    """LogScaler is a BasePreprocessor which standardizes the data by taking the log and then removing the mean and scaling to unit variance."""
 
     @log_arguments
     def __init__(self, with_mean=True, with_std=True):
@@ -275,7 +279,7 @@ def fit_transform(self, X):
         return self.obj.fit_transform(X.reshape(-1, 1)).reshape(-1)
 
     def transform(self, X):
-        """ Returns the transformed values of the dataset as a numpy array.
+        """Returns the transformed values of the dataset as a numpy array.
 
         Parameters:
             X (np.ndarray): the dataset
@@ -287,7 +291,7 @@ def transform(self, X):
         return self.obj.transform(X.reshape(-1, 1)).reshape(-1)
 
     def inverse_transform(self, X):
-        """ Returns the original values from the transformed values.
+        """Returns the original values from the transformed values.
 
         Parameters:
             X (np.ndarray): the transformed values
@@ -296,7 +300,11 @@ def inverse_transform(self, X):
             The original values from the transformed values
         """
         X = np.array(X)
-        return 10 ** (self.obj.inverse_transform(X.reshape(-1, 1)).reshape(-1)) - self.mean - 1e-3
+        return (
+            10 ** (self.obj.inverse_transform(X.reshape(-1, 1)).reshape(-1))
+            - self.mean
+            - 1e-3
+        )
 
     def _save(self):
         d = super()._save()
@@ -367,10 +375,17 @@ def preprocess(self, X, y, fit=False):
 
     def visualize_parameters_ipynb(self):
         from urllib.parse import quote
+
         from IPython.display import IFrame, display
 
         parameters = parameterize(self)
-        display(IFrame(f"https://oas.oloren.ai/jsonvis?json={quote(json.dumps(parameters))}", width=800, height=500))
+        display(
+            IFrame(
+                f"https://oas.oloren.ai/jsonvis?json={quote(json.dumps(parameters))}",
+                width=800,
+                height=500,
+            )
+        )
 
     @abstractmethod
     def _fit(self, X_train, y_train: np.ndarray) -> None:
@@ -388,7 +403,9 @@ def fit(
         self,
         X_train: Union[pd.DataFrame, np.ndarray],
         y_train: Union[pd.Series, list, np.ndarray],
-        valid: Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]] = None,
+        valid: Tuple[
+            Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]
+        ] = None,
         error_model: BaseErrorModel = None,
     ):
         """Calls the _fit method of the model to fit the model on the provided dataset.
@@ -399,6 +416,8 @@ def fit(
             valid (Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]]): Optional validation data, which can be used as with methods like early stopping and model averaging.
             error_model (BaseErrorModel): Optional error model, which can be used to predict confidence intervals.
         """
+        y_train = np.array(y_train)
+
         X_train_original = X_train
         y_train_original = y_train
 
@@ -428,11 +447,15 @@ def fit(
                     y_valid = self.normalization.transform(y_valid)
                     y_valid = np.array(y_valid)
             if (np.isreal(y_train.any()) == False) & (str(y_train).find("%") == -1):
-                logging.warning(" Dataset contains non-numeric values, which is incompatible with regression.")
+                logging.warning(
+                    " Dataset contains non-numeric values, which is incompatible with regression."
+                )
                 y_train.str.replace(r"[^0-9.%]", "", regex=True)
             logging.info("Regression dataset detected")
         else:
-            logging.error("Setting was not defined as either auto, regression, or classification.")
+            logging.error(
+                "Setting was not defined as either auto, regression, or classification."
+            )
             sys.exit(1)
 
         # Train model
@@ -457,8 +480,12 @@ def calibrate(self, X_valid, y_valid):
 
         from sklearn.calibration import calibration_curve
 
-        prob_pred, prob_true = calibration_curve(y_valid, y_pred_valid, n_bins=min(len(y_valid), 10))
-        prob_pred, prob_true = zip(*[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1])
+        prob_pred, prob_true = calibration_curve(
+            y_valid, y_pred_valid, n_bins=min(len(y_valid), 10)
+        )
+        prob_pred, prob_true = zip(
+            *[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1]
+        )
         prob_true = np.array(prob_true)
         prob_true = np.log(prob_true / (1 - prob_true))
         prob_pred = np.array(prob_pred)
@@ -470,7 +497,9 @@ def fit_class(
         self,
         X_train: Union[pd.DataFrame, np.ndarray],
         y_train: Union[pd.Series, list, np.ndarray],
-        valid: Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]] = None,
+        valid: Tuple[
+            Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]
+        ] = None,
     ):
         from sklearn.model_selection import KFold
         from tqdm import tqdm
@@ -492,8 +521,12 @@ def fit_class(
 
         from sklearn.calibration import calibration_curve
 
-        prob_pred, prob_true = calibration_curve(y_true, y_pred, n_bins=min(len(y_true), 10))
-        prob_pred, prob_true = zip(*[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1])
+        prob_pred, prob_true = calibration_curve(
+            y_true, y_pred, n_bins=min(len(y_true), 10)
+        )
+        prob_pred, prob_true = zip(
+            *[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1]
+        )
         prob_true = np.array(prob_true)
         prob_true = np.log(prob_true / (1 - prob_true))
         prob_pred = np.array(prob_pred)
@@ -503,7 +536,11 @@ def fit_class(
         self.fit(X_train, y_train)
 
     def _unnormalize(self, Y):
-        if self.normalization == "zscore" and hasattr(self, "ymean") and hasattr(self, "ystd"):
+        if (
+            self.normalization == "zscore"
+            and hasattr(self, "ymean")
+            and hasattr(self, "ystd")
+        ):
             result = Y * self.ystd + self.ymean
         elif issubclass(type(self.normalization), BasePreprocessor):
             result = self.normalization.inverse_transform(Y)
@@ -563,11 +600,18 @@ def predict(
                 if return_ci:
                     result["ci"] = errors
                 if return_vis:
-                    from olorenchemengine.visualizations.visualization import VisualizeError
+                    from olorenchemengine.visualizations.visualization import (
+                        VisualizeError,
+                    )
 
                     ci = self.error_model.quantile
                     result["vis"] = [
-                        VisualizeError(self.error_model.y_train, float(result["predicted"][i]), float(errors[i]), ci=ci)
+                        VisualizeError(
+                            self.error_model.y_train,
+                            float(result["predicted"][i]),
+                            float(errors[i]),
+                            ci=ci,
+                        )
                         for i in range(len(errors))
                     ]
             return result
@@ -589,7 +633,7 @@ def fit_cv(
         ci: float = 0.8,
         scoring: str = None,
     ):
-        """ Trains a production-ready model.
+        """Trains a production-ready model.
 
         This method trains the model on the entire dataset. It also performs an
         intermediate cross-validation step over dataset to both generate test
@@ -629,8 +673,8 @@ def fit_cv(
             elif self.setting == "classification":
                 scoring = "Average Precision"
 
-        from sklearn.model_selection import KFold
         from sklearn.calibration import calibration_curve
+        from sklearn.model_selection import KFold
 
         kf = KFold(n_splits=n_splits)
 
@@ -643,7 +687,9 @@ def fit_cv(
 
             if hasattr(self, "error_model") and self.setting == "regression":
                 y_pred_test = np.array(y_pred_test).flatten()
-                error_model = type(self.error_model)(*self.error_model.args, **self.error_model.kwargs)
+                error_model = type(self.error_model)(
+                    *self.error_model.args, **self.error_model.kwargs
+                )
                 error_model.build(model, X_train, y_train)
                 if isinstance(error_model, BaseAggregateErrorModel):
                     error_model.train_cv()
@@ -657,8 +703,12 @@ def fit_cv(
                     pred = np.concatenate((pred, y_pred_test))
                     true = np.concatenate((true, y_test))
             elif self.setting == "classification":
-                prob_pred, prob_true = calibration_curve(y_test, y_pred_test, n_bins=min(len(y_test), 10))
-                prob_pred, prob_true = zip(*[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1])
+                prob_pred, prob_true = calibration_curve(
+                    y_test, y_pred_test, n_bins=min(len(y_test), 10)
+                )
+                prob_pred, prob_true = zip(
+                    *[x for x in zip(prob_pred, prob_true) if 0 < x[1] < 1]
+                )
                 prob_true = np.array(prob_true)
                 prob_true = np.log(prob_true / (1 - prob_true))
                 prob_pred = np.array(prob_pred)
@@ -673,10 +723,10 @@ def fit_cv(
 
         if self.setting == "regression":
             self.calibrator = LinearRegression()
-            self.calibrator.fit(pred.reshape(-1,1), true.reshape(-1,1))
-            pred = self.calibrator.predict(pred.reshape(-1,1)).reshape(-1)
+            self.calibrator.fit(pred.reshape(-1, 1), true.reshape(-1, 1))
+            pred = self.calibrator.predict(pred.reshape(-1, 1)).reshape(-1)
             true = true.reshape(-1)
-            residuals = np.abs(pred-true)
+            residuals = np.abs(pred - true)
             if hasattr(self, "error_model"):
                 self.error_model._fit(residuals, scores, quantile=ci)
                 self.em_status = "fitted"
@@ -708,9 +758,10 @@ def test(
             self.error_model.fit(X, y)
             self.em_status = "fitted"
 
-        import joblib
         import json
 
+        import joblib
+
         dataset_hash = joblib.hash(X) + joblib.hash(y)
 
         y_pred = self.predict(X)
@@ -723,12 +774,15 @@ def test(
         d = {}
 
         if self.setting == "classification":
-            d.update({k: float(v(y, y_pred)) for k, v in classification_metrics.items()})
+            d.update(
+                {k: float(v(y, y_pred)) for k, v in classification_metrics.items()}
+            )
         elif self.setting == "regression":
             d.update({k: float(v(y, y_pred)) for k, v in regression_metrics.items()})
 
         # LOGGING START
         import requests
+
         try:
             response = requests.get(
                 f"https://api.oloren.ai/firestore/log_model_performance",
@@ -761,7 +815,9 @@ def _save(self) -> dict:
             d.update({"error_model": saves(self.error_model)})
         if hasattr(self, "em_status"):
             d.update({"em_status": self.em_status})
-        if hasattr(self, "normalization") and issubclass(type(self.normalization), BasePreprocessor):
+        if hasattr(self, "normalization") and issubclass(
+            type(self.normalization), BasePreprocessor
+        ):
             d.update({"normalization": self.normalization._save()})
         return d
 
@@ -811,7 +867,6 @@ def upload_oas(self, fname: str = None):
         print(f"Access your model at: https://oas.oloren.ai/endpoint/?mid={out[1].id}")
 
 
-
 class MakeMultiClassModel(BaseModel):
     """
     Base class for extending the classification capabilities of BaseModel to more than two classes, e.g. classes {W,X,Y,Z}.
@@ -829,35 +884,43 @@ def __init__(self, individual_classifier: BaseModel) -> None:
         self.classifiers = []
 
     def _fit(self, X_train, y_train: np.ndarray) -> None:
-        """Empty method to correctly override the fit method of the parent BaseModel class.
-        """
+        """Empty method to correctly override the fit method of the parent BaseModel class."""
         pass
 
     def fit(
         self,
         X_train: Union[pd.DataFrame, np.ndarray],
         y_train: Union[pd.Series, list, np.ndarray],
-        valid: Tuple[Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]] = None,
+        valid: Tuple[
+            Union[pd.DataFrame, np.ndarray], Union[pd.Series, list, np.ndarray]
+        ] = None,
     ):
 
         # Outputs a warning if, at a glance, it doesn't make sense for the dataset to be a classification dataset. The 0.1 can be adjusted.
         if ((np.unique(y_train)).size / (y_train.size)) > 0.1:
-            logging.warning(" # of unique values / # of total datapoints is above the default threshold of 0.1.")
+            logging.warning(
+                " # of unique values / # of total datapoints is above the default threshold of 0.1."
+            )
 
-        self.sorted_classes = np.unique(y_train)  # Sorts the unique values in the array. Ex. [Z,X,Y,Z,X,X,Y] -> [X,Y,Z]
+        self.sorted_classes = np.unique(
+            y_train
+        )  # Sorts the unique values in the array. Ex. [Z,X,Y,Z,X,X,Y] -> [X,Y,Z]
         for classes in self.sorted_classes:
             y_train_class = (y_train == classes).astype(
                 int
             )  # Sets the currently iterating class (ex. X) to 1 and all others (ex. Y,Z) to 0, converting the dataset to a binary classification problem.
-            self.model = self.classifier.copy()  # Creates a copy of the user-specified model (ex. Random Forest)
+            self.model = (
+                self.classifier.copy()
+            )  # Creates a copy of the user-specified model (ex. Random Forest)
             self.model.fit(
                 X_train, y_train_class
             )  # Calls the BaseModel fit method, since we now just have a binary classification dataset.
-            self.classifiers.append(self.model)  # Adds the results of the fitted model to an array for later analysis.
+            self.classifiers.append(
+                self.model
+            )  # Adds the results of the fitted model to an array for later analysis.
 
     def _predict(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series]):
-        """Empty method to correctly override the predict method of the parent BaseModel class.
-        """
+        """Empty method to correctly override the predict method of the parent BaseModel class."""
         pass
 
     def predict(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series]):
@@ -876,14 +939,20 @@ def predict(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series]):
             predictions[y_class] = classifier.predict(X)
 
         predictions = pd.DataFrame(predictions)
-        predictions = (predictions.T / predictions.sum(axis=1)).T  # normalizes outputs between 0 and 1
+        predictions = (
+            predictions.T / predictions.sum(axis=1)
+        ).T  # normalizes outputs between 0 and 1
 
         return predictions
 
     def _save(self):
         d = super()._save()
-        d.update({"classifiers": [saves(classifier) for classifier in self.classifiers],
-                  "sorted_classes": self.sorted_classes})
+        d.update(
+            {
+                "classifiers": [saves(classifier) for classifier in self.classifiers],
+                "sorted_classes": self.sorted_classes,
+            }
+        )
 
     def _load(self, d):
         super()._load(d)
@@ -901,7 +970,9 @@ class BaseSKLearnModel(BaseModel):
     """
 
     @log_arguments
-    def __init__(self, representation, regression_model, classification_model, log=True, **kwargs):
+    def __init__(
+        self, representation, regression_model, classification_model, log=True, **kwargs
+    ):
         self.representation = representation
         self.regression_model = regression_model
         self.classification_model = classification_model
@@ -930,7 +1001,9 @@ def _predict(self, X):
     def _save(self):
         d = super()._save()
         if self.setting == "classification":
-            d.update({"model": self.classification_model._save(), "setting": self.setting})
+            d.update(
+                {"model": self.classification_model._save(), "setting": self.setting}
+            )
         else:
             d.update({"model": self.regression_model._save(), "setting": self.setting})
         return d
@@ -944,6 +1017,7 @@ def _load(self, d):
             self.regression_model._load(d["model"])
             self.model = self.regression_model
 
+
 class BaseErrorModel(BaseClass):
     """Base class for error models.
 
@@ -982,7 +1056,11 @@ def build(
         self.y_pred_train = np.array(self.model.predict(self.X_train)).flatten()
 
     @abstractmethod
-    def calculate(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y_pred: np.ndarray,) -> np.ndarray:
+    def calculate(
+        self,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+        y_pred: np.ndarray,
+    ) -> np.ndarray:
         """To be implemented by the child class; calculates confidence scores from inputs.
 
         Args:
@@ -994,7 +1072,12 @@ def calculate(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y_pred:
         """
         pass
 
-    def fit(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y: Union[np.ndarray, list, pd.Series], **kwargs):
+    def fit(
+        self,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+        y: Union[np.ndarray, list, pd.Series],
+        **kwargs,
+    ):
         """Fits confidence scores to an external dataset
 
         Args:
@@ -1007,7 +1090,12 @@ def fit(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y: Union[np.n
         scores = np.array(self.calculate(X, y_pred))
         self._fit(residuals, scores, **kwargs)
 
-    def fit_cv(self, n_splits: int = 10, test_size_range: Union[tuple, list] = (0.05, 0.5), **kwargs):
+    def fit_cv(
+        self,
+        n_splits: int = 10,
+        test_size_range: Union[tuple, list] = (0.05, 0.5),
+        **kwargs,
+    ):
         """Fits confidence scores to the training dataset via cross validation.
 
         Args:
@@ -1079,8 +1167,15 @@ def _fit(
             y = np.array(y)
         elif method == "qbin":
             bin_labels = pd.qcut(scores, bins, labels=False, duplicates="drop")
-            X = np.array([np.mean(scores[bin_labels == i]) for i in range(max(bin_labels) + 1)])
-            y = np.array([pd.Series(residuals[bin_labels == i]).quantile(quantile) for i in range(max(bin_labels) + 1)])
+            X = np.array(
+                [np.mean(scores[bin_labels == i]) for i in range(max(bin_labels) + 1)]
+            )
+            y = np.array(
+                [
+                    pd.Series(residuals[bin_labels == i]).quantile(quantile)
+                    for i in range(max(bin_labels) + 1)
+                ]
+            )
         elif method == "roll":
             results = list(zip(scores, residuals))
             scores, residuals = zip(*sorted(results))
@@ -1097,8 +1192,8 @@ def _fit(
         funcs = [
             lambda x, a, b: a * np.exp(b * x),
             lambda x, a, b: a * np.log(x) + b,
-            lambda x, a, b: a * x ** b,
-            lambda x, a, b: a * x + b
+            lambda x, a, b: a * x**b,
+            lambda x, a, b: a * x + b,
         ]
         min_mse = None
         for func in funcs:
@@ -1116,7 +1211,9 @@ def _fit(
                 opt_popt = popt
                 min_mse = mse
         X_min, X_max = np.min(X), np.max(X)
-        self.reg = np.vectorize(lambda x: opt_func(max(min(x, X_max), X_min), *opt_popt))
+        self.reg = np.vectorize(
+            lambda x: opt_func(max(min(x, X_max), X_min), *opt_popt)
+        )
 
         plt.xlabel(self.__class__.__name__)
         plt.ylabel("Absolute Error")
@@ -1136,7 +1233,10 @@ def _fit(
         self.min_per_bin = min_per_bin
         self.filename = filename
 
-    def score(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series],) -> np.ndarray:
+    def score(
+        self,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+    ) -> np.ndarray:
         """Calculates confidence scores on a dataset.
 
         Args:
@@ -1165,7 +1265,9 @@ def copy(self) -> BaseErrorModel:
 
     def _save(self) -> dict:
         d = {}
-        if hasattr(self, "model") and not self is getattr(self.model, "error_model", None):
+        if hasattr(self, "model") and not self is getattr(
+            self.model, "error_model", None
+        ):
             d.update({"model": saves(self.model)})
         if hasattr(self, "X_train"):
             d.update({"X_train": self.X_train})
@@ -1199,7 +1301,16 @@ def _load(self, d) -> None:
             min_per_bin = d["min_per_bin"]
             filename = d["filename"]
 
-            self._fit(residuals, scores, method=method, bins=bins, window=window, quantile=quantile, min_per_bin=min_per_bin, filename=filename)
+            self._fit(
+                residuals,
+                scores,
+                method=method,
+                bins=bins,
+                window=window,
+                quantile=quantile,
+                min_per_bin=min_per_bin,
+                filename=filename,
+            )
 
 
 class BaseAggregateErrorModel(BaseErrorModel):
@@ -1210,7 +1321,9 @@ class BaseAggregateErrorModel(BaseErrorModel):
     """
 
     @log_arguments
-    def __init__(self, error_models: Union[BaseErrorModel, List[BaseErrorModel]], **kwargs):
+    def __init__(
+        self, error_models: Union[BaseErrorModel, List[BaseErrorModel]], **kwargs
+    ):
         if not isinstance(error_models, list):
             error_models = [error_models]
         self.error_models = error_models
@@ -1226,7 +1339,11 @@ def build(
             error_model.build(model, X, y)
         super().build(model, X, y)
 
-    def train(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y: Union[np.ndarray, list, pd.Series]):
+    def train(
+        self,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+        y: Union[np.ndarray, list, pd.Series],
+    ):
         """Trains the aggregate model to a validation dataset.
 
         Args:
@@ -1236,10 +1353,14 @@ def train(self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y: Union[np
         y_pred = np.array(self.model.predict(X)).flatten()
         residuals = np.abs(np.array(y) - y_pred)
 
-        scores = np.transpose(np.stack([self.get_scores(em, X, y_pred) for em in self.error_models]))
+        scores = np.transpose(
+            np.stack([self.get_scores(em, X, y_pred) for em in self.error_models])
+        )
         self._train(scores, residuals)
 
-    def train_cv(self, n_splits: int = 9, test_size_range: Union[tuple, list] = (0.1, 0.9)):
+    def train_cv(
+        self, n_splits: int = 9, test_size_range: Union[tuple, list] = (0.1, 0.9)
+    ):
         """Trains the aggregate model to the training dataset via cross validation.
 
         Args:
@@ -1279,7 +1400,10 @@ def train_cv(self, n_splits: int = 9, test_size_range: Union[tuple, list] = (0.1
         self._train(scores, residuals)
 
     def calculate(
-        self, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y_pred: np.ndarray, **kwargs
+        self,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+        y_pred: np.ndarray,
+        **kwargs,
     ) -> np.ndarray:
         """Computes aggregate error model score from inputs.
 
@@ -1294,7 +1418,10 @@ def calculate(
         return self._predict(scores, **kwargs)
 
     def get_scores(
-        self, em: BaseErrorModel, X: Union[pd.DataFrame, np.ndarray, list, pd.Series], y_pred: np.ndarray,
+        self,
+        em: BaseErrorModel,
+        X: Union[pd.DataFrame, np.ndarray, list, pd.Series],
+        y_pred: np.ndarray,
     ) -> np.ndarray:
         """Gets confidence scores from a model, inputs, and keyword arguments.
 
diff --git a/src/olorenchemengine/basics.py b/src/olorenchemengine/basics.py
index 963cde60..6e6b5349 100644
--- a/src/olorenchemengine/basics.py
+++ b/src/olorenchemengine/basics.py
@@ -1,26 +1,46 @@
 """ Machine learning algorithms for use with molecular vector representations and features from experimental data.
 """
 
+import io
+import json
 from tempfile import NamedTemporaryFile
+
 import joblib
-import io
+import numpy as np
+import pandas as pd
+import scipy
+from pandas.api.types import is_numeric_dtype
+from rdkit import Chem
 from sklearn.model_selection import RandomizedSearchCV
 
 import olorenchemengine as oce
-from olorenchemengine.representations import BaseCompoundVecRepresentation, BaseVecRepresentation
+from olorenchemengine.representations import (
+    BaseCompoundVecRepresentation,
+    BaseVecRepresentation,
+)
+
 from .base_class import *
-from rdkit import Chem
 
-import pandas as pd
-import numpy as np
-import scipy
-import json
+try:
+    from pytorch_lightning import LightningModule
 
-from pandas.api.types import is_numeric_dtype
+    lm_imports = (LightningModule, BaseClass)
+except:
+    lm_imports = (BaseClass,)
+
+if isinstance(RandomizedSearchCV, MagicMock):
+    rscvmodel_extends = (BaseEstimator,)
+else:
+    rscvmodel_extends = (
+        RandomizedSearchCV,
+        BaseEstimator,
+    )
 
-class RandomizedSearchCVModel(RandomizedSearchCV, BaseEstimator):
+
+class RandomizedSearchCVModel(*rscvmodel_extends):
 
     """Wrapper class for RandomizedSearchCV"""
+
     def __init__(self, *args, **kwargs):
         self.obj = None
         super().__init__(*args, **kwargs)
@@ -34,6 +54,7 @@ def predict(self, *args, **kwargs):
             raise ValueError("Model not fitted")
         return self.obj.predict(*args, **kwargs)
 
+
 class RandomForestRegressor(BaseEstimator):
 
     """Wrapper for sklearn RandomForestRegressor"""
@@ -65,9 +86,9 @@ class SVC(BaseEstimator):
 
     @log_arguments
     def __init__(self, *args, **kwargs):
-        from sklearn.svm import SVC
         from sklearn.pipeline import make_pipeline
         from sklearn.preprocessing import StandardScaler
+        from sklearn.svm import SVC
 
         self.obj = make_pipeline(StandardScaler(), SVC(*args, **kwargs))
 
@@ -78,21 +99,22 @@ class SVR(BaseEstimator):
 
     @log_arguments
     def __init__(self, *args, **kwargs):
-        from sklearn.svm import SVR
         from sklearn.pipeline import make_pipeline
         from sklearn.preprocessing import StandardScaler
+        from sklearn.svm import SVR
 
         self.obj = make_pipeline(StandardScaler(), SVR(*args, **kwargs))
 
+
 class KBestLinearRegression(BaseEstimator):
 
     """Selects the K-best features and then does linear regression"""
 
     @log_arguments
     def __init__(self, k=1, *args, **lwargs):
-        from sklearn.pipeline import Pipeline
-        from sklearn.linear_model import LinearRegression
         from sklearn.feature_selection import SelectKBest, mutual_info_regression
+        from sklearn.linear_model import LinearRegression
+        from sklearn.pipeline import Pipeline
 
         self.obj = Pipeline(
             [
@@ -101,6 +123,7 @@ def __init__(self, k=1, *args, **lwargs):
             ]
         )
 
+
 class LogisticRegression(BaseEstimator):
 
     """Wrapper for sklearn LogisticRegression"""
@@ -111,8 +134,9 @@ def __init__(self, *args, **kwargs):
 
         self.obj = LogisticRegression(*args, **kwargs)
 
+
 class RandomForestModel(BaseSKLearnModel):
-    """ Random forest model
+    """Random forest model
 
         Parameters:
             n_estimators (int): The number of trees in the forest.
@@ -146,15 +170,18 @@ def __init__(
     ):
         if not class_weight is None:
             try:
-                class_weight = {int(k):float(v) for k,v in json.loads(class_weight).items()}
+                class_weight = {
+                    int(k): float(v) for k, v in json.loads(class_weight).items()
+                }
             except:
-                 class_weight = None
-        regressor = RandomForestRegressor(n_estimators = n_estimators,
+                class_weight = None
+        regressor = RandomForestRegressor(
+            n_estimators=n_estimators,
             max_features=max_features,
             min_samples_split=4,
             min_samples_leaf=4,
-            max_depth = max_depth,
-            n_jobs=-1
+            max_depth=max_depth,
+            n_jobs=-1,
         )
         classifier = RandomForestClassifier(
             max_features=max_features,
@@ -169,8 +196,9 @@ def __init__(
         )
         super().__init__(representation, regressor, classifier, log=False, **kwargs)
 
+
 class AutoRandomForestModel(BaseSKLearnModel, BaseObject):
-    """ RandomForestModel where parameters have automatically tuned hyperparameters
+    """RandomForestModel where parameters have automatically tuned hyperparameters
 
         Parameters:
             representation (str): The representation to use for the model.
@@ -193,10 +221,9 @@ class AutoRandomForestModel(BaseSKLearnModel, BaseObject):
     def __init__(
         self, representation, n_iter=100, scoring=None, verbose=2, cv=5, **kwargs
     ):
-        from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
-        regressor = self.autofit(
-            RandomForestRegressor(), n_iter, cv, scoring, verbose
-        )
+        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+        regressor = self.autofit(RandomForestRegressor(), n_iter, cv, scoring, verbose)
         classifier = self.autofit(
             RandomForestClassifier(), n_iter, cv, scoring, verbose
         )
@@ -215,7 +242,7 @@ def autofit(self, model, n_iter, cv, scoring, verbose):
 
         Returns:
             model (sklearn model): The tuned model
-    """
+        """
 
         params = {
             "n_estimators": [50, 100, 200, 500, 1000],
@@ -235,6 +262,8 @@ def autofit(self, model, n_iter, cv, scoring, verbose):
             scoring=scoring,
             verbose=verbose,
         )
+
+
 class BaseMLPClassifier(BaseEstimator):
 
     """Wrapper for sklearn MLP"""
@@ -245,6 +274,7 @@ def __init__(self, *args, **kwargs):
 
         self.obj = MLPClassifier(*args, **kwargs)
 
+
 class BaseMLPRegressor(BaseEstimator):
 
     """Wrapper for sklearn MLP"""
@@ -255,8 +285,9 @@ def __init__(self, *args, **kwargs):
 
         self.obj = MLPRegressor(*args, **kwargs)
 
+
 class SklearnMLP(BaseSKLearnModel):
-    """ MLP Model based on sklearn implementation
+    """MLP Model based on sklearn implementation
 
         Parameters:
             representation (BaseVecRepresentation): The representation to use for the model.
@@ -279,48 +310,66 @@ class SklearnMLP(BaseSKLearnModel):
     model.predict(test['Drug'])
     ------------------------------
     """
+
     @log_arguments
-    def __init__(self, representation,
-        hidden_layer_sizes=[100], activation="relu", solver="adam",
-        alpha=0.0001, batch_size=32, learning_rate="constant",
-        learning_rate_init=0.001, power_t=0.5, max_iter=200, log=True, **kwargs):
+    def __init__(
+        self,
+        representation,
+        hidden_layer_sizes=[100],
+        activation="relu",
+        solver="adam",
+        alpha=0.0001,
+        batch_size=32,
+        learning_rate="constant",
+        learning_rate_init=0.001,
+        power_t=0.5,
+        max_iter=200,
+        log=True,
+        **kwargs
+    ):
 
         regressor = BaseMLPRegressor(
-            hidden_layer_sizes=hidden_layer_sizes, activation=activation, learning_rate=learning_rate,
-            solver=solver, alpha=alpha, batch_size=batch_size, learning_rate_init=learning_rate_init,
-            power_t=power_t, max_iter=max_iter)
+            hidden_layer_sizes=hidden_layer_sizes,
+            activation=activation,
+            learning_rate=learning_rate,
+            solver=solver,
+            alpha=alpha,
+            batch_size=batch_size,
+            learning_rate_init=learning_rate_init,
+            power_t=power_t,
+            max_iter=max_iter,
+        )
         classifier = BaseMLPClassifier(
-            hidden_layer_sizes=hidden_layer_sizes, activation=activation, learning_rate=learning_rate,
-            solver=solver, alpha=alpha, batch_size=batch_size, learning_rate_init=learning_rate_init,
-            power_t=power_t, max_iter=max_iter)
+            hidden_layer_sizes=hidden_layer_sizes,
+            activation=activation,
+            learning_rate=learning_rate,
+            solver=solver,
+            alpha=alpha,
+            batch_size=batch_size,
+            learning_rate_init=learning_rate_init,
+            power_t=power_t,
+            max_iter=max_iter,
+        )
         super().__init__(representation, regressor, classifier, log=False, **kwargs)
 
-try:
-    import pytorch_lightning as pl
-except ImportError:
-    oce.mock_imports(globals(), "pl")
 
-try:
-    import torch.nn as nn
-    import torch
-except ImportError:
-    oce.mock_imports(globals(), "nn", "torch")
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
 
-class plMLP(pl.LightningModule):
 
+class plMLP(pl.LightningModule):
     def __init__(self, hidden_layer_sizes, norm_layer, activation_layer, setting):
         super().__init__()
         hidden_layer_sizes = hidden_layer_sizes + [1]
         layers = []
-        for i in range(len(hidden_layer_sizes)-1):
-            layers.append(nn.Linear(hidden_layer_sizes[i], hidden_layer_sizes[i+1]))
-            if not activation_layer is None and i < len(hidden_layer_sizes)-2:
+        for i in range(len(hidden_layer_sizes) - 1):
+            layers.append(nn.Linear(hidden_layer_sizes[i], hidden_layer_sizes[i + 1]))
+            if not activation_layer is None and i < len(hidden_layer_sizes) - 2:
                 layers.append(activation_layer())
-            if not norm_layer is None and i < len(hidden_layer_sizes)-2:
+            if not norm_layer is None and i < len(hidden_layer_sizes) - 2:
                 layers.append(norm_layer())
-        self.layers = nn.Sequential(
-          *layers
-        )
+        self.layers = nn.Sequential(*layers)
         if setting == "classification":
             self.layers.append(nn.Sigmoid())
             self.loss = nn.CrossEntropyLoss()
@@ -335,15 +384,16 @@ def training_step(self, batch, batch_idx):
         x = x.view(x.size(0), -1)
         y_hat = self.layers(x)
         loss = self.loss(y_hat, y)
-        self.log('train_loss', loss)
+        self.log("train_loss", loss)
         return loss
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)
         return optimizer
 
+
 class TorchMLP(BaseModel):
-    """ MLP Model based on torch implementation
+    """MLP Model based on torch implementation
 
         Parameters:
             representation (BaseVecRepresentation): The representation to use for the model.
@@ -363,12 +413,20 @@ class TorchMLP(BaseModel):
     """
 
     @log_arguments
-    def __init__(self, representation, hidden_layer_sizes=[100],
-            norm_layer: str = None, activation_layer: str = None, dropout: float = 0.0,
-            log=True, **kwargs):
+    def __init__(
+        self,
+        representation,
+        hidden_layer_sizes=[100],
+        norm_layer: str = None,
+        activation_layer: str = None,
+        dropout: float = 0.0,
+        log=True,
+        **kwargs
+    ):
         self.representation = representation
         self.hidden_layer_sizes = hidden_layer_sizes
         import torch.nn as nn
+
         if activation_layer == "leakyrelu":
             self.activation_layer = nn.LeakyReLU
         elif activation_layer == "relu":
@@ -388,6 +446,7 @@ def __init__(self, representation, hidden_layer_sizes=[100],
             self.norm_layer = None
 
         self.dropout = dropout
+
         super().__init__(representation, log=False)
 
     def preprocess(self, X, y, fit=False):
@@ -401,20 +460,28 @@ def _fit(self, X, y):
         import torch
         import torch.nn as nn
 
-        self.network = plMLP(hidden_layer_sizes=[len(X[0])] + self.hidden_layer_sizes,
+        self.network = plMLP(
+            hidden_layer_sizes=[len(X[0])] + self.hidden_layer_sizes,
             norm_layer=self.norm_layer,
             activation_layer=self.activation_layer,
-            setting = self.setting)
+            setting=self.setting,
+        )
+
+        from torch.utils.data import DataLoader, TensorDataset
 
-        from torch.utils.data import TensorDataset, DataLoader
         y = np.array(y.tolist())
         device = torch.device(oce.CONFIG["MAP_LOCATION"])
-        dataset = TensorDataset(torch.from_numpy(X).float().to(device), torch.from_numpy(y).float().to(device))
+        dataset = TensorDataset(
+            torch.from_numpy(X).float().to(device),
+            torch.from_numpy(y).float().to(device),
+        )
         loader = DataLoader(dataset, batch_size=32)
 
         self.trainer = pl.Trainer(
             accelerator="auto",
-            devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
+            devices=1
+            if torch.cuda.is_available()
+            else None,  # limiting got iPython runs
             max_epochs=10,
         )
 
@@ -422,29 +489,39 @@ def _fit(self, X, y):
 
     def _predict(self, X):
         import torch
+
         # from torch.utils.data import TensorDataset, DataLoader
         device = torch.device(oce.CONFIG["MAP_LOCATION"])
         # dataset = TensorDataset(torch.from_numpy(X).float().to(device))
         # loader = DataLoader(dataset, batch_size=1)
         # return self.trainer.predict(self.network, loader)
         self.network.to(device)
-        return self.network(torch.from_numpy(X).float().to(device)).detach().cpu().numpy()
+        return (
+            self.network(torch.from_numpy(X).float().to(device)).detach().cpu().numpy()
+        )
 
     def _save(self) -> str:
         d = super()._save()
+        if not hasattr(self, "network"):
+            return d
         buffer = io.BytesIO()
         import torch
+
         torch.save(self.network, buffer)
         d.update({"save": buffer.getvalue()})
         return d
 
     def _load(self, d):
         super()._load(d)
+        if not "save" in d:
+            return
         import torch
+
         self.network = torch.load(io.BytesIO(d["save"]))
 
+
 class XGBoostModel(BaseSKLearnModel, BaseObject):
-    """ XGBoost model
+    """XGBoost model
 
         Parameters:
             representation (str): The representation to use for the model
@@ -465,23 +542,43 @@ class XGBoostModel(BaseSKLearnModel, BaseObject):
 
     @log_arguments
     def __init__(
-        self, representation, n_estimators=2000, max_depth=7, subsample=0.5,max_leaves=5,
-            learning_rate=0.05, colsample_bytree=0.8, **kwargs
+        self,
+        representation,
+        n_estimators=2000,
+        max_depth=7,
+        subsample=0.5,
+        max_leaves=5,
+        learning_rate=0.05,
+        colsample_bytree=0.8,
+        **kwargs
     ):
-        from xgboost import XGBRegressor, XGBClassifier
+        from xgboost import XGBClassifier, XGBRegressor
 
-        regressor = XGBRegressor(tree_method="gpu_hist", n_estimators=n_estimators,
-            max_depth=max_depth, max_leaves=max_leaves, learning_rate=learning_rate,
-            subsample=subsample, colsample_bytree=colsample_bytree)
+        regressor = XGBRegressor(
+            tree_method="gpu_hist",
+            n_estimators=n_estimators,
+            max_depth=max_depth,
+            max_leaves=max_leaves,
+            learning_rate=learning_rate,
+            subsample=subsample,
+            colsample_bytree=colsample_bytree,
+        )
 
-        classifier = XGBClassifier(tree_method="gpu_hist", n_estimators=n_estimators,
-            max_depth=max_depth, max_leaves=max_leaves, learning_rate=learning_rate,
-            subsample=subsample, colsample_bytree=colsample_bytree)
+        classifier = XGBClassifier(
+            tree_method="gpu_hist",
+            n_estimators=n_estimators,
+            max_depth=max_depth,
+            max_leaves=max_leaves,
+            learning_rate=learning_rate,
+            subsample=subsample,
+            colsample_bytree=colsample_bytree,
+        )
 
         super().__init__(representation, regressor, classifier, log=False, **kwargs)
 
+
 class ZWK_XGBoostModel(BaseSKLearnModel, BaseObject):
-    """ XGBoost model from https://github.com/smu-tao-group/ADMET_XGBoost
+    """XGBoost model from https://github.com/smu-tao-group/ADMET_XGBoost
 
         Parameters:
             representation (str): The representation to use for the model.
@@ -504,15 +601,16 @@ class ZWK_XGBoostModel(BaseSKLearnModel, BaseObject):
     def __init__(
         self, representation, n_iter=100, scoring=None, verbose=2, cv=5, **kwargs
     ):
-        from xgboost import XGBRegressor, XGBClassifier
+        from xgboost import XGBClassifier, XGBRegressor
 
-        if scoring == 'spearman':
+        if scoring == "spearman":
             from scipy.stats import spearmanr
             from sklearn.metrics import make_scorer
+
             def spearman_loss_func(y_true, y_pred):
-                """spearman metric
-                """
+                """spearman metric"""
                 return spearmanr(y_true, y_pred)[0]
+
             spearman = make_scorer(spearman_loss_func)
             scoring = spearman
 
@@ -560,9 +658,10 @@ def autofit(self, model, n_iter, cv, scoring, verbose):
             verbose=verbose,
         )
 
+
 class SupportVectorMachine(BaseSKLearnModel):
 
-    """ Support vector machine
+    """Support vector machine
 
         Parameters:
             representations (str): The representation to use for the model.
@@ -637,7 +736,7 @@ def predict(self, X):
 
 class KNN(BaseSKLearnModel):
 
-    """ KNN model
+    """KNN model
 
         Parameters:
             representations (str): The representation to use for the model.
@@ -661,7 +760,7 @@ def __init__(self, representation, **kwargs):
 
 class MLP(BaseSKLearnModel):
 
-    """ MLP model
+    """MLP model
 
         Parameters:
             representation (BaseVecRepresentation): The representation to use for the model.
@@ -681,13 +780,14 @@ class MLP(BaseSKLearnModel):
     model.fit(train['Drug'], train['Y'])
     model.predict(test['Drug'])
     ------------------------------
-     """
+    """
 
-    def preprocess(self, X, y, fit = False):
+    def preprocess(self, X, y, fit=False):
         if self.representation is None:
             return np.array(X)
         else:
-            return np.array(self.representation.convert(X, fit = fit))
+            return np.array(self.representation.convert(X, fit=fit))
+
     # specify parameters
     @log_arguments
     def __init__(
@@ -802,7 +902,7 @@ def __init__(self, config="full", reg="lr", **kwargs):
             self.reg = None
         super().__init__(log=False, **kwargs)
 
-    def preprocess(self, X, y, fit = False):
+    def preprocess(self, X, y, fit=False):
         """This method is used to preprocess the data before training.
 
         Parameters:
@@ -840,8 +940,8 @@ def _fit(self, X_train, y_train):
                     t0 = trans[0](df2["x"])
                     t1 = trans[1](df2["y"])
                     A, B, r, p, se = scipy.stats.linregress(t0, t1)
-                    if r ** 2 > r2max:
-                        r2max = r ** 2
+                    if r**2 > r2max:
+                        r2max = r**2
                         best = trans[3]
                         self.state[c] = (i, A, B, r2max)
                 print(c)
@@ -878,14 +978,14 @@ def _load(self, d):
 
 class FeaturesClassification(BaseModel):
 
-    """ FeaturesClassification uses machine learning models to classify features
-        based on their experimental data
+    """FeaturesClassification uses machine learning models to classify features
+    based on their experimental data
 
-        Attributes:
-            obj (sklearn.base.BaseEstimator): Machine learning model to use.
+    Attributes:
+        obj (sklearn.base.BaseEstimator): Machine learning model to use.
 
-        Parameters:
-            config (str): Configuration to use for the model."""
+    Parameters:
+        config (str): Configuration to use for the model."""
 
     @log_arguments
     def __init__(self, config="lineardiscriminant"):
diff --git a/src/olorenchemengine/benchmarks.py b/src/olorenchemengine/benchmarks.py
index 71c2364d..084e034f 100644
--- a/src/olorenchemengine/benchmarks.py
+++ b/src/olorenchemengine/benchmarks.py
@@ -3,17 +3,18 @@
 and for methods development
 """
 
-import olorenchemengine as oce
-from olorenchemengine import *
-
 from typing import *
+
 import pandas as pd
 from tqdm import tqdm
 
+import olorenchemengine as oce
+from olorenchemengine import *
 from olorenchemengine.base_class import *
 from olorenchemengine.manager import *
 from olorenchemengine.visualizations import *
 
+
 class BenchmarkDatasets(BaseVisualization):
     """
     Track and visualize the performance of models on user-defined datasets.
@@ -41,13 +42,21 @@ class BenchmarkDatasets(BaseVisualization):
     """
 
     @log_arguments
-    def __init__(self, datasets: List[BaseDataset], log = True):
-        super().__init__(log = False)
+    def __init__(self, datasets: List[BaseDataset], log=True):
+        super().__init__(log=False)
         self.packages = ["plotly"]
 
         self.datasets = datasets
-        self.metrics = ["ROC-AUC" if detect_setting(dataset.train_dataset[1]) == "classification" else "Root Mean Squared Error" for dataset in datasets]
-        self.managers = [ModelManager(dataset, metrics=[metric]) for dataset, metric in zip(self.datasets, self.metrics)]
+        self.metrics = [
+            "ROC-AUC"
+            if detect_setting(dataset.train_dataset[1]) == "classification"
+            else "Root Mean Squared Error"
+            for dataset in datasets
+        ]
+        self.managers = [
+            ModelManager(dataset, metrics=[metric])
+            for dataset, metric in zip(self.datasets, self.metrics)
+        ]
         self.model_names = []
         self.model_params = []
 
@@ -78,15 +87,22 @@ def get_data(self):
             benchmark_data = benchmark_data.append(model_database)
 
             if direction == "higher":
-                best_row = model_database[model_database[metric] == model_database[metric].max()].iloc[0]
+                best_row = model_database[
+                    model_database[metric] == model_database[metric].max()
+                ].iloc[0]
             elif direction == "lower":
-                best_row = model_database[model_database[metric] == model_database[metric].min()].iloc[0]
+                best_row = model_database[
+                    model_database[metric] == model_database[metric].min()
+                ].iloc[0]
 
             d[best_row["Model Name"]] = d[best_row["Model Name"]] + 1
 
         out = dict()
         out["model_names"] = self.model_names
-        out["model_params"] = [str(mp).replace("\n", "<br>") for mn, mp in zip(self.model_names, self.model_params)]
+        out["model_params"] = [
+            str(mp).replace("\n", "<br>")
+            for mn, mp in zip(self.model_names, self.model_params)
+        ]
         out["model_counts"] = [d[name] for name in self.model_names]
         return out
 
@@ -104,6 +120,7 @@ def _load(self, d):
         self.model_names = d["model_names"]
         self.model_params = d["model_params"]
 
+
 class BenchmarkTDC(BenchmarkDatasets):
     """
     Track and visualize the performance of models on selected ADME datasets from the Therapeutics Data Commons (TDC). https://tdcommons.ai/
@@ -133,69 +150,77 @@ class BenchmarkTDC(BenchmarkDatasets):
     """
 
     @log_arguments
-    def __init__(self, datasets: List[str] = ['all'], log = True):
+    def __init__(self, datasets: List[str] = ["all"], log=True):
         self.packages = ["plotly"]
         from tdc.benchmark_group import admet_group
 
         self.metrics_dict = {
-            'HIA_Hou':'ROC-AUC',
-            'Pgp_Broccatelli':'ROC-AUC',
-            'Bioavailability_Ma':'ROC-AUC',
-            'BBB_Martins':'ROC-AUC',
-            'hERG':'ROC-AUC',
-            'AMES':'ROC-AUC',
-            'DILI':'ROC-AUC',
-            'CYP3A4_Substrate_CarbonMangels':'ROC-AUC',
-            'Lipophilicity_AstraZeneca':'Mean Absolute Error',
-            'Solubility_AqSolDB':'Mean Absolute Error',
-            'LD50_Zhu':'Mean Absolute Error',
-            'Caco2_Wang':'Mean Absolute Error',
-            'PPBR_AZ':'Mean Absolute Error',
-            'CYP2C9_Veith':'Average Precision',
-            'CYP2D6_Veith':'Average Precision',
-            'CYP3A4_Veith':'Average Precision',
-            'CYP2C9_Substrate_CarbonMangels':'Average Precision',
-            'CYP2D6_Substrate_CarbonMangels':'Average Precision',
-            'Half_Life_Obach':'Spearman',
-            'Clearance_Hepatocyte_AZ':'Spearman',
-            'Clearance_Microsome_AZ':'Spearman',
-            'VDss_Lombardo':'Spearman',
+            "HIA_Hou": "ROC-AUC",
+            "Pgp_Broccatelli": "ROC-AUC",
+            "Bioavailability_Ma": "ROC-AUC",
+            "BBB_Martins": "ROC-AUC",
+            "hERG": "ROC-AUC",
+            "AMES": "ROC-AUC",
+            "DILI": "ROC-AUC",
+            "CYP3A4_Substrate_CarbonMangels": "ROC-AUC",
+            "Lipophilicity_AstraZeneca": "Mean Absolute Error",
+            "Solubility_AqSolDB": "Mean Absolute Error",
+            "LD50_Zhu": "Mean Absolute Error",
+            "Caco2_Wang": "Mean Absolute Error",
+            "PPBR_AZ": "Mean Absolute Error",
+            "CYP2C9_Veith": "Average Precision",
+            "CYP2D6_Veith": "Average Precision",
+            "CYP3A4_Veith": "Average Precision",
+            "CYP2C9_Substrate_CarbonMangels": "Average Precision",
+            "CYP2D6_Substrate_CarbonMangels": "Average Precision",
+            "Half_Life_Obach": "Spearman",
+            "Clearance_Hepatocyte_AZ": "Spearman",
+            "Clearance_Microsome_AZ": "Spearman",
+            "VDss_Lombardo": "Spearman",
         }
 
-        if datasets == ['all']:
+        if datasets == ["all"]:
             datasets = list(self.metrics_dict.keys())
 
-        group = admet_group(path = 'data/')
+        group = admet_group(path="data/")
         datasets_wrapped = []
         self.metrics = []
         for name in datasets:
             benchmark = group.get(name)
-            test = benchmark['test']
-            train, valid = group.get_train_valid_split(benchmark = name, split_type = 'default', seed = 0)
-            train['split'] = 'train'
-            test['split'] = 'test'
+            test = benchmark["test"]
+            train, valid = group.get_train_valid_split(
+                benchmark=name, split_type="default", seed=0
+            )
+            train["split"] = "train"
+            test["split"] = "test"
             self.metrics.append(self.metrics_dict[name])
 
             data = pd.concat([train, test]).to_csv()
             datasets_wrapped.append(
-                BaseDataset(data = data, structure_col = "Drug", property_col = "Y")
+                BaseDataset(data=data, structure_col="Drug", property_col="Y")
             )
             i = datasets.index(name)
             datasets_wrapped[i].train = train
             datasets_wrapped[i].test = test
-            datasets_wrapped[i].data = pd.concat([datasets_wrapped[i].train, datasets_wrapped[i].test])
-            
+            datasets_wrapped[i].data = pd.concat(
+                [datasets_wrapped[i].train, datasets_wrapped[i].test]
+            )
+
         self.dataset_names = datasets
         self.datasets = datasets_wrapped
-        self.managers = [ModelManager(dataset, metrics=[metric]) for dataset, metric in zip(self.datasets, self.metrics)]
+        self.managers = [
+            ModelManager(dataset, metrics=[metric])
+            for dataset, metric in zip(self.datasets, self.metrics)
+        ]
         self.model_names = []
         self.model_params = []
 
-    '''Returns DataFrame of perfomance metric for each dataset and model benchmarked.'''
+    """Returns DataFrame of perfomance metric for each dataset and model benchmarked."""
+
     def get_performance(self):
-        out = pd.DataFrame(data=None, columns=['Dataset', 'Metric'])
-        out['Dataset'] = self.dataset_names
-        out['Metric'] = [self.metrics_dict[name] for name in self.dataset_names]
+        out = pd.DataFrame(data=None, columns=["Dataset", "Metric"])
+        out["Dataset"] = self.dataset_names
+        out["Metric"] = [self.metrics_dict[name] for name in self.dataset_names]
         for name in self.model_names:
             out[name] = np.nan
         for manager in self.managers:
@@ -203,9 +228,12 @@ def get_performance(self):
             manager_data = manager.model_database
             for i, row in manager_data.iterrows():
                 index = i
-                out.loc[out['Dataset'] == dataset_name, row['Model Name']] = row[self.metrics_dict[dataset_name]]
+                out.loc[out["Dataset"] == dataset_name, row["Model Name"]] = row[
+                    self.metrics_dict[dataset_name]
+                ]
         return out
 
+
 class BenchmarkMolNet(BenchmarkDatasets):
     """
     Track and visualize the performance of models on selected datasets from the MoleculeNet collection. https://moleculenet.org/
@@ -241,8 +269,14 @@ class BenchmarkMolNet(BenchmarkDatasets):
     """
 
     @log_arguments
-    def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "default", 
-            eval="valid", log = True):
+    def __init__(
+        self,
+        datasets: List[str] = ["all"],
+        file_path=None,
+        mode="default",
+        eval="valid",
+        log=True,
+    ):
         self.mode = mode
         self.eval = eval
         if "default" in self.mode or "small" in self.mode:
@@ -259,7 +293,7 @@ def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "defa
                 "delaney_esol": "https://storage.googleapis.com/oloren-public-data/MoleculeNet/load_delaney.csv",
                 "freesolv": "https://storage.googleapis.com/oloren-public-data/MoleculeNet/load_freesolv.csv",
                 "lipo": "https://storage.googleapis.com/oloren-public-data/MoleculeNet/load_lipo.csv",
-                "bace_regression": "https://storage.googleapis.com/oloren-public-data/MoleculeNet/load_bace_regression.csv"
+                "bace_regression": "https://storage.googleapis.com/oloren-public-data/MoleculeNet/load_bace_regression.csv",
             }
         elif "geognn" in self.mode:
             self.datasets_loaded = {
@@ -274,11 +308,13 @@ def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "defa
                 "lipo": "https://storage.googleapis.com/oloren-public-data/MoleculeNet2/Lipophilicity.csv",
             }
         else:
-            assert False, "Invalid mode, should be either 'default', 'small' or 'geognn' with optional attribute small"
+            assert (
+                False
+            ), "Invalid mode, should be either 'default', 'small' or 'geognn' with optional attribute small"
 
         self.datasets_types = {
             "bace_classification": "classification",
-            'bace': 'classification',
+            "bace": "classification",
             "bbbp": "classification",
             "clintox": "classification",
             "hiv": "classification",
@@ -290,34 +326,839 @@ def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "defa
             "delaney_esol": "regression",
             "freesolv": "regression",
             "lipo": "regression",
-            "bace_regression": "regression"
+            "bace_regression": "regression",
         }
 
         self.dataset_property_cols = {
             "bace_classification": "Class",
-            'bace': 'Class',
+            "bace": "Class",
             "bbbp": "p_np",
-            "clintox": ["FDA_APPROVED","CT_TOX"],
+            "clintox": ["FDA_APPROVED", "CT_TOX"],
             "hiv": "HIV_active",
-            "muv": ["MUV-466","MUV-548","MUV-600","MUV-644","MUV-652","MUV-689","MUV-692","MUV-712","MUV-713","MUV-733","MUV-737","MUV-810","MUV-832","MUV-846","MUV-852","MUV-858","MUV-859"],
-            "pcba": ["PCBA-1030","PCBA-1379","PCBA-1452","PCBA-1454","PCBA-1457","PCBA-1458","PCBA-1460","PCBA-1461","PCBA-1468","PCBA-1469","PCBA-1471","PCBA-1479","PCBA-1631","PCBA-1634","PCBA-1688","PCBA-1721","PCBA-2100","PCBA-2101","PCBA-2147","PCBA-2242","PCBA-2326","PCBA-2451","PCBA-2517","PCBA-2528","PCBA-2546","PCBA-2549","PCBA-2551","PCBA-2662","PCBA-2675","PCBA-2676","PCBA-411","PCBA-463254","PCBA-485281","PCBA-485290","PCBA-485294","PCBA-485297","PCBA-485313","PCBA-485314","PCBA-485341","PCBA-485349","PCBA-485353","PCBA-485360","PCBA-485364","PCBA-485367","PCBA-492947","PCBA-493208","PCBA-504327","PCBA-504332","PCBA-504333","PCBA-504339","PCBA-504444","PCBA-504466","PCBA-504467","PCBA-504706","PCBA-504842","PCBA-504845","PCBA-504847","PCBA-504891","PCBA-540276","PCBA-540317","PCBA-588342","PCBA-588453","PCBA-588456","PCBA-588579","PCBA-588590","PCBA-588591","PCBA-588795","PCBA-588855","PCBA-602179","PCBA-602233","PCBA-602310","PCBA-602313","PCBA-602332","PCBA-624170","PCBA-624171","PCBA-624173","PCBA-624202","PCBA-624246","PCBA-624287","PCBA-624288","PCBA-624291","PCBA-624296","PCBA-624297","PCBA-624417","PCBA-651635","PCBA-651644","PCBA-651768","PCBA-651965","PCBA-652025","PCBA-652104","PCBA-652105","PCBA-652106","PCBA-686970","PCBA-686978","PCBA-686979","PCBA-720504","PCBA-720532","PCBA-720542","PCBA-720551","PCBA-720553","PCBA-720579","PCBA-720580","PCBA-720707","PCBA-720708","PCBA-720709","PCBA-720711","PCBA-743255","PCBA-743266","PCBA-875","PCBA-881","PCBA-883","PCBA-884","PCBA-885","PCBA-887","PCBA-891","PCBA-899","PCBA-902","PCBA-903","PCBA-904","PCBA-912","PCBA-914","PCBA-915","PCBA-924","PCBA-925","PCBA-926","PCBA-927","PCBA-938","PCBA-995"],
-            "sider": ["Hepatobiliary disorders","Metabolism and nutrition disorders","Product issues","Eye disorders","Investigations","Musculoskeletal and connective tissue disorders","Gastrointestinal disorders","Social circumstances","Immune system disorders","Reproductive system and breast disorders","Neoplasms benign, malignant and unspecified (incl cysts and polyps)","General disorders and administration site conditions","Endocrine disorders","Surgical and medical procedures","Vascular disorders","Blood and lymphatic system disorders","Skin and subcutaneous tissue disorders","Congenital, familial and genetic disorders","Infections and infestations","Respiratory, thoracic and mediastinal disorders","Psychiatric disorders","Renal and urinary disorders","Pregnancy, puerperium and perinatal conditions","Ear and labyrinth disorders","Cardiac disorders","Nervous system disorders","Injury, poisoning and procedural complications"],
-            "tox21": ["NR-AR","NR-AR-LBD","NR-AhR","NR-Aromatase","NR-ER","NR-ER-LBD","NR-PPAR-gamma","SR-ARE","SR-ATAD5","SR-HSE","SR-MMP","SR-p53"],
-            "toxcast": ["ACEA_T47D_80hr_Negative","ACEA_T47D_80hr_Positive","APR_HepG2_CellCycleArrest_24h_dn","APR_HepG2_CellCycleArrest_24h_up","APR_HepG2_CellCycleArrest_72h_dn","APR_HepG2_CellLoss_24h_dn","APR_HepG2_CellLoss_72h_dn","APR_HepG2_MicrotubuleCSK_24h_dn","APR_HepG2_MicrotubuleCSK_24h_up","APR_HepG2_MicrotubuleCSK_72h_dn","APR_HepG2_MicrotubuleCSK_72h_up","APR_HepG2_MitoMass_24h_dn","APR_HepG2_MitoMass_24h_up","APR_HepG2_MitoMass_72h_dn","APR_HepG2_MitoMass_72h_up","APR_HepG2_MitoMembPot_1h_dn","APR_HepG2_MitoMembPot_24h_dn","APR_HepG2_MitoMembPot_72h_dn","APR_HepG2_MitoticArrest_24h_up","APR_HepG2_MitoticArrest_72h_up","APR_HepG2_NuclearSize_24h_dn","APR_HepG2_NuclearSize_72h_dn","APR_HepG2_NuclearSize_72h_up","APR_HepG2_OxidativeStress_24h_up","APR_HepG2_OxidativeStress_72h_up","APR_HepG2_StressKinase_1h_up","APR_HepG2_StressKinase_24h_up","APR_HepG2_StressKinase_72h_up","APR_HepG2_p53Act_24h_up","APR_HepG2_p53Act_72h_up","APR_Hepat_Apoptosis_24hr_up","APR_Hepat_Apoptosis_48hr_up","APR_Hepat_CellLoss_24hr_dn","APR_Hepat_CellLoss_48hr_dn","APR_Hepat_DNADamage_24hr_up","APR_Hepat_DNADamage_48hr_up","APR_Hepat_DNATexture_24hr_up","APR_Hepat_DNATexture_48hr_up","APR_Hepat_MitoFxnI_1hr_dn","APR_Hepat_MitoFxnI_24hr_dn","APR_Hepat_MitoFxnI_48hr_dn","APR_Hepat_NuclearSize_24hr_dn","APR_Hepat_NuclearSize_48hr_dn","APR_Hepat_Steatosis_24hr_up","APR_Hepat_Steatosis_48hr_up","ATG_AP_1_CIS_dn","ATG_AP_1_CIS_up","ATG_AP_2_CIS_dn","ATG_AP_2_CIS_up","ATG_AR_TRANS_dn","ATG_AR_TRANS_up","ATG_Ahr_CIS_dn","ATG_Ahr_CIS_up","ATG_BRE_CIS_dn","ATG_BRE_CIS_up","ATG_CAR_TRANS_dn","ATG_CAR_TRANS_up","ATG_CMV_CIS_dn","ATG_CMV_CIS_up","ATG_CRE_CIS_dn","ATG_CRE_CIS_up","ATG_C_EBP_CIS_dn","ATG_C_EBP_CIS_up","ATG_DR4_LXR_CIS_dn","ATG_DR4_LXR_CIS_up","ATG_DR5_CIS_dn","ATG_DR5_CIS_up","ATG_E2F_CIS_dn","ATG_E2F_CIS_up","ATG_EGR_CIS_up","ATG_ERE_CIS_dn","ATG_ERE_CIS_up","ATG_ERRa_TRANS_dn","ATG_ERRg_TRANS_dn","ATG_ERRg_TRANS_up","ATG_ERa_TRANS_up","ATG_E_Box_CIS_dn","ATG_E_Box_CIS_up","ATG_Ets_CIS_dn","ATG_Ets_CIS_up","ATG_FXR_TRANS_up","ATG_FoxA2_CIS_dn","ATG_FoxA2_CIS_up","ATG_FoxO_CIS_dn","ATG_FoxO_CIS_up","ATG_GAL4_TRANS_dn","ATG_GATA_CIS_dn","ATG_GATA_CIS_up","ATG_GLI_CIS_dn","ATG_GLI_CIS_up","ATG_GRE_CIS_dn","ATG_GRE_CIS_up","ATG_GR_TRANS_dn","ATG_GR_TRANS_up","ATG_HIF1a_CIS_dn","ATG_HIF1a_CIS_up","ATG_HNF4a_TRANS_dn","ATG_HNF4a_TRANS_up","ATG_HNF6_CIS_dn","ATG_HNF6_CIS_up","ATG_HSE_CIS_dn","ATG_HSE_CIS_up","ATG_IR1_CIS_dn","ATG_IR1_CIS_up","ATG_ISRE_CIS_dn","ATG_ISRE_CIS_up","ATG_LXRa_TRANS_dn","ATG_LXRa_TRANS_up","ATG_LXRb_TRANS_dn","ATG_LXRb_TRANS_up","ATG_MRE_CIS_up","ATG_M_06_TRANS_up","ATG_M_19_CIS_dn","ATG_M_19_TRANS_dn","ATG_M_19_TRANS_up","ATG_M_32_CIS_dn","ATG_M_32_CIS_up","ATG_M_32_TRANS_dn","ATG_M_32_TRANS_up","ATG_M_61_TRANS_up","ATG_Myb_CIS_dn","ATG_Myb_CIS_up","ATG_Myc_CIS_dn","ATG_Myc_CIS_up","ATG_NFI_CIS_dn","ATG_NFI_CIS_up","ATG_NF_kB_CIS_dn","ATG_NF_kB_CIS_up","ATG_NRF1_CIS_dn","ATG_NRF1_CIS_up","ATG_NRF2_ARE_CIS_dn","ATG_NRF2_ARE_CIS_up","ATG_NURR1_TRANS_dn","ATG_NURR1_TRANS_up","ATG_Oct_MLP_CIS_dn","ATG_Oct_MLP_CIS_up","ATG_PBREM_CIS_dn","ATG_PBREM_CIS_up","ATG_PPARa_TRANS_dn","ATG_PPARa_TRANS_up","ATG_PPARd_TRANS_up","ATG_PPARg_TRANS_up","ATG_PPRE_CIS_dn","ATG_PPRE_CIS_up","ATG_PXRE_CIS_dn","ATG_PXRE_CIS_up","ATG_PXR_TRANS_dn","ATG_PXR_TRANS_up","ATG_Pax6_CIS_up","ATG_RARa_TRANS_dn","ATG_RARa_TRANS_up","ATG_RARb_TRANS_dn","ATG_RARb_TRANS_up","ATG_RARg_TRANS_dn","ATG_RARg_TRANS_up","ATG_RORE_CIS_dn","ATG_RORE_CIS_up","ATG_RORb_TRANS_dn","ATG_RORg_TRANS_dn","ATG_RORg_TRANS_up","ATG_RXRa_TRANS_dn","ATG_RXRa_TRANS_up","ATG_RXRb_TRANS_dn","ATG_RXRb_TRANS_up","ATG_SREBP_CIS_dn","ATG_SREBP_CIS_up","ATG_STAT3_CIS_dn","ATG_STAT3_CIS_up","ATG_Sox_CIS_dn","ATG_Sox_CIS_up","ATG_Sp1_CIS_dn","ATG_Sp1_CIS_up","ATG_TAL_CIS_dn","ATG_TAL_CIS_up","ATG_TA_CIS_dn","ATG_TA_CIS_up","ATG_TCF_b_cat_CIS_dn","ATG_TCF_b_cat_CIS_up","ATG_TGFb_CIS_dn","ATG_TGFb_CIS_up","ATG_THRa1_TRANS_dn","ATG_THRa1_TRANS_up","ATG_VDRE_CIS_dn","ATG_VDRE_CIS_up","ATG_VDR_TRANS_dn","ATG_VDR_TRANS_up","ATG_XTT_Cytotoxicity_up","ATG_Xbp1_CIS_dn","ATG_Xbp1_CIS_up","ATG_p53_CIS_dn","ATG_p53_CIS_up","BSK_3C_Eselectin_down","BSK_3C_HLADR_down","BSK_3C_ICAM1_down","BSK_3C_IL8_down","BSK_3C_MCP1_down","BSK_3C_MIG_down","BSK_3C_Proliferation_down","BSK_3C_SRB_down","BSK_3C_Thrombomodulin_down","BSK_3C_Thrombomodulin_up","BSK_3C_TissueFactor_down","BSK_3C_TissueFactor_up","BSK_3C_VCAM1_down","BSK_3C_Vis_down","BSK_3C_uPAR_down","BSK_4H_Eotaxin3_down","BSK_4H_MCP1_down","BSK_4H_Pselectin_down","BSK_4H_Pselectin_up","BSK_4H_SRB_down","BSK_4H_VCAM1_down","BSK_4H_VEGFRII_down","BSK_4H_uPAR_down","BSK_4H_uPAR_up","BSK_BE3C_HLADR_down","BSK_BE3C_IL1a_down","BSK_BE3C_IP10_down","BSK_BE3C_MIG_down","BSK_BE3C_MMP1_down","BSK_BE3C_MMP1_up","BSK_BE3C_PAI1_down","BSK_BE3C_SRB_down","BSK_BE3C_TGFb1_down","BSK_BE3C_tPA_down","BSK_BE3C_uPAR_down","BSK_BE3C_uPAR_up","BSK_BE3C_uPA_down","BSK_CASM3C_HLADR_down","BSK_CASM3C_IL6_down","BSK_CASM3C_IL6_up","BSK_CASM3C_IL8_down","BSK_CASM3C_LDLR_down","BSK_CASM3C_LDLR_up","BSK_CASM3C_MCP1_down","BSK_CASM3C_MCP1_up","BSK_CASM3C_MCSF_down","BSK_CASM3C_MCSF_up","BSK_CASM3C_MIG_down","BSK_CASM3C_Proliferation_down","BSK_CASM3C_Proliferation_up","BSK_CASM3C_SAA_down","BSK_CASM3C_SAA_up","BSK_CASM3C_SRB_down","BSK_CASM3C_Thrombomodulin_down","BSK_CASM3C_Thrombomodulin_up","BSK_CASM3C_TissueFactor_down","BSK_CASM3C_VCAM1_down","BSK_CASM3C_VCAM1_up","BSK_CASM3C_uPAR_down","BSK_CASM3C_uPAR_up","BSK_KF3CT_ICAM1_down","BSK_KF3CT_IL1a_down","BSK_KF3CT_IP10_down","BSK_KF3CT_IP10_up","BSK_KF3CT_MCP1_down","BSK_KF3CT_MCP1_up","BSK_KF3CT_MMP9_down","BSK_KF3CT_SRB_down","BSK_KF3CT_TGFb1_down","BSK_KF3CT_TIMP2_down","BSK_KF3CT_uPA_down","BSK_LPS_CD40_down","BSK_LPS_Eselectin_down","BSK_LPS_Eselectin_up","BSK_LPS_IL1a_down","BSK_LPS_IL1a_up","BSK_LPS_IL8_down","BSK_LPS_IL8_up","BSK_LPS_MCP1_down","BSK_LPS_MCSF_down","BSK_LPS_PGE2_down","BSK_LPS_PGE2_up","BSK_LPS_SRB_down","BSK_LPS_TNFa_down","BSK_LPS_TNFa_up","BSK_LPS_TissueFactor_down","BSK_LPS_TissueFactor_up","BSK_LPS_VCAM1_down","BSK_SAg_CD38_down","BSK_SAg_CD40_down","BSK_SAg_CD69_down","BSK_SAg_Eselectin_down","BSK_SAg_Eselectin_up","BSK_SAg_IL8_down","BSK_SAg_IL8_up","BSK_SAg_MCP1_down","BSK_SAg_MIG_down","BSK_SAg_PBMCCytotoxicity_down","BSK_SAg_PBMCCytotoxicity_up","BSK_SAg_Proliferation_down","BSK_SAg_SRB_down","BSK_hDFCGF_CollagenIII_down","BSK_hDFCGF_EGFR_down","BSK_hDFCGF_EGFR_up","BSK_hDFCGF_IL8_down","BSK_hDFCGF_IP10_down","BSK_hDFCGF_MCSF_down","BSK_hDFCGF_MIG_down","BSK_hDFCGF_MMP1_down","BSK_hDFCGF_MMP1_up","BSK_hDFCGF_PAI1_down","BSK_hDFCGF_Proliferation_down","BSK_hDFCGF_SRB_down","BSK_hDFCGF_TIMP1_down","BSK_hDFCGF_VCAM1_down","CEETOX_H295R_11DCORT_dn","CEETOX_H295R_ANDR_dn","CEETOX_H295R_CORTISOL_dn","CEETOX_H295R_DOC_dn","CEETOX_H295R_DOC_up","CEETOX_H295R_ESTRADIOL_dn","CEETOX_H295R_ESTRADIOL_up","CEETOX_H295R_ESTRONE_dn","CEETOX_H295R_ESTRONE_up","CEETOX_H295R_OHPREG_up","CEETOX_H295R_OHPROG_dn","CEETOX_H295R_OHPROG_up","CEETOX_H295R_PROG_up","CEETOX_H295R_TESTO_dn","CLD_ABCB1_48hr","CLD_ABCG2_48hr","CLD_CYP1A1_24hr","CLD_CYP1A1_48hr","CLD_CYP1A1_6hr","CLD_CYP1A2_24hr","CLD_CYP1A2_48hr","CLD_CYP1A2_6hr","CLD_CYP2B6_24hr","CLD_CYP2B6_48hr","CLD_CYP2B6_6hr","CLD_CYP3A4_24hr","CLD_CYP3A4_48hr","CLD_CYP3A4_6hr","CLD_GSTA2_48hr","CLD_SULT2A_24hr","CLD_SULT2A_48hr","CLD_UGT1A1_24hr","CLD_UGT1A1_48hr","NCCT_HEK293T_CellTiterGLO","NCCT_QuantiLum_inhib_2_dn","NCCT_QuantiLum_inhib_dn","NCCT_TPO_AUR_dn","NCCT_TPO_GUA_dn","NHEERL_ZF_144hpf_TERATOSCORE_up",
-                        "NVS_ADME_hCYP19A1","NVS_ADME_hCYP1A1","NVS_ADME_hCYP1A2","NVS_ADME_hCYP2A6","NVS_ADME_hCYP2B6","NVS_ADME_hCYP2C19","NVS_ADME_hCYP2C9","NVS_ADME_hCYP2D6","NVS_ADME_hCYP3A4","NVS_ADME_hCYP4F12","NVS_ADME_rCYP2C12","NVS_ENZ_hAChE","NVS_ENZ_hAMPKa1","NVS_ENZ_hAurA","NVS_ENZ_hBACE","NVS_ENZ_hCASP5","NVS_ENZ_hCK1D","NVS_ENZ_hDUSP3","NVS_ENZ_hES","NVS_ENZ_hElastase","NVS_ENZ_hFGFR1","NVS_ENZ_hGSK3b","NVS_ENZ_hMMP1","NVS_ENZ_hMMP13","NVS_ENZ_hMMP2","NVS_ENZ_hMMP3","NVS_ENZ_hMMP7","NVS_ENZ_hMMP9","NVS_ENZ_hPDE10","NVS_ENZ_hPDE4A1","NVS_ENZ_hPDE5","NVS_ENZ_hPI3Ka","NVS_ENZ_hPTEN","NVS_ENZ_hPTPN11","NVS_ENZ_hPTPN12","NVS_ENZ_hPTPN13","NVS_ENZ_hPTPN9","NVS_ENZ_hPTPRC","NVS_ENZ_hSIRT1","NVS_ENZ_hSIRT2","NVS_ENZ_hTrkA","NVS_ENZ_hVEGFR2","NVS_ENZ_oCOX1","NVS_ENZ_oCOX2","NVS_ENZ_rAChE","NVS_ENZ_rCNOS","NVS_ENZ_rMAOAC","NVS_ENZ_rMAOAP","NVS_ENZ_rMAOBC","NVS_ENZ_rMAOBP","NVS_ENZ_rabI2C","NVS_GPCR_bAdoR_NonSelective","NVS_GPCR_bDR_NonSelective","NVS_GPCR_g5HT4","NVS_GPCR_gH2","NVS_GPCR_gLTB4","NVS_GPCR_gLTD4","NVS_GPCR_gMPeripheral_NonSelective","NVS_GPCR_gOpiateK","NVS_GPCR_h5HT2A","NVS_GPCR_h5HT5A","NVS_GPCR_h5HT6","NVS_GPCR_h5HT7","NVS_GPCR_hAT1","NVS_GPCR_hAdoRA1","NVS_GPCR_hAdoRA2a","NVS_GPCR_hAdra2A","NVS_GPCR_hAdra2C","NVS_GPCR_hAdrb1","NVS_GPCR_hAdrb2","NVS_GPCR_hAdrb3","NVS_GPCR_hDRD1","NVS_GPCR_hDRD2s","NVS_GPCR_hDRD4.4","NVS_GPCR_hH1","NVS_GPCR_hLTB4_BLT1","NVS_GPCR_hM1","NVS_GPCR_hM2","NVS_GPCR_hM3","NVS_GPCR_hM4","NVS_GPCR_hNK2","NVS_GPCR_hOpiate_D1","NVS_GPCR_hOpiate_mu","NVS_GPCR_hTXA2","NVS_GPCR_p5HT2C","NVS_GPCR_r5HT1_NonSelective","NVS_GPCR_r5HT_NonSelective","NVS_GPCR_rAdra1B","NVS_GPCR_rAdra1_NonSelective","NVS_GPCR_rAdra2_NonSelective","NVS_GPCR_rAdrb_NonSelective","NVS_GPCR_rNK1","NVS_GPCR_rNK3","NVS_GPCR_rOpiate_NonSelective","NVS_GPCR_rOpiate_NonSelectiveNa","NVS_GPCR_rSST","NVS_GPCR_rTRH","NVS_GPCR_rV1","NVS_GPCR_rabPAF","NVS_GPCR_rmAdra2B","NVS_IC_hKhERGCh","NVS_IC_rCaBTZCHL","NVS_IC_rCaDHPRCh_L","NVS_IC_rNaCh_site2","NVS_LGIC_bGABARa1","NVS_LGIC_h5HT3","NVS_LGIC_hNNR_NBungSens","NVS_LGIC_rGABAR_NonSelective","NVS_LGIC_rNNR_BungSens","NVS_MP_hPBR","NVS_MP_rPBR","NVS_NR_bER","NVS_NR_bPR","NVS_NR_cAR","NVS_NR_hAR","NVS_NR_hCAR_Antagonist","NVS_NR_hER","NVS_NR_hFXR_Agonist","NVS_NR_hFXR_Antagonist","NVS_NR_hGR","NVS_NR_hPPARa","NVS_NR_hPPARg","NVS_NR_hPR","NVS_NR_hPXR","NVS_NR_hRAR_Antagonist","NVS_NR_hRARa_Agonist","NVS_NR_hTRa_Antagonist","NVS_NR_mERa","NVS_NR_rAR","NVS_NR_rMR","NVS_OR_gSIGMA_NonSelective","NVS_TR_gDAT","NVS_TR_hAdoT","NVS_TR_hDAT","NVS_TR_hNET","NVS_TR_hSERT","NVS_TR_rNET","NVS_TR_rSERT","NVS_TR_rVMAT2","OT_AR_ARELUC_AG_1440","OT_AR_ARSRC1_0480","OT_AR_ARSRC1_0960","OT_ER_ERaERa_0480","OT_ER_ERaERa_1440","OT_ER_ERaERb_0480","OT_ER_ERaERb_1440","OT_ER_ERbERb_0480","OT_ER_ERbERb_1440","OT_ERa_EREGFP_0120","OT_ERa_EREGFP_0480","OT_FXR_FXRSRC1_0480","OT_FXR_FXRSRC1_1440","OT_NURR1_NURR1RXRa_0480","OT_NURR1_NURR1RXRa_1440","TOX21_ARE_BLA_Agonist_ch1","TOX21_ARE_BLA_Agonist_ch2",
-                        "TOX21_ARE_BLA_agonist_ratio","TOX21_ARE_BLA_agonist_viability","TOX21_AR_BLA_Agonist_ch1","TOX21_AR_BLA_Agonist_ch2","TOX21_AR_BLA_Agonist_ratio","TOX21_AR_BLA_Antagonist_ch1","TOX21_AR_BLA_Antagonist_ch2","TOX21_AR_BLA_Antagonist_ratio","TOX21_AR_BLA_Antagonist_viability","TOX21_AR_LUC_MDAKB2_Agonist","TOX21_AR_LUC_MDAKB2_Antagonist","TOX21_AR_LUC_MDAKB2_Antagonist2","TOX21_AhR_LUC_Agonist","TOX21_Aromatase_Inhibition","TOX21_AutoFluor_HEK293_Cell_blue","TOX21_AutoFluor_HEK293_Media_blue","TOX21_AutoFluor_HEPG2_Cell_blue","TOX21_AutoFluor_HEPG2_Cell_green","TOX21_AutoFluor_HEPG2_Media_blue","TOX21_AutoFluor_HEPG2_Media_green","TOX21_ELG1_LUC_Agonist","TOX21_ERa_BLA_Agonist_ch1","TOX21_ERa_BLA_Agonist_ch2","TOX21_ERa_BLA_Agonist_ratio","TOX21_ERa_BLA_Antagonist_ch1","TOX21_ERa_BLA_Antagonist_ch2","TOX21_ERa_BLA_Antagonist_ratio","TOX21_ERa_BLA_Antagonist_viability","TOX21_ERa_LUC_BG1_Agonist","TOX21_ERa_LUC_BG1_Antagonist","TOX21_ESRE_BLA_ch1","TOX21_ESRE_BLA_ch2","TOX21_ESRE_BLA_ratio","TOX21_ESRE_BLA_viability","TOX21_FXR_BLA_Antagonist_ch1","TOX21_FXR_BLA_Antagonist_ch2","TOX21_FXR_BLA_agonist_ch2","TOX21_FXR_BLA_agonist_ratio","TOX21_FXR_BLA_antagonist_ratio","TOX21_FXR_BLA_antagonist_viability","TOX21_GR_BLA_Agonist_ch1","TOX21_GR_BLA_Agonist_ch2","TOX21_GR_BLA_Agonist_ratio","TOX21_GR_BLA_Antagonist_ch2","TOX21_GR_BLA_Antagonist_ratio","TOX21_GR_BLA_Antagonist_viability","TOX21_HSE_BLA_agonist_ch1","TOX21_HSE_BLA_agonist_ch2","TOX21_HSE_BLA_agonist_ratio","TOX21_HSE_BLA_agonist_viability","TOX21_MMP_ratio_down","TOX21_MMP_ratio_up","TOX21_MMP_viability","TOX21_NFkB_BLA_agonist_ch1","TOX21_NFkB_BLA_agonist_ch2","TOX21_NFkB_BLA_agonist_ratio","TOX21_NFkB_BLA_agonist_viability","TOX21_PPARd_BLA_Agonist_viability","TOX21_PPARd_BLA_Antagonist_ch1","TOX21_PPARd_BLA_agonist_ch1","TOX21_PPARd_BLA_agonist_ch2","TOX21_PPARd_BLA_agonist_ratio","TOX21_PPARd_BLA_antagonist_ratio","TOX21_PPARd_BLA_antagonist_viability","TOX21_PPARg_BLA_Agonist_ch1","TOX21_PPARg_BLA_Agonist_ch2","TOX21_PPARg_BLA_Agonist_ratio","TOX21_PPARg_BLA_Antagonist_ch1","TOX21_PPARg_BLA_antagonist_ratio","TOX21_PPARg_BLA_antagonist_viability","TOX21_TR_LUC_GH3_Agonist","TOX21_TR_LUC_GH3_Antagonist","TOX21_VDR_BLA_Agonist_viability","TOX21_VDR_BLA_Antagonist_ch1","TOX21_VDR_BLA_agonist_ch2","TOX21_VDR_BLA_agonist_ratio","TOX21_VDR_BLA_antagonist_ratio","TOX21_VDR_BLA_antagonist_viability","TOX21_p53_BLA_p1_ch1","TOX21_p53_BLA_p1_ch2","TOX21_p53_BLA_p1_ratio","TOX21_p53_BLA_p1_viability","TOX21_p53_BLA_p2_ch1","TOX21_p53_BLA_p2_ch2","TOX21_p53_BLA_p2_ratio","TOX21_p53_BLA_p2_viability","TOX21_p53_BLA_p3_ch1","TOX21_p53_BLA_p3_ch2","TOX21_p53_BLA_p3_ratio","TOX21_p53_BLA_p3_viability","TOX21_p53_BLA_p4_ch1","TOX21_p53_BLA_p4_ch2","TOX21_p53_BLA_p4_ratio","TOX21_p53_BLA_p4_viability","TOX21_p53_BLA_p5_ch1","TOX21_p53_BLA_p5_ch2","TOX21_p53_BLA_p5_ratio","TOX21_p53_BLA_p5_viability","Tanguay_ZF_120hpf_AXIS_up","Tanguay_ZF_120hpf_ActivityScore","Tanguay_ZF_120hpf_BRAI_up","Tanguay_ZF_120hpf_CFIN_up","Tanguay_ZF_120hpf_CIRC_up","Tanguay_ZF_120hpf_EYE_up","Tanguay_ZF_120hpf_JAW_up","Tanguay_ZF_120hpf_MORT_up","Tanguay_ZF_120hpf_OTIC_up","Tanguay_ZF_120hpf_PE_up","Tanguay_ZF_120hpf_PFIN_up","Tanguay_ZF_120hpf_PIG_up","Tanguay_ZF_120hpf_SNOU_up","Tanguay_ZF_120hpf_SOMI_up","Tanguay_ZF_120hpf_SWIM_up","Tanguay_ZF_120hpf_TRUN_up","Tanguay_ZF_120hpf_TR_up","Tanguay_ZF_120hpf_YSE_up"],
+            "muv": [
+                "MUV-466",
+                "MUV-548",
+                "MUV-600",
+                "MUV-644",
+                "MUV-652",
+                "MUV-689",
+                "MUV-692",
+                "MUV-712",
+                "MUV-713",
+                "MUV-733",
+                "MUV-737",
+                "MUV-810",
+                "MUV-832",
+                "MUV-846",
+                "MUV-852",
+                "MUV-858",
+                "MUV-859",
+            ],
+            "pcba": [
+                "PCBA-1030",
+                "PCBA-1379",
+                "PCBA-1452",
+                "PCBA-1454",
+                "PCBA-1457",
+                "PCBA-1458",
+                "PCBA-1460",
+                "PCBA-1461",
+                "PCBA-1468",
+                "PCBA-1469",
+                "PCBA-1471",
+                "PCBA-1479",
+                "PCBA-1631",
+                "PCBA-1634",
+                "PCBA-1688",
+                "PCBA-1721",
+                "PCBA-2100",
+                "PCBA-2101",
+                "PCBA-2147",
+                "PCBA-2242",
+                "PCBA-2326",
+                "PCBA-2451",
+                "PCBA-2517",
+                "PCBA-2528",
+                "PCBA-2546",
+                "PCBA-2549",
+                "PCBA-2551",
+                "PCBA-2662",
+                "PCBA-2675",
+                "PCBA-2676",
+                "PCBA-411",
+                "PCBA-463254",
+                "PCBA-485281",
+                "PCBA-485290",
+                "PCBA-485294",
+                "PCBA-485297",
+                "PCBA-485313",
+                "PCBA-485314",
+                "PCBA-485341",
+                "PCBA-485349",
+                "PCBA-485353",
+                "PCBA-485360",
+                "PCBA-485364",
+                "PCBA-485367",
+                "PCBA-492947",
+                "PCBA-493208",
+                "PCBA-504327",
+                "PCBA-504332",
+                "PCBA-504333",
+                "PCBA-504339",
+                "PCBA-504444",
+                "PCBA-504466",
+                "PCBA-504467",
+                "PCBA-504706",
+                "PCBA-504842",
+                "PCBA-504845",
+                "PCBA-504847",
+                "PCBA-504891",
+                "PCBA-540276",
+                "PCBA-540317",
+                "PCBA-588342",
+                "PCBA-588453",
+                "PCBA-588456",
+                "PCBA-588579",
+                "PCBA-588590",
+                "PCBA-588591",
+                "PCBA-588795",
+                "PCBA-588855",
+                "PCBA-602179",
+                "PCBA-602233",
+                "PCBA-602310",
+                "PCBA-602313",
+                "PCBA-602332",
+                "PCBA-624170",
+                "PCBA-624171",
+                "PCBA-624173",
+                "PCBA-624202",
+                "PCBA-624246",
+                "PCBA-624287",
+                "PCBA-624288",
+                "PCBA-624291",
+                "PCBA-624296",
+                "PCBA-624297",
+                "PCBA-624417",
+                "PCBA-651635",
+                "PCBA-651644",
+                "PCBA-651768",
+                "PCBA-651965",
+                "PCBA-652025",
+                "PCBA-652104",
+                "PCBA-652105",
+                "PCBA-652106",
+                "PCBA-686970",
+                "PCBA-686978",
+                "PCBA-686979",
+                "PCBA-720504",
+                "PCBA-720532",
+                "PCBA-720542",
+                "PCBA-720551",
+                "PCBA-720553",
+                "PCBA-720579",
+                "PCBA-720580",
+                "PCBA-720707",
+                "PCBA-720708",
+                "PCBA-720709",
+                "PCBA-720711",
+                "PCBA-743255",
+                "PCBA-743266",
+                "PCBA-875",
+                "PCBA-881",
+                "PCBA-883",
+                "PCBA-884",
+                "PCBA-885",
+                "PCBA-887",
+                "PCBA-891",
+                "PCBA-899",
+                "PCBA-902",
+                "PCBA-903",
+                "PCBA-904",
+                "PCBA-912",
+                "PCBA-914",
+                "PCBA-915",
+                "PCBA-924",
+                "PCBA-925",
+                "PCBA-926",
+                "PCBA-927",
+                "PCBA-938",
+                "PCBA-995",
+            ],
+            "sider": [
+                "Hepatobiliary disorders",
+                "Metabolism and nutrition disorders",
+                "Product issues",
+                "Eye disorders",
+                "Investigations",
+                "Musculoskeletal and connective tissue disorders",
+                "Gastrointestinal disorders",
+                "Social circumstances",
+                "Immune system disorders",
+                "Reproductive system and breast disorders",
+                "Neoplasms benign, malignant and unspecified (incl cysts and polyps)",
+                "General disorders and administration site conditions",
+                "Endocrine disorders",
+                "Surgical and medical procedures",
+                "Vascular disorders",
+                "Blood and lymphatic system disorders",
+                "Skin and subcutaneous tissue disorders",
+                "Congenital, familial and genetic disorders",
+                "Infections and infestations",
+                "Respiratory, thoracic and mediastinal disorders",
+                "Psychiatric disorders",
+                "Renal and urinary disorders",
+                "Pregnancy, puerperium and perinatal conditions",
+                "Ear and labyrinth disorders",
+                "Cardiac disorders",
+                "Nervous system disorders",
+                "Injury, poisoning and procedural complications",
+            ],
+            "tox21": [
+                "NR-AR",
+                "NR-AR-LBD",
+                "NR-AhR",
+                "NR-Aromatase",
+                "NR-ER",
+                "NR-ER-LBD",
+                "NR-PPAR-gamma",
+                "SR-ARE",
+                "SR-ATAD5",
+                "SR-HSE",
+                "SR-MMP",
+                "SR-p53",
+            ],
+            "toxcast": [
+                "ACEA_T47D_80hr_Negative",
+                "ACEA_T47D_80hr_Positive",
+                "APR_HepG2_CellCycleArrest_24h_dn",
+                "APR_HepG2_CellCycleArrest_24h_up",
+                "APR_HepG2_CellCycleArrest_72h_dn",
+                "APR_HepG2_CellLoss_24h_dn",
+                "APR_HepG2_CellLoss_72h_dn",
+                "APR_HepG2_MicrotubuleCSK_24h_dn",
+                "APR_HepG2_MicrotubuleCSK_24h_up",
+                "APR_HepG2_MicrotubuleCSK_72h_dn",
+                "APR_HepG2_MicrotubuleCSK_72h_up",
+                "APR_HepG2_MitoMass_24h_dn",
+                "APR_HepG2_MitoMass_24h_up",
+                "APR_HepG2_MitoMass_72h_dn",
+                "APR_HepG2_MitoMass_72h_up",
+                "APR_HepG2_MitoMembPot_1h_dn",
+                "APR_HepG2_MitoMembPot_24h_dn",
+                "APR_HepG2_MitoMembPot_72h_dn",
+                "APR_HepG2_MitoticArrest_24h_up",
+                "APR_HepG2_MitoticArrest_72h_up",
+                "APR_HepG2_NuclearSize_24h_dn",
+                "APR_HepG2_NuclearSize_72h_dn",
+                "APR_HepG2_NuclearSize_72h_up",
+                "APR_HepG2_OxidativeStress_24h_up",
+                "APR_HepG2_OxidativeStress_72h_up",
+                "APR_HepG2_StressKinase_1h_up",
+                "APR_HepG2_StressKinase_24h_up",
+                "APR_HepG2_StressKinase_72h_up",
+                "APR_HepG2_p53Act_24h_up",
+                "APR_HepG2_p53Act_72h_up",
+                "APR_Hepat_Apoptosis_24hr_up",
+                "APR_Hepat_Apoptosis_48hr_up",
+                "APR_Hepat_CellLoss_24hr_dn",
+                "APR_Hepat_CellLoss_48hr_dn",
+                "APR_Hepat_DNADamage_24hr_up",
+                "APR_Hepat_DNADamage_48hr_up",
+                "APR_Hepat_DNATexture_24hr_up",
+                "APR_Hepat_DNATexture_48hr_up",
+                "APR_Hepat_MitoFxnI_1hr_dn",
+                "APR_Hepat_MitoFxnI_24hr_dn",
+                "APR_Hepat_MitoFxnI_48hr_dn",
+                "APR_Hepat_NuclearSize_24hr_dn",
+                "APR_Hepat_NuclearSize_48hr_dn",
+                "APR_Hepat_Steatosis_24hr_up",
+                "APR_Hepat_Steatosis_48hr_up",
+                "ATG_AP_1_CIS_dn",
+                "ATG_AP_1_CIS_up",
+                "ATG_AP_2_CIS_dn",
+                "ATG_AP_2_CIS_up",
+                "ATG_AR_TRANS_dn",
+                "ATG_AR_TRANS_up",
+                "ATG_Ahr_CIS_dn",
+                "ATG_Ahr_CIS_up",
+                "ATG_BRE_CIS_dn",
+                "ATG_BRE_CIS_up",
+                "ATG_CAR_TRANS_dn",
+                "ATG_CAR_TRANS_up",
+                "ATG_CMV_CIS_dn",
+                "ATG_CMV_CIS_up",
+                "ATG_CRE_CIS_dn",
+                "ATG_CRE_CIS_up",
+                "ATG_C_EBP_CIS_dn",
+                "ATG_C_EBP_CIS_up",
+                "ATG_DR4_LXR_CIS_dn",
+                "ATG_DR4_LXR_CIS_up",
+                "ATG_DR5_CIS_dn",
+                "ATG_DR5_CIS_up",
+                "ATG_E2F_CIS_dn",
+                "ATG_E2F_CIS_up",
+                "ATG_EGR_CIS_up",
+                "ATG_ERE_CIS_dn",
+                "ATG_ERE_CIS_up",
+                "ATG_ERRa_TRANS_dn",
+                "ATG_ERRg_TRANS_dn",
+                "ATG_ERRg_TRANS_up",
+                "ATG_ERa_TRANS_up",
+                "ATG_E_Box_CIS_dn",
+                "ATG_E_Box_CIS_up",
+                "ATG_Ets_CIS_dn",
+                "ATG_Ets_CIS_up",
+                "ATG_FXR_TRANS_up",
+                "ATG_FoxA2_CIS_dn",
+                "ATG_FoxA2_CIS_up",
+                "ATG_FoxO_CIS_dn",
+                "ATG_FoxO_CIS_up",
+                "ATG_GAL4_TRANS_dn",
+                "ATG_GATA_CIS_dn",
+                "ATG_GATA_CIS_up",
+                "ATG_GLI_CIS_dn",
+                "ATG_GLI_CIS_up",
+                "ATG_GRE_CIS_dn",
+                "ATG_GRE_CIS_up",
+                "ATG_GR_TRANS_dn",
+                "ATG_GR_TRANS_up",
+                "ATG_HIF1a_CIS_dn",
+                "ATG_HIF1a_CIS_up",
+                "ATG_HNF4a_TRANS_dn",
+                "ATG_HNF4a_TRANS_up",
+                "ATG_HNF6_CIS_dn",
+                "ATG_HNF6_CIS_up",
+                "ATG_HSE_CIS_dn",
+                "ATG_HSE_CIS_up",
+                "ATG_IR1_CIS_dn",
+                "ATG_IR1_CIS_up",
+                "ATG_ISRE_CIS_dn",
+                "ATG_ISRE_CIS_up",
+                "ATG_LXRa_TRANS_dn",
+                "ATG_LXRa_TRANS_up",
+                "ATG_LXRb_TRANS_dn",
+                "ATG_LXRb_TRANS_up",
+                "ATG_MRE_CIS_up",
+                "ATG_M_06_TRANS_up",
+                "ATG_M_19_CIS_dn",
+                "ATG_M_19_TRANS_dn",
+                "ATG_M_19_TRANS_up",
+                "ATG_M_32_CIS_dn",
+                "ATG_M_32_CIS_up",
+                "ATG_M_32_TRANS_dn",
+                "ATG_M_32_TRANS_up",
+                "ATG_M_61_TRANS_up",
+                "ATG_Myb_CIS_dn",
+                "ATG_Myb_CIS_up",
+                "ATG_Myc_CIS_dn",
+                "ATG_Myc_CIS_up",
+                "ATG_NFI_CIS_dn",
+                "ATG_NFI_CIS_up",
+                "ATG_NF_kB_CIS_dn",
+                "ATG_NF_kB_CIS_up",
+                "ATG_NRF1_CIS_dn",
+                "ATG_NRF1_CIS_up",
+                "ATG_NRF2_ARE_CIS_dn",
+                "ATG_NRF2_ARE_CIS_up",
+                "ATG_NURR1_TRANS_dn",
+                "ATG_NURR1_TRANS_up",
+                "ATG_Oct_MLP_CIS_dn",
+                "ATG_Oct_MLP_CIS_up",
+                "ATG_PBREM_CIS_dn",
+                "ATG_PBREM_CIS_up",
+                "ATG_PPARa_TRANS_dn",
+                "ATG_PPARa_TRANS_up",
+                "ATG_PPARd_TRANS_up",
+                "ATG_PPARg_TRANS_up",
+                "ATG_PPRE_CIS_dn",
+                "ATG_PPRE_CIS_up",
+                "ATG_PXRE_CIS_dn",
+                "ATG_PXRE_CIS_up",
+                "ATG_PXR_TRANS_dn",
+                "ATG_PXR_TRANS_up",
+                "ATG_Pax6_CIS_up",
+                "ATG_RARa_TRANS_dn",
+                "ATG_RARa_TRANS_up",
+                "ATG_RARb_TRANS_dn",
+                "ATG_RARb_TRANS_up",
+                "ATG_RARg_TRANS_dn",
+                "ATG_RARg_TRANS_up",
+                "ATG_RORE_CIS_dn",
+                "ATG_RORE_CIS_up",
+                "ATG_RORb_TRANS_dn",
+                "ATG_RORg_TRANS_dn",
+                "ATG_RORg_TRANS_up",
+                "ATG_RXRa_TRANS_dn",
+                "ATG_RXRa_TRANS_up",
+                "ATG_RXRb_TRANS_dn",
+                "ATG_RXRb_TRANS_up",
+                "ATG_SREBP_CIS_dn",
+                "ATG_SREBP_CIS_up",
+                "ATG_STAT3_CIS_dn",
+                "ATG_STAT3_CIS_up",
+                "ATG_Sox_CIS_dn",
+                "ATG_Sox_CIS_up",
+                "ATG_Sp1_CIS_dn",
+                "ATG_Sp1_CIS_up",
+                "ATG_TAL_CIS_dn",
+                "ATG_TAL_CIS_up",
+                "ATG_TA_CIS_dn",
+                "ATG_TA_CIS_up",
+                "ATG_TCF_b_cat_CIS_dn",
+                "ATG_TCF_b_cat_CIS_up",
+                "ATG_TGFb_CIS_dn",
+                "ATG_TGFb_CIS_up",
+                "ATG_THRa1_TRANS_dn",
+                "ATG_THRa1_TRANS_up",
+                "ATG_VDRE_CIS_dn",
+                "ATG_VDRE_CIS_up",
+                "ATG_VDR_TRANS_dn",
+                "ATG_VDR_TRANS_up",
+                "ATG_XTT_Cytotoxicity_up",
+                "ATG_Xbp1_CIS_dn",
+                "ATG_Xbp1_CIS_up",
+                "ATG_p53_CIS_dn",
+                "ATG_p53_CIS_up",
+                "BSK_3C_Eselectin_down",
+                "BSK_3C_HLADR_down",
+                "BSK_3C_ICAM1_down",
+                "BSK_3C_IL8_down",
+                "BSK_3C_MCP1_down",
+                "BSK_3C_MIG_down",
+                "BSK_3C_Proliferation_down",
+                "BSK_3C_SRB_down",
+                "BSK_3C_Thrombomodulin_down",
+                "BSK_3C_Thrombomodulin_up",
+                "BSK_3C_TissueFactor_down",
+                "BSK_3C_TissueFactor_up",
+                "BSK_3C_VCAM1_down",
+                "BSK_3C_Vis_down",
+                "BSK_3C_uPAR_down",
+                "BSK_4H_Eotaxin3_down",
+                "BSK_4H_MCP1_down",
+                "BSK_4H_Pselectin_down",
+                "BSK_4H_Pselectin_up",
+                "BSK_4H_SRB_down",
+                "BSK_4H_VCAM1_down",
+                "BSK_4H_VEGFRII_down",
+                "BSK_4H_uPAR_down",
+                "BSK_4H_uPAR_up",
+                "BSK_BE3C_HLADR_down",
+                "BSK_BE3C_IL1a_down",
+                "BSK_BE3C_IP10_down",
+                "BSK_BE3C_MIG_down",
+                "BSK_BE3C_MMP1_down",
+                "BSK_BE3C_MMP1_up",
+                "BSK_BE3C_PAI1_down",
+                "BSK_BE3C_SRB_down",
+                "BSK_BE3C_TGFb1_down",
+                "BSK_BE3C_tPA_down",
+                "BSK_BE3C_uPAR_down",
+                "BSK_BE3C_uPAR_up",
+                "BSK_BE3C_uPA_down",
+                "BSK_CASM3C_HLADR_down",
+                "BSK_CASM3C_IL6_down",
+                "BSK_CASM3C_IL6_up",
+                "BSK_CASM3C_IL8_down",
+                "BSK_CASM3C_LDLR_down",
+                "BSK_CASM3C_LDLR_up",
+                "BSK_CASM3C_MCP1_down",
+                "BSK_CASM3C_MCP1_up",
+                "BSK_CASM3C_MCSF_down",
+                "BSK_CASM3C_MCSF_up",
+                "BSK_CASM3C_MIG_down",
+                "BSK_CASM3C_Proliferation_down",
+                "BSK_CASM3C_Proliferation_up",
+                "BSK_CASM3C_SAA_down",
+                "BSK_CASM3C_SAA_up",
+                "BSK_CASM3C_SRB_down",
+                "BSK_CASM3C_Thrombomodulin_down",
+                "BSK_CASM3C_Thrombomodulin_up",
+                "BSK_CASM3C_TissueFactor_down",
+                "BSK_CASM3C_VCAM1_down",
+                "BSK_CASM3C_VCAM1_up",
+                "BSK_CASM3C_uPAR_down",
+                "BSK_CASM3C_uPAR_up",
+                "BSK_KF3CT_ICAM1_down",
+                "BSK_KF3CT_IL1a_down",
+                "BSK_KF3CT_IP10_down",
+                "BSK_KF3CT_IP10_up",
+                "BSK_KF3CT_MCP1_down",
+                "BSK_KF3CT_MCP1_up",
+                "BSK_KF3CT_MMP9_down",
+                "BSK_KF3CT_SRB_down",
+                "BSK_KF3CT_TGFb1_down",
+                "BSK_KF3CT_TIMP2_down",
+                "BSK_KF3CT_uPA_down",
+                "BSK_LPS_CD40_down",
+                "BSK_LPS_Eselectin_down",
+                "BSK_LPS_Eselectin_up",
+                "BSK_LPS_IL1a_down",
+                "BSK_LPS_IL1a_up",
+                "BSK_LPS_IL8_down",
+                "BSK_LPS_IL8_up",
+                "BSK_LPS_MCP1_down",
+                "BSK_LPS_MCSF_down",
+                "BSK_LPS_PGE2_down",
+                "BSK_LPS_PGE2_up",
+                "BSK_LPS_SRB_down",
+                "BSK_LPS_TNFa_down",
+                "BSK_LPS_TNFa_up",
+                "BSK_LPS_TissueFactor_down",
+                "BSK_LPS_TissueFactor_up",
+                "BSK_LPS_VCAM1_down",
+                "BSK_SAg_CD38_down",
+                "BSK_SAg_CD40_down",
+                "BSK_SAg_CD69_down",
+                "BSK_SAg_Eselectin_down",
+                "BSK_SAg_Eselectin_up",
+                "BSK_SAg_IL8_down",
+                "BSK_SAg_IL8_up",
+                "BSK_SAg_MCP1_down",
+                "BSK_SAg_MIG_down",
+                "BSK_SAg_PBMCCytotoxicity_down",
+                "BSK_SAg_PBMCCytotoxicity_up",
+                "BSK_SAg_Proliferation_down",
+                "BSK_SAg_SRB_down",
+                "BSK_hDFCGF_CollagenIII_down",
+                "BSK_hDFCGF_EGFR_down",
+                "BSK_hDFCGF_EGFR_up",
+                "BSK_hDFCGF_IL8_down",
+                "BSK_hDFCGF_IP10_down",
+                "BSK_hDFCGF_MCSF_down",
+                "BSK_hDFCGF_MIG_down",
+                "BSK_hDFCGF_MMP1_down",
+                "BSK_hDFCGF_MMP1_up",
+                "BSK_hDFCGF_PAI1_down",
+                "BSK_hDFCGF_Proliferation_down",
+                "BSK_hDFCGF_SRB_down",
+                "BSK_hDFCGF_TIMP1_down",
+                "BSK_hDFCGF_VCAM1_down",
+                "CEETOX_H295R_11DCORT_dn",
+                "CEETOX_H295R_ANDR_dn",
+                "CEETOX_H295R_CORTISOL_dn",
+                "CEETOX_H295R_DOC_dn",
+                "CEETOX_H295R_DOC_up",
+                "CEETOX_H295R_ESTRADIOL_dn",
+                "CEETOX_H295R_ESTRADIOL_up",
+                "CEETOX_H295R_ESTRONE_dn",
+                "CEETOX_H295R_ESTRONE_up",
+                "CEETOX_H295R_OHPREG_up",
+                "CEETOX_H295R_OHPROG_dn",
+                "CEETOX_H295R_OHPROG_up",
+                "CEETOX_H295R_PROG_up",
+                "CEETOX_H295R_TESTO_dn",
+                "CLD_ABCB1_48hr",
+                "CLD_ABCG2_48hr",
+                "CLD_CYP1A1_24hr",
+                "CLD_CYP1A1_48hr",
+                "CLD_CYP1A1_6hr",
+                "CLD_CYP1A2_24hr",
+                "CLD_CYP1A2_48hr",
+                "CLD_CYP1A2_6hr",
+                "CLD_CYP2B6_24hr",
+                "CLD_CYP2B6_48hr",
+                "CLD_CYP2B6_6hr",
+                "CLD_CYP3A4_24hr",
+                "CLD_CYP3A4_48hr",
+                "CLD_CYP3A4_6hr",
+                "CLD_GSTA2_48hr",
+                "CLD_SULT2A_24hr",
+                "CLD_SULT2A_48hr",
+                "CLD_UGT1A1_24hr",
+                "CLD_UGT1A1_48hr",
+                "NCCT_HEK293T_CellTiterGLO",
+                "NCCT_QuantiLum_inhib_2_dn",
+                "NCCT_QuantiLum_inhib_dn",
+                "NCCT_TPO_AUR_dn",
+                "NCCT_TPO_GUA_dn",
+                "NHEERL_ZF_144hpf_TERATOSCORE_up",
+                "NVS_ADME_hCYP19A1",
+                "NVS_ADME_hCYP1A1",
+                "NVS_ADME_hCYP1A2",
+                "NVS_ADME_hCYP2A6",
+                "NVS_ADME_hCYP2B6",
+                "NVS_ADME_hCYP2C19",
+                "NVS_ADME_hCYP2C9",
+                "NVS_ADME_hCYP2D6",
+                "NVS_ADME_hCYP3A4",
+                "NVS_ADME_hCYP4F12",
+                "NVS_ADME_rCYP2C12",
+                "NVS_ENZ_hAChE",
+                "NVS_ENZ_hAMPKa1",
+                "NVS_ENZ_hAurA",
+                "NVS_ENZ_hBACE",
+                "NVS_ENZ_hCASP5",
+                "NVS_ENZ_hCK1D",
+                "NVS_ENZ_hDUSP3",
+                "NVS_ENZ_hES",
+                "NVS_ENZ_hElastase",
+                "NVS_ENZ_hFGFR1",
+                "NVS_ENZ_hGSK3b",
+                "NVS_ENZ_hMMP1",
+                "NVS_ENZ_hMMP13",
+                "NVS_ENZ_hMMP2",
+                "NVS_ENZ_hMMP3",
+                "NVS_ENZ_hMMP7",
+                "NVS_ENZ_hMMP9",
+                "NVS_ENZ_hPDE10",
+                "NVS_ENZ_hPDE4A1",
+                "NVS_ENZ_hPDE5",
+                "NVS_ENZ_hPI3Ka",
+                "NVS_ENZ_hPTEN",
+                "NVS_ENZ_hPTPN11",
+                "NVS_ENZ_hPTPN12",
+                "NVS_ENZ_hPTPN13",
+                "NVS_ENZ_hPTPN9",
+                "NVS_ENZ_hPTPRC",
+                "NVS_ENZ_hSIRT1",
+                "NVS_ENZ_hSIRT2",
+                "NVS_ENZ_hTrkA",
+                "NVS_ENZ_hVEGFR2",
+                "NVS_ENZ_oCOX1",
+                "NVS_ENZ_oCOX2",
+                "NVS_ENZ_rAChE",
+                "NVS_ENZ_rCNOS",
+                "NVS_ENZ_rMAOAC",
+                "NVS_ENZ_rMAOAP",
+                "NVS_ENZ_rMAOBC",
+                "NVS_ENZ_rMAOBP",
+                "NVS_ENZ_rabI2C",
+                "NVS_GPCR_bAdoR_NonSelective",
+                "NVS_GPCR_bDR_NonSelective",
+                "NVS_GPCR_g5HT4",
+                "NVS_GPCR_gH2",
+                "NVS_GPCR_gLTB4",
+                "NVS_GPCR_gLTD4",
+                "NVS_GPCR_gMPeripheral_NonSelective",
+                "NVS_GPCR_gOpiateK",
+                "NVS_GPCR_h5HT2A",
+                "NVS_GPCR_h5HT5A",
+                "NVS_GPCR_h5HT6",
+                "NVS_GPCR_h5HT7",
+                "NVS_GPCR_hAT1",
+                "NVS_GPCR_hAdoRA1",
+                "NVS_GPCR_hAdoRA2a",
+                "NVS_GPCR_hAdra2A",
+                "NVS_GPCR_hAdra2C",
+                "NVS_GPCR_hAdrb1",
+                "NVS_GPCR_hAdrb2",
+                "NVS_GPCR_hAdrb3",
+                "NVS_GPCR_hDRD1",
+                "NVS_GPCR_hDRD2s",
+                "NVS_GPCR_hDRD4.4",
+                "NVS_GPCR_hH1",
+                "NVS_GPCR_hLTB4_BLT1",
+                "NVS_GPCR_hM1",
+                "NVS_GPCR_hM2",
+                "NVS_GPCR_hM3",
+                "NVS_GPCR_hM4",
+                "NVS_GPCR_hNK2",
+                "NVS_GPCR_hOpiate_D1",
+                "NVS_GPCR_hOpiate_mu",
+                "NVS_GPCR_hTXA2",
+                "NVS_GPCR_p5HT2C",
+                "NVS_GPCR_r5HT1_NonSelective",
+                "NVS_GPCR_r5HT_NonSelective",
+                "NVS_GPCR_rAdra1B",
+                "NVS_GPCR_rAdra1_NonSelective",
+                "NVS_GPCR_rAdra2_NonSelective",
+                "NVS_GPCR_rAdrb_NonSelective",
+                "NVS_GPCR_rNK1",
+                "NVS_GPCR_rNK3",
+                "NVS_GPCR_rOpiate_NonSelective",
+                "NVS_GPCR_rOpiate_NonSelectiveNa",
+                "NVS_GPCR_rSST",
+                "NVS_GPCR_rTRH",
+                "NVS_GPCR_rV1",
+                "NVS_GPCR_rabPAF",
+                "NVS_GPCR_rmAdra2B",
+                "NVS_IC_hKhERGCh",
+                "NVS_IC_rCaBTZCHL",
+                "NVS_IC_rCaDHPRCh_L",
+                "NVS_IC_rNaCh_site2",
+                "NVS_LGIC_bGABARa1",
+                "NVS_LGIC_h5HT3",
+                "NVS_LGIC_hNNR_NBungSens",
+                "NVS_LGIC_rGABAR_NonSelective",
+                "NVS_LGIC_rNNR_BungSens",
+                "NVS_MP_hPBR",
+                "NVS_MP_rPBR",
+                "NVS_NR_bER",
+                "NVS_NR_bPR",
+                "NVS_NR_cAR",
+                "NVS_NR_hAR",
+                "NVS_NR_hCAR_Antagonist",
+                "NVS_NR_hER",
+                "NVS_NR_hFXR_Agonist",
+                "NVS_NR_hFXR_Antagonist",
+                "NVS_NR_hGR",
+                "NVS_NR_hPPARa",
+                "NVS_NR_hPPARg",
+                "NVS_NR_hPR",
+                "NVS_NR_hPXR",
+                "NVS_NR_hRAR_Antagonist",
+                "NVS_NR_hRARa_Agonist",
+                "NVS_NR_hTRa_Antagonist",
+                "NVS_NR_mERa",
+                "NVS_NR_rAR",
+                "NVS_NR_rMR",
+                "NVS_OR_gSIGMA_NonSelective",
+                "NVS_TR_gDAT",
+                "NVS_TR_hAdoT",
+                "NVS_TR_hDAT",
+                "NVS_TR_hNET",
+                "NVS_TR_hSERT",
+                "NVS_TR_rNET",
+                "NVS_TR_rSERT",
+                "NVS_TR_rVMAT2",
+                "OT_AR_ARELUC_AG_1440",
+                "OT_AR_ARSRC1_0480",
+                "OT_AR_ARSRC1_0960",
+                "OT_ER_ERaERa_0480",
+                "OT_ER_ERaERa_1440",
+                "OT_ER_ERaERb_0480",
+                "OT_ER_ERaERb_1440",
+                "OT_ER_ERbERb_0480",
+                "OT_ER_ERbERb_1440",
+                "OT_ERa_EREGFP_0120",
+                "OT_ERa_EREGFP_0480",
+                "OT_FXR_FXRSRC1_0480",
+                "OT_FXR_FXRSRC1_1440",
+                "OT_NURR1_NURR1RXRa_0480",
+                "OT_NURR1_NURR1RXRa_1440",
+                "TOX21_ARE_BLA_Agonist_ch1",
+                "TOX21_ARE_BLA_Agonist_ch2",
+                "TOX21_ARE_BLA_agonist_ratio",
+                "TOX21_ARE_BLA_agonist_viability",
+                "TOX21_AR_BLA_Agonist_ch1",
+                "TOX21_AR_BLA_Agonist_ch2",
+                "TOX21_AR_BLA_Agonist_ratio",
+                "TOX21_AR_BLA_Antagonist_ch1",
+                "TOX21_AR_BLA_Antagonist_ch2",
+                "TOX21_AR_BLA_Antagonist_ratio",
+                "TOX21_AR_BLA_Antagonist_viability",
+                "TOX21_AR_LUC_MDAKB2_Agonist",
+                "TOX21_AR_LUC_MDAKB2_Antagonist",
+                "TOX21_AR_LUC_MDAKB2_Antagonist2",
+                "TOX21_AhR_LUC_Agonist",
+                "TOX21_Aromatase_Inhibition",
+                "TOX21_AutoFluor_HEK293_Cell_blue",
+                "TOX21_AutoFluor_HEK293_Media_blue",
+                "TOX21_AutoFluor_HEPG2_Cell_blue",
+                "TOX21_AutoFluor_HEPG2_Cell_green",
+                "TOX21_AutoFluor_HEPG2_Media_blue",
+                "TOX21_AutoFluor_HEPG2_Media_green",
+                "TOX21_ELG1_LUC_Agonist",
+                "TOX21_ERa_BLA_Agonist_ch1",
+                "TOX21_ERa_BLA_Agonist_ch2",
+                "TOX21_ERa_BLA_Agonist_ratio",
+                "TOX21_ERa_BLA_Antagonist_ch1",
+                "TOX21_ERa_BLA_Antagonist_ch2",
+                "TOX21_ERa_BLA_Antagonist_ratio",
+                "TOX21_ERa_BLA_Antagonist_viability",
+                "TOX21_ERa_LUC_BG1_Agonist",
+                "TOX21_ERa_LUC_BG1_Antagonist",
+                "TOX21_ESRE_BLA_ch1",
+                "TOX21_ESRE_BLA_ch2",
+                "TOX21_ESRE_BLA_ratio",
+                "TOX21_ESRE_BLA_viability",
+                "TOX21_FXR_BLA_Antagonist_ch1",
+                "TOX21_FXR_BLA_Antagonist_ch2",
+                "TOX21_FXR_BLA_agonist_ch2",
+                "TOX21_FXR_BLA_agonist_ratio",
+                "TOX21_FXR_BLA_antagonist_ratio",
+                "TOX21_FXR_BLA_antagonist_viability",
+                "TOX21_GR_BLA_Agonist_ch1",
+                "TOX21_GR_BLA_Agonist_ch2",
+                "TOX21_GR_BLA_Agonist_ratio",
+                "TOX21_GR_BLA_Antagonist_ch2",
+                "TOX21_GR_BLA_Antagonist_ratio",
+                "TOX21_GR_BLA_Antagonist_viability",
+                "TOX21_HSE_BLA_agonist_ch1",
+                "TOX21_HSE_BLA_agonist_ch2",
+                "TOX21_HSE_BLA_agonist_ratio",
+                "TOX21_HSE_BLA_agonist_viability",
+                "TOX21_MMP_ratio_down",
+                "TOX21_MMP_ratio_up",
+                "TOX21_MMP_viability",
+                "TOX21_NFkB_BLA_agonist_ch1",
+                "TOX21_NFkB_BLA_agonist_ch2",
+                "TOX21_NFkB_BLA_agonist_ratio",
+                "TOX21_NFkB_BLA_agonist_viability",
+                "TOX21_PPARd_BLA_Agonist_viability",
+                "TOX21_PPARd_BLA_Antagonist_ch1",
+                "TOX21_PPARd_BLA_agonist_ch1",
+                "TOX21_PPARd_BLA_agonist_ch2",
+                "TOX21_PPARd_BLA_agonist_ratio",
+                "TOX21_PPARd_BLA_antagonist_ratio",
+                "TOX21_PPARd_BLA_antagonist_viability",
+                "TOX21_PPARg_BLA_Agonist_ch1",
+                "TOX21_PPARg_BLA_Agonist_ch2",
+                "TOX21_PPARg_BLA_Agonist_ratio",
+                "TOX21_PPARg_BLA_Antagonist_ch1",
+                "TOX21_PPARg_BLA_antagonist_ratio",
+                "TOX21_PPARg_BLA_antagonist_viability",
+                "TOX21_TR_LUC_GH3_Agonist",
+                "TOX21_TR_LUC_GH3_Antagonist",
+                "TOX21_VDR_BLA_Agonist_viability",
+                "TOX21_VDR_BLA_Antagonist_ch1",
+                "TOX21_VDR_BLA_agonist_ch2",
+                "TOX21_VDR_BLA_agonist_ratio",
+                "TOX21_VDR_BLA_antagonist_ratio",
+                "TOX21_VDR_BLA_antagonist_viability",
+                "TOX21_p53_BLA_p1_ch1",
+                "TOX21_p53_BLA_p1_ch2",
+                "TOX21_p53_BLA_p1_ratio",
+                "TOX21_p53_BLA_p1_viability",
+                "TOX21_p53_BLA_p2_ch1",
+                "TOX21_p53_BLA_p2_ch2",
+                "TOX21_p53_BLA_p2_ratio",
+                "TOX21_p53_BLA_p2_viability",
+                "TOX21_p53_BLA_p3_ch1",
+                "TOX21_p53_BLA_p3_ch2",
+                "TOX21_p53_BLA_p3_ratio",
+                "TOX21_p53_BLA_p3_viability",
+                "TOX21_p53_BLA_p4_ch1",
+                "TOX21_p53_BLA_p4_ch2",
+                "TOX21_p53_BLA_p4_ratio",
+                "TOX21_p53_BLA_p4_viability",
+                "TOX21_p53_BLA_p5_ch1",
+                "TOX21_p53_BLA_p5_ch2",
+                "TOX21_p53_BLA_p5_ratio",
+                "TOX21_p53_BLA_p5_viability",
+                "Tanguay_ZF_120hpf_AXIS_up",
+                "Tanguay_ZF_120hpf_ActivityScore",
+                "Tanguay_ZF_120hpf_BRAI_up",
+                "Tanguay_ZF_120hpf_CFIN_up",
+                "Tanguay_ZF_120hpf_CIRC_up",
+                "Tanguay_ZF_120hpf_EYE_up",
+                "Tanguay_ZF_120hpf_JAW_up",
+                "Tanguay_ZF_120hpf_MORT_up",
+                "Tanguay_ZF_120hpf_OTIC_up",
+                "Tanguay_ZF_120hpf_PE_up",
+                "Tanguay_ZF_120hpf_PFIN_up",
+                "Tanguay_ZF_120hpf_PIG_up",
+                "Tanguay_ZF_120hpf_SNOU_up",
+                "Tanguay_ZF_120hpf_SOMI_up",
+                "Tanguay_ZF_120hpf_SWIM_up",
+                "Tanguay_ZF_120hpf_TRUN_up",
+                "Tanguay_ZF_120hpf_TR_up",
+                "Tanguay_ZF_120hpf_YSE_up",
+            ],
             "delaney_esol": "measured log solubility in mols per litre",
             "freesolv": "y",
             "lipo": "exp",
-            "bace_regression": "pIC50"
+            "bace_regression": "pIC50",
         }
 
-        if datasets == ['all']:
+        if datasets == ["all"]:
             datasets = list(self.datasets_loaded.keys())
 
         dfs = {}
 
         from os import path
+
         if "geognn" in mode:
             base_path = path.join(path.expanduser("~"), f".oce/moleculenet2/")
         else:
@@ -332,14 +1173,19 @@ def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "defa
                 if not path.exists(base_path):
                     os.mkdir(base_path)
                 df = pd.read_csv(self.datasets_loaded[name])
-                df.to_csv(dataset_path, index = False)
-            if 'split' in df.columns:
-                df["split"] = df['split'].str.lower()
+                df.to_csv(dataset_path, index=False)
+            if "split" in df.columns:
+                df["split"] = df["split"].str.lower()
             dfs[name] = df
 
         self.dataset_names = datasets
         self.dfs = dfs
-        self.metrics = {name: "ROC-AUC" if self.datasets_types[name] == "classification" else "Root Mean Squared Error" for name in self.dataset_names}
+        self.metrics = {
+            name: "ROC-AUC"
+            if self.datasets_types[name] == "classification"
+            else "Root Mean Squared Error"
+            for name in self.dataset_names
+        }
         self.model_names = []
         self.model_params = []
         self.model_database = pd.DataFrame()
@@ -350,7 +1196,8 @@ def __init__(self, datasets: List[str] = ['all'], file_path = None, mode = "defa
                 d = pickle.load(f)
                 self._load(d["instance_save"])
 
-    '''Runs the benchmark dataset on each of the desired models, including datasets with multi-task requirements.'''
+    """Runs the benchmark dataset on each of the desired models, including datasets with multi-task requirements."""
+
     def run(self, models: Union[BaseModel, List[BaseModel]]):
         if not isinstance(models, list):
             models = [models]
@@ -368,12 +1215,14 @@ def run(self, models: Union[BaseModel, List[BaseModel]]):
                     if len(property_cols) > 5:
                         import random
                         from random import sample
+
                         random.seed(42)
                         property_cols = sample(property_cols, 5)
                 elif "single" in self.mode:
                     if len(property_cols) > 1:
                         import random
                         from random import sample
+
                         random.seed(42)
                         property_cols = sample(property_cols, 1)
 
@@ -381,16 +1230,23 @@ def run(self, models: Union[BaseModel, List[BaseModel]]):
                     structure_col = "smiles"
                 elif "mol" in df.columns:
                     structure_col = "mol"
-                dataset = oce.BaseDataset(data = df.to_csv(), structure_col = structure_col) + oce.CleanStructures()
+                dataset = (
+                    oce.BaseDataset(data=df.to_csv(), structure_col=structure_col)
+                    + oce.CleanStructures()
+                )
 
                 if "geognn" in self.mode:
-                    dataset = dataset + oce.gg_ScaffoldSplit(split_proportions = [0.8, 0.1, 0.1])
+                    dataset = dataset + oce.gg_ScaffoldSplit(
+                        split_proportions=[0.8, 0.1, 0.1]
+                    )
                 else:
-                    dataset = dataset + oce.ScaffoldSplit(split_proportions = [0.8, 0.1, 0.1])
+                    dataset = dataset + oce.ScaffoldSplit(
+                        split_proportions=[0.8, 0.1, 0.1]
+                    )
 
                 for property_col in tqdm(property_cols):
                     print(f"Running property {property_col}")
-                    
+
                     dataset2 = dataset.copy()
                     dataset2.property_col = property_col
                     dataset2 = dataset2 + oce.CleanStructures()
@@ -404,10 +1260,17 @@ def run(self, models: Union[BaseModel, List[BaseModel]]):
 
                     l.append(results[self.metrics[name]])
 
-                self.model_database = self.model_database.append({"Model Name": model.name,
-                    "Model Parameters": oce.pretty_params_str(model),
-                    "Evaluation": self.eval,
-                    "Dataset": name, "Metric": np.mean(l), "Full Metrics List": np.array(l)}, ignore_index=True)
+                self.model_database = self.model_database.append(
+                    {
+                        "Model Name": model.name,
+                        "Model Parameters": oce.pretty_params_str(model),
+                        "Evaluation": self.eval,
+                        "Dataset": name,
+                        "Metric": np.mean(l),
+                        "Full Metrics List": np.array(l),
+                    },
+                    ignore_index=True,
+                )
                 if not self.file_path is None:
                     oce.save(self, self.file_path)
 
@@ -417,4 +1280,5 @@ def _save(self):
 
     def _load(self, d):
         import io
+
         self.model_database = pd.read_csv(io.StringIO(d["model_database"]))
diff --git a/src/olorenchemengine/dataset.py b/src/olorenchemengine/dataset.py
index 0fada25c..b23065fd 100644
--- a/src/olorenchemengine/dataset.py
+++ b/src/olorenchemengine/dataset.py
@@ -1,6 +1,7 @@
-from typing_extensions import Self
 import warnings
 
+from typing_extensions import Self
+
 warnings.filterwarnings("ignore")
 
 from time import sleep
@@ -9,8 +10,9 @@
 
 from .base_class import *
 
+
 class BaseDataset(BaseClass):
-    """ BaseDataset for all dataset objects
+    """BaseDataset for all dataset objects
 
     BaseDataset holds its data in a Pandas DataFrame.
 
@@ -25,18 +27,20 @@ class BaseDataset(BaseClass):
     @log_arguments
     def __init__(
         self,
-        name: str =None,
-        data: str =None,
+        name: str = None,
+        data: str = None,
         structure_col: str = None,
-        property_col: str =None,
-        feature_cols: list =[],
+        property_col: str = None,
+        feature_cols: list = [],
         date_col: str = None,
         log=True,
-        **kwargs
+        **kwargs,
     ):
 
         if name is None:
-            code = base64.urlsafe_b64encode(hashlib.md5(str(data).encode("utf-8")).digest())[:16].decode("utf-8")
+            code = base64.urlsafe_b64encode(
+                hashlib.md5(str(data).encode("utf-8")).digest()
+            )[:16].decode("utf-8")
             self.name = f"Dataset-{code}"
         else:
             self.name = name
@@ -59,51 +63,57 @@ def __init__(
 
     @property
     def entire_dataset(self):
-        """ Returns the entire dataset
+        """Returns the entire dataset
 
         Returns:
             pd.DataFrame: The entire dataset
-            """
+        """
 
         return self.data[self.input_cols], self.data[self.property_col]
 
     @property
     def entire_dataset_split(self):
-        """ Returns a tuple of three elements where the first is the input train data,
+        """Returns a tuple of three elements where the first is the input train data,
         the second is the input validation data, and the third is the input test data
 
         Returns:
             tuple: (train_data, val_data, test_data)
         """
-        assert "split" in self.data.columns, "Dataset not split yet, please use a splitter or define the split column"
+        assert (
+            "split" in self.data.columns
+        ), "Dataset not split yet, please use a splitter or define the split column"
         return self.data["split"]
 
     @property
     def train_dataset(self):
-        """ Returns the train dataset
-        """
-        assert "split" in self.data.columns, "Dataset not split yet, please use a splitter or define the split column"
+        """Returns the train dataset"""
+        # assert "split" in list(
+        #     self.data.columns
+        # ), f"Dataset not split yet, please use a splitter or define the split column {list(self.data.columns)}"
         train = self.data[self.data["split"] == "train"]
         return train[self.input_cols], train[self.property_col]
 
     @property
     def valid_dataset(self):
-        """ Gives a tuple of two elements where the first is the input val data and
+        """Gives a tuple of two elements where the first is the input val data and
         the second is the property of interest
 
         Returns:
             pd.DataFrame: The validation data
 
         """
-        assert "split" in self.data.columns, "Dataset not split yet, please use a splitter or define the split column"
+        assert (
+            "split" in self.data.columns
+        ), "Dataset not split yet, please use a splitter or define the split column"
         valid = self.data[self.data["split"] == "valid"]
         return valid[self.input_cols], valid[self.property_col]
 
     @property
     def trainval_dataset(self):
-        """ Returns the train and validation dataset
-        """
-        assert "split" in self.data.columns, "Dataset not split yet, please use a splitter or define the split column"
+        """Returns the train and validation dataset"""
+        assert (
+            "split" in self.data.columns
+        ), "Dataset not split yet, please use a splitter or define the split column"
         train = self.data[self.data["split"] == "train"]
         val = self.data[self.data["split"] == "valid"]
         trainval = pd.concat([train, val])
@@ -111,23 +121,29 @@ def trainval_dataset(self):
 
     @property
     def test_dataset(self):
-        """ Gives a tuple of two elements where the first is the input test data and
+        """Gives a tuple of two elements where the first is the input test data and
         the second is the property of interest
 
         Returns:
             pd.DataFrame: The test data
 
         """
-        assert "split" in self.data.columns, "Dataset not split yet, please use a splitter or define the split column"
+        assert (
+            "split" in self.data.columns
+        ), "Dataset not split yet, please use a splitter or define the split column"
         test = self.data[self.data["split"] == "test"]
         return test[self.input_cols], test[self.property_col]
 
     @property
     def size(self):
-        return (len(self.train_dataset[0]), len(self.valid_dataset[0]), len(self.test_dataset[0]))
+        return (
+            len(self.train_dataset[0]),
+            len(self.valid_dataset[0]),
+            len(self.test_dataset[0]),
+        )
 
     def transform(self, dataset: Self):
-        """ Combines this dataset with the passed dataset object"""
+        """Combines this dataset with the passed dataset object"""
         # Rename passed dataset to homogenize with this dataset
 
     def _save(self):
@@ -139,8 +155,9 @@ def _load(self, d):
 
         self.data = pd.read_csv(StringIO(d["data_save"]))
 
+
 class BaseDatasetTransform(BaseClass):
-    """ Applies a transformation onto the inputted BaseDataset.
+    """Applies a transformation onto the inputted BaseDataset.
 
     Transformation applied as defined in the abstract method transform.
 
@@ -150,7 +167,7 @@ class BaseDatasetTransform(BaseClass):
 
     @abstractmethod
     def transform(self, dataset: BaseDataset) -> BaseDataset:
-        """ Applies a transformation onto the inputted BaseDataset.
+        """Applies a transformation onto the inputted BaseDataset.
 
         Parameters:
         dataset (BaseDataset): The dataset to transform.
@@ -163,13 +180,16 @@ def _save(self):
     def _load(self, d):
         pass
 
+
 def func(self: BaseDataset, other: BaseDatasetTransform) -> BaseDataset:
     return other.transform(self)
 
+
 BaseDataset.__add__ = func
 
+
 class DatasetFromCSV(BaseDataset):
-    """ DatasetFromFile for all dataset objects
+    """DatasetFromFile for all dataset objects
 
     Parameters:
         file_path (str): Relative or absolute to a local CSV file
@@ -182,12 +202,13 @@ def __init__(self, file_path, log=True, **kwargs):
         super().__init__(data=self.data.to_csv(), log=False, **kwargs)
 
 
-import requests
 import io
 
+import requests
+
 
 class DatasetFromCDDSearch(BaseDataset):
-    """ Dataset for retreiving data from CDD via a saved search.
+    """Dataset for retreiving data from CDD via a saved search.
 
     Requires a CDD Token to be set.
 
@@ -197,7 +218,9 @@ class DatasetFromCDDSearch(BaseDataset):
         update (bool): Whether or no to update the cached dataset by redoing the CDD search"""
 
     @log_arguments
-    def __init__(self, search_id, cache_file_path=None, update=True, log=True, **kwargs):
+    def __init__(
+        self, search_id, cache_file_path=None, update=True, log=True, **kwargs
+    ):
 
         if cache_file_path is None:
             cache_file_path = f"{search_id}.csv"
@@ -212,13 +235,15 @@ def __init__(self, search_id, cache_file_path=None, update=True, log=True, **kwa
         super().__init__(data=self.data.to_csv(), log=False, **kwargs)
 
     def run_saved_search(self, search_id):
-        """  Uses a CDD Token (passed as search_id) to search saved datasets to find and return its
+        """Uses a CDD Token (passed as search_id) to search saved datasets to find and return its
         related dataset export id.
 
         Parameters:
         search_id (str): The ID of the saved CDD search to use.
         """
-        base_url = f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        base_url = (
+            f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        )
         headers = {"X-CDD-token": f"{oce.CONFIG['CDD_TOKEN']}"}
         url = base_url + f"searches/{search_id}"
 
@@ -233,7 +258,9 @@ def check_export_status(self, export_id):
         export_id (str): The unique export ID of the dataset searched for
 
         """
-        base_url = f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        base_url = (
+            f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        )
         headers = {"X-CDD-token": f"{oce.CONFIG['CDD_TOKEN']}"}
         url = base_url + f"export_progress/{export_id}"
 
@@ -248,7 +275,9 @@ def get_export(self, export_id):
         export_id (str): The unique export ID of the dataset searched for
 
         """
-        base_url = f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        base_url = (
+            f"https://app.collaborativedrug.com/api/v1/vaults/{oce.CONFIG['VAULT_ID']}/"
+        )
         headers = {"X-CDD-token": f"{oce.CONFIG['CDD_TOKEN']}"}
         url = base_url + f"exports/{export_id}"
 
@@ -257,7 +286,7 @@ def get_export(self, export_id):
         return pd.read_csv(data_stream)
 
     def get_dataset_cdd_saved_search(self, search_id):
-        """ Uses a CDD Token (passed as search_id) to search saved datasets to find and return its
+        """Uses a CDD Token (passed as search_id) to search saved datasets to find and return its
         related dataset export id. Using the export id, it then checks the export status and
         returns the dataset's data in CSV format.
 
@@ -269,7 +298,7 @@ def get_dataset_cdd_saved_search(self, search_id):
         status = "new"
         while True:
             print(f"Export status is {status}, checking in {2**i} seconds...")
-            sleep(2 ** i)
+            sleep(2**i)
             status = self.check_export_status(export_id)
             if status == "finished":
                 print("Export ready!")
@@ -277,10 +306,12 @@ def get_dataset_cdd_saved_search(self, search_id):
             i += 1
         return self.get_export(export_id)
 
+
 from rdkit import Chem
 
+
 class CleanStructures(BaseDatasetTransform):
-    """ CleanStructures creates a new dataset from the original dataset by removing
+    """CleanStructures creates a new dataset from the original dataset by removing
     structures that are not valid.
 
     Parameters:
@@ -289,13 +320,17 @@ class CleanStructures(BaseDatasetTransform):
 
     def transform(self, dataset: BaseDataset, dropna_property: bool = True, **kwargs):
         cols = dataset.data.columns.tolist()
+
         def try_clean(s):
             try:
                 return Chem.MolToSmiles(Chem.MolFromSmiles(s))
             except:
                 return None
+
         start_num = len(dataset.data)
-        dataset.data[dataset.structure_col] = dataset.data[dataset.structure_col].apply(lambda x: try_clean(x))
+        dataset.data[dataset.structure_col] = dataset.data[dataset.structure_col].apply(
+            lambda x: try_clean(x)
+        )
         dataset.data = dataset.data.dropna(subset=[dataset.structure_col])
         if hasattr(dataset, "property_col") and not dataset.property_col is None:
             dataset.data = dataset.data.dropna(subset=[dataset.property_col])
@@ -303,8 +338,9 @@ def try_clean(s):
         print(f"{start_num - len(dataset.data)} structure(s) were removed.")
         return dataset
 
+
 class Discretize(BaseDatasetTransform):
-    """ Discretize creates a new dataset from the original dataset by discretizing
+    """Discretize creates a new dataset from the original dataset by discretizing
     the property column.
 
     Parameters:
@@ -318,11 +354,14 @@ def __init__(self, prop_cutoff: float, dir: str = "larger", log=True, **kwargs):
         self.prop_cutoff = prop_cutoff
 
     def transform(self, dataset: BaseDataset, **kwargs):
-        dataset.data[dataset.property_col] = [1 if x > self.prop_cutoff else 0 for x in dataset.data[dataset.property_col]]
+        dataset.data[dataset.property_col] = [
+            1 if x > self.prop_cutoff else 0 for x in dataset.data[dataset.property_col]
+        ]
         return dataset
 
+
 class OneHotEncode(BaseDatasetTransform):
-    """ This one hot encodes a given feature column
+    """This one hot encodes a given feature column
 
     Parameters:
         feature_col (str): The feature column to one hot encode.
@@ -336,4 +375,4 @@ def transform(self, dataset: BaseDataset, **kwargs):
         one_hot = pd.get_dummies(dataset.data[self.feature_col])
         dataset.data = dataset.data.join(one_hot)
         dataset.feature_cols += one_hot.columns.tolist()
-        return dataset
\ No newline at end of file
+        return dataset
diff --git a/src/olorenchemengine/ensemble.py b/src/olorenchemengine/ensemble.py
index e9880ec8..be6333f5 100644
--- a/src/olorenchemengine/ensemble.py
+++ b/src/olorenchemengine/ensemble.py
@@ -1,8 +1,9 @@
 """ Ensembling methods to combine `BaseModel`s to create better, combined models.
 """
 from .base_class import *
-from .representations import *
 from .basics import *
+from .representations import *
+
 
 def get_oof(self, model, X, y, kf):
     # https://www.kaggle.com/code/arthurtok/introduction-to-ensembling-stacking-in-python/notebook
@@ -19,8 +20,9 @@ def get_oof(self, model, X, y, kf):
 
     return oof_train
 
+
 class Averager(BaseModel):
-    """ Averager averages the predictions of multiple models for an ensembled prediction.
+    """Averager averages the predictions of multiple models for an ensembled prediction.
 
     Parameters:
         models (list): list of BaseModel objects to be averaged.
@@ -38,9 +40,10 @@ class Averager(BaseModel):
     model.predict(test['Drug'])
     ------------------------------
     """
+
     @log_arguments
-    def __init__(self, models: List[BaseModel], n: int =1, log: bool = True, **kwargs):
-        """ BaseStacking constructor
+    def __init__(self, models: List[BaseModel], n: int = 1, log: bool = True, **kwargs):
+        """BaseStacking constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
@@ -58,8 +61,8 @@ def __init__(self, models: List[BaseModel], n: int =1, log: bool = True, **kwarg
                 self.models.append(model)
         super().__init__(log=False, **kwargs)
 
-    def preprocess(self, X, y, fit = False):
-        """ Preprocesses the data for the model.
+    def preprocess(self, X, y, fit=False):
+        """Preprocesses the data for the model.
 
         Parameters:
             X (pd.DataFrame): Dataframe of features.
@@ -69,7 +72,7 @@ def preprocess(self, X, y, fit = False):
             X (pd.DataFrame): Dataframe of features."""
         return X
 
-    def _fit(self, X, y, valid = None):
+    def _fit(self, X, y, valid=None):
         for model in self.models:
             if issubclass(type(model), BaseModel):
                 model.fit(X, y)
@@ -77,7 +80,7 @@ def _fit(self, X, y, valid = None):
     def _predict(self, X):
         results = np.zeros(X.shape[0])
         for model in self.models:
-            results += model.predict(X)/len(self.models)
+            results += model.predict(X) / len(self.models)
         return results
 
     def _save(self):
@@ -93,9 +96,10 @@ def _load(self, d):
         for i, model in enumerate(self.models):
             model._load(d["model_save_dict"][i])
 
+
 class BaseStacking(BaseModel):
 
-    """ BaseStacking stacks the predictions of models for an ensembled prediction.
+    """BaseStacking stacks the predictions of models for an ensembled prediction.
 
     Parameters:
         models (List[BaseModel]): list of models to use for the learners to be stacked together.
@@ -105,9 +109,17 @@ class BaseStacking(BaseModel):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel], stacker_model: BaseModel, n: int =1, log: bool = True,
-        oof = False, nfolds = 5, **kwargs):
-        """ BaseStacking constructor
+    def __init__(
+        self,
+        models: List[BaseModel],
+        stacker_model: BaseModel,
+        n: int = 1,
+        log: bool = True,
+        oof=False,
+        nfolds=5,
+        **kwargs
+    ):
+        """BaseStacking constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
@@ -128,7 +140,7 @@ def __init__(self, models: List[BaseModel], stacker_model: BaseModel, n: int =1,
         self.nfolds = nfolds
         super().__init__(log=False, **kwargs)
 
-    def preprocess(self, X, y, fit = False):
+    def preprocess(self, X, y, fit=False):
         return X
 
     def featurize(self, X):
@@ -156,7 +168,7 @@ def featurize(self, X):
             data = data.reshape(-1, 1)
         return data
 
-    def _fit(self, X, y, valid = None):
+    def _fit(self, X, y, valid=None):
         if not self.oof:
             for model in self.models:
                 if issubclass(type(model), BaseModel):
@@ -170,7 +182,8 @@ def _fit(self, X, y, valid = None):
                 self.stacker_model.fit(data, y_valid)
         else:
             from sklearn.model_selection import KFold
-            kf = KFold(n_splits = self.nfolds)
+
+            kf = KFold(n_splits=self.nfolds)
 
             data = []
             for model in self.models:
@@ -215,9 +228,10 @@ def _load(self, d):
             model._load(d["model_save_dict"][i])
         self.stacker_model._load(d["stacker_model"])
 
+
 class BestStacker(BaseStacking):
 
-    """ BestStacker is a stacking method that uses the best model from a collection of models to make an ensembled prediction.
+    """BestStacker is a stacking method that uses the best model from a collection of models to make an ensembled prediction.
 
     Parameters:
         models (List[BaseModel]): list of models to use for the learners to be stacked together.
@@ -238,8 +252,10 @@ class BestStacker(BaseStacking):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel], n: int = 1, k: int = 1, log: bool = True):
-        """ BestStacker constructor
+    def __init__(
+        self, models: List[BaseModel], n: int = 1, k: int = 1, log: bool = True
+    ):
+        """BestStacker constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
@@ -248,9 +264,10 @@ def __init__(self, models: List[BaseModel], n: int = 1, k: int = 1, log: bool =
 
         super().__init__(models, KBestLinearRegression(k=k), n=n, log=False)
 
+
 class LinearRegressionStacker(BaseStacking):
 
-    """ LinearRegressionStacker is a stacking method that uses linear regression on the predictions
+    """LinearRegressionStacker is a stacking method that uses linear regression on the predictions
         from a collection of models to make an ensembled prediction.
 
         Parameters:
@@ -272,21 +289,20 @@ class LinearRegressionStacker(BaseStacking):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel],  n: int = 1, log: bool = True):
-        """ LinearRegressionStacker constructor
+    def __init__(self, models: List[BaseModel], n: int = 1, log: bool = True):
+        """LinearRegressionStacker constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
             n (int, optional): Number of times to repeat the given models. Defaults to 1.
             log (bool, optional): Whether or not to log the arguments of this constructor,
         """
-        super().__init__(models, LinearRegression(),
-        n=n, log=False)
+        super().__init__(models, LinearRegression(), n=n, log=False)
 
 
 class SKLearnStacker(BaseSKLearnModel):
 
-    """ SKLearnStacker is a stacking method that uses a sklearn-like models to make an ensembled prediction.
+    """SKLearnStacker is a stacking method that uses a sklearn-like models to make an ensembled prediction.
 
     Attributes:
         reg_stack: a sklearn-like model to use for stacking the models for regression tasks.
@@ -296,9 +312,15 @@ class SKLearnStacker(BaseSKLearnModel):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel], regression_stacker_model: BaseEstimator,
-        classification_stacker_model: BaseEstimator, n: int = 1, log: bool = True):
-        """ SKLearnStacker constructor
+    def __init__(
+        self,
+        models: List[BaseModel],
+        regression_stacker_model: BaseEstimator,
+        classification_stacker_model: BaseEstimator,
+        n: int = 1,
+        log: bool = True,
+    ):
+        """SKLearnStacker constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
@@ -309,14 +331,12 @@ def __init__(self, models: List[BaseModel], regression_stacker_model: BaseEstima
                 should be True only for unnested classes. Defaults to True."""
         self.reg_stack = BaseStacking(models, regression_stacker_model, n=n)
         self.class_stack = BaseStacking(models, classification_stacker_model, n=n)
-        super().__init__(None,
-            self.reg_stack,
-            self.class_stack,
-            log = False)
+        super().__init__(None, self.reg_stack, self.class_stack, log=False)
+
 
 class RFStacker(SKLearnStacker):
 
-    """ RFStacker is a subclass of SKLearnStacker that uses a random forest models to make an ensembled prediction.
+    """RFStacker is a subclass of SKLearnStacker that uses a random forest models to make an ensembled prediction.
 
         Parameters:
             models (List[BaseModel]): list of models to use for the learners to be stacked together.
@@ -340,9 +360,14 @@ class RFStacker(SKLearnStacker):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel], n_estimators: int = 100,
-        n: int = 1, log: bool = True):
-        """ RFStacker constructor
+    def __init__(
+        self,
+        models: List[BaseModel],
+        n_estimators: int = 100,
+        n: int = 1,
+        log: bool = True,
+    ):
+        """RFStacker constructor
 
         Args:
             models (List[BaseModel]): List of models to use for the learners to be stacked together.
@@ -350,29 +375,28 @@ def __init__(self, models: List[BaseModel], n_estimators: int = 100,
             n (int, optional): Number of times to repeat the given models. Defaults to 1.
             log (bool, optional): Whether or not to log the arguments of this constructor,
         """
-        rf_regressor = RandomForestRegressor(n_estimators = n_estimators,
+        rf_regressor = RandomForestRegressor(
+            n_estimators=n_estimators,
             max_features="log2",
             min_samples_split=4,
-            min_samples_leaf=4)
+            min_samples_leaf=4,
+        )
         rf_classifier = RandomForestClassifier(
             max_features="log2",
             min_samples_split=4,
             min_samples_leaf=4,
-            n_estimators = n_estimators,
+            n_estimators=n_estimators,
             n_jobs=-1,
-            criterion='entropy',
-            class_weight="balanced_subsample"
-            )
+            criterion="entropy",
+            class_weight="balanced_subsample",
+        )
+
+        super().__init__(models, rf_regressor, rf_classifier, n=n, log=False)
 
-        super().__init__(models,
-            rf_regressor,
-            rf_classifier,
-            n=n,
-            log = False)
 
 class MLPStacker(SKLearnStacker):
 
-    """ MLPStacker is a subclass of SKLearnStacker that uses a multi-layer perceptron model to make an ensembled prediction.
+    """MLPStacker is a subclass of SKLearnStacker that uses a multi-layer perceptron model to make an ensembled prediction.
 
         Parameters:
             models (List[BaseModel]): list of models to use for the learners to be stacked together.
@@ -404,16 +428,43 @@ class MLPStacker(SKLearnStacker):
     """
 
     @log_arguments
-    def __init__(self, models, layer_dims = [2048,512,128], activation = "tanh", epochs = 100, batch_size = 16, verbose = 0, n=1,log = True):
-
-        super().__init__(models,
-            TorchMLP(None, layer_dims= layer_dims, activation = activation, epochs=epochs, batch_size=batch_size, verbose = 0),
-            TorchMLP(None, layer_dims= layer_dims, activation = activation, epochs=epochs, batch_size=batch_size, verbose = 0),
+    def __init__(
+        self,
+        models,
+        layer_dims=[2048, 512, 128],
+        activation="tanh",
+        epochs=100,
+        batch_size=16,
+        verbose=0,
+        n=1,
+        log=True,
+    ):
+
+        super().__init__(
+            models,
+            TorchMLP(
+                None,
+                layer_dims=layer_dims,
+                activation=activation,
+                epochs=epochs,
+                batch_size=batch_size,
+                verbose=0,
+            ),
+            TorchMLP(
+                None,
+                layer_dims=layer_dims,
+                activation=activation,
+                epochs=epochs,
+                batch_size=batch_size,
+                verbose=0,
+            ),
             n=n,
-            log = False)
+            log=False,
+        )
+
 
 class BaseBoosting(BaseModel):
-    """ BaseBoosting uses models in a gradient boosting fashion to create an ensembled model.
+    """BaseBoosting uses models in a gradient boosting fashion to create an ensembled model.
 
     Parameters:
         models (List[BaseModel]): list of models to use for the learners to be stacked together.
@@ -435,9 +486,16 @@ class BaseBoosting(BaseModel):
     ------------------------------
     """
 
-
     @log_arguments
-    def __init__(self, models: List[BaseModel], n: int =1, oof=False, nfolds = 5, log: bool =True, **kwargs):
+    def __init__(
+        self,
+        models: List[BaseModel],
+        n: int = 1,
+        oof=False,
+        nfolds=5,
+        log: bool = True,
+        **kwargs
+    ):
         """
 
         Args:
@@ -459,12 +517,13 @@ def __init__(self, models: List[BaseModel], n: int =1, oof=False, nfolds = 5, lo
         self.oof = oof
         self.nfolds = nfolds
 
-        super().__init__(log = False, **kwargs)
+        super().__init__(log=False, **kwargs)
 
     def _fit(self, X_train, y_train):
         if self.oof:
             from sklearn.model_selection import KFold
-            kf = KFold(n_splits = self.nfolds)
+
+            kf = KFold(n_splits=self.nfolds)
 
         for i, model in enumerate(self.models):
             if self.oof:
@@ -474,47 +533,51 @@ def _fit(self, X_train, y_train):
                 model.fit(X_train, y_train)
                 y_pred = np.array(model.predict(X_train)).flatten()
             y_train = y_train - y_pred
-    
+
     def _error_calc(self, residuals):
         initial_residual, other_residuals = residuals[0], residuals[1:]
         y = np.array(([initial_residual]))
         for index in range(len(other_residuals)):
             curr_residual = other_residuals[index]
             total_residual = y[-1] - curr_residual
-            y = np.append(y, [total_residual], axis = 0)
-        
-        stdevs = list(np.std(y, axis = 1))
+            y = np.append(y, [total_residual], axis=0)
+
+        stdevs = list(np.std(y, axis=1))
         return stdevs
 
-    def _waterfall(self, y_data = None, normalize = False):
+    def _waterfall(self, y_data=None, normalize=False):
         self.stdevs = self._error_calc(self.residuals)
-        if y_data is not None: 
+        if y_data is not None:
             naive_mean = np.average(y_data)
-            naive_residuals = [(y - naive_mean)/np.std(y_data) for y in y_data] #normalized so they're in the same format as predict's input
-            if normalize:    
-                naive_residuals = naive_residuals 
-            else: 
+            naive_residuals = [
+                (y - naive_mean) / np.std(y_data) for y in y_data
+            ]  # normalized so they're in the same format as predict's input
+            if normalize:
+                naive_residuals = naive_residuals
+            else:
                 naive_residuals = self._unnormalize(np.array(naive_residuals))
             self.residuals.insert(0, naive_residuals)
-            self.stdevs.insert(0, np.std(naive_residuals))  
+            self.stdevs.insert(0, np.std(naive_residuals))
         return self.stdevs
 
-    def _predict(self, X, waterfall = False, normalize = False):
-        if waterfall: 
-            self.residuals = [] 
+    def _predict(self, X, waterfall=False, normalize=False):
+        if waterfall:
+            self.residuals = []
         y = np.zeros((len(X)))
         for model in self.models:
             prediction = np.array(model.predict(X)).flatten()
             y = y + prediction
             if waterfall:
-                if normalize: 
+                if normalize:
                     self.residuals.append(prediction)
                 else:
                     self.residuals.append(self._unnormalize(prediction))
         if self.setting == "classification":
             y = np.clip(y, 0, 1)
-            if waterfall: 
-                self.residuals = [np.clip(residual, 0, 1) for residual in self.residuals]
+            if waterfall:
+                self.residuals = [
+                    np.clip(residual, 0, 1) for residual in self.residuals
+                ]
         return y.flatten()
 
     def _save(self):
@@ -530,8 +593,9 @@ def _load(self, d):
         for i, model in enumerate(self.models):
             model._load(d["model_save_dict"][i])
 
+
 class Resample1(BaseModel):
-    """ Sample from imbalanced dataset. Take all compounds from smaller class and
+    """Sample from imbalanced dataset. Take all compounds from smaller class and
     then sample an equal number from the larger class.
 
     Parameters:
@@ -552,24 +616,30 @@ class Resample1(BaseModel):
     """
 
     @log_arguments
-    def __init__(self, model: BaseModel, log = True):
+    def __init__(self, model: BaseModel, log=True):
         self.model = model
-        super().__init__(log = False)
+        super().__init__(log=False)
 
     def _fit(self, X_train, y_train):
         self.model.fit(X_train, y_train)
 
     def fit(self, X_train, y_train):
         y_train = np.array(y_train)
-        assert len(np.unique(y_train)) == 2, "Resample1 can only be used on binary classification tasks"
+        assert (
+            len(np.unique(y_train)) == 2
+        ), "Resample1 can only be used on binary classification tasks"
         v1 = np.unique(y_train)[0]
         v2 = np.unique(y_train)[1]
         ix1 = np.where(y_train == v1)[0]
         ix2 = np.where(y_train == v2)[0]
         if len(ix1) > len(ix2):
-            ix = np.concatenate((ix2, np.random.choice(ix1, size = len(ix2), replace = False)))
+            ix = np.concatenate(
+                (ix2, np.random.choice(ix1, size=len(ix2), replace=False))
+            )
         else:
-            ix = np.concatenate((ix1, np.random.choice(ix2, size = len(ix1), replace = False)))
+            ix = np.concatenate(
+                (ix1, np.random.choice(ix2, size=len(ix1), replace=False))
+            )
         if isinstance(X_train, list):
             X_train = [X_train[i] for i in ix]
         elif isinstance(X_train, pd.Series) or isinstance(X_train, pd.DataFrame):
@@ -589,21 +659,24 @@ def _save(self):
     def _load(self, d):
         self.model._load(d["model_save"])
 
+
 class Resample2(BaseModel):
-    """ Sample from imbalanced dataset. Take all compounds from smaller class and
+    """Sample from imbalanced dataset. Take all compounds from smaller class and
     then sample an equal number from the larger class."""
 
     @log_arguments
-    def __init__(self, model: BaseModel, log = True):
+    def __init__(self, model: BaseModel, log=True):
         self.model = model
-        super().__init__(log = False)
+        super().__init__(log=False)
 
     def _fit(self, X_train, y_train):
         self.model.fit(X_train, y_train)
 
     def fit(self, X_train, y_train):
         y_train = np.array(y_train)
-        assert len(np.unique(y_train)) == 2, "Resample2 can only be used on binary classification tasks"
+        assert (
+            len(np.unique(y_train)) == 2
+        ), "Resample2 can only be used on binary classification tasks"
         v1 = np.unique(y_train)[0]
         v2 = np.unique(y_train)[1]
         ix1 = np.where(y_train == v1)[0]
@@ -611,9 +684,21 @@ def fit(self, X_train, y_train):
         print(len(ix1))
         print(len(ix2))
         if len(ix1) > len(ix2):
-            ix = np.concatenate((ix1, ix2, np.random.choice(ix2, size = len(ix1) - len(ix2), replace = True)))
+            ix = np.concatenate(
+                (
+                    ix1,
+                    ix2,
+                    np.random.choice(ix2, size=len(ix1) - len(ix2), replace=True),
+                )
+            )
         else:
-            ix = np.concatenate((ix1, ix2, np.random.choice(ix1, size = len(ix2) - len(ix1), replace = True)))
+            ix = np.concatenate(
+                (
+                    ix1,
+                    ix2,
+                    np.random.choice(ix1, size=len(ix2) - len(ix1), replace=True),
+                )
+            )
         if isinstance(X_train, list):
             X_train = [X_train[i] for i in ix]
         elif isinstance(X_train, pd.Series) or isinstance(X_train, pd.DataFrame):
@@ -636,7 +721,7 @@ def _load(self, d):
 
 
 class ResampleAdaboost(BaseBoosting):
-    """ ResampleAdaBoost performs the Adaboost with sampling weighting being done
+    """ResampleAdaBoost performs the Adaboost with sampling weighting being done
     via resampling of the dataset to create an ensembled model.
 
     Parameters:
@@ -664,7 +749,16 @@ class ResampleAdaboost(BaseBoosting):
     """
 
     @log_arguments
-    def __init__(self, models: List[BaseModel], n: int =1, factor: int =8, size: int = None, equation: str ="abs",log: bool=True, **kwargs):
+    def __init__(
+        self,
+        models: List[BaseModel],
+        n: int = 1,
+        factor: int = 8,
+        size: int = None,
+        equation: str = "abs",
+        log: bool = True,
+        **kwargs
+    ):
         """_summary_
 
         Args:
@@ -690,15 +784,15 @@ def _fit(self, X_train, y_train):
         y_train_df = pd.DataFrame(y_train)
         y_train = np.array(y_train)
 
-        w = np.array([1/len(X_train)] * len(X_train))
+        w = np.array([1 / len(X_train)] * len(X_train))
         self.betas = list()
 
         i = 0
         for model in self.models:
-            i+=1
+            i += 1
 
-            p = w/np.sum(w)
-            idxs = np.random.choice(len(X_train), size=size,replace=True, p=p)
+            p = w / np.sum(w)
+            idxs = np.random.choice(len(X_train), size=size, replace=True, p=p)
             X_train_temp = X_train.iloc[idxs, :]
             y_train_temp = y_train_df.iloc[idxs, :].to_numpy().flatten()
             model.fit(X_train_temp, y_train_temp)
@@ -707,17 +801,17 @@ def _fit(self, X_train, y_train):
 
             D = np.max(res)
             if self.equation == "abs":
-                loss = res/D
+                loss = res / D
             elif self.equation == "sqr":
-                loss = res**2/D**2
+                loss = res**2 / D**2
             else:
-                loss = 1 - np.exp(-res/D)
+                loss = 1 - np.exp(-res / D)
 
-            loss = loss/np.max(loss)
-            avg_loss = np.sum(loss*p)
-            beta = avg_loss/(1-avg_loss)
+            loss = loss / np.max(loss)
+            avg_loss = np.sum(loss * p)
+            beta = avg_loss / (1 - avg_loss)
 
-            w = w * np.power(beta, 1-loss)
+            w = w * np.power(beta, 1 - loss)
             self.betas.append(beta)
 
         self.betas = np.array(self.betas)
@@ -730,15 +824,20 @@ def _predict(self, X):
 
         predictions = list()
 
-        predictions = np.dot(ys.transpose(), np.log10(1/self.betas))/np.sum(np.log10(1/self.betas))
+        predictions = np.dot(ys.transpose(), np.log10(1 / self.betas)) / np.sum(
+            np.log10(1 / self.betas)
+        )
 
         return np.array(predictions)
 
     def _save(self):
         d = super()._save()
+        if not hasattr(self, "betas"):
+            return d
         d.update({"betas": self.betas})
         return d
 
     def _load(self, d):
         super()._load(d)
-        self.betas = d["betas"]
\ No newline at end of file
+        if "betas" in d:
+            self.betas = d["betas"]
diff --git a/src/olorenchemengine/external/ChemProp/main.py b/src/olorenchemengine/external/ChemProp/main.py
index e1168670..30800d7f 100644
--- a/src/olorenchemengine/external/ChemProp/main.py
+++ b/src/olorenchemengine/external/ChemProp/main.py
@@ -4,33 +4,28 @@
 
 """
 
-from olorenchemengine.representations import AtomFeaturizer, BondFeaturizer, TorchGeometricGraph
-from olorenchemengine.base_class import BaseModel, log_arguments, QuantileTransformer
-
-from olorenchemengine.internal import mock_imports
-
-try:
-    from torch_geometric.loader import DataLoader
-    from torch_geometric.data import Data, Dataset
-    from torch_geometric.data.data import size_repr
-    from torch_geometric.nn import global_mean_pool
-    from torch_scatter import scatter_sum
-except ImportError:
-    mock_imports(globals(), "DataLoader", "Data", "Dataset", "size_repr", "global_mean_pool", "scatter_sum")
-
-
-try:
-    import torch
-    import torch.nn as nn
-    import torch.utils.data
-except ImportError:
-    mock_imports(globals(), "torch", "nn")
+import io
 
 import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.data
 from rdkit import Chem
-import io
+from torch_geometric.data import Data, Dataset
+from torch_geometric.data.data import size_repr
+from torch_geometric.loader import DataLoader
+from torch_geometric.nn import global_mean_pool
+from torch_scatter import scatter_sum
 from tqdm import tqdm
 
+from olorenchemengine.base_class import BaseModel, QuantileTransformer, log_arguments
+from olorenchemengine.internal import mock_imports
+from olorenchemengine.representations import (
+    AtomFeaturizer,
+    BondFeaturizer,
+    TorchGeometricGraph,
+)
+
 
 class RevIndexedData(Data):
     def __init__(self, orig):
@@ -66,6 +61,7 @@ def __repr__(self):
             info = [size_repr(key, item, indent=2) for key, item in self]
             return "{}(\n{}\n)".format(cls, ",\n".join(info))
 
+
 class RevIndexedDataset(Dataset):
     def __init__(self, orig):
         super(RevIndexedDataset, self).__init__()
@@ -77,13 +73,15 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.dataset)
 
+
 def onek_encoding_unk(value, choices):
     encoding = [0] * (len(choices) + 1)
     index = choices.index(value) if value in choices else -1
     encoding[index] = 1
     return encoding
 
-class ChemProp_AF (AtomFeaturizer):
+
+class ChemProp_AF(AtomFeaturizer):
     @property
     def length(self):
         return 151
@@ -91,7 +89,7 @@ def length(self):
     def convert(self, atom: Chem.Atom) -> np.ndarray:
 
         ATOM_FEATURES = {
-            "atomic_num": list(range(1,119)),
+            "atomic_num": list(range(1, 119)),
             "degree": [0, 1, 2, 3, 4, 5],
             "formal_charge": [-1, -2, 1, 2, 0],
             "chiral_tag": [0, 1, 2, 3],
@@ -118,9 +116,10 @@ def convert(self, atom: Chem.Atom) -> np.ndarray:
         )  # scaled to about the same range as other features
         return np.array(features)
 
-class ChemProp_BF (BondFeaturizer):
+
+class ChemProp_BF(BondFeaturizer):
     @property
-    def length (self):
+    def length(self):
         return 14
 
     def convert(self, bond: Chem.Bond) -> np.ndarray:
@@ -140,10 +139,11 @@ def convert(self, bond: Chem.Bond) -> np.ndarray:
             fbond += onek_encoding_unk(int(bond.GetStereo()), list(range(6)))
         return np.array(fbond)
 
-class ChemPropDataLoading (TorchGeometricGraph):
+
+class ChemPropDataLoading(TorchGeometricGraph):
     @log_arguments
     def __init__(self):
-        super().__init__(ChemProp_AF(), ChemProp_BF(), log = False)
+        super().__init__(ChemProp_AF(), ChemProp_BF(), log=False)
 
     def _convert(self, smiles, y=None, **kwargs):
         data = super()._convert(smiles, y=y, **kwargs)
@@ -161,16 +161,19 @@ def initialize_weights(model: nn.Module) -> None:
         else:
             nn.init.xavier_normal_(param)
 
+
 def directed_mp(message, edge_index, revedge_index):
     m = scatter_sum(message, edge_index[1], dim=0)
     m_all = m[edge_index[0]]
     m_rev = message[revedge_index]
     return m_all - m_rev
 
+
 def aggregate_at_nodes(num_nodes, message, edge_index):
     m = scatter_sum(message, edge_index[1], dim=0, dim_size=num_nodes)
     return m[torch.arange(num_nodes)]
 
+
 #### Key Message-Passing algorithm in the model ####
 class DMPNNEncoder(nn.Module):
     def __init__(self, hidden_size, node_fdim, edge_fdim, depth=3):
@@ -210,6 +213,7 @@ def forward(self, data):
         # readout: pyg global pooling
         return global_mean_pool(node_attr, batch)
 
+
 #### General ChemProp model training method, to be used with _fit in the ChemPropModel() class implementation. ####
 def train(config, loader, setting, device=torch.device("cpu")):
     criterion = config["loss"]
@@ -218,7 +222,7 @@ def train(config, loader, setting, device=torch.device("cpu")):
     scheduler = config["scheduler"]
 
     if setting == "classification":
-        criterion = nn.BCEWithLogitsLoss(reduction = "none")
+        criterion = nn.BCEWithLogitsLoss(reduction="none")
     else:
         criterion = nn.MSELoss()
 
@@ -229,14 +233,15 @@ def train(config, loader, setting, device=torch.device("cpu")):
 
         optimizer.zero_grad()
         out = model(batch)
-        loss = criterion(torch.squeeze(out,1), batch.y.float())
+        loss = criterion(torch.squeeze(out, 1), batch.y.float())
         loss.sum().backward()
 
         optimizer.step()
         scheduler.step()
 
+
 #### General ChemProp model prediction method, to be used with _predict in the ChemPropModel() class implementation. ####
-def predict(config, loader, setting = "classification", device=torch.device("cpu")):
+def predict(config, loader, setting="classification", device=torch.device("cpu")):
     model = config["model"]
 
     model = model.to(device)
@@ -254,6 +259,7 @@ def predict(config, loader, setting = "classification", device=torch.device("cpu
     predictions = np.concatenate([y.cpu().numpy() for y in y_pred])
     return predictions.flatten()
 
+
 class ChemPropModel(BaseModel):
     """ChemProp is the model presented in `Analyzing Learned Molecular Representations for Property Prediction <hhttps://doi.org/10.1021/acs.jcim.9b00237>` _.
     Here, we adapt its PyTorch Geometric implementation as in the `Github repository <https://github.com/itakigawa/pyg_chemprop>`_
@@ -277,15 +283,17 @@ class ChemPropModel(BaseModel):
     """
 
     @log_arguments
-    def __init__(self,
+    def __init__(
+        self,
         dropout_rate: float = 0.0,
         epochs: int = 3,
         batch_size: int = 50,
         lr: float = 1e-3,
         hidden_size: int = 300,
-        depth: int =3,
-        map_location = "cuda:0",
-         **kwargs):
+        depth: int = 3,
+        map_location="cuda:0",
+        **kwargs
+    ):
 
         self.representation = ChemPropDataLoading()
 
@@ -321,7 +329,7 @@ def __init__(self,
 
     def preprocess(self, X, y, **kwargs):
         if y is None:
-            y = [None]*len(X)
+            y = [None] * len(X)
         return self.representation.convert(X, ys=y)
 
     def _fit(self, X, y, **kwargs):
@@ -338,7 +346,9 @@ def _fit(self, X, y, **kwargs):
             "scheduler": self.scheduler,
         }
 
-        if (len(np.unique(y, return_counts=True)) > 2 and (((np.unique(y)).size / (y.size)) > 0.1)):
+        if len(np.unique(y, return_counts=True)) > 2 and (
+            ((np.unique(y)).size / (y.size)) > 0.1
+        ):
             self.setting = "regression"
 
         for epoch in range(self.epochs):
@@ -367,4 +377,4 @@ def _save(self) -> str:
 
     def _load(self, d):
         super()._load(d)
-        self.model = torch.load(io.BytesIO(d["save"]))
\ No newline at end of file
+        self.model = torch.load(io.BytesIO(d["save"]))
diff --git a/src/olorenchemengine/external/GINNetwork/main.py b/src/olorenchemengine/external/GINNetwork/main.py
index 130c3ab8..6e472342 100644
--- a/src/olorenchemengine/external/GINNetwork/main.py
+++ b/src/olorenchemengine/external/GINNetwork/main.py
@@ -6,13 +6,15 @@
 """
 from olorenchemengine.base_class import log_arguments
 from olorenchemengine.gnn import BaseLightningModule
+
 from .operations import GINNetwork
 
+
 class GINModel(BaseLightningModule):
-    
+
     """
     GINModel class from `GINEPLUS GitHub repository <https://github.com/RBrossard/GINEPLUS>`.
-    
+
     Parameters:
         task_type (str): Type of task to perform. Options are 'regression' or 'classification'. This
             task_type will automatically be set by BaseTorchGeometricModel and usually
@@ -26,19 +28,32 @@ class GINModel(BaseLightningModule):
         conv_radius (int): Radius of convolution.
         optim (str): Optimizer to use. Options are 'adam'.
     """
+
     @log_arguments
-    def __init__(self, task_type = "classification", hidden: str = 100, lr: float = 0.001,
-                 layers=3, dropout=0.5, virtual_node=False,
-                 conv_radius=3, conv_type='gin+',optim="adam",**kwargs):
-        assert conv_type in ['gcn', 'gin', 'gin+', 'naivegin+']
+    def __init__(
+        self,
+        task_type="classification",
+        hidden: str = 100,
+        lr: float = 0.001,
+        layers=3,
+        dropout=0.5,
+        virtual_node=False,
+        conv_radius=3,
+        conv_type="gin+",
+        optim="adam",
+        **kwargs
+    ):
+        assert conv_type in ["gcn", "gin", "gin+", "naivegin+"]
         super().__init__(optim=optim)
 
         # Network
         out_dim = 1
-        self.network = GINNetwork(hidden=hidden,
-                                         out_dim=out_dim,
-                                         layers=layers,
-                                         dropout=dropout,
-                                         virtual_node=virtual_node,
-                                         k=conv_radius,
-                                         conv_type=conv_type)
\ No newline at end of file
+        self.network = GINNetwork(
+            hidden=hidden,
+            out_dim=out_dim,
+            layers=layers,
+            dropout=dropout,
+            virtual_node=virtual_node,
+            k=conv_radius,
+            conv_type=conv_type,
+        )
diff --git a/src/olorenchemengine/external/GINNetwork/operations.py b/src/olorenchemengine/external/GINNetwork/operations.py
index dff461df..45e202e7 100644
--- a/src/olorenchemengine/external/GINNetwork/operations.py
+++ b/src/olorenchemengine/external/GINNetwork/operations.py
@@ -1,22 +1,14 @@
-from olorenchemengine.internal import mock_imports
-
-try:
-    import torch
-    from torch import nn as nn
-    from torch.nn import functional as F
-except:
-    mock_imports(globals(), "torch", "nn", "F")
-
-try:
-    from torch_geometric.nn import MessagePassing
-    from torch_geometric import nn as nng
-    from torch_sparse import SparseTensor, coalesce
-    from torch_scatter import scatter_add
-except ImportError:
-    mock_imports(globals(), "MessagePassing", "nng", "SparseTensor", "coalesce", "scatter_add")
+from copy import copy
 
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+from torch_geometric import nn as nng
+from torch_geometric.nn import MessagePassing
+from torch_scatter import scatter_add
+from torch_sparse import SparseTensor, coalesce
 
-from copy import copy
+from olorenchemengine.internal import mock_imports
 
 
 class MLP(nn.Module):
@@ -25,7 +17,7 @@ def __init__(self, in_dim, out_dim):
         self.main = [
             nn.Linear(in_dim, 2 * in_dim),
             nn.BatchNorm1d(2 * in_dim),
-            nn.ReLU()
+            nn.ReLU(),
         ]
         self.main.append(nn.Linear(2 * in_dim, out_dim))
         self.main = nn.Sequential(*self.main)
@@ -40,7 +32,7 @@ def forward(self, x):
 class OGBMolEmbedding(nn.Module):
     def __init__(self, dim, embed_edge=True, x_as_list=False):
         super().__init__()
-        from ogb.graphproppred.mol_encoder import BondEncoder, AtomEncoder
+        from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
 
         self.atom_embedding = AtomEncoder(emb_dim=dim)
         if embed_edge:
@@ -49,14 +41,14 @@ def __init__(self, dim, embed_edge=True, x_as_list=False):
 
     def forward(self, data):
         data = new(data)
-        #try:
+        # try:
         data.x = self.atom_embedding(data.x)
         if self.x_as_list:
             data.x = [data.x]
-        if hasattr(self, 'edge_embedding'):
+        if hasattr(self, "edge_embedding"):
             data.edge_attr = self.edge_embedding(data.edge_attr)
         return data
-        #except Exception as e:
+        # except Exception as e:
         #    print(e)
         #    print(data.smiles)
 
@@ -80,24 +72,18 @@ def __init__(self, dim, conv_type="gin"):
         super().__init__()
         self.conv_type = conv_type
         if "gin" in conv_type:
-            self.mlp = nn.Sequential(
-                MLP(dim, dim),
-                nn.BatchNorm1d(dim),
-                nn.ReLU()
-            )
+            self.mlp = nn.Sequential(MLP(dim, dim), nn.BatchNorm1d(dim), nn.ReLU())
         elif "gcn" in conv_type:
             self.W0 = nn.Linear(dim, dim)
             self.W1 = nn.Linear(dim, dim)
-            self.nl_bn = nn.Sequential(
-                nn.BatchNorm1d(dim),
-                nn.ReLU()
-            )
+            self.nl_bn = nn.Sequential(nn.BatchNorm1d(dim), nn.ReLU())
         else:
-            raise NotImplementedError('Unrecognised model conv : {}'.format(conv_type))
+            raise NotImplementedError("Unrecognised model conv : {}".format(conv_type))
 
     def forward(self, virtual_node, embeddings, batch_vector):
         if batch_vector.size(0) > 0:  # ...or the operation will crash for empty graphs
             from torch_geometric.nn import global_add_pool
+
             G = global_add_pool(embeddings, batch_vector)
         else:
             G = torch.zeros_like(virtual_node)
@@ -108,25 +94,39 @@ def forward(self, virtual_node, embeddings, batch_vector):
             virtual_node = self.W0(virtual_node) + self.W1(G)
             virtual_node = self.nl_bn(virtual_node)
         else:
-            raise NotImplementedError('Unrecognised model conv : {}'.format(self.conv_type))
+            raise NotImplementedError(
+                "Unrecognised model conv : {}".format(self.conv_type)
+            )
         return virtual_node
 
 
 class ConvBlock(nn.Module):
-    def __init__(self, dim, dropout=0.5, activation=F.relu, virtual_node=False, virtual_node_agg=True,
-                 k=4, last_layer=False, conv_type='gin', edge_embedding=None):
+    def __init__(
+        self,
+        dim,
+        dropout=0.5,
+        activation=F.relu,
+        virtual_node=False,
+        virtual_node_agg=True,
+        k=4,
+        last_layer=False,
+        conv_type="gin",
+        edge_embedding=None,
+    ):
         super().__init__()
         self.edge_embed = edge_embedding
         self.conv_type = conv_type
-        if conv_type == 'gin+':
+        if conv_type == "gin+":
             self.conv = GINEPLUS(MLP(dim, dim), dim, k=k)
-        elif conv_type == 'naivegin+':
+        elif conv_type == "naivegin+":
             self.conv = NAIVEGINEPLUS(MLP(dim, dim), dim, k=k)
-        elif conv_type == 'gin':
+        elif conv_type == "gin":
             from torch_geometric.nn import GINEConv
+
             self.conv = GINEConv(MLP(dim, dim), train_eps=True)
-        elif conv_type == 'gcn':
+        elif conv_type == "gcn":
             from torch_geometric.nn import GCNConv
+
             self.conv = GCNConv(dim, dim)
         self.norm = nn.BatchNorm1d(dim)
         self.act = activation or nn.Identity()
@@ -145,17 +145,17 @@ def forward(self, data):
         mhei, d = data.multihop_edge_index, data.distance
         h = x
         if self.virtual_node:
-            if self.conv_type == 'gin+':
+            if self.conv_type == "gin+":
                 h[0] = h[0] + data.virtual_node[b]
             else:
                 h = h + data.virtual_node[b]
-        if self.conv_type == 'gin':
+        if self.conv_type == "gin":
             H = self.conv(h, ei, edge_attr=self.edge_embed(ea))
-        elif self.conv_type == 'gcn':
+        elif self.conv_type == "gcn":
             H = self.conv(h, ei)
         else:
             H = self.conv(h, mhei, d, self.edge_embed(ea))
-        if self.conv_type == 'gin+':
+        if self.conv_type == "gin+":
             h = H[0]
         else:
             h = H
@@ -168,7 +168,7 @@ def forward(self, data):
             v = self.vn_aggregator(data.virtual_node, h, b)
             v = F.dropout(v, self.dropout_ratio, training=self.training)
             data.virtual_node = v
-        if self.conv_type == 'gin+':
+        if self.conv_type == "gin+":
             H[0] = h
             h = H
         data.x = h
@@ -179,7 +179,12 @@ class GlobalPool(nn.Module):
     def __init__(self, fun, cat_size=False, cat_candidates=False):
         super().__init__()
         self.cat_size = cat_size
-        from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+        from torch_geometric.nn import (
+            global_add_pool,
+            global_max_pool,
+            global_mean_pool,
+        )
+
         if fun.lower() == "add":
             self.fun = global_add_pool
         elif fun.lower() == "mean":
@@ -191,48 +196,71 @@ def __init__(self, fun, cat_size=False, cat_candidates=False):
     def forward(self, batch):
         x, b = batch.x, batch.batch
         pooled = self.fun(x, b, size=batch.num_graphs)
-        from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+        from torch_geometric.nn import (
+            global_add_pool,
+            global_max_pool,
+            global_mean_pool,
+        )
 
         if self.cat_size:
-            sizes = global_add_pool(torch.ones(x.size(0), 1).type_as(x), b, size=batch.num_graphs)
+            sizes = global_add_pool(
+                torch.ones(x.size(0), 1).type_as(x), b, size=batch.num_graphs
+            )
             pooled = torch.cat([pooled, sizes], dim=1)
         if self.cat_candidates:
             ei = batch.edge_index
             mask = batch.edge_attr == 3
-            candidates = scatter_add(x[ei[0, mask]], b[ei[0, mask]], dim=0, dim_size=batch.num_graphs)
+            candidates = scatter_add(
+                x[ei[0, mask]], b[ei[0, mask]], dim=0, dim_size=batch.num_graphs
+            )
             pooled = torch.cat([pooled, candidates], dim=1)
         return pooled
 
 
 class GINNetwork(nn.Module):
-    def __init__(self, hidden=100, out_dim=128, layers=3, dropout=0.5, virtual_node=False,
-                 k=4, conv_type='gin+'):
+    def __init__(
+        self,
+        hidden=100,
+        out_dim=128,
+        layers=3,
+        dropout=0.5,
+        virtual_node=False,
+        k=4,
+        conv_type="gin+",
+    ):
         super().__init__()
         self.k = k
         self.conv_type = conv_type
-        from ogb.graphproppred.mol_encoder import BondEncoder, AtomEncoder
-        convs = [ConvBlock(hidden,
-                           dropout=dropout,
-                           virtual_node=virtual_node,
-                           k=min(i + 1, k),
-                           conv_type=conv_type,
-                           edge_embedding=BondEncoder(emb_dim=hidden))
-                 for i in range(layers - 1)]
-        convs.append(ConvBlock(hidden,
-                               dropout=dropout,
-                               virtual_node=virtual_node,
-                               virtual_node_agg=False,  # on last layer, use but do not update virtual node
-                               last_layer=True,
-                               k=min(layers, k),
-                               conv_type=conv_type,
-                               edge_embedding=BondEncoder(emb_dim=hidden)))
+        from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
+
+        convs = [
+            ConvBlock(
+                hidden,
+                dropout=dropout,
+                virtual_node=virtual_node,
+                k=min(i + 1, k),
+                conv_type=conv_type,
+                edge_embedding=BondEncoder(emb_dim=hidden),
+            )
+            for i in range(layers - 1)
+        ]
+        convs.append(
+            ConvBlock(
+                hidden,
+                dropout=dropout,
+                virtual_node=virtual_node,
+                virtual_node_agg=False,  # on last layer, use but do not update virtual node
+                last_layer=True,
+                k=min(layers, k),
+                conv_type=conv_type,
+                edge_embedding=BondEncoder(emb_dim=hidden),
+            )
+        )
         self.main = nn.Sequential(
-            OGBMolEmbedding(hidden, embed_edge=False, x_as_list=(conv_type == 'gin+')),
-            *convs)
-        self.aggregate = nn.Sequential(
-            GlobalPool('mean'),
-            nn.Linear(hidden, out_dim)
+            OGBMolEmbedding(hidden, embed_edge=False, x_as_list=(conv_type == "gin+")),
+            *convs
         )
+        self.aggregate = nn.Sequential(GlobalPool("mean"), nn.Linear(hidden, out_dim))
         self.virtual_node = virtual_node
         if self.virtual_node:
             self.v0 = nn.Parameter(torch.zeros(1, hidden), requires_grad=True)
@@ -242,7 +270,7 @@ def forward(self, data):
         if self.virtual_node:
             data.virtual_node = self.v0.expand(data.num_graphs, self.v0.shape[-1])
         g = self.main(data)
-        if self.conv_type == 'gin+':
+        if self.conv_type == "gin+":
             g.x = g.x[0]
         return self.aggregate(g)
 
@@ -267,7 +295,9 @@ def make_multihop_edges(data, k):
         return data
 
     # Get the distance 0
-    multihop_edge_index = torch.arange(0, N, dtype=data.edge_index[0].dtype, device=data.x.device)
+    multihop_edge_index = torch.arange(
+        0, N, dtype=data.edge_index[0].dtype, device=data.x.device
+    )
     distance = torch.zeros_like(multihop_edge_index)
     multihop_edge_index = multihop_edge_index.unsqueeze(0).repeat(2, 1)
 
@@ -275,12 +305,18 @@ def make_multihop_edges(data, k):
     multihop_edge_index = torch.cat((multihop_edge_index, data.edge_index), dim=1)
     distance = torch.cat((distance, torch.ones_like(data.edge_index[0])), dim=0)
 
-    A = SparseTensor(row=data.edge_index[0], col=data.edge_index[1], value=torch.ones_like(data.edge_index[0]).float(),
-                     sparse_sizes=(N, N), is_sorted=False)
+    A = SparseTensor(
+        row=data.edge_index[0],
+        col=data.edge_index[1],
+        value=torch.ones_like(data.edge_index[0]).float(),
+        sparse_sizes=(N, N),
+        is_sorted=False,
+    )
     Ad = A  # A to the power d
 
     # Get larger distances
     from torch_sparse import matmul
+
     for d in range(2, k + 1):
         Ad = matmul(Ad, A)
         row, col, v = Ad.coo()
@@ -290,7 +326,9 @@ def make_multihop_edges(data, k):
         distance = torch.cat((distance, d_edge_attr), dim=0)
 
     # remove dupicate, keep only shortest distance
-    multihop_edge_index, distance = coalesce(multihop_edge_index, distance, N, N, op='min')
+    multihop_edge_index, distance = coalesce(
+        multihop_edge_index, distance, N, N, op="min"
+    )
 
     data.multihop_edge_index = multihop_edge_index
     data.distance = distance
@@ -300,7 +338,7 @@ def make_multihop_edges(data, k):
 
 class NAIVEGINEPLUS(MessagePassing):
     def __init__(self, fun, dim, k=4, **kwargs):
-        super().__init__(aggr='add')
+        super().__init__(aggr="add")
         self.k = k
         self.nn = fun
         self.eps = nn.Parameter(torch.zeros(k + 1, dim), requires_grad=True)
@@ -310,9 +348,13 @@ def forward(self, x, multihop_edge_index, distance, edge_attr):
         result = (1 + self.eps[0]) * x
         for i in range(self.k):
             if i == 0:
-                out = self.propagate(multihop_edge_index[:, distance == i + 1], edge_attr=edge_attr, x=x)
+                out = self.propagate(
+                    multihop_edge_index[:, distance == i + 1], edge_attr=edge_attr, x=x
+                )
             else:
-                out = self.propagate(multihop_edge_index[:, distance == i + 1], edge_attr=None, x=x)
+                out = self.propagate(
+                    multihop_edge_index[:, distance == i + 1], edge_attr=None, x=x
+                )
             result += (1 + self.eps[i + 1]) * out
         result = self.nn(result)
         return result
@@ -324,12 +366,14 @@ def message(self, x_j, edge_attr):
             return F.relu(x_j)
 
     def __repr__(self):
-        return '{}(nn={}, k={})'.format(self.__class__.__name__, self.nn, self.eps.size(0))
+        return "{}(nn={}, k={})".format(
+            self.__class__.__name__, self.nn, self.eps.size(0)
+        )
 
 
 class GINEPLUS(MessagePassing):
     def __init__(self, fun, dim, k=4, **kwargs):
-        super().__init__(aggr='add')
+        super().__init__(aggr="add")
         self.k = k
         self.nn = fun
         self.eps = nn.Parameter(torch.zeros(k + 1, dim), requires_grad=True)
@@ -343,9 +387,13 @@ def forward(self, XX, multihop_edge_index, distance, edge_attr):
             if i >= self.k:
                 break
             if i == 0:
-                out = self.propagate(multihop_edge_index[:, distance == i + 1], edge_attr=edge_attr, x=x)
+                out = self.propagate(
+                    multihop_edge_index[:, distance == i + 1], edge_attr=edge_attr, x=x
+                )
             else:
-                out = self.propagate(multihop_edge_index[:, distance == i + 1], edge_attr=None, x=x)
+                out = self.propagate(
+                    multihop_edge_index[:, distance == i + 1], edge_attr=None, x=x
+                )
             result += (1 + self.eps[i + 1]) * out
         result = self.nn(result)
         return [result] + XX
@@ -358,4 +406,4 @@ def message(self, x_j, edge_attr):
 
 
 def new(data):
-    return copy(data)
\ No newline at end of file
+    return copy(data)
diff --git a/src/olorenchemengine/external/HondaSTRep/main.py b/src/olorenchemengine/external/HondaSTRep/main.py
index c2b60a31..2109dfc0 100644
--- a/src/olorenchemengine/external/HondaSTRep/main.py
+++ b/src/olorenchemengine/external/HondaSTRep/main.py
@@ -1,22 +1,20 @@
+from typing import Any, List, Union
+
 import numpy as np
 import pandas as pd
-
-from typing import List, Union, Any
+import torch
+from torch.utils.data import DataLoader
 
 import olorenchemengine as oce
-from olorenchemengine.base_class import log_arguments, BaseModel
-from olorenchemengine.representations import BaseVecRepresentation, SMILESRepresentation
+from olorenchemengine.base_class import BaseModel, log_arguments
 from olorenchemengine.internal import download_public_file
-from .operations import WordVocab, TrfmSeq2seq, Seq2seqDataset
+from olorenchemengine.representations import BaseVecRepresentation, SMILESRepresentation
+
+from .operations import Seq2seqDataset, TrfmSeq2seq, WordVocab
 
-try:
-    from torch.utils.data import DataLoader
-    import torch
-except ImportError:
-    oce.mock_imports(globals(), "DataLoader", "torch")
 
 class HondaSTRep(BaseVecRepresentation):
-    """ HondaSTRep is an implementation of the molecular representation provided
+    """HondaSTRep is an implementation of the molecular representation provided
     by Honda et al. in `SMILES Transformer: Pre-trained Molecular Fingerprint for Low Data Drug Discovery
     <https://arxiv.org/abs/1911.04738>`."""
 
@@ -32,15 +30,27 @@ def __init__(self):
 
     def _convert(self, s: str) -> np.ndarray:
         seq2seqdataset = Seq2seqDataset([s], self.vocab)
-        loader = DataLoader(seq2seqdataset, batch_size=32, shuffle=False, num_workers=oce.CONFIG["NUM_WORKERS"])
+        loader = DataLoader(
+            seq2seqdataset,
+            batch_size=32,
+            shuffle=False,
+            num_workers=oce.CONFIG["NUM_WORKERS"],
+        )
         output = self.trfm.encode(next(iter(loader)))
         return output
 
-    def _convert_list(self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None) -> List[np.ndarray]:
+    def _convert_list(
+        self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None
+    ) -> List[np.ndarray]:
         if isinstance(smiles_list, pd.Series):
             smiles_list = smiles_list.tolist()
         seq2seqdataset = Seq2seqDataset(smiles_list, self.vocab)
-        loader = DataLoader(seq2seqdataset, batch_size=32, shuffle=False, num_workers=oce.CONFIG["NUM_WORKERS"])
+        loader = DataLoader(
+            seq2seqdataset,
+            batch_size=32,
+            shuffle=False,
+            num_workers=oce.CONFIG["NUM_WORKERS"],
+        )
         output = None
         for b, sm in enumerate(loader):
             sm = torch.t(sm)
@@ -50,4 +60,4 @@ def _convert_list(self, smiles_list: List[str], ys: List[Union[int, float, np.nu
                 output = encoding
             else:
                 output = np.concatenate([output, encoding])
-        return output
\ No newline at end of file
+        return output
diff --git a/src/olorenchemengine/external/HondaSTRep/operations.py b/src/olorenchemengine/external/HondaSTRep/operations.py
index f586fe3c..500eaf03 100644
--- a/src/olorenchemengine/external/HondaSTRep/operations.py
+++ b/src/olorenchemengine/external/HondaSTRep/operations.py
@@ -29,17 +29,28 @@
 import argparse
 import pickle
 from collections import Counter
+
 import olorenchemengine as oce
 from olorenchemengine.internal import mock_imports
 
+
 class TorchVocab(object):
     """
     :property freqs: collections.Counter, コーパス中の単語の出現頻度を保持するオブジェクト
     :property stoi: collections.defaultdict, string → id の対応を示す辞書
     :property itos: collections.defaultdict, id → string の対応を示す辞書
     """
-    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],
-                 vectors=None, unk_init=None, vectors_cache=None):
+
+    def __init__(
+        self,
+        counter,
+        max_size=None,
+        min_freq=1,
+        specials=["<pad>", "<oov>"],
+        vectors=None,
+        unk_init=None,
+        vectors_cache=None,
+    ):
         """
         :param counter: collections.Counter, データ中に含まれる単語の頻度を計測するためのcounter
         :param max_size: int, vocabularyの最大のサイズ. Noneの場合は最大値なし. defaultはNone
@@ -109,7 +120,12 @@ def __init__(self, counter, max_size=None, min_freq=1):
         self.eos_index = 2
         self.sos_index = 3
         self.mask_index = 4
-        super().__init__(counter, specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"], max_size=max_size, min_freq=min_freq)
+        super().__init__(
+            counter,
+            specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"],
+            max_size=max_size,
+            min_freq=min_freq,
+        )
 
     # override用
     def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:
@@ -120,7 +136,7 @@ def from_seq(self, seq, join=False, with_pad=False):
         pass
 
     @staticmethod
-    def load_vocab(vocab_path: str) -> 'Vocab':
+    def load_vocab(vocab_path: str) -> "Vocab":
         with open(vocab_path, "rb") as f:
             return pickle.load(f)
 
@@ -144,7 +160,9 @@ def __init__(self, texts, max_size=None, min_freq=1):
                 counter[word] += 1
         super().__init__(counter, max_size=max_size, min_freq=min_freq)
 
-    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
+    def to_seq(
+        self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False
+    ):
         if isinstance(sentence, str):
             sentence = sentence.split()
 
@@ -167,11 +185,11 @@ def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_le
         return (seq, origin_seq_len) if with_len else seq
 
     def from_seq(self, seq, join=False, with_pad=False):
-        words = [self.itos[idx]
-                 if idx < len(self.itos)
-                 else "<%d>" % idx
-                 for idx in seq
-                 if not with_pad or idx != self.pad_index]
+        words = [
+            self.itos[idx] if idx < len(self.itos) else "<%d>" % idx
+            for idx in seq
+            if not with_pad or idx != self.pad_index
+        ]
 
         return " ".join(words) if join else words
 
@@ -180,22 +198,17 @@ def load_vocab(vocab_path: str):
         with open(vocab_path, "rb") as f:
             return pickle.load(f)
 
-import numpy as np
-import pandas as pd
-
-try:
-    import torch
-    from torch import nn
-    from torch import optim
-    from torch.autograd import Variable
-    from torch.nn import functional as F
-    from torch.utils.data import DataLoader
-except ImportError:
-    mock_imports(globals(), "torch", "nn", "optim", "Variable", "F", "DataLoader")
 
+import math
 
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn, optim
+from torch.autograd import Variable
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
 from tqdm import tqdm
-import math
 
 PAD = 0
 UNK = 1
@@ -203,26 +216,30 @@ def load_vocab(vocab_path: str):
 SOS = 3
 MASK = 4
 
+
 class PositionalEncoding(nn.Module):
     "Implement the PE function. No batch support?"
+
     def __init__(self, d_model, dropout, max_len=5000):
         super(PositionalEncoding, self).__init__()
         self.dropout = nn.Dropout(p=dropout)
 
         # Compute the positional encodings once in log space.
-        pe = torch.zeros(max_len, d_model) # (T,H)
-        position = torch.arange(0., max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, d_model)  # (T,H)
+        position = torch.arange(0.0, max_len).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0.0, d_model, 2) * -(math.log(10000.0) / d_model)
+        )
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
         pe = pe.unsqueeze(0)
-        self.register_buffer('pe', pe)
+        self.register_buffer("pe", pe)
 
     def forward(self, x):
-        x = x + Variable(self.pe[:, :x.size(1)],
-                         requires_grad=False)
+        x = x + Variable(self.pe[:, : x.size(1)], requires_grad=False)
         return self.dropout(x)
 
+
 class TrfmSeq2seq(nn.Module):
     def __init__(self, in_size, hidden_size, out_size, n_layers, dropout=0.1):
         super(TrfmSeq2seq, self).__init__()
@@ -234,8 +251,13 @@ def __init__(self, in_size, hidden_size, out_size, n_layers, dropout=0.1):
         self.pe = PositionalEncoding(hidden_size, dropout)
         self.pe.to(oce.CONFIG["DEVICE"])
 
-        self.trfm = nn.Transformer(d_model=hidden_size, nhead=4,
-            num_encoder_layers=n_layers, num_decoder_layers=n_layers, dim_feedforward=hidden_size)
+        self.trfm = nn.Transformer(
+            d_model=hidden_size,
+            nhead=4,
+            num_encoder_layers=n_layers,
+            num_decoder_layers=n_layers,
+            dim_feedforward=hidden_size,
+        )
         self.trfm.to(oce.CONFIG["DEVICE"])
 
         self.out = nn.Linear(hidden_size, out_size)
@@ -244,180 +266,209 @@ def __init__(self, in_size, hidden_size, out_size, n_layers, dropout=0.1):
     def forward(self, src):
         # src: (T,B)
         embedded = self.embed(src)  # (T,B,H)
-        embedded = self.pe(embedded) # (T,B,H)
-        hidden = self.trfm(embedded, embedded) # (T,B,H)
-        out = self.out(hidden) # (T,B,V)
-        out = F.log_softmax(out, dim=2) # (T,B,V)
-        return out # (T,B,V)
+        embedded = self.pe(embedded)  # (T,B,H)
+        hidden = self.trfm(embedded, embedded)  # (T,B,H)
+        out = self.out(hidden)  # (T,B,V)
+        out = F.log_softmax(out, dim=2)  # (T,B,V)
+        return out  # (T,B,V)
 
     def _encode(self, src):
         # src: (T,B)
         embedded = self.embed(src)  # (T,B,H)
-        embedded = self.pe(embedded) # (T,B,H)
+        embedded = self.pe(embedded)  # (T,B,H)
         output = embedded
         for i in range(self.trfm.encoder.num_layers - 1):
             output = self.trfm.encoder.layers[i](output, None)  # (T,B,H)
         penul = output.detach().cpu().numpy()
         output = self.trfm.encoder.layers[-1](output, None)  # (T,B,H)
         if self.trfm.encoder.norm:
-            output = self.trfm.encoder.norm(output) # (T,B,H)
+            output = self.trfm.encoder.norm(output)  # (T,B,H)
         output = output.detach().cpu().numpy()
         # mean, max, first*2
-        return np.hstack([np.mean(output, axis=0), np.max(output, axis=0), output[0,:,:], penul[0,:,:] ]) # (B,4H)
+        return np.hstack(
+            [
+                np.mean(output, axis=0),
+                np.max(output, axis=0),
+                output[0, :, :],
+                penul[0, :, :],
+            ]
+        )  # (B,4H)
 
     def encode(self, src):
         # src: (T,B)
         batch_size = src.shape[1]
-        if batch_size<=100:
+        if batch_size <= 100:
             return self._encode(src)
-        else: # Batch is too large to load
-            print('There are {:d} molecules. It will take a little time.'.format(batch_size))
-            st,ed = 0,100
-            out = self._encode(src[:,st:ed]) # (B,4H)
-            while ed<batch_size:
+        else:  # Batch is too large to load
+            print(
+                "There are {:d} molecules. It will take a little time.".format(
+                    batch_size
+                )
+            )
+            st, ed = 0, 100
+            out = self._encode(src[:, st:ed])  # (B,4H)
+            while ed < batch_size:
                 st += 100
                 ed += 100
-                out = np.concatenate([out, self._encode(src[:,st:ed])], axis=0)
+                out = np.concatenate([out, self._encode(src[:, st:ed])], axis=0)
             return out
 
-try:
-    from torch.utils.data import Dataset
-except ImportError:
-    mock_imports(globals(), "Dataset")
+
+from torch.utils.data import Dataset
+
 
 def split(sm):
-    '''
+    """
     function: Split SMILES into words. Care for Cl, Br, Si, Se, Na etc.
     input: A SMILES
     output: A string with space between words
-    '''
+    """
     arr = []
     i = 0
-    while i < len(sm)-1:
-        if not sm[i] in ['%', 'C', 'B', 'S', 'N', 'R', 'X', 'L', 'A', 'M', \
-                        'T', 'Z', 's', 't', 'H', '+', '-', 'K', 'F']:
+    while i < len(sm) - 1:
+        if not sm[i] in [
+            "%",
+            "C",
+            "B",
+            "S",
+            "N",
+            "R",
+            "X",
+            "L",
+            "A",
+            "M",
+            "T",
+            "Z",
+            "s",
+            "t",
+            "H",
+            "+",
+            "-",
+            "K",
+            "F",
+        ]:
             arr.append(sm[i])
             i += 1
-        elif sm[i]=='%':
-            arr.append(sm[i:i+3])
+        elif sm[i] == "%":
+            arr.append(sm[i : i + 3])
             i += 3
-        elif sm[i]=='C' and sm[i+1]=='l':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "C" and sm[i + 1] == "l":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='C' and sm[i+1]=='a':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "C" and sm[i + 1] == "a":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='C' and sm[i+1]=='u':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "C" and sm[i + 1] == "u":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='B' and sm[i+1]=='r':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "B" and sm[i + 1] == "r":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='B' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "B" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='B' and sm[i+1]=='a':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "B" and sm[i + 1] == "a":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='B' and sm[i+1]=='i':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "B" and sm[i + 1] == "i":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='S' and sm[i+1]=='i':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "S" and sm[i + 1] == "i":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='S' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "S" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='S' and sm[i+1]=='r':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "S" and sm[i + 1] == "r":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='N' and sm[i+1]=='a':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "N" and sm[i + 1] == "a":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='N' and sm[i+1]=='i':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "N" and sm[i + 1] == "i":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='R' and sm[i+1]=='b':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "R" and sm[i + 1] == "b":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='R' and sm[i+1]=='a':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "R" and sm[i + 1] == "a":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='X' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "X" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='L' and sm[i+1]=='i':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "L" and sm[i + 1] == "i":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='A' and sm[i+1]=='l':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "A" and sm[i + 1] == "l":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='A' and sm[i+1]=='s':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "A" and sm[i + 1] == "s":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='A' and sm[i+1]=='g':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "A" and sm[i + 1] == "g":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='A' and sm[i+1]=='u':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "A" and sm[i + 1] == "u":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='M' and sm[i+1]=='g':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "M" and sm[i + 1] == "g":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='M' and sm[i+1]=='n':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "M" and sm[i + 1] == "n":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='T' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "T" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='Z' and sm[i+1]=='n':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "Z" and sm[i + 1] == "n":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='s' and sm[i+1]=='i':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "s" and sm[i + 1] == "i":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='s' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "s" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='t' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "t" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='H' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "H" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='+' and sm[i+1]=='2':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "+" and sm[i + 1] == "2":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='+' and sm[i+1]=='3':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "+" and sm[i + 1] == "3":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='+' and sm[i+1]=='4':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "+" and sm[i + 1] == "4":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='-' and sm[i+1]=='2':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "-" and sm[i + 1] == "2":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='-' and sm[i+1]=='3':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "-" and sm[i + 1] == "3":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='-' and sm[i+1]=='4':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "-" and sm[i + 1] == "4":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='K' and sm[i+1]=='r':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "K" and sm[i + 1] == "r":
+            arr.append(sm[i : i + 2])
             i += 2
-        elif sm[i]=='F' and sm[i+1]=='e':
-            arr.append(sm[i:i+2])
+        elif sm[i] == "F" and sm[i + 1] == "e":
+            arr.append(sm[i : i + 2])
             i += 2
         else:
             arr.append(sm[i])
             i += 1
-    if i == len(sm)-1:
+    if i == len(sm) - 1:
         arr.append(sm[i])
-    return ' '.join(arr)
+    return " ".join(arr)
 
-class Seq2seqDataset(Dataset):
 
+class Seq2seqDataset(Dataset):
     def __init__(self, smiles, vocab, seq_len=220):
         self.smiles = smiles
         self.vocab = vocab
@@ -431,6 +482,6 @@ def __getitem__(self, item):
         sm = split(sm).split()
         content = [self.vocab.stoi.get(token, self.vocab.unk_index) for token in sm]
         X = [self.vocab.sos_index] + content + [self.vocab.eos_index]
-        padding = [self.vocab.pad_index]*(self.seq_len - len(X))
+        padding = [self.vocab.pad_index] * (self.seq_len - len(X))
         X.extend(padding)
-        return torch.tensor(X)
\ No newline at end of file
+        return torch.tensor(X)
diff --git a/src/olorenchemengine/external/MolCLR/main.py b/src/olorenchemengine/external/MolCLR/main.py
index 61ad44ed..53295fec 100644
--- a/src/olorenchemengine/external/MolCLR/main.py
+++ b/src/olorenchemengine/external/MolCLR/main.py
@@ -1,44 +1,49 @@
-import olorenchemengine as oce
-
-from olorenchemengine.base_class import log_arguments, BaseModel
-from olorenchemengine.representations import AtomFeaturizer, BondFeaturizer, TorchGeometricGraph, SMILESRepresentation, BaseVecRepresentation
-from olorenchemengine.internal import download_public_file
-
+import numpy as np
 from rdkit import Chem
 from rdkit.Chem.rdchem import BondType as BT
-
-try:
-    from torch_geometric.data import DataLoader
-except ImportError:
-    oce.mock_imports(globals(), "DataLoader")
-
-import numpy as np
+from torch_geometric.data import DataLoader
 from tqdm import tqdm
 
-ATOM_LIST = list(range(1,119))
+import olorenchemengine as oce
+from olorenchemengine.base_class import BaseModel, log_arguments
+from olorenchemengine.internal import download_public_file
+from olorenchemengine.representations import (
+    AtomFeaturizer,
+    BaseVecRepresentation,
+    BondFeaturizer,
+    SMILESRepresentation,
+    TorchGeometricGraph,
+)
+
+ATOM_LIST = list(range(1, 119))
 CHIRALITY_LIST = [
     Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
     Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
     Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
-    Chem.rdchem.ChiralType.CHI_OTHER
+    Chem.rdchem.ChiralType.CHI_OTHER,
 ]
 BOND_LIST = [BT.SINGLE, BT.DOUBLE, BT.TRIPLE, BT.AROMATIC]
 BONDDIR_LIST = [
     Chem.rdchem.BondDir.NONE,
     Chem.rdchem.BondDir.ENDUPRIGHT,
-    Chem.rdchem.BondDir.ENDDOWNRIGHT
+    Chem.rdchem.BondDir.ENDDOWNRIGHT,
 ]
 
+
 class MolCLR_AF(AtomFeaturizer):
     @property
     def length(self):
         return 2
 
     def convert(self, atom: Chem.Atom):
-        atom_feature = [ATOM_LIST.index(atom.GetAtomicNum()), CHIRALITY_LIST.index(atom.GetChiralTag())]
+        atom_feature = [
+            ATOM_LIST.index(atom.GetAtomicNum()),
+            CHIRALITY_LIST.index(atom.GetChiralTag()),
+        ]
         x = np.array(atom_feature)
         return x
 
+
 class MolCLR_BF(BondFeaturizer):
     @property
     def length(self):
@@ -46,29 +51,40 @@ def length(self):
 
     def convert(self, bond: Chem.Bond):
         edge_feature = [
-                BOND_LIST.index(bond.GetBondType()),
-                BONDDIR_LIST.index(bond.GetBondDir())
-            ]
+            BOND_LIST.index(bond.GetBondType()),
+            BONDDIR_LIST.index(bond.GetBondDir()),
+        ]
         x = np.array(edge_feature)
         return x
 
+
 class MolCLR_PYG(TorchGeometricGraph):
     @log_arguments
     def __init__(self):
-        super().__init__(MolCLR_AF(), MolCLR_BF(), log = False)
+        super().__init__(MolCLR_AF(), MolCLR_BF(), log=False)
+
 
 class MolCLR(BaseModel):
 
-    model_config = {"num_layer": 5,      # number of graph conv layers
-            "emb_dim": 300,                  # embedding dimension in graph conv layers
-            "feat_dim": 512,                 # output feature dimention
-            "drop_ratio": 0.3,               # dropout ratio
-            "pool": "mean"}
+    model_config = {
+        "num_layer": 5,  # number of graph conv layers
+        "emb_dim": 300,  # embedding dimension in graph conv layers
+        "feat_dim": 512,  # output feature dimention
+        "drop_ratio": 0.3,  # dropout ratio
+        "pool": "mean",
+    }
 
     @log_arguments
-    def __init__(self, model_type = "ginet", epochs = 100, batch_size = 32,
-                 init_lr = 0.0005, init_base_lr = 0.0001, weight_decay = 1e-6,
-                 **kwargs):
+    def __init__(
+        self,
+        model_type="ginet",
+        epochs=100,
+        batch_size=32,
+        init_lr=0.0005,
+        init_base_lr=0.0001,
+        weight_decay=1e-6,
+        **kwargs,
+    ):
         self.model_type = model_type
 
         self.config = {
@@ -76,7 +92,7 @@ def __init__(self, model_type = "ginet", epochs = 100, batch_size = 32,
             "batch_size": batch_size,
             "init_lr": init_lr,
             "init_base_lr": init_base_lr,
-            "weight_decay": weight_decay
+            "weight_decay": weight_decay,
         }
 
         self.representation = MolCLR_PYG()
@@ -85,50 +101,86 @@ def __init__(self, model_type = "ginet", epochs = 100, batch_size = 32,
 
     def preprocess(self, X, y, **kwargs):
         if y is None:
-            y = [None]*len(X)
+            y = [None] * len(X)
         return self.representation.convert(X, ys=y)
 
-    def _fit(self, X, y = None, **kwargs):
+    def _fit(self, X, y=None, **kwargs):
 
         import torch
-        from torch import nn
         import torch.nn.functional as F
-        from torch.utils.tensorboard import SummaryWriter
+        from sklearn.metrics import (
+            mean_absolute_error,
+            mean_squared_error,
+            roc_auc_score,
+        )
+        from torch import nn
         from torch.optim.lr_scheduler import CosineAnnealingLR
-        from sklearn.metrics import roc_auc_score, mean_squared_error, mean_absolute_error
+        from torch.utils.tensorboard import SummaryWriter
 
-        if self.setting == 'regression':
+        if self.setting == "regression":
             self.criterion = nn.MSELoss()
-        elif self.setting == 'classification':
-            self.criterion = nn.BCEWithLogitsLoss(reduction = "none")
+        elif self.setting == "classification":
+            self.criterion = nn.BCEWithLogitsLoss(reduction="none")
 
         if self.model_type == "ginet":
             from .model import GINet
+
             self.model = GINet(self.setting, **self.model_config)
             save_path = download_public_file("MolCLR/pretrained_gin.pth")
         elif self.model_type == "gcn":
             from .model import GCN
+
             self.model = GCN(self.setting, **self.model_config)
             save_path = download_public_file("MolCLR/pretrained_gcn.pth")
 
-        self.model.load_my_state_dict(torch.load(save_path, map_location = oce.CONFIG["MAP_LOCATION"]))
+        self.model.load_my_state_dict(
+            torch.load(save_path, map_location=oce.CONFIG["MAP_LOCATION"])
+        )
         self.model.to(oce.CONFIG["DEVICE"])
 
         layer_list = []
         for name, param in self.model.named_parameters():
-            if 'pred_lin' in name:
+            if "pred_lin" in name:
                 layer_list.append(name)
 
-        params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] in layer_list, self.model.named_parameters()))))
-        base_params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in layer_list, self.model.named_parameters()))))
+        params = list(
+            map(
+                lambda x: x[1],
+                list(
+                    filter(
+                        lambda kv: kv[0] in layer_list, self.model.named_parameters()
+                    )
+                ),
+            )
+        )
+        base_params = list(
+            map(
+                lambda x: x[1],
+                list(
+                    filter(
+                        lambda kv: kv[0] not in layer_list,
+                        self.model.named_parameters(),
+                    )
+                ),
+            )
+        )
 
         optimizer = torch.optim.Adam(
-            [{'params': base_params, 'lr': self.config['init_base_lr']}, {'params': params}],
-            self.config['init_lr'], weight_decay=self.config['weight_decay']
+            [
+                {"params": base_params, "lr": self.config["init_base_lr"]},
+                {"params": params},
+            ],
+            self.config["init_lr"],
+            weight_decay=self.config["weight_decay"],
         )
 
-        train_loader = DataLoader(X, batch_size=self.config["batch_size"], shuffle=True, num_workers=oce.CONFIG["NUM_WORKERS"])
-        for epoch_counter in range(self.config['epochs']):
+        train_loader = DataLoader(
+            X,
+            batch_size=self.config["batch_size"],
+            shuffle=True,
+            num_workers=oce.CONFIG["NUM_WORKERS"],
+        )
+        for epoch_counter in range(self.config["epochs"]):
             print(f"Epoch {epoch_counter+1}/{self.config['epochs']}")
             for bn, batch in enumerate(train_loader):
 
@@ -136,7 +188,7 @@ def _fit(self, X, y = None, **kwargs):
                 _, pred = self.model(batch)
                 y = batch.y.view(pred.shape).to(torch.float64)
 
-                #Loss matrix
+                # Loss matrix
                 loss_mat = self.criterion(pred.double(), y)
 
                 optimizer.zero_grad()
@@ -156,7 +208,12 @@ def _step(self, model, data, n_iter):
     def _predict(self, X, **kwargs):
         import torch
 
-        loader = DataLoader(X, batch_size=self.config["batch_size"], shuffle=False, num_workers=oce.CONFIG["NUM_WORKERS"])
+        loader = DataLoader(
+            X,
+            batch_size=self.config["batch_size"],
+            shuffle=False,
+            num_workers=oce.CONFIG["NUM_WORKERS"],
+        )
         self.model.eval()
         y_pred = []
 
@@ -169,18 +226,28 @@ def _predict(self, X, **kwargs):
 
         return np.vstack(y_pred).flatten()
 
+
 class MolCLRVecRep(BaseVecRepresentation):
 
-    model_config = {"num_layer": 5,      # number of graph conv layers
-            "emb_dim": 300,                  # embedding dimension in graph conv layers
-            "feat_dim": 512,                 # output feature dimention
-            "drop_ratio": 0.3,               # dropout ratio
-            "pool": "mean"}
+    model_config = {
+        "num_layer": 5,  # number of graph conv layers
+        "emb_dim": 300,  # embedding dimension in graph conv layers
+        "feat_dim": 512,  # output feature dimention
+        "drop_ratio": 0.3,  # dropout ratio
+        "pool": "mean",
+    }
 
     @log_arguments
-    def __init__(self, model_type = "ginet", epochs = 100, batch_size = 32,
-                 init_lr = 0.0005, init_base_lr = 0.0001, weight_decay = 1e-6,
-                 **kwargs):
+    def __init__(
+        self,
+        model_type="ginet",
+        epochs=100,
+        batch_size=32,
+        init_lr=0.0005,
+        init_base_lr=0.0001,
+        weight_decay=1e-6,
+        **kwargs,
+    ):
         self.model_type = model_type
 
         self.config = {
@@ -188,26 +255,34 @@ def __init__(self, model_type = "ginet", epochs = 100, batch_size = 32,
             "batch_size": batch_size,
             "init_lr": init_lr,
             "init_base_lr": init_base_lr,
-            "weight_decay": weight_decay
+            "weight_decay": weight_decay,
         }
 
         import torch
-        from torch import nn
         import torch.nn.functional as F
-        from torch.utils.tensorboard import SummaryWriter
+        from sklearn.metrics import (
+            mean_absolute_error,
+            mean_squared_error,
+            roc_auc_score,
+        )
+        from torch import nn
         from torch.optim.lr_scheduler import CosineAnnealingLR
-        from sklearn.metrics import roc_auc_score, mean_squared_error, mean_absolute_error
+        from torch.utils.tensorboard import SummaryWriter
 
         if self.model_type == "ginet":
             from .model import GINet
+
             self.model = GINet("regression", **self.model_config)
             save_path = download_public_file("MolCLR/pretrained_gin.pth")
         elif self.model_type == "gcn":
             from .model import GCN
+
             self.model = GCN("regression", **self.model_config)
             save_path = download_public_file("MolCLR/pretrained_gcn.pth")
 
-        self.model.load_my_state_dict(torch.load(save_path, map_location = oce.CONFIG["MAP_LOCATION"]))
+        self.model.load_my_state_dict(
+            torch.load(save_path, map_location=oce.CONFIG["MAP_LOCATION"])
+        )
         self.model.to(oce.CONFIG["DEVICE"])
 
     def _convert(self, X, **kwargs):
@@ -217,7 +292,12 @@ def convert(self, X, **kwargs):
         X = MolCLR_PYG().convert(X, ys=None)
         import torch
 
-        loader = DataLoader(X, batch_size=self.config["batch_size"], shuffle=True, num_workers=oce.CONFIG["NUM_WORKERS"])
+        loader = DataLoader(
+            X,
+            batch_size=self.config["batch_size"],
+            shuffle=True,
+            num_workers=oce.CONFIG["NUM_WORKERS"],
+        )
         self.model.eval()
         y_pred = []
 
@@ -228,4 +308,4 @@ def convert(self, X, **kwargs):
                 rep, _ = self.model(batch)
             y_pred.append(rep.cpu().detach().numpy())
 
-        return np.concatenate(y_pred)
\ No newline at end of file
+        return np.concatenate(y_pred)
diff --git a/src/olorenchemengine/external/MolCLR/model.py b/src/olorenchemengine/external/MolCLR/model.py
index 0360a60c..6a192656 100644
--- a/src/olorenchemengine/external/MolCLR/model.py
+++ b/src/olorenchemengine/external/MolCLR/model.py
@@ -1,35 +1,35 @@
 import math
 
 import torch
-from torch import nn
 import torch.nn.functional as F
-from torch.nn import Parameter
-from torch.nn import Linear, LayerNorm, ReLU
-
 import torch_sparse
-from torch_sparse import SparseTensor, matmul, fill_diag, sum as sparsesum, mul
-from torch_scatter import scatter
-from torch_scatter import scatter_add
-
-from torch_geometric.nn import MessagePassing, GCNConv
+from torch import nn
+from torch.nn import LayerNorm, Linear, Parameter, ReLU
+from torch_geometric.nn import (
+    GCNConv,
+    MessagePassing,
+    global_add_pool,
+    global_max_pool,
+    global_mean_pool,
+)
 from torch_geometric.utils import add_self_loops, degree, softmax
-from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
-
 from torch_geometric.utils.num_nodes import maybe_num_nodes
+from torch_scatter import scatter, scatter_add
+from torch_sparse import SparseTensor, fill_diag, matmul, mul
+from torch_sparse import sum as sparsesum
 
-num_atom_type = 119 # including the extra mask tokens
+num_atom_type = 119  # including the extra mask tokens
 num_chirality_tag = 3
 
-num_bond_type = 5 # including aromatic and self-loop edge
-num_bond_direction = 3 
+num_bond_type = 5  # including aromatic and self-loop edge
+num_bond_direction = 3
+
 
 class GINEConv(MessagePassing):
     def __init__(self, emb_dim):
         super(GINEConv, self).__init__()
         self.mlp = nn.Sequential(
-            nn.Linear(emb_dim, 2*emb_dim), 
-            nn.ReLU(), 
-            nn.Linear(2*emb_dim, emb_dim)
+            nn.Linear(emb_dim, 2 * emb_dim), nn.ReLU(), nn.Linear(2 * emb_dim, emb_dim)
         )
         self.edge_embedding1 = nn.Embedding(num_bond_type, emb_dim)
         self.edge_embedding2 = nn.Embedding(num_bond_direction, emb_dim)
@@ -43,12 +43,13 @@ def forward(self, x, edge_index, edge_attr):
 
         # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 # bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
         edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
 
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + \
-            self.edge_embedding2(edge_attr[:,1])
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
         return self.propagate(edge_index, x=x, edge_attr=edge_embeddings)
 
@@ -69,9 +70,17 @@ class GINet(nn.Module):
     Output:
         node representations
     """
-    def __init__(self, 
-        task='classification', num_layer=5, emb_dim=300, feat_dim=512, 
-        drop_ratio=0, pool='mean', pred_n_layer=2, pred_act='softplus'
+
+    def __init__(
+        self,
+        task="classification",
+        num_layer=5,
+        emb_dim=300,
+        feat_dim=512,
+        drop_ratio=0,
+        pool="mean",
+        pred_n_layer=2,
+        pred_act="softplus",
     ):
         super(GINet, self).__init__()
         self.num_layer = num_layer
@@ -95,54 +104,52 @@ def __init__(self,
         for layer in range(num_layer):
             self.batch_norms.append(nn.BatchNorm1d(emb_dim))
 
-        if pool == 'mean':
+        if pool == "mean":
             self.pool = global_mean_pool
-        elif pool == 'max':
+        elif pool == "max":
             self.pool = global_max_pool
-        elif pool == 'add':
+        elif pool == "add":
             self.pool = global_add_pool
         self.feat_lin = nn.Linear(self.emb_dim, self.feat_dim)
 
-        if self.task == 'classification':
+        if self.task == "classification":
             out_dim = 1
-        elif self.task == 'regression':
+        elif self.task == "regression":
             out_dim = 1
-        
+
         self.pred_n_layer = max(1, pred_n_layer)
 
-        if pred_act == 'relu':
+        if pred_act == "relu":
             pred_head = [
-                nn.Linear(self.feat_dim, self.feat_dim//2), 
-                nn.ReLU(inplace=True)
+                nn.Linear(self.feat_dim, self.feat_dim // 2),
+                nn.ReLU(inplace=True),
             ]
             for _ in range(self.pred_n_layer - 1):
-                pred_head.extend([
-                    nn.Linear(self.feat_dim//2, self.feat_dim//2), 
-                    nn.ReLU(inplace=True),
-                ])
-            pred_head.append(nn.Linear(self.feat_dim//2, out_dim))
-        elif pred_act == 'softplus':
-            pred_head = [
-                nn.Linear(self.feat_dim, self.feat_dim//2), 
-                nn.Softplus()
-            ]
+                pred_head.extend(
+                    [
+                        nn.Linear(self.feat_dim // 2, self.feat_dim // 2),
+                        nn.ReLU(inplace=True),
+                    ]
+                )
+            pred_head.append(nn.Linear(self.feat_dim // 2, out_dim))
+        elif pred_act == "softplus":
+            pred_head = [nn.Linear(self.feat_dim, self.feat_dim // 2), nn.Softplus()]
             for _ in range(self.pred_n_layer - 1):
-                pred_head.extend([
-                    nn.Linear(self.feat_dim//2, self.feat_dim//2), 
-                    nn.Softplus()
-                ])
+                pred_head.extend(
+                    [nn.Linear(self.feat_dim // 2, self.feat_dim // 2), nn.Softplus()]
+                )
         else:
-            raise ValueError('Undefined activation function')
-        
-        pred_head.append(nn.Linear(self.feat_dim//2, out_dim))
+            raise ValueError("Undefined activation function")
+
+        pred_head.append(nn.Linear(self.feat_dim // 2, out_dim))
         self.pred_head = nn.Sequential(*pred_head)
 
     def forward(self, data):
         x = data.x
         edge_index = data.edge_index
         edge_attr = data.edge_attr
-        
-        h = self.x_embedding1(x[:,0]) + self.x_embedding2(x[:,1])
+
+        h = self.x_embedding1(x[:, 0]) + self.x_embedding2(x[:, 1])
 
         for layer in range(self.num_layer):
             h = self.gnns[layer](h, edge_index, edge_attr)
@@ -154,7 +161,7 @@ def forward(self, data):
 
         h = self.pool(h, data.batch)
         h = self.feat_lin(h)
-        
+
         return h, self.pred_head(h)
 
     def load_my_state_dict(self, state_dict):
@@ -167,15 +174,16 @@ def load_my_state_dict(self, state_dict):
                 param = param.data
             own_state[name].copy_(param)
 
+
 def gcn_norm(edge_index, num_nodes=None):
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
-    edge_weight = torch.ones((edge_index.size(1), ), device=edge_index.device)
+    edge_weight = torch.ones((edge_index.size(1),), device=edge_index.device)
 
     row, col = edge_index[0], edge_index[1]
     deg = scatter_add(edge_weight, col, dim=0, dim_size=num_nodes)
     deg_inv_sqrt = deg.pow_(-0.5)
-    deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
+    deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float("inf"), 0)
     return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
 
 
@@ -208,11 +216,13 @@ def forward(self, x, edge_index, edge_attr):
 
         # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 #bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
         edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
-        
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])
+
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
         edge_index, __ = gcn_norm(edge_index)
 
@@ -235,7 +245,15 @@ def message_and_aggregate(self, adj_t, x):
 
 
 class GCN(nn.Module):
-    def __init__(self, task='classification', num_layer=5, emb_dim=300, feat_dim=256, drop_ratio=0, pool='mean'):
+    def __init__(
+        self,
+        task="classification",
+        num_layer=5,
+        emb_dim=300,
+        feat_dim=256,
+        drop_ratio=0,
+        pool="mean",
+    ):
         super(GCN, self).__init__()
         self.num_layer = num_layer
         self.emb_dim = emb_dim
@@ -262,28 +280,28 @@ def __init__(self, task='classification', num_layer=5, emb_dim=300, feat_dim=256
         for layer in range(num_layer):
             self.batch_norms.append(nn.BatchNorm1d(emb_dim))
 
-        if pool == 'mean':
+        if pool == "mean":
             self.pool = global_mean_pool
-        elif pool == 'add':
+        elif pool == "add":
             self.pool = global_add_pool
-        elif pool == 'max':
+        elif pool == "max":
             self.pool = global_max_pool
         else:
-            raise ValueError('Not defined pooling!')
-        
+            raise ValueError("Not defined pooling!")
+
         self.feat_lin = nn.Linear(self.emb_dim, self.feat_dim)
 
-        if self.task == 'classification':
+        if self.task == "classification":
             self.pred_head = nn.Sequential(
-                nn.Linear(self.feat_dim, self.feat_dim//2), 
+                nn.Linear(self.feat_dim, self.feat_dim // 2),
                 nn.Softplus(),
-                nn.Linear(self.feat_dim//2, 1)
+                nn.Linear(self.feat_dim // 2, 1),
             )
-        elif self.task == 'regression':
+        elif self.task == "regression":
             self.pred_head = nn.Sequential(
-                nn.Linear(self.feat_dim, self.feat_dim//2), 
+                nn.Linear(self.feat_dim, self.feat_dim // 2),
                 nn.Softplus(),
-                nn.Linear(self.feat_dim//2, 1)
+                nn.Linear(self.feat_dim // 2, 1),
             )
 
     def forward(self, data):
@@ -291,7 +309,7 @@ def forward(self, data):
         edge_index = data.edge_index
         edge_attr = data.edge_attr
 
-        h = self.x_embedding1(x[:,0]) + self.x_embedding2(x[:,1])
+        h = self.x_embedding1(x[:, 0]) + self.x_embedding2(x[:, 1])
 
         for layer in range(self.num_layer):
             h = self.gnns[layer](h, edge_index, edge_attr)
@@ -314,4 +332,4 @@ def load_my_state_dict(self, state_dict):
             if isinstance(param, nn.parameter.Parameter):
                 # backwards compatibility for serialized parameters
                 param = param.data
-            own_state[name].copy_(param)
\ No newline at end of file
+            own_state[name].copy_(param)
diff --git a/src/olorenchemengine/external/SPGNN/main.py b/src/olorenchemengine/external/SPGNN/main.py
index 8a6d5941..fa39ca2f 100644
--- a/src/olorenchemengine/external/SPGNN/main.py
+++ b/src/olorenchemengine/external/SPGNN/main.py
@@ -3,99 +3,106 @@
 `GitHub repository <https://github.com/snap-stanford/pretrain-gnns>`_
 """
 
-from olorenchemengine.internal import mock_imports
-
-try:
-    from torch_geometric.data import DataLoader
-except ImportError:
-    mock_imports(globals(), "DataLoader")
-
-try:
-    import torch
-    import torch.nn as nn
-    import torch.optim as optim
-except ImportError:
-    mock_imports(globals(), "torch", "nn", "optim")
-
+import io
+import os
+from typing import Dict, List, Optional, Tuple, Union
 
-from tqdm import tqdm
 import numpy as np
-
-from .model import GNN_graphpred, GNN, global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from rdkit import Chem
 from sklearn.metrics import roc_auc_score
+from torch_geometric.data import DataLoader
+from tqdm import tqdm
 
 import olorenchemengine as oce
-from olorenchemengine.representations import AtomFeaturizer, BaseRepresentation, BondFeaturizer, TorchGeometricGraph, BaseVecRepresentation
-
-import os
-import io
-from typing import Union, List, Tuple, Dict, Optional
-
-from olorenchemengine.base_class import BaseModel, log_arguments, QuantileTransformer
-from olorenchemengine.internal import download_public_file
-
-from rdkit import Chem
+from olorenchemengine.base_class import BaseModel, QuantileTransformer, log_arguments
+from olorenchemengine.internal import download_public_file, mock_imports
+from olorenchemengine.representations import (
+    AtomFeaturizer,
+    BaseRepresentation,
+    BaseVecRepresentation,
+    BondFeaturizer,
+    TorchGeometricGraph,
+)
+
+from .model import (
+    GNN,
+    GlobalAttention,
+    GNN_graphpred,
+    Set2Set,
+    global_add_pool,
+    global_max_pool,
+    global_mean_pool,
+)
 
 allowable_features = {
-    'possible_atomic_num_list' : list(range(1, 119)),
-    'possible_formal_charge_list' : [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5],
-    'possible_chirality_list' : [
+    "possible_atomic_num_list": list(range(1, 119)),
+    "possible_formal_charge_list": [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5],
+    "possible_chirality_list": [
         Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
         Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
         Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
-        Chem.rdchem.ChiralType.CHI_OTHER
+        Chem.rdchem.ChiralType.CHI_OTHER,
     ],
-    'possible_hybridization_list' : [
+    "possible_hybridization_list": [
         Chem.rdchem.HybridizationType.S,
-        Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,
-        Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.SP3D,
-        Chem.rdchem.HybridizationType.SP3D2, Chem.rdchem.HybridizationType.UNSPECIFIED
+        Chem.rdchem.HybridizationType.SP,
+        Chem.rdchem.HybridizationType.SP2,
+        Chem.rdchem.HybridizationType.SP3,
+        Chem.rdchem.HybridizationType.SP3D,
+        Chem.rdchem.HybridizationType.SP3D2,
+        Chem.rdchem.HybridizationType.UNSPECIFIED,
     ],
-    'possible_numH_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8],
-    'possible_implicit_valence_list' : [0, 1, 2, 3, 4, 5, 6],
-    'possible_degree_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
-    'possible_bonds' : [
+    "possible_numH_list": [0, 1, 2, 3, 4, 5, 6, 7, 8],
+    "possible_implicit_valence_list": [0, 1, 2, 3, 4, 5, 6],
+    "possible_degree_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+    "possible_bonds": [
         Chem.rdchem.BondType.SINGLE,
         Chem.rdchem.BondType.DOUBLE,
         Chem.rdchem.BondType.TRIPLE,
-        Chem.rdchem.BondType.AROMATIC
+        Chem.rdchem.BondType.AROMATIC,
     ],
-    'possible_bond_dirs' : [ # only for double bond stereo information
+    "possible_bond_dirs": [  # only for double bond stereo information
         Chem.rdchem.BondDir.NONE,
         Chem.rdchem.BondDir.ENDUPRIGHT,
-        Chem.rdchem.BondDir.ENDDOWNRIGHT
-    ]
+        Chem.rdchem.BondDir.ENDDOWNRIGHT,
+    ],
 }
 
+
 class SPGNN_AF(AtomFeaturizer):
     @property
     def length(self):
         return 2
 
     def convert(self, atom: Chem.Atom):
-        atom_feature = [allowable_features['possible_atomic_num_list'].index(
-            atom.GetAtomicNum())] + [allowable_features[
-            'possible_chirality_list'].index(atom.GetChiralTag())]
+        atom_feature = [
+            allowable_features["possible_atomic_num_list"].index(atom.GetAtomicNum())
+        ] + [allowable_features["possible_chirality_list"].index(atom.GetChiralTag())]
         x = np.array(atom_feature)
         return x
 
+
 class SPGNN_BF(BondFeaturizer):
     @property
     def length(self):
         return 2
 
     def convert(self, bond: Chem.Bond):
-        edge_feature = [allowable_features['possible_bonds'].index(
-                bond.GetBondType())] + [allowable_features[
-                                            'possible_bond_dirs'].index(
-                bond.GetBondDir())]
+        edge_feature = [
+            allowable_features["possible_bonds"].index(bond.GetBondType())
+        ] + [allowable_features["possible_bond_dirs"].index(bond.GetBondDir())]
         x = np.array(edge_feature)
         return x
 
+
 class SPGNN_PYG(TorchGeometricGraph):
     @log_arguments
     def __init__(self):
-        super().__init__(SPGNN_AF(), SPGNN_BF(), log = False)
+        super().__init__(SPGNN_AF(), SPGNN_BF(), log=False)
+
 
 class SPGNNVecRep(BaseVecRepresentation):
     """SPGNN_REP gives the output of the model presented in `Strategies for pre-training graph neural networks <https://arxiv.org/abs/1905.12265>`_
@@ -108,7 +115,8 @@ class SPGNNVecRep(BaseVecRepresentation):
         model_type (str): Type of model to use; default: "contextpred"
     """
 
-    available_pretrained_models = ["contextpred",
+    available_pretrained_models = [
+        "contextpred",
         "edgepred",
         "infomax",
         "masking",
@@ -119,27 +127,32 @@ class SPGNNVecRep(BaseVecRepresentation):
         "supervised",
         "gat_supervised_contextpred",
         "gat_supervised",
-        "gat_contextpred"]
+        "gat_contextpred",
+    ]
 
     @log_arguments
-    def __init__(self, model_type = "contextpred",
-        batch_size = 32,
-        epochs = 100,
-        lr = 0.001,
-        lr_scale = 1,
-        decay = 0,
-        num_layer = 5,
-        emb_dim = 300,
-        dropout_ratio = 0.5,
-        graph_pooling = "mean",
-        JK = "last",
-        gnn_type = "gin", **kwargs):
+    def __init__(
+        self,
+        model_type="contextpred",
+        batch_size=32,
+        epochs=100,
+        lr=0.001,
+        lr_scale=1,
+        decay=0,
+        num_layer=5,
+        emb_dim=300,
+        dropout_ratio=0.5,
+        graph_pooling="mean",
+        JK="last",
+        gnn_type="gin",
+        **kwargs,
+    ):
 
         self.model_type = model_type
 
         if "gat" in model_type:
             gnn_type = "gat"
-            emb_dim= 300
+            emb_dim = 300
 
         self.representation = SPGNN_PYG()
 
@@ -149,11 +162,13 @@ def __init__(self, model_type = "contextpred",
         self.map_location = oce.CONFIG["MAP_LOCATION"]
         self.device = oce.CONFIG["DEVICE"]
 
-        self.model = GNN(num_layer, emb_dim, JK, dropout_ratio, gnn_type = gnn_type)
+        self.model = GNN(num_layer, emb_dim, JK, dropout_ratio, gnn_type=gnn_type)
 
         input_model_file = download_public_file(f"SPGNN_saves/{self.model_type}.pth")
 
-        self.model.load_state_dict(torch.load(input_model_file, map_location=self.map_location))
+        self.model.load_state_dict(
+            torch.load(input_model_file, map_location=self.map_location)
+        )
         self.model.eval()
         self.model.to(self.device)
 
@@ -165,9 +180,11 @@ def __init__(self, model_type = "contextpred",
             self.pool = global_max_pool
         elif graph_pooling == "attention":
             if self.JK == "concat":
-                self.pool = GlobalAttention(gate_nn = torch.nn.Linear((self.num_layer + 1) * emb_dim, 1))
+                self.pool = GlobalAttention(
+                    gate_nn=torch.nn.Linear((self.num_layer + 1) * emb_dim, 1)
+                )
             else:
-                self.pool = GlobalAttention(gate_nn = torch.nn.Linear(emb_dim, 1))
+                self.pool = GlobalAttention(gate_nn=torch.nn.Linear(emb_dim, 1))
         elif graph_pooling[:-1] == "set2set":
             set2set_iter = int(graph_pooling[-1])
             if self.JK == "concat":
@@ -177,13 +194,16 @@ def __init__(self, model_type = "contextpred",
         else:
             raise ValueError("Invalid graph pooling type.")
 
-    def _convert(self, smiles: str, y: Union[int, float, np.number] = None) -> np.ndarray:
+    def _convert(
+        self, smiles: str, y: Union[int, float, np.number] = None
+    ) -> np.ndarray:
         assert Exception, "Please directly use convert method"
 
     def convert(self, smiles, **kwargs):
         X = self.representation.convert(smiles)
-        loader = DataLoader(X, batch_size=self.batch_size,
-            shuffle=False, num_workers=self.num_workers)
+        loader = DataLoader(
+            X, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
+        )
 
         self.model.eval()
         y_pred = []
@@ -198,9 +218,10 @@ def convert(self, smiles, **kwargs):
             y_pred.append(pred)
         return np.concatenate(y_pred)
 
+
 def train(model, device, loader, optimizer, setting):
     if setting == "classification":
-        criterion = nn.BCEWithLogitsLoss(reduction = "none")
+        criterion = nn.BCEWithLogitsLoss(reduction="none")
     else:
         criterion = nn.MSELoss()
     model.train()
@@ -210,7 +231,7 @@ def train(model, device, loader, optimizer, setting):
         pred = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch)
         y = batch.y.view(pred.shape).to(torch.float64)
 
-        #Loss matrix
+        # Loss matrix
         loss_mat = criterion(pred.double(), y)
 
         optimizer.zero_grad()
@@ -220,7 +241,8 @@ def train(model, device, loader, optimizer, setting):
 
         optimizer.step()
 
-def predict(model, device, loader, setting = "classification"):
+
+def predict(model, device, loader, setting="classification"):
     model.eval()
     y_pred = []
 
@@ -232,11 +254,12 @@ def predict(model, device, loader, setting = "classification"):
 
         y_pred.append(pred)
     if setting == "classification":
-        y_pred = nn.Sigmoid()(torch.cat(y_pred, dim = 0)).cpu().numpy()
+        y_pred = nn.Sigmoid()(torch.cat(y_pred, dim=0)).cpu().numpy()
     else:
-        y_pred = torch.cat(y_pred, dim = 0).cpu().numpy()
+        y_pred = torch.cat(y_pred, dim=0).cpu().numpy()
     return y_pred.flatten()
 
+
 class SPGNN(BaseModel):
     """SPGNN is the model presented in `Strategies for pre-training graph neural networks <https://arxiv.org/abs/1905.12265>`_
     `GitHub repository <https://github.com/snap-stanford/pretrain-gnns>`_
@@ -248,7 +271,8 @@ class SPGNN(BaseModel):
         model_type (str): Type of model to use; default: "contextpred"
     """
 
-    available_pretrained_models = ["contextpred",
+    available_pretrained_models = [
+        "contextpred",
         "edgepred",
         "infomax",
         "masking",
@@ -259,27 +283,32 @@ class SPGNN(BaseModel):
         "supervised",
         "gat_supervised_contextpred",
         "gat_supervised",
-        "gat_contextpred"]
+        "gat_contextpred",
+    ]
 
     @log_arguments
-    def __init__(self, model_type = "contextpred",
-        batch_size = 32,
-        epochs = 100,
-        lr = 0.001,
-        lr_scale = 1,
-        decay = 0,
-        num_layer = 5,
-        emb_dim = 300,
-        dropout_ratio = 0.5,
-        graph_pooling = "mean",
-        JK = "last",
-        gnn_type = "gin", **kwargs):
+    def __init__(
+        self,
+        model_type="contextpred",
+        batch_size=32,
+        epochs=100,
+        lr=0.001,
+        lr_scale=1,
+        decay=0,
+        num_layer=5,
+        emb_dim=300,
+        dropout_ratio=0.5,
+        graph_pooling="mean",
+        JK="last",
+        gnn_type="gin",
+        **kwargs,
+    ):
 
         self.model_type = model_type
 
         if "gat" in model_type:
             gnn_type = "gat"
-            emb_dim= 300
+            emb_dim = 300
         self.representation = SPGNN_PYG()
 
         self.batch_size = batch_size
@@ -289,13 +318,15 @@ def __init__(self, model_type = "contextpred",
         self.map_location = oce.CONFIG["MAP_LOCATION"]
         self.device = oce.CONFIG["DEVICE"]
 
-        self.model = GNN_graphpred(num_layer,
+        self.model = GNN_graphpred(
+            num_layer,
             emb_dim,
             1,
-            JK = JK,
-            drop_ratio = dropout_ratio,
-            graph_pooling = graph_pooling,
-            gnn_type = gnn_type)
+            JK=JK,
+            drop_ratio=dropout_ratio,
+            graph_pooling=graph_pooling,
+            gnn_type=gnn_type,
+        )
 
         input_model_file = download_public_file(f"SPGNN_saves/{self.model_type}.pth")
 
@@ -308,25 +339,33 @@ def __init__(self, model_type = "contextpred",
         model_param_group = []
         model_param_group.append({"params": self.model.gnn.parameters()})
         if graph_pooling == "attention":
-            model_param_group.append({"params": self.model.pool.parameters(), "lr":lr*lr_scale})
-        model_param_group.append({"params": self.model.graph_pred_linear.parameters(), "lr":lr* lr_scale})
+            model_param_group.append(
+                {"params": self.model.pool.parameters(), "lr": lr * lr_scale}
+            )
+        model_param_group.append(
+            {"params": self.model.graph_pred_linear.parameters(), "lr": lr * lr_scale}
+        )
 
         self.optimizer = optim.Adam(model_param_group, lr=lr, weight_decay=decay)
         self.sigmoid = nn.Sigmoid()
 
     def preprocess(self, X, y, **kwargs):
         if y is None:
-            y = [None]*len(X)
+            y = [None] * len(X)
         return self.representation.convert(X, ys=y)
 
     def _fit(self, X, y, **kwargs):
-        dataloader = DataLoader(X, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)
+        dataloader = DataLoader(
+            X, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers
+        )
         for epoch in range(self.epochs):
             train(self.model, self.device, dataloader, self.optimizer, self.setting)
 
     def _predict(self, X, **kwargs):
-        dataloader = DataLoader(X, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers)
-        return predict(self.model, self.device, dataloader, setting = self.setting)
+        dataloader = DataLoader(
+            X, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
+        )
+        return predict(self.model, self.device, dataloader, setting=self.setting)
 
     def _save(self) -> str:
         d = super()._save()
@@ -341,4 +380,4 @@ def _load(self, d):
 
     @classmethod
     def AllInstances(cls):
-        return [cls(model_type=mt) for mt in cls.available_pretrained_models]
\ No newline at end of file
+        return [cls(model_type=mt) for mt in cls.available_pretrained_models]
diff --git a/src/olorenchemengine/external/SPGNN/model.py b/src/olorenchemengine/external/SPGNN/model.py
index 0abd51d5..0ce573b7 100644
--- a/src/olorenchemengine/external/SPGNN/model.py
+++ b/src/olorenchemengine/external/SPGNN/model.py
@@ -1,25 +1,29 @@
 """ Defines the model architecture of SPGNN.
 """
 
+import torch
+import torch.nn.functional as F
+from torch_geometric.nn import (
+    GlobalAttention,
+    MessagePassing,
+    Set2Set,
+    global_add_pool,
+    global_max_pool,
+    global_mean_pool,
+)
+from torch_geometric.nn.inits import glorot, zeros
+from torch_geometric.utils import add_self_loops, degree, softmax
+from torch_scatter import scatter_add
+
 from olorenchemengine.internal import mock_imports
 
-try:
-    import torch
-    from torch_geometric.nn import MessagePassing
-    from torch_geometric.utils import add_self_loops, degree, softmax
-    from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set
-    import torch.nn.functional as F
-    from torch_scatter import scatter_add
-    from torch_geometric.nn.inits import glorot, zeros
-except:
-    mock_imports(globals(), "torch", "MessagePassing", "add_self_loops", "degree", "softmax", "global_add_pool", "global_mean_pool", "global_max_pool", "GlobalAttention", "Set2Set", "F", "scatter_add", "glorot", "zeros")
-
-num_atom_type = 120 #including the extra mask tokens
+num_atom_type = 120  # including the extra mask tokens
 num_chirality_tag = 3
 
-num_bond_type = 6 #including aromatic and self-loop edge, and extra masked tokens
+num_bond_type = 6  # including aromatic and self-loop edge, and extra masked tokens
 num_bond_direction = 3
 
+
 class GINConv(MessagePassing):
     """
     Extension of GIN aggregation to incorporate edge information by concatenation.
@@ -29,10 +33,15 @@ class GINConv(MessagePassing):
 
     See https://arxiv.org/abs/1810.00826
     """
-    def __init__(self, emb_dim, aggr = "add"):
+
+    def __init__(self, emb_dim, aggr="add"):
         super(GINConv, self).__init__(aggr=aggr)
-        #multi-layer perceptron
-        self.mlp = torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim), torch.nn.ReLU(), torch.nn.Linear(2*emb_dim, emb_dim))
+        # multi-layer perceptron
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(emb_dim, 2 * emb_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(2 * emb_dim, emb_dim),
+        )
         self.edge_embedding1 = torch.nn.Embedding(num_bond_type, emb_dim)
         self.edge_embedding2 = torch.nn.Embedding(num_bond_direction, emb_dim)
 
@@ -41,16 +50,18 @@ def __init__(self, emb_dim, aggr = "add"):
         self.aggr = aggr
 
     def forward(self, x, edge_index, edge_attr):
-        #add self loops in the edge space
-        edge_index, _ = add_self_loops(edge_index, num_nodes = x.size(0))
+        # add self loops in the edge space
+        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
 
-        #add features corresponding to self-loop edges.
+        # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 #bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
-        edge_attr = torch.cat((edge_attr, self_loop_attr), dim = 0)
+        edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
 
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
         return self.propagate(edge_index=edge_index, x=x, edge_attr=edge_embeddings)
 
@@ -62,8 +73,7 @@ def update(self, aggr_out):
 
 
 class GCNConv(MessagePassing):
-
-    def __init__(self, emb_dim, aggr = "add"):
+    def __init__(self, emb_dim, aggr="add"):
         super(GCNConv, self).__init__()
 
         self.emb_dim = emb_dim
@@ -78,40 +88,44 @@ def __init__(self, emb_dim, aggr = "add"):
 
     def norm(self, edge_index, num_nodes, dtype):
         ### assuming that self-loops have been already added in edge_index
-        edge_weight = torch.ones((edge_index.size(1), ), dtype=dtype,
-                                     device=edge_index.device)
+        edge_weight = torch.ones(
+            (edge_index.size(1),), dtype=dtype, device=edge_index.device
+        )
         row, col = edge_index
         deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
         deg_inv_sqrt = deg.pow(-0.5)
-        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
+        deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0
 
         return deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
 
-
     def forward(self, x, edge_index, edge_attr):
-        #add self loops in the edge space
-        edge_index = add_self_loops(edge_index, num_nodes = x.size(0))
+        # add self loops in the edge space
+        edge_index = add_self_loops(edge_index, num_nodes=x.size(0))
 
-        #add features corresponding to self-loop edges.
+        # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 #bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
-        edge_attr = torch.cat((edge_attr, self_loop_attr), dim = 0)
+        edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
 
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
         norm = self.norm(edge_index, x.size(0), x.dtype)
 
         x = self.linear(x)
 
-        return self.propagate(self.aggr, edge_index, x=x, edge_attr=edge_embeddings, norm = norm)
+        return self.propagate(
+            self.aggr, edge_index, x=x, edge_attr=edge_embeddings, norm=norm
+        )
 
     def message(self, x_j, edge_attr, norm):
         return norm.view(-1, 1) * (x_j + edge_attr)
 
 
 class GATConv(MessagePassing):
-    def __init__(self, emb_dim, heads=2, negative_slope=0.2, aggr = "add"):
+    def __init__(self, emb_dim, heads=2, negative_slope=0.2, aggr="add"):
         super(GATConv, self).__init__()
 
         self.aggr = aggr
@@ -138,25 +152,27 @@ def reset_parameters(self):
         zeros(self.bias)
 
     def forward(self, x, edge_index, edge_attr):
-        #add self loops in the edge space
-        edge_index, _ = add_self_loops(edge_index, num_nodes = x.size(0))
+        # add self loops in the edge space
+        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
 
-        #add features corresponding to self-loop edges.
+        # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 #bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
-        edge_attr = torch.cat((edge_attr, self_loop_attr), dim = 0)
+        edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
 
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
-        #x = self.weight_linear(x).view(-1, self.heads, self.emb_dim)
+        # x = self.weight_linear(x).view(-1, self.heads, self.emb_dim)
         x = self.weight_linear(x)
-        #return self.propagate(self.aggr, edge_index[0], x=x, edge_attr=edge_embeddings)
-        return self.propagate( edge_index, x=x, edge_attr=edge_embeddings)
+        # return self.propagate(self.aggr, edge_index[0], x=x, edge_attr=edge_embeddings)
+        return self.propagate(edge_index, x=x, edge_attr=edge_embeddings)
 
     def message(self, edge_index, x_i, x_j, edge_attr):
-        x_i=x_i.view(-1, self.heads, self.emb_dim)
-        x_j=x_j.view(-1, self.heads, self.emb_dim)
+        x_i = x_i.view(-1, self.heads, self.emb_dim)
+        x_j = x_j.view(-1, self.heads, self.emb_dim)
         edge_attr = edge_attr.view(-1, self.heads, self.emb_dim)
         x_j += edge_attr
 
@@ -166,7 +182,7 @@ def message(self, edge_index, x_i, x_j, edge_attr):
         alpha = softmax(alpha, edge_index[0])
 
         out = x_j * alpha.view(-1, self.heads, 1)
-        out = out.view(-1,self.heads*self.emb_dim)
+        out = out.view(-1, self.heads * self.emb_dim)
         return out
 
     def update(self, aggr_out):
@@ -176,7 +192,7 @@ def update(self, aggr_out):
 
 
 class GraphSAGEConv(MessagePassing):
-    def __init__(self, emb_dim, aggr = "mean"):
+    def __init__(self, emb_dim, aggr="mean"):
         super(GraphSAGEConv, self).__init__()
 
         self.emb_dim = emb_dim
@@ -190,16 +206,18 @@ def __init__(self, emb_dim, aggr = "mean"):
         self.aggr = aggr
 
     def forward(self, x, edge_index, edge_attr):
-        #add self loops in the edge space
-        edge_index = add_self_loops(edge_index, num_nodes = x.size(0))
+        # add self loops in the edge space
+        edge_index = add_self_loops(edge_index, num_nodes=x.size(0))
 
-        #add features corresponding to self-loop edges.
+        # add features corresponding to self-loop edges.
         self_loop_attr = torch.zeros(x.size(0), 2)
-        self_loop_attr[:,0] = 4 #bond type for self-loop edge
+        self_loop_attr[:, 0] = 4  # bond type for self-loop edge
         self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
-        edge_attr = torch.cat((edge_attr, self_loop_attr), dim = 0)
+        edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0)
 
-        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
 
         x = self.linear(x)
 
@@ -209,8 +227,7 @@ def message(self, x_j, edge_attr):
         return x_j + edge_attr
 
     def update(self, aggr_out):
-        return F.normalize(aggr_out, p = 2, dim = -1)
-
+        return F.normalize(aggr_out, p=2, dim=-1)
 
 
 class GNN(torch.nn.Module):
@@ -226,7 +243,8 @@ class GNN(torch.nn.Module):
     Output:
         node representations
     """
-    def __init__(self, num_layer, emb_dim, JK = "last", drop_ratio = 0, gnn_type = "gin"):
+
+    def __init__(self, num_layer, emb_dim, JK="last", drop_ratio=0, gnn_type="gin"):
         super(GNN, self).__init__()
         self.num_layer = num_layer
         self.drop_ratio = drop_ratio
@@ -245,7 +263,7 @@ def __init__(self, num_layer, emb_dim, JK = "last", drop_ratio = 0, gnn_type = "
         self.gnns = torch.nn.ModuleList()
         for layer in range(num_layer):
             if gnn_type == "gin":
-                self.gnns.append(GINConv(emb_dim, aggr = "add"))
+                self.gnns.append(GINConv(emb_dim, aggr="add"))
             elif gnn_type == "gcn":
                 self.gnns.append(GCNConv(emb_dim))
             elif gnn_type == "gat":
@@ -258,7 +276,7 @@ def __init__(self, num_layer, emb_dim, JK = "last", drop_ratio = 0, gnn_type = "
         for layer in range(num_layer):
             self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))
 
-    #def forward(self, x, edge_index, edge_attr):
+    # def forward(self, x, edge_index, edge_attr):
     def forward(self, *argv):
         if len(argv) == 3:
             x, edge_index, edge_attr = argv[0], argv[1], argv[2]
@@ -268,31 +286,31 @@ def forward(self, *argv):
         else:
             raise ValueError("unmatched number of arguments.")
 
-        x = self.x_embedding1(x[:,0]) + self.x_embedding2(x[:,1])
+        x = self.x_embedding1(x[:, 0]) + self.x_embedding2(x[:, 1])
 
         h_list = [x]
         for layer in range(self.num_layer):
             h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
             h = self.batch_norms[layer](h)
-            #h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
+            # h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
             if layer == self.num_layer - 1:
-                #remove relu for the last layer
-                h = F.dropout(h, self.drop_ratio, training = self.training)
+                # remove relu for the last layer
+                h = F.dropout(h, self.drop_ratio, training=self.training)
             else:
-                h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
+                h = F.dropout(F.relu(h), self.drop_ratio, training=self.training)
             h_list.append(h)
 
         ### Different implementations of Jk-concat
         if self.JK == "concat":
-            node_representation = torch.cat(h_list, dim = 1)
+            node_representation = torch.cat(h_list, dim=1)
         elif self.JK == "last":
             node_representation = h_list[-1]
         elif self.JK == "max":
             h_list = [h.unsqueeze_(0) for h in h_list]
-            node_representation = torch.max(torch.cat(h_list, dim = 0), dim = 0)[0]
+            node_representation = torch.max(torch.cat(h_list, dim=0), dim=0)[0]
         elif self.JK == "sum":
             h_list = [h.unsqueeze_(0) for h in h_list]
-            node_representation = torch.sum(torch.cat(h_list, dim = 0), dim = 0)[0]
+            node_representation = torch.sum(torch.cat(h_list, dim=0), dim=0)[0]
 
         return node_representation
 
@@ -312,7 +330,17 @@ class GNN_graphpred(torch.nn.Module):
     See https://arxiv.org/abs/1810.00826
     JK-net: https://arxiv.org/abs/1806.03536
     """
-    def __init__(self, num_layer, emb_dim, num_tasks, JK = "last", drop_ratio = 0, graph_pooling = "mean", gnn_type = "gin"):
+
+    def __init__(
+        self,
+        num_layer,
+        emb_dim,
+        num_tasks,
+        JK="last",
+        drop_ratio=0,
+        graph_pooling="mean",
+        gnn_type="gin",
+    ):
         super(GNN_graphpred, self).__init__()
         self.num_layer = num_layer
         self.drop_ratio = drop_ratio
@@ -323,9 +351,9 @@ def __init__(self, num_layer, emb_dim, num_tasks, JK = "last", drop_ratio = 0, g
         if self.num_layer < 2:
             raise ValueError("Number of GNN layers must be greater than 1.")
 
-        self.gnn = GNN(num_layer, emb_dim, JK, drop_ratio, gnn_type = gnn_type)
+        self.gnn = GNN(num_layer, emb_dim, JK, drop_ratio, gnn_type=gnn_type)
 
-        #Different kind of graph pooling
+        # Different kind of graph pooling
         if graph_pooling == "sum":
             self.pool = global_add_pool
         elif graph_pooling == "mean":
@@ -334,9 +362,11 @@ def __init__(self, num_layer, emb_dim, num_tasks, JK = "last", drop_ratio = 0, g
             self.pool = global_max_pool
         elif graph_pooling == "attention":
             if self.JK == "concat":
-                self.pool = GlobalAttention(gate_nn = torch.nn.Linear((self.num_layer + 1) * emb_dim, 1))
+                self.pool = GlobalAttention(
+                    gate_nn=torch.nn.Linear((self.num_layer + 1) * emb_dim, 1)
+                )
             else:
-                self.pool = GlobalAttention(gate_nn = torch.nn.Linear(emb_dim, 1))
+                self.pool = GlobalAttention(gate_nn=torch.nn.Linear(emb_dim, 1))
         elif graph_pooling[:-1] == "set2set":
             set2set_iter = int(graph_pooling[-1])
             if self.JK == "concat":
@@ -346,18 +376,22 @@ def __init__(self, num_layer, emb_dim, num_tasks, JK = "last", drop_ratio = 0, g
         else:
             raise ValueError("Invalid graph pooling type.")
 
-        #For graph-level binary classification
+        # For graph-level binary classification
         if graph_pooling[:-1] == "set2set":
             self.mult = 2
         else:
             self.mult = 1
 
         if self.JK == "concat":
-            self.graph_pred_linear = torch.nn.Linear(self.mult * (self.num_layer + 1) * self.emb_dim, self.num_tasks)
+            self.graph_pred_linear = torch.nn.Linear(
+                self.mult * (self.num_layer + 1) * self.emb_dim, self.num_tasks
+            )
         else:
-            self.graph_pred_linear = torch.nn.Linear(self.mult * self.emb_dim, self.num_tasks)
+            self.graph_pred_linear = torch.nn.Linear(
+                self.mult * self.emb_dim, self.num_tasks
+            )
 
-    def from_pretrained(self, model_file, map_location = "cpu"):
+    def from_pretrained(self, model_file, map_location="cpu"):
         self.gnn.load_state_dict(torch.load(model_file, map_location=map_location))
 
     def forward(self, *argv):
@@ -365,10 +399,15 @@ def forward(self, *argv):
             x, edge_index, edge_attr, batch = argv[0], argv[1], argv[2], argv[3]
         elif len(argv) == 1:
             data = argv[0]
-            x, edge_index, edge_attr, batch = data.x, data.edge_index, data.edge_attr, data.batch
+            x, edge_index, edge_attr, batch = (
+                data.x,
+                data.edge_index,
+                data.edge_attr,
+                data.batch,
+            )
         else:
             raise ValueError("unmatched number of arguments.")
 
         node_representation = self.gnn(x, edge_index, edge_attr)
 
-        return self.graph_pred_linear(self.pool(node_representation, batch))
\ No newline at end of file
+        return self.graph_pred_linear(self.pool(node_representation, batch))
diff --git a/src/olorenchemengine/external/__init__.py b/src/olorenchemengine/external/__init__.py
index 2502ce97..d452ec20 100644
--- a/src/olorenchemengine/external/__init__.py
+++ b/src/olorenchemengine/external/__init__.py
@@ -1,9 +1,9 @@
+from .ChemProp.main import ChemPropModel
 from .GINNetwork.main import GINModel
-from .SPGNN.main import SPGNN, SPGNNVecRep
+from .HondaSTRep.main import HondaSTRep
+from .HondaSTRep.operations import WordVocab
 from .mol2vec.main import Mol2Vec
+from .MolCLR.main import MolCLR, MolCLRVecRep
 from .piCalculax import calc_pI
+from .SPGNN.main import SPGNN, SPGNNVecRep
 from .stoned import *
-from .ChemProp.main import ChemPropModel
-from .HondaSTRep.main import HondaSTRep
-from .HondaSTRep.operations import WordVocab
-from .MolCLR.main import MolCLR, MolCLRVecRep
\ No newline at end of file
diff --git a/src/olorenchemengine/external/mol2vec/main.py b/src/olorenchemengine/external/mol2vec/main.py
index 304cc077..50da87a7 100644
--- a/src/olorenchemengine/external/mol2vec/main.py
+++ b/src/olorenchemengine/external/mol2vec/main.py
@@ -1,40 +1,46 @@
 import os
+from typing import Any, List, Union
 
 import numpy as np
-
-from typing import List, Union, Any
 from rdkit import Chem
 
 import olorenchemengine as oce
-
 from olorenchemengine.base_class import log_arguments
-from olorenchemengine.representations import BaseVecRepresentation
 from olorenchemengine.internal import download_public_file
-from .operations import MolSentence, sentences2vec, mol2alt_sentence, DfVec
+from olorenchemengine.representations import BaseVecRepresentation
+
+from .operations import DfVec, MolSentence, mol2alt_sentence, sentences2vec
 
-class Mol2Vec(BaseVecRepresentation):
 
+class Mol2Vec(BaseVecRepresentation):
     @log_arguments
     def __init__(self):
-        print("Requires gensim installation https://github.com/RaRe-Technologies/gensim")
+        print(
+            "Requires gensim installation https://github.com/RaRe-Technologies/gensim"
+        )
 
         model_path = download_public_file("mol2vec/model_300dim.pkl")
-        
+
         oce.import_or_install("gensim")
-        
+
         from gensim.models import word2vec
-        
+
         self.model = word2vec.Word2Vec.load(model_path)
         super().__init__(log=False)
 
     def _convert(self, s: str) -> np.ndarray:
         mol = Chem.MolFromSmiles(s)
-        sentence = MolSentence(mol2alt_sentence(mol,1))
-        X = DfVec(sentences2vec([sentence], self.model, unseen='UNK')[0]).vec
+        sentence = MolSentence(mol2alt_sentence(mol, 1))
+        X = DfVec(sentences2vec([sentence], self.model, unseen="UNK")[0]).vec
         return np.array(X)
 
-    def _convert_list(self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None) -> List[np.ndarray]:
+    def _convert_list(
+        self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None
+    ) -> List[np.ndarray]:
         mols = [Chem.MolFromSmiles(s) for s in smiles_list]
         sentences = [MolSentence(mol2alt_sentence(mol, 1)) for mol in mols]
-        X = [DfVec(sentence).vec for sentence in sentences2vec(sentences, self.model, unseen='UNK')]
-        return np.array(X)
\ No newline at end of file
+        X = [
+            DfVec(sentence).vec
+            for sentence in sentences2vec(sentences, self.model, unseen="UNK")
+        ]
+        return np.array(X)
diff --git a/src/olorenchemengine/external/mol2vec/operations.py b/src/olorenchemengine/external/mol2vec/operations.py
index 5b902e4a..c79c1250 100644
--- a/src/olorenchemengine/external/mol2vec/operations.py
+++ b/src/olorenchemengine/external/mol2vec/operations.py
@@ -35,27 +35,29 @@
 ==============================
 w. PR for gensim 4.0 from tonyreina and samoturk
 """
-from tqdm import tqdm
+import timeit
+
 import numpy as np
 import pandas as pd
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.Chem import PandasTools
-import timeit
 from joblib import Parallel, delayed
+from rdkit import Chem
+from rdkit.Chem import AllChem, PandasTools
+from tqdm import tqdm
+
 
 class DfVec(object):
     """
     Helper class to store vectors in a pandas DataFrame
-    
-    Parameters  
-    ---------- 
+
+    Parameters
+    ----------
     vec: np.array
     """
+
     def __init__(self, vec):
         self.vec = vec
         if type(self.vec) != np.ndarray:
-            raise TypeError('numpy.ndarray expected, got %s' % type(self.vec))
+            raise TypeError("numpy.ndarray expected, got %s" % type(self.vec))
 
     def __str__(self):
         return "%s dimensional vector" % str(self.vec.shape)
@@ -69,18 +71,18 @@ def __len__(self):
 
 
 class MolSentence:
-    """Class for storing mol sentences in pandas DataFrame
-    """
+    """Class for storing mol sentences in pandas DataFrame"""
+
     def __init__(self, sentence):
         self.sentence = sentence
         if type(self.sentence[0]) != str:
-            raise TypeError('List with strings expected')
+            raise TypeError("List with strings expected")
 
     def __len__(self):
         return len(self.sentence)
 
     def __str__(self):  # String representation
-        return 'MolSentence with %i words' % len(self.sentence)
+        return "MolSentence with %i words" % len(self.sentence)
 
     __repr__ = __str__  # Default representation
 
@@ -100,7 +102,6 @@ def __iter__(self):  # Iterate over words (for word in MolSentence:...)
     _repr_html_ = __str__
 
 
-
 def mol2sentence(mol, radius):
 
     """Calculates ECFP (Morgan fingerprint) and returns identifiers of substructures as 'sentence' (string).
@@ -108,11 +109,11 @@ def mol2sentence(mol, radius):
     combined.
     NOTE: Words are ALWAYS reordered according to atom order in the input mol object.
     NOTE: Due to the way how Morgan FPs are generated, number of identifiers at each radius is smaller
-    
+
     Parameters
     ----------
     mol : rdkit.Chem.rdchem.Mol
-    radius : float 
+    radius : float
         Fingerprint radius
     Returns
     -------
@@ -123,25 +124,29 @@ def mol2sentence(mol, radius):
     """
     radii = list(range(int(radius) + 1))
     info = {}
-    _ = AllChem.GetMorganFingerprint(mol, radius, bitInfo=info)  # info: dictionary identifier, atom_idx, radius
+    _ = AllChem.GetMorganFingerprint(
+        mol, radius, bitInfo=info
+    )  # info: dictionary identifier, atom_idx, radius
 
     mol_atoms = [a.GetIdx() for a in mol.GetAtoms()]
     dict_atoms = {x: {r: None for r in radii} for x in mol_atoms}
 
     for element in info:
         for atom_idx, radius_at in info[element]:
-            dict_atoms[atom_idx][radius_at] = element  # {atom number: {fp radius: identifier}}
+            dict_atoms[atom_idx][
+                radius_at
+            ] = element  # {atom number: {fp radius: identifier}}
 
     # iterate over all atoms and radii
     identifier_sentences = []
-    
+
     for r in radii:  # iterate over radii to get one sentence per radius
         identifiers = []
         for atom in dict_atoms:  # iterate over atoms
             # get one sentence per radius
             identifiers.append(dict_atoms[atom][r])
         identifier_sentences.append(list(map(str, [x for x in identifiers if x])))
-    
+
     # merge identifiers alternating radius to sentence: atom 0 radius0, atom 0 radius 1, etc.
     identifiers_alt = []
     for atom in dict_atoms:  # iterate over atoms
@@ -160,13 +165,13 @@ def mol2alt_sentence(mol, radius):
     combined.
     NOTE: Words are ALWAYS reordered according to atom order in the input mol object.
     NOTE: Due to the way how Morgan FPs are generated, number of identifiers at each radius is smaller
-    
+
     Parameters
     ----------
     mol : rdkit.Chem.rdchem.Mol
-    radius : float 
+    radius : float
         Fingerprint radius
-    
+
     Returns
     -------
     list
@@ -175,14 +180,18 @@ def mol2alt_sentence(mol, radius):
     """
     radii = list(range(int(radius) + 1))
     info = {}
-    _ = AllChem.GetMorganFingerprint(mol, radius, bitInfo=info)  # info: dictionary identifier, atom_idx, radius
+    _ = AllChem.GetMorganFingerprint(
+        mol, radius, bitInfo=info
+    )  # info: dictionary identifier, atom_idx, radius
 
     mol_atoms = [a.GetIdx() for a in mol.GetAtoms()]
     dict_atoms = {x: {r: None for r in radii} for x in mol_atoms}
 
     for element in info:
         for atom_idx, radius_at in info[element]:
-            dict_atoms[atom_idx][radius_at] = element  # {atom number: {fp radius: identifier}}
+            dict_atoms[atom_idx][
+                radius_at
+            ] = element  # {atom number: {fp radius: identifier}}
 
     # merge identifiers alternating radius to sentence: atom 0 radius0, atom 0 radius 1, etc.
     identifiers_alt = []
@@ -196,8 +205,7 @@ def mol2alt_sentence(mol, radius):
 
 
 def _parallel_job(mol, r):
-    """Helper function for joblib jobs
-    """
+    """Helper function for joblib jobs"""
     if mol is not None:
         smiles = Chem.MolToSmiles(mol)
         mol = Chem.MolFromSmiles(smiles)
@@ -210,13 +218,13 @@ def _read_smi(file_name):
         line = file_name.readline()
         if not line:
             break
-        yield Chem.MolFromSmiles(line.split('\t')[0])
+        yield Chem.MolFromSmiles(line.split("\t")[0])
 
 
-def generate_corpus(in_file, out_file, r, sentence_type='alt', n_jobs=1):
+def generate_corpus(in_file, out_file, r, sentence_type="alt", n_jobs=1):
 
     """Generates corpus file from sdf
-    
+
     Parameters
     ----------
     in_file : str
@@ -226,8 +234,8 @@ def generate_corpus(in_file, out_file, r, sentence_type='alt', n_jobs=1):
     r : int
         Radius of morgan fingerprint
     sentence_type : str
-        Options:    'all' - generates all corpus files for all types of sentences, 
-                    'alt' - generates a corpus file with only combined alternating sentence, 
+        Options:    'all' - generates all corpus files for all types of sentences,
+                    'alt' - generates a corpus file with only combined alternating sentence,
                     'individual' - generates corpus files for each radius
     n_jobs : int
         Number of cores to use (only 'alt' sentence type is parallelized)
@@ -236,49 +244,54 @@ def generate_corpus(in_file, out_file, r, sentence_type='alt', n_jobs=1):
     """
 
     # File type detection
-    in_split = in_file.split('.')
-    if in_split[-1].lower() not in ['sdf', 'smi', 'ism', 'gz']:
-        raise ValueError('File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)')
+    in_split = in_file.split(".")
+    if in_split[-1].lower() not in ["sdf", "smi", "ism", "gz"]:
+        raise ValueError("File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)")
     gzipped = False
-    if in_split[-1].lower() == 'gz':
+    if in_split[-1].lower() == "gz":
         gzipped = True
-        if in_split[-2].lower() not in ['sdf', 'smi', 'ism']:
-            raise ValueError('File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)')
+        if in_split[-2].lower() not in ["sdf", "smi", "ism"]:
+            raise ValueError(
+                "File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)"
+            )
 
     file_handles = []
-    
+
     # write only files which contain corpus
-    if (sentence_type == 'individual') or (sentence_type == 'all'):
-        
-        f1 = open(out_file+'_r0.corpus', "w")
-        f2 = open(out_file+'_r1.corpus', "w")
+    if (sentence_type == "individual") or (sentence_type == "all"):
+
+        f1 = open(out_file + "_r0.corpus", "w")
+        f2 = open(out_file + "_r1.corpus", "w")
         file_handles.append(f1)
         file_handles.append(f2)
 
-    if (sentence_type == 'alt') or (sentence_type == 'all'):
+    if (sentence_type == "alt") or (sentence_type == "all"):
         f3 = open(out_file, "w")
         file_handles.append(f3)
-    
+
     if gzipped:
         import gzip
-        if in_split[-2].lower() == 'sdf':
-            mols_file = gzip.open(in_file, mode='r')
+
+        if in_split[-2].lower() == "sdf":
+            mols_file = gzip.open(in_file, mode="r")
             suppl = Chem.ForwardSDMolSupplier(mols_file)
         else:
-            mols_file = gzip.open(in_file, mode='rt')
+            mols_file = gzip.open(in_file, mode="rt")
             suppl = _read_smi(mols_file)
     else:
-        if in_split[-1].lower() == 'sdf':
+        if in_split[-1].lower() == "sdf":
             suppl = Chem.ForwardSDMolSupplier(in_file)
         else:
-            mols_file = open(in_file, mode='rt')
+            mols_file = open(in_file, mode="rt")
             suppl = _read_smi(mols_file)
 
-    if sentence_type == 'alt':  # This can run parallelized
-        result = Parallel(n_jobs=n_jobs, verbose=1)(delayed(_parallel_job)(mol, r) for mol in suppl)
+    if sentence_type == "alt":  # This can run parallelized
+        result = Parallel(n_jobs=n_jobs, verbose=1)(
+            delayed(_parallel_job)(mol, r) for mol in suppl
+        )
         for i, line in enumerate(result):
-            f3.write(str(line) + '\n')
-        print('% molecules successfully processed.')
+            f3.write(str(line) + "\n")
+        print("% molecules successfully processed.")
 
     else:
         for mol in suppl:
@@ -292,12 +305,12 @@ def generate_corpus(in_file, out_file, r, sentence_type='alt', n_jobs=1):
                 alternating_sentence_r0r1 = " ".join(alternating_sentence)
 
                 if len(smiles) != 0:
-                    if (sentence_type == 'individual') or (sentence_type == 'all'):
-                        f1.write(str(identifier_sentence_r0)+'\n')
-                        f2.write(str(identifier_sentence_r1)+'\n')
+                    if (sentence_type == "individual") or (sentence_type == "all"):
+                        f1.write(str(identifier_sentence_r0) + "\n")
+                        f2.write(str(identifier_sentence_r1) + "\n")
 
-                    if (sentence_type == 'alt') or (sentence_type == 'all'):
-                        f3.write(str(alternating_sentence_r0r1)+'\n')
+                    if (sentence_type == "alt") or (sentence_type == "all"):
+                        f3.write(str(alternating_sentence_r0r1) + "\n")
 
     for fh in file_handles:
         fh.close()
@@ -311,7 +324,7 @@ def _read_corpus(file_name):
         yield line.split()
 
 
-def insert_unk(corpus, out_corpus, threshold=3, uncommon='UNK'):
+def insert_unk(corpus, out_corpus, threshold=3, uncommon="UNK"):
     """Handling of uncommon "words" (i.e. identifiers). It finds all least common identifiers (defined by threshold) and
     replaces them by 'uncommon' string.
     Parameters
@@ -330,7 +343,7 @@ def insert_unk(corpus, out_corpus, threshold=3, uncommon='UNK'):
     # Find least common identifiers in corpus
     f = open(corpus)
     unique = {}
-    for i, x in tqdm(enumerate(_read_corpus(f)), desc='Counting identifiers in corpus'):
+    for i, x in tqdm(enumerate(_read_corpus(f)), desc="Counting identifiers in corpus"):
         for identifier in x:
             if identifier not in unique:
                 unique[identifier] = 1
@@ -341,8 +354,8 @@ def insert_unk(corpus, out_corpus, threshold=3, uncommon='UNK'):
     f.close()
 
     f = open(corpus)
-    fw = open(out_corpus, mode='w')
-    for line in tqdm(_read_corpus(f), total=n_lines, desc='Inserting %s' % uncommon):
+    fw = open(out_corpus, mode="w")
+    for line in tqdm(_read_corpus(f), total=n_lines, desc="Inserting %s" % uncommon):
         intersection = set(line) & least_common
         if len(intersection) > 0:
             new_line = []
@@ -351,18 +364,26 @@ def insert_unk(corpus, out_corpus, threshold=3, uncommon='UNK'):
                     new_line.append(uncommon)
                 else:
                     new_line.append(item)
-            fw.write(" ".join(new_line) + '\n')
+            fw.write(" ".join(new_line) + "\n")
         else:
-            fw.write(" ".join(line) + '\n')
+            fw.write(" ".join(line) + "\n")
     f.close()
     fw.close()
 
 
-def train_word2vec_model(infile_name, outfile_name=None, vector_size=100, window=10, min_count=3, n_jobs=1,
-                         method='skip-gram', **kwargs):
+def train_word2vec_model(
+    infile_name,
+    outfile_name=None,
+    vector_size=100,
+    window=10,
+    min_count=3,
+    n_jobs=1,
+    method="skip-gram",
+    **kwargs
+):
     """Trains word2vec (Mol2vec, ProtVec) model on corpus file extracted from molecule/protein sequences.
     The corpus file is treated as LineSentence corpus (one sentence = one line, words separated by whitespaces)
-    
+
     Parameters
     ----------
     infile_name : str
@@ -379,36 +400,44 @@ def train_word2vec_model(infile_name, outfile_name=None, vector_size=100, window
         Number of cpu cores used for calculation
     method : str
         Method to use in model training. Options cbow and skip-gram, default: skip-gram)
-    
+
     Returns
     -------
     word2vec.Word2Vec
     """
-    if method.lower() == 'skip-gram':
+    if method.lower() == "skip-gram":
         sg = 1
-    elif method.lower() == 'cbow':
+    elif method.lower() == "cbow":
         sg = 0
     else:
-        raise ValueError('skip-gram or cbow are only valid options')
-  
+        raise ValueError("skip-gram or cbow are only valid options")
+
     start = timeit.default_timer()
 
     from gensim.models import word2vec
+
     corpus = word2vec.LineSentence(infile_name)
-    model = word2vec.Word2Vec(corpus, size=vector_size, window=window, min_count=min_count, workers=n_jobs, sg=sg,
-                              **kwargs)
+    model = word2vec.Word2Vec(
+        corpus,
+        size=vector_size,
+        window=window,
+        min_count=min_count,
+        workers=n_jobs,
+        sg=sg,
+        **kwargs
+    )
     if outfile_name:
         model.save(outfile_name)
-    
+
     stop = timeit.default_timer()
-    print('Runtime: ', round((stop - start)/60, 2), ' minutes')
+    print("Runtime: ", round((stop - start) / 60, 2), " minutes")
     return model
-    
-    
+
+
 def remove_salts_solvents(smiles, hac=3):
     """Remove solvents and ions have max 'hac' heavy atoms. This function removes any fragment in molecule that has
     number of heavy atoms <= "hac" and it might not be an actual solvent or salt
-    
+
     Parameters
     ----------
     smiles : str
@@ -425,14 +454,14 @@ def remove_salts_solvents(smiles, hac=3):
         mol = Chem.MolFromSmiles(str(el))
         if mol.GetNumHeavyAtoms() <= hac:
             save.append(mol)
-        
+
     return ".".join([Chem.MolToSmiles(x) for x in save])
 
 
 def sentences2vec(sentences, model, unseen=None):
     """Generate vectors for each sentence (list) in a list of sentences. Vector is simply a
     sum of vectors for individual words.
-    
+
     Parameters
     ----------
     sentences : list, array
@@ -447,6 +476,7 @@ def sentences2vec(sentences, model, unseen=None):
     np.array
     """
     import gensim
+
     GENSIM4 = gensim.__version__ > "4"  # GENSIM4 has some API changes from GENSIM 3
     if GENSIM4:
         keys = set(model.wv.key_to_index)
@@ -458,11 +488,26 @@ def sentences2vec(sentences, model, unseen=None):
 
     for sentence in sentences:
         if unseen:
-            vec.append(sum([model.wv.word_vec(y) if y in set(sentence) & keys
-                       else unseen_vec for y in sentence]))
+            vec.append(
+                sum(
+                    [
+                        model.wv.word_vec(y)
+                        if y in set(sentence) & keys
+                        else unseen_vec
+                        for y in sentence
+                    ]
+                )
+            )
         else:
-            vec.append(sum([model.wv.word_vec(y) for y in sentence 
-                            if y in set(sentence) & keys]))
+            vec.append(
+                sum(
+                    [
+                        model.wv.word_vec(y)
+                        for y in sentence
+                        if y in set(sentence) & keys
+                    ]
+                )
+            )
     return np.array(vec)
 
 
@@ -487,42 +532,53 @@ def featurize(in_file, out_file, model_path, r, uncommon=None):
     """
     # Load the model
     from gensim.models import word2vec
+
     word2vec_model = word2vec.Word2Vec.load(model_path)
     if uncommon:
         try:
             word2vec_model[uncommon]
         except KeyError:
-            raise KeyError('Selected word for uncommon: %s not in vocabulary' % uncommon)
+            raise KeyError(
+                "Selected word for uncommon: %s not in vocabulary" % uncommon
+            )
 
     # File type detection
-    in_split = in_file.split('.')
+    in_split = in_file.split(".")
     f_type = in_split[-1].lower()
-    if f_type not in ['sdf', 'smi', 'ism', 'gz']:
-        raise ValueError('File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)')
-    if f_type == 'gz':
-        if in_split[-2].lower() not in ['sdf', 'smi', 'ism']:
-            raise ValueError('File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)')
+    if f_type not in ["sdf", "smi", "ism", "gz"]:
+        raise ValueError("File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)")
+    if f_type == "gz":
+        if in_split[-2].lower() not in ["sdf", "smi", "ism"]:
+            raise ValueError(
+                "File extension not supported (sdf, smi, ism, sdf.gz, smi.gz)"
+            )
         else:
             f_type = in_split[-2].lower()
 
-    print('Loading molecules.')
-    if f_type == 'sdf':
+    print("Loading molecules.")
+    if f_type == "sdf":
         df = PandasTools.LoadSDF(in_file)
         print("Keeping only molecules that can be processed by RDKit.")
-        df = df[df['ROMol'].notnull()]
-        df['Smiles'] = df['ROMol'].map(Chem.MolToSmiles)
+        df = df[df["ROMol"].notnull()]
+        df["Smiles"] = df["ROMol"].map(Chem.MolToSmiles)
     else:
-        df = pd.read_csv(in_file, delimiter='\t', usecols=[0, 1], names=['Smiles', 'ID'])  # Assume <tab> separated
-        PandasTools.AddMoleculeColumnToFrame(df, smilesCol='Smiles')
+        df = pd.read_csv(
+            in_file, delimiter="\t", usecols=[0, 1], names=["Smiles", "ID"]
+        )  # Assume <tab> separated
+        PandasTools.AddMoleculeColumnToFrame(df, smilesCol="Smiles")
         print("Keeping only molecules that can be processed by RDKit.")
-        df = df[df['ROMol'].notnull()]
-        df['Smiles'] = df['ROMol'].map(Chem.MolToSmiles)  # Recreate SMILES
-
-    print('Featurizing molecules.')
-    df['mol-sentence'] = df.apply(lambda x: MolSentence(mol2alt_sentence(x['ROMol'], r)), axis=1)
-    vectors = sentences2vec(df['mol-sentence'], word2vec_model, unseen=uncommon)
-    df_vec = pd.DataFrame(vectors, columns=['mol2vec-%03i' % x for x in range(vectors.shape[1])])
+        df = df[df["ROMol"].notnull()]
+        df["Smiles"] = df["ROMol"].map(Chem.MolToSmiles)  # Recreate SMILES
+
+    print("Featurizing molecules.")
+    df["mol-sentence"] = df.apply(
+        lambda x: MolSentence(mol2alt_sentence(x["ROMol"], r)), axis=1
+    )
+    vectors = sentences2vec(df["mol-sentence"], word2vec_model, unseen=uncommon)
+    df_vec = pd.DataFrame(
+        vectors, columns=["mol2vec-%03i" % x for x in range(vectors.shape[1])]
+    )
     df_vec.index = df.index
     df = df.join(df_vec)
 
-    df.drop(['ROMol', 'mol-sentence'], axis=1).to_csv(out_file)
\ No newline at end of file
+    df.drop(["ROMol", "mol-sentence"], axis=1).to_csv(out_file)
diff --git a/src/olorenchemengine/external/piCalculax.py b/src/olorenchemengine/external/piCalculax.py
index 296905cf..58f7bed5 100644
--- a/src/olorenchemengine/external/piCalculax.py
+++ b/src/olorenchemengine/external/piCalculax.py
@@ -5,13 +5,12 @@
 For more information visit the github repository: https://github.com/EBjerrum/pICalculax
 """
 
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from numpy import unique, argmin
 from numpy import abs as np_abs
+from numpy import argmin, unique
 
 # RuleTables
 from rdkit import Chem
+from rdkit.Chem import AllChem
 
 # Format of ruletable tuple (Name, SMARTS, pKa (evt. as list), type (evt.as list) )
 # Type is 0 = acidic (charge below pKa)
@@ -23,196 +22,435 @@
 # Default pKa tables defined at the end of the file
 
 pkatable_grimsley = [
-    ('Tetramethylrhodamine #3',
-     Chem.MolFromSmiles('CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1'), [-10, 1.19, 3.59, 10],
-     [1, 1, 0, 1]),
-    ('6-hydroxy-3H-xanthen-3-one', Chem.MolFromSmarts('c2c3ccc(O)cc3oc3cc(=O)ccc2-3'), [2.94, 6.56], [1, 0]),
-    ('Anilin', Chem.MolFromSmarts('[$([NH2]c1ccccc1)]'), 4.87, 1),
-    ('N-acetyl Hydrazine', Chem.MolFromSmarts('[$([NH2]NC(=O))]'), 2.81, 1),  # Reaxys
-    ('Hydrazine', Chem.MolFromSmarts('[$([NH2]N)]'), 8.1, 1),  # Wikiped
-    ('N-acetyl Piperazine', Chem.MolFromSmarts('[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]'), 7.1, 1),  # Reaxys
-    ('Piperazine', Chem.MolFromSmarts('C1CNCCN1'), 9.73, 1),  # Wikiped
-    ('Pyridine', Chem.MolFromSmarts('n1ccccc1'), 5.25, 1),  # Wikiped
-    ('Sulfonate', Chem.MolFromSmarts('[$([OH]S(=O)=O)]'), -1.5, 0),  # Guess
-    ('Phosphonate', Chem.MolFromSmarts('[OH]P([OH])(=O)[#6]'), [1.5, 7.0], [0, -1]),  # Guess
-    ('7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine', Chem.MolFromSmarts('C1C=NC2=CN=C(N=C2O1)N'), 3.6, 1),  # PKA Guessed
-    ('Tyrosine', Chem.MolFromSmarts('[$([O-,OH]c1ccccc1)]'), 10.3, 0),
-    ('Histidine', Chem.MolFromSmarts(
-        '[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1'),
-     6.6, 1),
-    ('Histidine1', Chem.MolFromSmarts('[$(cC)]1cnc[nH]1'), 6.6, 1),
-    ('Histidine2', Chem.MolFromSmarts('[$(cC)]1c[nH]cn1'), 6.6, 1),
-    ('Glutamic Acid', Chem.MolFromSmarts('[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 4.2, 0),
-    ('Cysteine', Chem.MolFromSmarts('[$([S-,SH][CD2]C(N)C(=O))]'), 6.8, 0),
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#175])]'), 6.8, 0),  # Sidechain treated as normal atom
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#178])]'), 6.8, 0),  # Sidechain treated as normal atom
+    (
+        "Tetramethylrhodamine #3",
+        Chem.MolFromSmiles(
+            "CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1"
+        ),
+        [-10, 1.19, 3.59, 10],
+        [1, 1, 0, 1],
+    ),
+    (
+        "6-hydroxy-3H-xanthen-3-one",
+        Chem.MolFromSmarts("c2c3ccc(O)cc3oc3cc(=O)ccc2-3"),
+        [2.94, 6.56],
+        [1, 0],
+    ),
+    ("Anilin", Chem.MolFromSmarts("[$([NH2]c1ccccc1)]"), 4.87, 1),
+    ("N-acetyl Hydrazine", Chem.MolFromSmarts("[$([NH2]NC(=O))]"), 2.81, 1),  # Reaxys
+    ("Hydrazine", Chem.MolFromSmarts("[$([NH2]N)]"), 8.1, 1),  # Wikiped
+    (
+        "N-acetyl Piperazine",
+        Chem.MolFromSmarts("[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]"),
+        7.1,
+        1,
+    ),  # Reaxys
+    ("Piperazine", Chem.MolFromSmarts("C1CNCCN1"), 9.73, 1),  # Wikiped
+    ("Pyridine", Chem.MolFromSmarts("n1ccccc1"), 5.25, 1),  # Wikiped
+    ("Sulfonate", Chem.MolFromSmarts("[$([OH]S(=O)=O)]"), -1.5, 0),  # Guess
+    (
+        "Phosphonate",
+        Chem.MolFromSmarts("[OH]P([OH])(=O)[#6]"),
+        [1.5, 7.0],
+        [0, -1],
+    ),  # Guess
+    (
+        "7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine",
+        Chem.MolFromSmarts("C1C=NC2=CN=C(N=C2O1)N"),
+        3.6,
+        1,
+    ),  # PKA Guessed
+    ("Tyrosine", Chem.MolFromSmarts("[$([O-,OH]c1ccccc1)]"), 10.3, 0),
+    (
+        "Histidine",
+        Chem.MolFromSmarts(
+            "[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1"
+        ),
+        6.6,
+        1,
+    ),
+    ("Histidine1", Chem.MolFromSmarts("[$(cC)]1cnc[nH]1"), 6.6, 1),
+    ("Histidine2", Chem.MolFromSmarts("[$(cC)]1c[nH]cn1"), 6.6, 1),
+    (
+        "Glutamic Acid",
+        Chem.MolFromSmarts("[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        4.2,
+        0,
+    ),
+    ("Cysteine", Chem.MolFromSmarts("[$([S-,SH][CD2]C(N)C(=O))]"), 6.8, 0),
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#175])]"),
+        6.8,
+        0,
+    ),  # Sidechain treated as normal atom
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#178])]"),
+        6.8,
+        0,
+    ),  # Sidechain treated as normal atom
     # ('Aspartic Acid', Chem.MolFromSmarts('[$([CX3][CD2])](=O)[OX1H0-,OX2H1]'),3.90,0),
-    ('Aspartic Acid v2', Chem.MolFromSmarts('[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 3.5, 0),
-    ('Arginine', Chem.MolFromSmarts('[$(N[CD2][CD2])]C(=N)N'), 12.0, 1),
-    ('Lysine', Chem.MolFromSmarts('[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]'), 10.5, 1),
-    ('End Carboxylate (avg)', Chem.MolFromSmarts('[$([CX3]CN)](=O)[OX1H0-,OX2H1]'), 3.3, 0),
+    (
+        "Aspartic Acid v2",
+        Chem.MolFromSmarts("[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        3.5,
+        0,
+    ),
+    ("Arginine", Chem.MolFromSmarts("[$(N[CD2][CD2])]C(=N)N"), 12.0, 1),
+    ("Lysine", Chem.MolFromSmarts("[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]"), 10.5, 1),
+    (
+        "End Carboxylate (avg)",
+        Chem.MolFromSmarts("[$([CX3]CN)](=O)[OX1H0-,OX2H1]"),
+        3.3,
+        0,
+    ),
     # ('End Amine (avg)',Chem.MolFromSmarts('[$([NX3;H2,H1;!$(NC=O)]CC=O)]'),9.15,1),
-    ('End Amine (avg) PseudoAtom safe', Chem.MolFromSmarts(
-        '[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]'),
-     7.7, 1),
-    ('Acrydine', Chem.MolFromSmarts('c1c2ccccc2nc2ccccc21'), 6.15, 1),
-    ('Generic Carboxylic Acid', Chem.MolFromSmarts('[CX3](=O)[OX1H0-,OX2H1]'), 4.75, 0),
-    (
-    'Generic tert. Amine, not amide', Chem.MolFromSmarts('[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]'), 9.8,
-    1),
-    ('Generic Amine, not amide', Chem.MolFromSmarts(
-        '[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]'),
-     10., 1),
+    (
+        "End Amine (avg) PseudoAtom safe",
+        Chem.MolFromSmarts(
+            "[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]"
+        ),
+        7.7,
+        1,
+    ),
+    ("Acrydine", Chem.MolFromSmarts("c1c2ccccc2nc2ccccc21"), 6.15, 1),
+    ("Generic Carboxylic Acid", Chem.MolFromSmarts("[CX3](=O)[OX1H0-,OX2H1]"), 4.75, 0),
+    (
+        "Generic tert. Amine, not amide",
+        Chem.MolFromSmarts("[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]"),
+        9.8,
+        1,
+    ),
+    (
+        "Generic Amine, not amide",
+        Chem.MolFromSmarts(
+            "[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]"
+        ),
+        10.0,
+        1,
+    ),
 ]
 
 condensed_grimsley = [
-    ('Condensed Tyrosine', Chem.MolFromSmarts('[#189]'), 10.3, 0),
-    ('Condensed Histidine', Chem.MolFromSmarts('[#179]'), 6.6, 1),
-    ('Condensed Glutamic Acid', Chem.MolFromSmarts('[#177]'), 4.2, 0),
-    ('Condensed Aspartic Acid', Chem.MolFromSmarts('[#174]'), 3.5, 0),
-    ('Condensed Arginine', Chem.MolFromSmarts('[#172]'), 12.0, 1),
-    ('Condensed Lysine', Chem.MolFromSmarts('[#182]'), 10.5, 1),
+    ("Condensed Tyrosine", Chem.MolFromSmarts("[#189]"), 10.3, 0),
+    ("Condensed Histidine", Chem.MolFromSmarts("[#179]"), 6.6, 1),
+    ("Condensed Glutamic Acid", Chem.MolFromSmarts("[#177]"), 4.2, 0),
+    ("Condensed Aspartic Acid", Chem.MolFromSmarts("[#174]"), 3.5, 0),
+    ("Condensed Arginine", Chem.MolFromSmarts("[#172]"), 12.0, 1),
+    ("Condensed Lysine", Chem.MolFromSmarts("[#182]"), 10.5, 1),
     # ('Condensed AA end amino subst.', Chem.MolFromSmarts('[#0D1]'),9.15,1),
     # Can it somehow match up the wrong end? with a protection group??
-    ('Condensed AA end amino', Chem.MolFromSmarts(
-        '[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]'),
-     7.7, 1),
+    (
+        "Condensed AA end amino",
+        Chem.MolFromSmarts(
+            "[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]"
+        ),
+        7.7,
+        1,
+    ),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed cys end amino', Chem.MolFromSmarts('[#175;D2]'), 7.7, 1),
+    ("Condensed cys end amino", Chem.MolFromSmarts("[#175;D2]"), 7.7, 1),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed AA end Carboxylic', Chem.MolFromSmarts(
-        '[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]'),
-     3.3, 0)
+    (
+        "Condensed AA end Carboxylic",
+        Chem.MolFromSmarts(
+            "[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]"
+        ),
+        3.3,
+        0,
+    ),
 ]
 
 # IPC_peptide as refered in KOZLOWSKI, Lukasz P. "IPC - Isoelectric Point Calculator.". Biol. Direct. 2016, vol 11, p. 55.
 pkatable_ipc = [
-    ('Tetramethylrhodamine #3',
-     Chem.MolFromSmiles('CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1'), [-10, 1.19, 3.59, 10],
-     [1, 1, 0, 1]),
+    (
+        "Tetramethylrhodamine #3",
+        Chem.MolFromSmiles(
+            "CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1"
+        ),
+        [-10, 1.19, 3.59, 10],
+        [1, 1, 0, 1],
+    ),
     # Guess DOI: 10.1039/C0OB01045F
-    ('6-hydroxy-3H-xanthen-3-one', Chem.MolFromSmarts('c2c3ccc(O)cc3oc3cc(=O)ccc2-3'), [2.94, 6.56], [1, 0]),
-    ('Anilin', Chem.MolFromSmarts('[$([NH2]c1ccccc1)]'), 4.87, 1),
-    ('N-acetyl Hydrazine', Chem.MolFromSmarts('[$([NH2]NC(=O))]'), 2.81, 1),  # Reaxys
-    ('Hydrazine', Chem.MolFromSmarts('[$([NH2]N)]'), 8.1, 1),  # Wikiped
-    ('N-acetyl Piperazine', Chem.MolFromSmarts('[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]'), 7.1, 1),  # Reaxys
-    ('Piperazine', Chem.MolFromSmarts('C1CNCCN1'), 9.73, 1),  # Wikiped
-    ('Pyridine', Chem.MolFromSmarts('n1ccccc1'), 5.25, 1),  # Wikiped
-    ('Sulfonate', Chem.MolFromSmarts('[$([OH]S(=O)=O)]'), -1.5, 0),  # Guess
-    ('Phosphonate', Chem.MolFromSmarts('[OH]P([OH])(=O)[#6]'), [1.5, 7.0], [0, -1]),  # Guess
-    ('7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine', Chem.MolFromSmarts('C1C=NC2=CN=C(N=C2O1)N'), 3.6, 1),  # PKA Guessed
-    ('Tyrosine', Chem.MolFromSmarts('[$([O-,OH]c1ccccc1)]'), 10.071, 0),
-    ('Histidine', Chem.MolFromSmarts(
-        '[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1'),
-     6.018, 1),
-    ('Histidine1', Chem.MolFromSmarts('[$(cC)]1cnc[nH]1'), 6.018, 1),
-    ('Histidine2', Chem.MolFromSmarts('[$(cC)]1c[nH]cn1'), 6.018, 1),
-    ('Glutamic Acid', Chem.MolFromSmarts('[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 4.317, 0),
-    ('Cysteine', Chem.MolFromSmarts('[$([S-,SH][CD2]C(N)C(=O))]'), 8.297, 0),
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#175])]'), 8.297, 0),  # Sidechain treated as normal atom
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#178])]'), 8.297, 0),  # Sidechain treated as normal atom
+    (
+        "6-hydroxy-3H-xanthen-3-one",
+        Chem.MolFromSmarts("c2c3ccc(O)cc3oc3cc(=O)ccc2-3"),
+        [2.94, 6.56],
+        [1, 0],
+    ),
+    ("Anilin", Chem.MolFromSmarts("[$([NH2]c1ccccc1)]"), 4.87, 1),
+    ("N-acetyl Hydrazine", Chem.MolFromSmarts("[$([NH2]NC(=O))]"), 2.81, 1),  # Reaxys
+    ("Hydrazine", Chem.MolFromSmarts("[$([NH2]N)]"), 8.1, 1),  # Wikiped
+    (
+        "N-acetyl Piperazine",
+        Chem.MolFromSmarts("[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]"),
+        7.1,
+        1,
+    ),  # Reaxys
+    ("Piperazine", Chem.MolFromSmarts("C1CNCCN1"), 9.73, 1),  # Wikiped
+    ("Pyridine", Chem.MolFromSmarts("n1ccccc1"), 5.25, 1),  # Wikiped
+    ("Sulfonate", Chem.MolFromSmarts("[$([OH]S(=O)=O)]"), -1.5, 0),  # Guess
+    (
+        "Phosphonate",
+        Chem.MolFromSmarts("[OH]P([OH])(=O)[#6]"),
+        [1.5, 7.0],
+        [0, -1],
+    ),  # Guess
+    (
+        "7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine",
+        Chem.MolFromSmarts("C1C=NC2=CN=C(N=C2O1)N"),
+        3.6,
+        1,
+    ),  # PKA Guessed
+    ("Tyrosine", Chem.MolFromSmarts("[$([O-,OH]c1ccccc1)]"), 10.071, 0),
+    (
+        "Histidine",
+        Chem.MolFromSmarts(
+            "[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1"
+        ),
+        6.018,
+        1,
+    ),
+    ("Histidine1", Chem.MolFromSmarts("[$(cC)]1cnc[nH]1"), 6.018, 1),
+    ("Histidine2", Chem.MolFromSmarts("[$(cC)]1c[nH]cn1"), 6.018, 1),
+    (
+        "Glutamic Acid",
+        Chem.MolFromSmarts("[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        4.317,
+        0,
+    ),
+    ("Cysteine", Chem.MolFromSmarts("[$([S-,SH][CD2]C(N)C(=O))]"), 8.297, 0),
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#175])]"),
+        8.297,
+        0,
+    ),  # Sidechain treated as normal atom
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#178])]"),
+        8.297,
+        0,
+    ),  # Sidechain treated as normal atom
     # ('Aspartic Acid', Chem.MolFromSmarts('[$([CX3][CD2])](=O)[OX1H0-,OX2H1]'),3.90,0),
-    ('Aspartic Acid v2', Chem.MolFromSmarts('[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 3.887, 0),
-    ('Arginine', Chem.MolFromSmarts('[$(N[CD2][CD2])]C(=N)N'), 12.503, 1),
-    ('Lysine', Chem.MolFromSmarts('[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]'), 10.517, 1),
-    ('End Carboxylate (avg)', Chem.MolFromSmarts('[$([CX3]CN)](=O)[OX1H0-,OX2H1]'), 2.383, 0),
+    (
+        "Aspartic Acid v2",
+        Chem.MolFromSmarts("[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        3.887,
+        0,
+    ),
+    ("Arginine", Chem.MolFromSmarts("[$(N[CD2][CD2])]C(=N)N"), 12.503, 1),
+    (
+        "Lysine",
+        Chem.MolFromSmarts("[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]"),
+        10.517,
+        1,
+    ),
+    (
+        "End Carboxylate (avg)",
+        Chem.MolFromSmarts("[$([CX3]CN)](=O)[OX1H0-,OX2H1]"),
+        2.383,
+        0,
+    ),
     # ('End Amine (avg)',Chem.MolFromSmarts('[$([NX3;H2,H1;!$(NC=O)]CC=O)]'),9.15,1),
-    ('End Amine (avg) PseudoAtom safe', Chem.MolFromSmarts(
-        '[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]'),
-     9.564, 1),
-    ('Acrydine', Chem.MolFromSmarts('c1c2ccccc2nc2ccccc21'), 6.15, 1),
-    ('Generic Carboxylic Acid', Chem.MolFromSmarts('[CX3](=O)[OX1H0-,OX2H1]'), 4.75, 0),
-    (
-    'Generic tert. Amine, not amide', Chem.MolFromSmarts('[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]'), 9.8,
-    1),
-    ('Generic Amine, not amide', Chem.MolFromSmarts(
-        '[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]'),
-     10., 1),
+    (
+        "End Amine (avg) PseudoAtom safe",
+        Chem.MolFromSmarts(
+            "[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]"
+        ),
+        9.564,
+        1,
+    ),
+    ("Acrydine", Chem.MolFromSmarts("c1c2ccccc2nc2ccccc21"), 6.15, 1),
+    ("Generic Carboxylic Acid", Chem.MolFromSmarts("[CX3](=O)[OX1H0-,OX2H1]"), 4.75, 0),
+    (
+        "Generic tert. Amine, not amide",
+        Chem.MolFromSmarts("[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]"),
+        9.8,
+        1,
+    ),
+    (
+        "Generic Amine, not amide",
+        Chem.MolFromSmarts(
+            "[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]"
+        ),
+        10.0,
+        1,
+    ),
 ]
 
 condensed_ipc = [
-    ('Condensed Tyrosine', Chem.MolFromSmarts('[#189]'), 10.071, 0),
-    ('Condensed Histidine', Chem.MolFromSmarts('[#179]'), 6.018, 1),
-    ('Condensed Glutamic Acid', Chem.MolFromSmarts('[#177]'), 4.317, 0),
-    ('Condensed Aspartic Acid', Chem.MolFromSmarts('[#174]'), 3.887, 0),
-    ('Condensed Arginine', Chem.MolFromSmarts('[#172]'), 12.503, 1),
-    ('Condensed Lysine', Chem.MolFromSmarts('[#182]'), 10.517, 1),
+    ("Condensed Tyrosine", Chem.MolFromSmarts("[#189]"), 10.071, 0),
+    ("Condensed Histidine", Chem.MolFromSmarts("[#179]"), 6.018, 1),
+    ("Condensed Glutamic Acid", Chem.MolFromSmarts("[#177]"), 4.317, 0),
+    ("Condensed Aspartic Acid", Chem.MolFromSmarts("[#174]"), 3.887, 0),
+    ("Condensed Arginine", Chem.MolFromSmarts("[#172]"), 12.503, 1),
+    ("Condensed Lysine", Chem.MolFromSmarts("[#182]"), 10.517, 1),
     # ('Condensed AA end amino subst.', Chem.MolFromSmarts('[#0D1]'),9.15,1),
     # Can it somehow match up the wrong end? with a protection group??
-    ('Condensed AA end amino', Chem.MolFromSmarts(
-        '[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]'),
-     9.564, 1),
+    (
+        "Condensed AA end amino",
+        Chem.MolFromSmarts(
+            "[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]"
+        ),
+        9.564,
+        1,
+    ),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed cys end amino', Chem.MolFromSmarts('[#175;D2]'), 9.564, 1),
+    ("Condensed cys end amino", Chem.MolFromSmarts("[#175;D2]"), 9.564, 1),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed AA end Carboxylic', Chem.MolFromSmarts(
-        '[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]'),
-     2.383, 0)
+    (
+        "Condensed AA end Carboxylic",
+        Chem.MolFromSmarts(
+            "[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]"
+        ),
+        2.383,
+        0,
+    ),
 ]
 
 # EMBOSS iep. 2017, http://emboss.sourceforge.net/apps/release/6.6/emboss/apps/iep.html.
 pkatable_EMBOSS = [
-    ('Tetramethylrhodamine #3',
-     Chem.MolFromSmiles('CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1'), [-10, 1.19, 3.59, 10],
-     [1, 1, 0, 1]),
+    (
+        "Tetramethylrhodamine #3",
+        Chem.MolFromSmiles(
+            "CN(C)C1=C-C=C2-C(C4=C(C(=O)O)C=CC=C4)=C3-C=C-C(=[N+](C)C)C=C3OC2=C1"
+        ),
+        [-10, 1.19, 3.59, 10],
+        [1, 1, 0, 1],
+    ),
     # Guess DOI: 10.1039/C0OB01045F
-    ('6-hydroxy-3H-xanthen-3-one', Chem.MolFromSmarts('c2c3ccc(O)cc3oc3cc(=O)ccc2-3'), [2.94, 6.56], [1, 0]),
-    ('Anilin', Chem.MolFromSmarts('[$([NH2]c1ccccc1)]'), 4.87, 1),
-    ('N-acetyl Hydrazine', Chem.MolFromSmarts('[$([NH2]NC(=O))]'), 2.81, 1),  # Reaxys
-    ('Hydrazine', Chem.MolFromSmarts('[$([NH2]N)]'), 8.1, 1),  # Wikiped
-    ('N-acetyl Piperazine', Chem.MolFromSmarts('[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]'), 7.1, 1),  # Reaxys
-    ('Piperazine', Chem.MolFromSmarts('C1CNCCN1'), 9.73, 1),  # Wikiped
-    ('Pyridine', Chem.MolFromSmarts('n1ccccc1'), 5.25, 1),  # Wikiped
-    ('Sulfonate', Chem.MolFromSmarts('[$([OH]S(=O)=O)]'), -1.5, 0),  # Guess
-    ('Phosphonate', Chem.MolFromSmarts('[OH]P([OH])(=O)[#6]'), [1.5, 7.0], [0, -1]),  # Guess
-    ('7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine', Chem.MolFromSmarts('C1C=NC2=CN=C(N=C2O1)N'), 3.6, 1),  # PKA Guessed
-    ('Tyrosine', Chem.MolFromSmarts('[$([O-,OH]c1ccccc1)]'), 10.1, 0),
-    ('Histidine', Chem.MolFromSmarts(
-        '[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1'),
-     6.5, 1),
-    ('Histidine1', Chem.MolFromSmarts('[$(cC)]1cnc[nH]1'), 6.5, 1),
-    ('Histidine2', Chem.MolFromSmarts('[$(cC)]1c[nH]cn1'), 6.5, 1),
-    ('Glutamic Acid', Chem.MolFromSmarts('[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 4.1, 0),
-    ('Cysteine', Chem.MolFromSmarts('[$([S-,SH][CD2]C(N)C(=O))]'), 8.5, 0),
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#175])]'), 8.5, 0),  # Sidechain treated as normal atom
-    ('Condensed Cysteine', Chem.MolFromSmarts('[$([SH][#178])]'), 8.5, 0),  # Sidechain treated as normal atom
+    (
+        "6-hydroxy-3H-xanthen-3-one",
+        Chem.MolFromSmarts("c2c3ccc(O)cc3oc3cc(=O)ccc2-3"),
+        [2.94, 6.56],
+        [1, 0],
+    ),
+    ("Anilin", Chem.MolFromSmarts("[$([NH2]c1ccccc1)]"), 4.87, 1),
+    ("N-acetyl Hydrazine", Chem.MolFromSmarts("[$([NH2]NC(=O))]"), 2.81, 1),  # Reaxys
+    ("Hydrazine", Chem.MolFromSmarts("[$([NH2]N)]"), 8.1, 1),  # Wikiped
+    (
+        "N-acetyl Piperazine",
+        Chem.MolFromSmarts("[N+&!H0,NX3&H0;$(N1(C)CCN(C(=O)C)CC1)]"),
+        7.1,
+        1,
+    ),  # Reaxys
+    ("Piperazine", Chem.MolFromSmarts("C1CNCCN1"), 9.73, 1),  # Wikiped
+    ("Pyridine", Chem.MolFromSmarts("n1ccccc1"), 5.25, 1),  # Wikiped
+    ("Sulfonate", Chem.MolFromSmarts("[$([OH]S(=O)=O)]"), -1.5, 0),  # Guess
+    (
+        "Phosphonate",
+        Chem.MolFromSmarts("[OH]P([OH])(=O)[#6]"),
+        [1.5, 7.0],
+        [0, -1],
+    ),  # Guess
+    (
+        "7H-Pyrimido[4,5-b][1,4]oxazine-2,4-diamine",
+        Chem.MolFromSmarts("C1C=NC2=CN=C(N=C2O1)N"),
+        3.6,
+        1,
+    ),  # PKA Guessed
+    ("Tyrosine", Chem.MolFromSmarts("[$([O-,OH]c1ccccc1)]"), 10.1, 0),
+    (
+        "Histidine",
+        Chem.MolFromSmarts(
+            "[CH2X4][#6X3]1:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]:[$([#7X3H+,#7X2H0+0]:[#6X3H]:[#7X3H]),$([#7X3H])]:[#6X3H]1"
+        ),
+        6.5,
+        1,
+    ),
+    ("Histidine1", Chem.MolFromSmarts("[$(cC)]1cnc[nH]1"), 6.5, 1),
+    ("Histidine2", Chem.MolFromSmarts("[$(cC)]1c[nH]cn1"), 6.5, 1),
+    (
+        "Glutamic Acid",
+        Chem.MolFromSmarts("[$([CX3][CD2][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        4.1,
+        0,
+    ),
+    ("Cysteine", Chem.MolFromSmarts("[$([S-,SH][CD2]C(N)C(=O))]"), 8.5, 0),
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#175])]"),
+        8.5,
+        0,
+    ),  # Sidechain treated as normal atom
+    (
+        "Condensed Cysteine",
+        Chem.MolFromSmarts("[$([SH][#178])]"),
+        8.5,
+        0,
+    ),  # Sidechain treated as normal atom
     # ('Aspartic Acid', Chem.MolFromSmarts('[$([CX3][CD2])](=O)[OX1H0-,OX2H1]'),3.90,0),
-    ('Aspartic Acid v2', Chem.MolFromSmarts('[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]'), 3.9, 0),
-    ('Arginine', Chem.MolFromSmarts('[$(N[CD2][CD2])]C(=N)N'), 12.5, 1),
-    ('Lysine', Chem.MolFromSmarts('[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]'), 10.8, 1),
-    ('End Carboxylate (avg)', Chem.MolFromSmarts('[$([CX3]CN)](=O)[OX1H0-,OX2H1]'), 3.6, 0),
+    (
+        "Aspartic Acid v2",
+        Chem.MolFromSmarts("[$([CX3][CD2]C(N)C(=O))](=O)[OX1H0-,OX2H1]"),
+        3.9,
+        0,
+    ),
+    ("Arginine", Chem.MolFromSmarts("[$(N[CD2][CD2])]C(=N)N"), 12.5, 1),
+    ("Lysine", Chem.MolFromSmarts("[$([ND1][CD2][CD2][CD2][CD2]C(N)C(=O))]"), 10.8, 1),
+    (
+        "End Carboxylate (avg)",
+        Chem.MolFromSmarts("[$([CX3]CN)](=O)[OX1H0-,OX2H1]"),
+        3.6,
+        0,
+    ),
     # ('End Amine (avg)',Chem.MolFromSmarts('[$([NX3;H2,H1;!$(NC=O)]CC=O)]'),9.15,1),
-    ('End Amine (avg) PseudoAtom safe', Chem.MolFromSmarts(
-        '[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]'),
-     8.6, 1),
-    ('Acrydine', Chem.MolFromSmarts('c1c2ccccc2nc2ccccc21'), 6.15, 1),
-    ('Generic Carboxylic Acid', Chem.MolFromSmarts('[CX3](=O)[OX1H0-,OX2H1]'), 4.75, 0),
-    (
-    'Generic tert. Amine, not amide', Chem.MolFromSmarts('[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]'), 9.8,
-    1),
-    ('Generic Amine, not amide', Chem.MolFromSmarts(
-        '[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]'),
-     10., 1),
+    (
+        "End Amine (avg) PseudoAtom safe",
+        Chem.MolFromSmarts(
+            "[$([NX3;H2,H1;!$(NC=O);!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]CC=O)]"
+        ),
+        8.6,
+        1,
+    ),
+    ("Acrydine", Chem.MolFromSmarts("c1c2ccccc2nc2ccccc21"), 6.15, 1),
+    ("Generic Carboxylic Acid", Chem.MolFromSmarts("[CX3](=O)[OX1H0-,OX2H1]"), 4.75, 0),
+    (
+        "Generic tert. Amine, not amide",
+        Chem.MolFromSmarts("[N+&!H0&NX3,NX3;$(N(C)(C)C);!$(NC=[!#6]);!$(NC#[!#6])]"),
+        9.8,
+        1,
+    ),
+    (
+        "Generic Amine, not amide",
+        Chem.MolFromSmarts(
+            "[N+,NX3;!H0;!$(NC=[!#6]);!$(NC#[!#6])!$(N[#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193])]"
+        ),
+        10.0,
+        1,
+    ),
 ]
 
 condensed_EMBOSS = [
-    ('Condensed Tyrosine', Chem.MolFromSmarts('[#189]'), 10.1, 0),
-    ('Condensed Histidine', Chem.MolFromSmarts('[#179]'), 6.5, 1),
-    ('Condensed Glutamic Acid', Chem.MolFromSmarts('[#177]'), 4.4, 0),
-    ('Condensed Aspartic Acid', Chem.MolFromSmarts('[#174]'), 3.9, 0),
-    ('Condensed Arginine', Chem.MolFromSmarts('[#172]'), 12.5, 1),
-    ('Condensed Lysine', Chem.MolFromSmarts('[#182]'), 10.8, 1),
+    ("Condensed Tyrosine", Chem.MolFromSmarts("[#189]"), 10.1, 0),
+    ("Condensed Histidine", Chem.MolFromSmarts("[#179]"), 6.5, 1),
+    ("Condensed Glutamic Acid", Chem.MolFromSmarts("[#177]"), 4.4, 0),
+    ("Condensed Aspartic Acid", Chem.MolFromSmarts("[#174]"), 3.9, 0),
+    ("Condensed Arginine", Chem.MolFromSmarts("[#172]"), 12.5, 1),
+    ("Condensed Lysine", Chem.MolFromSmarts("[#182]"), 10.8, 1),
     # ('Condensed AA end amino subst.', Chem.MolFromSmarts('[#0D1]'),9.15,1),
     # Can it somehow match up the wrong end? with a protection group??
-    ('Condensed AA end amino', Chem.MolFromSmarts(
-        '[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]'),
-     8.6, 1),
+    (
+        "Condensed AA end amino",
+        Chem.MolFromSmarts(
+            "[#171,#172,#173,#174,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193;D1]"
+        ),
+        8.6,
+        1,
+    ),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed cys end amino', Chem.MolFromSmarts('[#175;D2]'), 8.6, 1),
+    ("Condensed cys end amino", Chem.MolFromSmarts("[#175;D2]"), 8.6, 1),
     # ('Condensed AA end Carboxylic Subst.', Chem.MolFromSmarts('[$([#0][OH])]'),2.15,0), #All acidic/basic should be replaced with glycine so they can be removed??
-    ('Condensed AA end Carboxylic', Chem.MolFromSmarts(
-        '[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]'),
-     3.6, 0)
+    (
+        "Condensed AA end Carboxylic",
+        Chem.MolFromSmarts(
+            "[$([#171,#172,#173,#174,#175,#176,#177,#178,#179,#180,#181,#182,#183,#184,#185,#186,#187,#188,#189,#190,#191,#192,#193][OH])]"
+        ),
+        3.6,
+        0,
+    ),
 ]
 
 # Defaut pKa_tables
@@ -222,48 +460,65 @@
 
 condensedtable = condensed_grimsley
 
-def find_pKas(mol, pkatable=pkatable, condensedtable=condensedtable, debug=False, returnindices=False):
+
+def find_pKas(
+    mol,
+    pkatable=pkatable,
+    condensedtable=condensedtable,
+    debug=False,
+    returnindices=False,
+):
     """This function finds pKa values in a supplied molecule by matching it up to SMARTS rules in a table
     Returns the list of pKa values and a list with the charge of the group at pH levels blow the pKa.
     Example: So histidine has a +1 charge at pH levels below the pKa of ~6"""
     chargelist = []
     pkalist = []
-    dummy = Chem.MolFromSmarts('[#178]')  # Used Gly instead of #0 or *, for compatibility with building blocks
+    dummy = Chem.MolFromSmarts(
+        "[#178]"
+    )  # Used Gly instead of #0 or *, for compatibility with building blocks
 
     if returnindices:
         atomindices = []
         for atom in mol.GetAtoms():
-            atom.SetProp('O_i', str(atom.GetIdx()))  # Save the Original Idx as a property
+            atom.SetProp(
+                "O_i", str(atom.GetIdx())
+            )  # Save the Original Idx as a property
 
     # Look for condensed AA atoms with potential pKa values. Substitute with Glycine
     for s in condensedtable:
         if mol.HasSubstructMatch(s[1]):
-            if debug: print
+            if debug:
+                print
             "Match: %s" % (s[0]),
             # How many??
             substructmatches = mol.GetSubstructMatches(s[1])
             n = len(substructmatches)
-            if debug: print
+            if debug:
+                print
             "times %s" % (n)
             pkalist += [s[2]] * n
             chargelist += [s[3]] * n
             if returnindices:
                 for match in substructmatches:
-                    atomindices += [[int(mol.GetAtomWithIdx(idx).GetProp('O_i')) for idx in
-                                     match]]  # Retrieve the original index of matches
+                    atomindices += [
+                        [int(mol.GetAtomWithIdx(idx).GetProp("O_i")) for idx in match]
+                    ]  # Retrieve the original index of matches
             mol = AllChem.ReplaceSubstructs(mol, s[1], dummy)[
-                0]  # Delete acidic/basic groups to prevent multiple matching
+                0
+            ]  # Delete acidic/basic groups to prevent multiple matching
 
     # Look for molecular fragments with known pKa's
     for s in pkatable:
         # print s[0]
         if mol.HasSubstructMatch(s[1]):
-            if debug: print
+            if debug:
+                print
             "Match: %s" % (s[0]),
             # How many??
             substructmatches = mol.GetSubstructMatches(s[1])
             n = len(substructmatches)
-            if debug: print
+            if debug:
+                print
             "times %s" % (n)
             if type(s[2]) == type(float()):
                 pkalist += [s[2]] * n
@@ -275,7 +530,8 @@ def find_pKas(mol, pkatable=pkatable, condensedtable=condensedtable, debug=False
             if returnindices:
                 for match in substructmatches:
                     atomindices += [
-                        [int(mol.GetAtomWithIdx(idx).GetProp('O_i')) for idx in match]]  # Retrieve the original i
+                        [int(mol.GetAtomWithIdx(idx).GetProp("O_i")) for idx in match]
+                    ]  # Retrieve the original i
             mol = AllChem.ReplaceSubstructs(mol, s[1], dummy)[0]
     if returnindices:
         return pkalist, chargelist, atomindices
@@ -313,21 +569,30 @@ def pI(pkalist, chargelist):
         # Find pI by simulation in the pH range 0 to 0
         chargecol = []
         for i in range(0, 1400):
-            ph = i / 100.
+            ph = i / 100.0
             chargecol.append(charge(ph, pkalist, chargelist))  # Calculate charge
-        pI = argmin(
-            np_abs(chargecol)) / 100.  # Simply taking the smallest absolute value, and dividing the index with 100
+        pI = (
+            argmin(np_abs(chargecol)) / 100.0
+        )  # Simply taking the smallest absolute value, and dividing the index with 100
     # print "pI %.1f"%pI
     return pI
 
-from olorenchemengine.representations import BaseCompoundVecRepresentation
+
 import numpy as np
+
 from olorenchemengine.base_class import log_arguments
+from olorenchemengine.representations import BaseCompoundVecRepresentation
+
+
 class calc_pI(BaseCompoundVecRepresentation):
     @log_arguments
     def __init__(self, ph=7.4, log=True, **kwargs):
-        self.ph=ph
-        super().__init__(names=["pI", f"Charge{self.ph}", f"TotalAbsCharge{self.ph}"], log=False, **kwargs)
+        self.ph = ph
+        super().__init__(
+            names=["pI", f"Charge{self.ph}", f"TotalAbsCharge{self.ph}"],
+            log=False,
+            **kwargs,
+        )
 
     def _convert(self, smiles, y=None):
         # find pKa values and charge class
@@ -335,4 +600,4 @@ def _convert(self, smiles, y=None):
         pkalist, chargelist = find_pKas(m)
         # Calculate pI
         pIpred = pI(pkalist, chargelist)
-        return np.array([pIpred, *charge(self.ph, pkalist, chargelist)])
\ No newline at end of file
+        return np.array([pIpred, *charge(self.ph, pkalist, chargelist)])
diff --git a/src/olorenchemengine/external/stoned.py b/src/olorenchemengine/external/stoned.py
index 1530de11..fe29429d 100644
--- a/src/olorenchemengine/external/stoned.py
+++ b/src/olorenchemengine/external/stoned.py
@@ -207,44 +207,45 @@
 #   3. Added condition to sanitize SMILES
 
 import numpy as np
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.DataStructs.cDataStructs import TanimotoSimilarity
-from rdkit.Chem import Mol
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem, Mol
 from rdkit.Chem.AtomPairs.Sheridan import GetBPFingerprint, GetBTFingerprint
 from rdkit.Chem.Pharm2D import Generate, Gobbi_Pharm2D
-
-from rdkit import RDLogger
+from rdkit.DataStructs.cDataStructs import TanimotoSimilarity
 
 import olorenchemengine as oce
 
 RDLogger.DisableLog("rdApp.*")
 
+
 def randomize_smiles(mol):
-    '''Returns a random (dearomatized) SMILES given an rdkit mol object of a molecule.
+    """Returns a random (dearomatized) SMILES given an rdkit mol object of a molecule.
     Parameters:
     mol (rdkit.Chem.rdchem.Mol) :  RdKit mol object (None if invalid smile string smi)
-    
+
     Returns:
     mol (rdkit.Chem.rdchem.Mol) : RdKit mol object  (None if invalid smile string smi)
-    '''
+    """
     if not mol:
         return None
 
     Chem.Kekulize(mol)
-    return Chem.MolToSmiles(mol, canonical=False, doRandom=True, isomericSmiles=False,  kekuleSmiles=True)
+    return Chem.MolToSmiles(
+        mol, canonical=False, doRandom=True, isomericSmiles=False, kekuleSmiles=True
+    )
+
 
 def sanitize_smiles(smi):
-    '''Return a canonical smile representation of smi
-    
+    """Return a canonical smile representation of smi
+
     Parameters:
-    smi (string) : smile string to be canonicalized 
-    
+    smi (string) : smile string to be canonicalized
+
     Returns:
     mol (rdkit.Chem.rdchem.Mol) : RdKit mol object                          (None if invalid smile string smi)
     smi_canon (string)          : Canonicalized smile representation of smi (None if invalid smile string smi)
-    conversion_successful (bool): True/False to indicate if conversion was  successful 
-    '''
+    conversion_successful (bool): True/False to indicate if conversion was  successful
+    """
     try:
         mol = Chem.MolFromSmiles(smi, sanitize=True)
         smi_canon = Chem.MolToSmiles(mol, isomericSmiles=False, canonical=True)
@@ -253,138 +254,158 @@ def sanitize_smiles(smi):
         return (mol, smi_canon, True)
     except:
         return (None, None, False)
-    
+
 
 def get_selfie_chars(selfie):
-    '''Obtain a list of all selfie characters in string selfie
-    
-    Parameters: 
-    selfie (string) : A selfie string - representing a molecule 
-    
-    Example: 
+    """Obtain a list of all selfie characters in string selfie
+
+    Parameters:
+    selfie (string) : A selfie string - representing a molecule
+
+    Example:
     >>> get_selfie_chars('[C][=C][C][=C][C][=C][Ring1][Branch1_1]')
     ['[C]', '[=C]', '[C]', '[=C]', '[C]', '[=C]', '[Ring1]', '[Branch1_1]']
-    
+
     Returns:
     chars_selfie: list of selfie characters present in molecule selfie
-    '''
-    chars_selfie = [] # A list of all SELFIE sybols from string selfie
-    while selfie != '':
-        chars_selfie.append(selfie[selfie.find('['): selfie.find(']')+1])
-        selfie = selfie[selfie.find(']')+1:]
+    """
+    chars_selfie = []  # A list of all SELFIE sybols from string selfie
+    while selfie != "":
+        chars_selfie.append(selfie[selfie.find("[") : selfie.find("]") + 1])
+        selfie = selfie[selfie.find("]") + 1 :]
     return chars_selfie
 
+
 def mutate_selfie(selfie, max_molecules_len, write_fail_cases=False, index=None):
-    '''Return a mutated selfie string (only one mutation on selfie is performed)
-    
-    Mutations are done until a valid molecule is obtained 
-    Rules of mutation: With a 33.3% propbabily, either: 
+    """Return a mutated selfie string (only one mutation on selfie is performed)
+
+    Mutations are done until a valid molecule is obtained
+    Rules of mutation: With a 33.3% propbabily, either:
         1. Add a random SELFIE character in the string
         2. Replace a random SELFIE character with another
         3. Delete a random character
-    
+
     Parameters:
-    selfie            (string)  : SELFIE string to be mutated 
+    selfie            (string)  : SELFIE string to be mutated
     max_molecules_len (int)     : Mutations of SELFIE string are allowed up to this length
     write_fail_cases  (bool)    : If true, failed mutations are recorded in "selfie_failure_cases.txt"
-    
+
     Returns:
     selfie_mutated    (string)  : Mutated SELFIE string
     smiles_canon      (string)  : canonical smile of mutated SELFIE string
-    '''
-    valid=False
+    """
+    valid = False
     fail_counter = 0
     chars_selfie = get_selfie_chars(selfie)
-    
+
     while not valid:
         fail_counter += 1
-        
+
         oce.import_or_install("selfies")
-        
+
         import selfies
 
-        alphabet = list(selfies.get_semantic_robust_alphabet()) # 34 SELFIE characters 
+        alphabet = list(selfies.get_semantic_robust_alphabet())  # 34 SELFIE characters
 
-        choice_ls = [1, 2, 3] # 1=Insert; 2=Replace; 3=Delete
+        choice_ls = [1, 2, 3]  # 1=Insert; 2=Replace; 3=Delete
         random_choice = np.random.choice(choice_ls, 1)[0]
-        
+
         # Insert a character in a Random Location
-        if random_choice == 1: 
+        if random_choice == 1:
             if index is None:
-                random_index = np.random.randint(len(chars_selfie)+1)
+                random_index = np.random.randint(len(chars_selfie) + 1)
             else:
                 random_index = index
             random_character = np.random.choice(alphabet, size=1)[0]
-            
-            selfie_mutated_chars = chars_selfie[:random_index] + [random_character] + chars_selfie[random_index:]
 
-        # Replace a random character 
+            selfie_mutated_chars = (
+                chars_selfie[:random_index]
+                + [random_character]
+                + chars_selfie[random_index:]
+            )
+
+        # Replace a random character
         elif random_choice == 2:
-            if index is None:                         
+            if index is None:
                 random_index = np.random.randint(len(chars_selfie))
             else:
                 random_index = index
             random_character = np.random.choice(alphabet, size=1)[0]
             if random_index == 0:
-                selfie_mutated_chars = [random_character] + chars_selfie[random_index+1:]
+                selfie_mutated_chars = [random_character] + chars_selfie[
+                    random_index + 1 :
+                ]
             else:
-                selfie_mutated_chars = chars_selfie[:random_index] + [random_character] + chars_selfie[random_index+1:]
-                
+                selfie_mutated_chars = (
+                    chars_selfie[:random_index]
+                    + [random_character]
+                    + chars_selfie[random_index + 1 :]
+                )
+
         # Delete a random character
         elif random_choice == 3:
-            if index is None: 
+            if index is None:
                 random_index = np.random.randint(len(chars_selfie))
             else:
                 random_index = index
             if random_index == 0:
-                selfie_mutated_chars = chars_selfie[random_index+1:]
+                selfie_mutated_chars = chars_selfie[random_index + 1 :]
             else:
-                selfie_mutated_chars = chars_selfie[:random_index] + chars_selfie[random_index+1:]
-                
-        else: 
-            raise Exception('Invalid Operation trying to be performed')
+                selfie_mutated_chars = (
+                    chars_selfie[:random_index] + chars_selfie[random_index + 1 :]
+                )
+
+        else:
+            raise Exception("Invalid Operation trying to be performed")
 
         selfie_mutated = "".join(x for x in selfie_mutated_chars)
         sf = "".join(x for x in chars_selfie)
-        
+
         try:
             smiles = decoder(selfie_mutated)
             mol, smiles_canon, done = sanitize_smiles(smiles)
-            if len(selfie_mutated_chars) > max_molecules_len or smiles_canon=="":
+            if len(selfie_mutated_chars) > max_molecules_len or smiles_canon == "":
                 done = False
             if done:
                 valid = True
             else:
                 valid = False
         except:
-            valid=False
+            valid = False
             if fail_counter > 1 and write_fail_cases == True:
                 f = open("selfie_failure_cases.txt", "a+")
-                f.write('Tried to mutate SELFIE: '+str(sf)+' To Obtain: '+str(selfie_mutated) + '\n')
+                f.write(
+                    "Tried to mutate SELFIE: "
+                    + str(sf)
+                    + " To Obtain: "
+                    + str(selfie_mutated)
+                    + "\n"
+                )
                 f.close()
-    
+
     return (selfie_mutated, smiles_canon)
 
-def get_mutated_SELFIES(selfies_ls, num_mutations): 
-    ''' Mutate all the SELFIES in 'selfies_ls' 'num_mutations' number of times. 
-    
+
+def get_mutated_SELFIES(selfies_ls, num_mutations):
+    """Mutate all the SELFIES in 'selfies_ls' 'num_mutations' number of times.
+
     Parameters:
-    selfies_ls   (list)  : A list of SELFIES 
+    selfies_ls   (list)  : A list of SELFIES
     num_mutations (int)  : number of mutations to perform on each SELFIES within 'selfies_ls'
-    
+
     Returns:
     selfies_ls   (list)  : A list of mutated SELFIES
-    
-    '''
-    for _ in range(num_mutations): 
+
+    """
+    for _ in range(num_mutations):
         selfie_ls_mut_ls = []
-        for str_ in selfies_ls: 
-            
+        for str_ in selfies_ls:
+
             str_chars = get_selfie_chars(str_)
             max_molecules_len = len(str_chars) + num_mutations
-            
+
             selfie_mutated, _ = mutate_selfie(str_, max_molecules_len)
             selfie_ls_mut_ls.append(selfie_mutated)
-        
+
         selfies_ls = selfie_ls_mut_ls.copy()
     return selfies_ls
diff --git a/src/olorenchemengine/gnn.py b/src/olorenchemengine/gnn.py
index e1826984..0cc36480 100644
--- a/src/olorenchemengine/gnn.py
+++ b/src/olorenchemengine/gnn.py
@@ -1,26 +1,22 @@
 """ Building blocks for graph neural networks.
 """
 
+import torch
+import torch.nn as nn
+from pytorch_lightning import LightningModule
+
 from .base_class import *
 from .representations import BaseRepresentation, TorchGeometricGraph
-from .internal import mock_imports
-
-try:
-    import torch
-    import torch.nn as nn
-except:
-    mock_imports(globals(), "torch", "nn")
-
-try:
 
-    from pytorch_lightning import LightningModule
-except:
-    LightningModule = object
+if isinstance(LightningModule, MagicMock):
+    lm_imports = (BaseClass,)
+else:
+    lm_imports = (LightningModule, BaseClass)
 
 
-class BaseLightningModule(BaseClass, LightningModule):
+class BaseLightningModule(*lm_imports):
 
-    """ BaseLightningModule allows for the use of a Pytorch Lightning module as a BaseClass to be incorporated into the framework.
+    """BaseLightningModule allows for the use of a Pytorch Lightning module as a BaseClass to be incorporated into the framework.
 
     Parameters:
         optim (str, optional): parameter describing what kind of optimizer to use. Defaults to "adam".
@@ -32,7 +28,7 @@ def __init__(self, optim: str = "adam", input_dimensions: Tuple = None):
         self.optim = optim
 
     def set_task_type(self, task_type, pos_weight=torch.tensor([1])):
-        """ Sets the task type for the model.
+        """Sets the task type for the model.
 
         Parameters:
             task_type (str): the task type to set the model to.
@@ -40,12 +36,14 @@ def set_task_type(self, task_type, pos_weight=torch.tensor([1])):
 
         """
         if task_type == "classification":
-            self.loss_fun = nn.BCEWithLogitsLoss(reduction="none", pos_weight=pos_weight)
+            self.loss_fun = nn.BCEWithLogitsLoss(
+                reduction="none", pos_weight=pos_weight
+            )
         else:
             self.loss_fun = nn.MSELoss()
 
     def loss(self, y_pred, y_true):
-        """ Calculate the loss for the model.
+        """Calculate the loss for the model.
 
         Parameters:
             y_pred (torch.tensor): the predictions for the model.
@@ -105,10 +103,10 @@ def configure_optimizers(self):
 
 class MultiSequential(nn.Sequential):
 
-    """ Helper class to allow for the use of nn.Sequential with multi-input torch.nn modules"""
+    """Helper class to allow for the use of nn.Sequential with multi-input torch.nn modules"""
 
     def forward(self, *input):
-        """ Forward pass of the module.
+        """Forward pass of the module.
 
         Parameters:
             *input (torch.tensor): the input data.
@@ -127,7 +125,7 @@ def forward(self, *input):
 
 class AttentiveFP(BaseLightningModule):
 
-    """ AttentiveFP is a wrapper for the PyTorch Geometric interpretation of https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959.
+    """AttentiveFP is a wrapper for the PyTorch Geometric interpretation of https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959.
 
     Parameters:
         hidden_channels (int, optional): the number of hidden channels to use in the model. Defaults to 4.
@@ -175,7 +173,7 @@ def __init__(
             self.activation = nn.ELU
 
     def create(self, dimensions):
-        """ Create the model.
+        """Create the model.
 
         Parameters:
             dimensions (list): the dimensions of the input data.
@@ -206,7 +204,10 @@ def create(self, dimensions):
                     dropout=self.dropout,
                 ),
                 *[
-                    nn.Sequential(nn.Linear(self.layer_dims[i], self.layer_dims[i + 1]), self.activation())
+                    nn.Sequential(
+                        nn.Linear(self.layer_dims[i], self.layer_dims[i + 1]),
+                        self.activation(),
+                    )
                     for i in range(len(self.layer_dims) - 1)
                 ],
                 nn.Linear(self.layer_dims[-1], 1)
@@ -229,7 +230,7 @@ def forward(self, data):
 
 class BaseTorchGeometricModel(BaseModel):
 
-    """ BaseTorchGeometricModel is a base class for models in the PyTorch Geometric framework.
+    """BaseTorchGeometricModel is a base class for models in the PyTorch Geometric framework.
 
     Parameters:
         network (nn.Module): The network to be used for the model.
@@ -292,7 +293,10 @@ def _fit(self, X_train, y_train):
         if len(values) == 2:
             if self.pos_weight == "balanced":
                 self.network.set_task_type(
-                    "classification", pos_weight=torch.tensor([(len(y_train) - sum(y_train)) / sum(y_train)])
+                    "classification",
+                    pos_weight=torch.tensor(
+                        [(len(y_train) - sum(y_train)) / sum(y_train)]
+                    ),
                 )
             else:
                 self.network.set_task_type("classification")
@@ -307,6 +311,7 @@ def _fit(self, X_train, y_train):
 
     def _predict(self, X):
         self.network.eval()
+        
         from torch_geometric.data import DataLoader as PyGDataLoader
 
         dataloader = PyGDataLoader(X, batch_size=self.batch_size, num_workers=oce.CONFIG["NUM_WORKERS"])
@@ -336,7 +341,7 @@ def _load(self, d):
 
 class TLFromCheckpoint(BaseLightningModule):
 
-    """ TLFromCheckpoint is a base class for transfer-learning from an OlorenVec PyTorch-lightning checkpoint.
+    """TLFromCheckpoint is a base class for transfer-learning from an OlorenVec PyTorch-lightning checkpoint.
 
     Parameters:
         model_path (str, option): The path to the PyTorch-lightning checkpoint. Ise
@@ -372,12 +377,24 @@ def __init__(
             logging.warn("Overriding map_location to cpu as no GPUs are available.")
 
         state_dict = OrderedDict(
-            [(k.replace("model.", ""), v) for k, v in torch.load(path, map_location=map_location)["state_dict"].items()]
+            [
+                (k.replace("model.", ""), v)
+                for k, v in torch.load(path, map_location=map_location)[
+                    "state_dict"
+                ].items()
+            ]
         )
 
         from olorenchemengine.pyg.gcn import GNN
 
-        self.A = GNN(gnn_type="gcn", num_tasks=num_tasks, num_layer=5, emb_dim=300, drop_ratio=0.5, virtual_node=False)
+        self.A = GNN(
+            gnn_type="gcn",
+            num_tasks=num_tasks,
+            num_layer=5,
+            emb_dim=300,
+            drop_ratio=0.5,
+            virtual_node=False,
+        )
         if not reset:
             self.A.load_state_dict(state_dict)
         self.B = nn.Sequential(
@@ -391,4 +408,3 @@ def __init__(
             )
         )
         self.network = nn.Sequential(OrderedDict([("A", self.A), ("B", self.B)]))
-
diff --git a/src/olorenchemengine/hyperparameters.py b/src/olorenchemengine/hyperparameters.py
index b5698f81..754ac826 100644
--- a/src/olorenchemengine/hyperparameters.py
+++ b/src/olorenchemengine/hyperparameters.py
@@ -14,21 +14,21 @@
 from abc import abstractproperty
 from typing import *
 
+from hyperopt import fmin, hp, tpe
+from hyperopt.pyll import scope
+
 from .base_class import *
-from .manager import *
 from .benchmarks import *
+from .manager import *
 
-from hyperopt import hp, fmin, tpe, hp
-from hyperopt.pyll import scope
 
 class Opt(BaseClass):
-
     @abstractproperty
     def get_hp(self):
         pass
 
     @log_arguments
-    def __init__(self, label, *args, use_int = False, **kwargs):
+    def __init__(self, label, *args, use_int=False, **kwargs):
         self.label = label
         if use_int:
             self.hp = scope.int(self.get_hp(label, *args, **kwargs))
@@ -41,60 +41,61 @@ def _load(self, d):
     def _save(self):
         return {}
 
-class OptChoice(Opt):
 
+class OptChoice(Opt):
     @property
     def get_hp(self):
         return hp.choice
 
-class OptRandInt(Opt):
 
+class OptRandInt(Opt):
     @property
     def get_hp(self):
         return hp.randint
 
-class OptUniform(Opt):
 
+class OptUniform(Opt):
     @property
     def get_hp(self):
         return hp.uniform
 
-class OptQUniform(Opt):
 
+class OptQUniform(Opt):
     @property
     def get_hp(self):
         return hp.quniform
 
-class OptLogUniform(Opt):
 
+class OptLogUniform(Opt):
     @property
     def get_hp(self):
         return hp.loguniform
 
-class OptQLogUniform(Opt):
 
+class OptQLogUniform(Opt):
     @property
     def get_hp(self):
         return hp.qloguniform
 
-class OptQNormal(Opt):
 
+class OptQNormal(Opt):
     @property
     def get_hp(self):
         return hp.qnormal
 
-class OptLogNormal(Opt):
 
+class OptLogNormal(Opt):
     @property
     def get_hp(self):
         return hp.lognormal
 
-class OptQLogNormal(Opt):
 
+class OptQLogNormal(Opt):
     @property
     def get_hp(self):
         return hp.qnormal
 
+
 def index_hyperparameters(object: BaseClass) -> dict:
     """
     Returns a dictionary of hyperparameters for the model.
@@ -111,13 +112,17 @@ def index_hyperparameters(object: BaseClass) -> dict:
             d_ = index_hyperparameters(arg)
             d.update(d_)
             if not len(d) == count + len(d_):
-                raise Exception(f"Hyperparameter indexing failed, overlapping keys in {d_}")
+                raise Exception(
+                    f"Hyperparameter indexing failed, overlapping keys in {d_}"
+                )
         for k, v in object.kwargs.items():
             count = len(d)
             d_ = index_hyperparameters(v)
             d.update(d_)
             if not len(d) == count + len(d_):
-                raise Exception(f"Hyperparameter indexing failed, overlapping keys in {d_}")
+                raise Exception(
+                    f"Hyperparameter indexing failed, overlapping keys in {d_}"
+                )
         return d
     elif issubclass(type(object), dict):
         d = {}
@@ -126,13 +131,17 @@ def index_hyperparameters(object: BaseClass) -> dict:
             d_ = index_hyperparameters(arg)
             d.update(d_)
             if not len(d) == count + len(d_):
-                raise Exception(f"Hyperparameter indexing failed, overlapping keys in {d_}")
+                raise Exception(
+                    f"Hyperparameter indexing failed, overlapping keys in {d_}"
+                )
         for k, v in object["kwargs"].items():
             count = len(d)
             d_ = index_hyperparameters(v)
             d.update(d_)
             if not len(d) == count + len(d_):
-                raise Exception(f"Hyperparameter indexing failed, overlapping keys in {d_}")
+                raise Exception(
+                    f"Hyperparameter indexing failed, overlapping keys in {d_}"
+                )
         return d
     elif issubclass(type(object), list):
         d = {}
@@ -141,7 +150,9 @@ def index_hyperparameters(object: BaseClass) -> dict:
             d_ = index_hyperparameters(x)
             d.update(d_)
             if not len(d) == count + len(d_):
-                raise Exception(f"Hyperparameter indexing failed, overlapping keys in {d_}")
+                raise Exception(
+                    f"Hyperparameter indexing failed, overlapping keys in {d_}"
+                )
         return d
     elif (
         object is None
@@ -154,21 +165,44 @@ def index_hyperparameters(object: BaseClass) -> dict:
         print(object)
         raise ValueError
 
+
 def load_hyperparameters_(object: BaseClass, hyperparameter_dictionary: dict) -> dict:
     if issubclass(type(object), Opt):
-        return load_hyperparameters_(hyperparameter_dictionary[object.label], hyperparameter_dictionary)
+        return load_hyperparameters_(
+            hyperparameter_dictionary[object.label], hyperparameter_dictionary
+        )
 
     if issubclass(type(object), BaseClass):
         return {
             **{"BC_class_name": type(object).__name__},
-            **{"args": [load_hyperparameters_(arg, hyperparameter_dictionary) for arg in object.args]},
-            **{"kwargs": {k: load_hyperparameters_(v, hyperparameter_dictionary) for k, v in object.kwargs.items()}},
+            **{
+                "args": [
+                    load_hyperparameters_(arg, hyperparameter_dictionary)
+                    for arg in object.args
+                ]
+            },
+            **{
+                "kwargs": {
+                    k: load_hyperparameters_(v, hyperparameter_dictionary)
+                    for k, v in object.kwargs.items()
+                }
+            },
         }
     elif issubclass(type(object), dict):
         return {
             **{"BC_class_name": object["BC_class_name"]},
-            **{"args": [load_hyperparameters_(arg, hyperparameter_dictionary) for arg in object["args"]]},
-            **{"kwargs": {k: load_hyperparameters_(v, hyperparameter_dictionary) for k, v in object["kwargs"].items()}},
+            **{
+                "args": [
+                    load_hyperparameters_(arg, hyperparameter_dictionary)
+                    for arg in object["args"]
+                ]
+            },
+            **{
+                "kwargs": {
+                    k: load_hyperparameters_(v, hyperparameter_dictionary)
+                    for k, v in object["kwargs"].items()
+                }
+            },
         }
     elif (
         object is None
@@ -183,11 +217,17 @@ def load_hyperparameters_(object: BaseClass, hyperparameter_dictionary: dict) ->
         print(object)
         raise ValueError
 
+
 def load_hyperparameters(object: BaseClass, hyperparameter_dictionary: dict) -> dict:
     mp = load_hyperparameters_(object, hyperparameter_dictionary)
     return oce.create_BC(mp)
 
-def optimize(model: Union[BaseModel, dict], runner: Union[BaseModelManager, Callable], max_evals = 3):
+
+def optimize(
+    model: Union[BaseModel, dict],
+    runner: Union[BaseModelManager, Callable],
+    max_evals=3,
+):
     hyperparameter_index = index_hyperparameters(model)
 
     def objective(hyperparameter_dictionary, model = model, runner = runner):
@@ -199,8 +239,10 @@ def objective(hyperparameter_dictionary, model = model, runner = runner):
             metric = runner(model)
         return metric
 
-    best = fmin(fn=objective,
+    best = fmin(
+        fn=objective,
         space=hyperparameter_index,
         algo=tpe.suggest,
-        max_evals=max_evals,)
-    return best
\ No newline at end of file
+        max_evals=max_evals,
+    )
+    return best
diff --git a/src/olorenchemengine/internal.py b/src/olorenchemengine/internal.py
index d4aafc07..b8398b60 100644
--- a/src/olorenchemengine/internal.py
+++ b/src/olorenchemengine/internal.py
@@ -1,33 +1,44 @@
 import base64
+import contextlib
 import hashlib
 import inspect
 import json
 import os
 import pickle
+import subprocess
+import sys
+import tempfile
 from abc import ABC, abstractmethod
 from typing import Callable, Union
-import gcsfs
-import subprocess
+from unittest.mock import MagicMock
+
+import numpy as np
+import pandas as pd
 import pyrebase  # Default pyrebase is pyrebase3 which won't work. Need to install pyrebase4 (pip install pyrebase4)
 from google.cloud.firestore import Client
 from google.oauth2.credentials import Credentials
-import sys
-from unittest.mock import MagicMock
 
 import olorenchemengine
 
-sys.modules["olorenautoml"] = olorenchemengine # important for backwards compatibility of some models
+sys.modules[
+    "olorenautoml"
+] = olorenchemengine  # important for backwards compatibility of some models
+
 
 def mock_imports(g, *args):
     for arg in args:
         g[arg] = MagicMock()
 
+
 def all_subclasses(cls):
-    """ Helper function to return all subclasses of class"""
-    return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in all_subclasses(c)])
+    """Helper function to return all subclasses of class"""
+    return set(cls.__subclasses__()).union(
+        [s for c in cls.__subclasses__() for s in all_subclasses(c)]
+    )
+
 
 class OASConnector:
-    """ Class which links oce to OAS and can move data between them using Firestore and an API"""
+    """Class which links oce to OAS and can move data between them using Firestore and an API"""
 
     def __init__(self) -> None:
         self.storage = None
@@ -43,7 +54,9 @@ def authenticate(self):
 
         # Use the google verify password REST api to authenticate and generate user tokens
         def log_in_first_time(token):
-            request_url = f"{FIREBASE_REST_API}:signInWithCustomToken?key={FIREBASE_API_KEY}"
+            request_url = (
+                f"{FIREBASE_REST_API}:signInWithCustomToken?key={FIREBASE_API_KEY}"
+            )
             headers = {"content-type": "application/json; charset=UTF-8"}
             data = json.dumps({"token": token, "returnSecureToken": True})
             resp = requests.post(request_url, headers=headers, data=data)
@@ -57,9 +70,13 @@ def log_in_first_time(token):
             return resp.json()
 
         def get_id_token_from_refresh(refresh_token):
-            request_url = f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_API_KEY}"
+            request_url = (
+                f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_API_KEY}"
+            )
             headers = {"content-type": "application/json; charset=UTF-8"}
-            data = json.dumps({"refresh_token": refresh_token, "grant_type": "refresh_token"})
+            data = json.dumps(
+                {"refresh_token": refresh_token, "grant_type": "refresh_token"}
+            )
             resp = requests.post(request_url, headers=headers, data=data)
             try:
                 resp.raise_for_status()
@@ -105,16 +122,18 @@ def get_id_token_from_refresh(refresh_token):
             "appId": "1:602366687071:web:f35531f7142084b86a28ea",
             "measurementId": "G-MVNK4DPQ60",
         }
-        firebase = pyrebase.initialize_app(config)
+        self.firebase = pyrebase.initialize_app(config)
 
-        auth = firebase.auth()
-        user = auth.refresh(response["refreshToken"])
+        # self.auth = self.firebase.auth()
+        # response = self.auth.refresh(response["refreshToken"])
 
-        self.storage = firebase.storage()
+        self.storage = self.firebase.storage()
         self.logging_db = Client("oloren-ai", creds)
         self.uid = response["uid"]
         self.uid_token = response["idToken"]
 
+        return self.uid_token
+
     def upload_vis(self, visualization):
         self.authenticate()
 
@@ -130,7 +149,6 @@ def upload_vis(self, visualization):
 
         return out
 
-
     def upload_model(self, model, model_name):
         self.authenticate()
 
@@ -139,15 +157,17 @@ def upload_model(self, model, model_name):
         with tempfile.NamedTemporaryFile() as tmp:
             pickle.dump(saves(model), tmp)
             tmp.flush()
-            self.storage.child(self.uid + "/models/" + model_name + ".pkl").put(tmp.name, self.uid_token)
+            response = self.storage.child(f"{self.uid}/models/{model_name}.pkl").put(
+                tmp.name, self.uid_token
+            )
 
-        out = self.logging_db.collection("models").add(
+        return self.logging_db.collection("models").add(
             {
                 "uid": self.uid,
                 "did": None,
                 "source": "oce",
                 "data": {
-                    "file_path": "gs://oloren-ai.appspot.com/" + self.uid + "/models/" + model_name + ".pkl",
+                    "file_path": f"gs://oloren-ai.appspot.com/{self.uid}/models/{model_name}.pkl",
                     "name": model_name,
                     "description": pretty_params_str(model),
                     "parameters": json.dumps(parameterize(model)),
@@ -157,12 +177,12 @@ def upload_model(self, model, model_name):
             }
         )
 
-        return out
 
 oas_connector = OASConnector()
 
 ignored_kwargs = ["map_location", "num_workers"]
 
+
 def download_public_file(path, redownload=False):
     """Download a public file from Oloren's Public Storage, and returns the contents.
 
@@ -172,23 +192,29 @@ def download_public_file(path, redownload=False):
 
     local_path = os.path.join(os.path.expanduser("~"), ".oce", path)
 
-    if os.path.exists(local_path) and not redownload: return local_path
+    if os.path.exists(local_path) and not redownload:
+        return local_path
 
     if not os.path.exists(os.path.dirname(local_path)):
         os.makedirs(os.path.dirname(local_path))
 
     print(f"Downloading {path}...")
-    fs = gcsfs.GCSFileSystem()
-    with fs.open(f'gs://oloren-public-data/{path}', 'rb') as f:
-        with open(local_path, 'wb') as out:
-            out.write(f.read())
+    import urllib.request
 
+    urllib.request.urlretrieve(
+        f"https://storage.googleapis.com/oloren-public-data/{path}", local_path
+    )
     return local_path
 
 
 def get_default_args(func):
     signature = inspect.signature(func)
-    return {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}
+    return {
+        k: v.default
+        for k, v in signature.parameters.items()
+        if v.default is not inspect.Parameter.empty
+    }
+
 
 def log_arguments(func: Callable[..., None]) -> Callable[..., None]:
     """
@@ -203,23 +229,407 @@ def log_arguments(func: Callable[..., None]) -> Callable[..., None]:
     """
 
     def wrapper(self, *args, **kwargs):
-
-        if "log" in kwargs.keys() and not kwargs["log"]:
-            pass
-        else:
+        if "log" not in kwargs.keys() or kwargs["log"]:
             kwds = get_default_args(func)
             for k, v in kwds.items():
-                if not k in kwargs.keys():
-                    kwargs.update({k: v})
+                if k not in kwargs:
+                    kwargs[k] = v
             self.args = args
-            self.kwargs = {k: v for k, v in kwargs.items() if not k in ignored_kwargs}
-        return func(self, *args, **kwargs)
+            self.kwargs = {k: v for k, v in kwargs.items() if k not in ignored_kwargs}
+        if _runtime.is_local:
+            return func(self, *args, **kwargs)
+        import uuid
+
+        REMOTE_ID = str(uuid.uuid4())
+
+        if "BaseRemoteSymbol" in str(type(self)):
+            self.REMOTE_ID = REMOTE_ID
+            return func(self, *args, **kwargs)
+
+        _runtime.add_instruction(
+            {
+                "type": "CREATE",
+                "REMOTE_ID": REMOTE_ID,
+                "parameters": parameterize(self),
+            }
+        )
+        self.REMOTE_ID = REMOTE_ID
 
     wrapper.__wrapped__ = func
 
     return wrapper
 
-class BaseClass(ABC):
+
+def _create_BC_if_necessary(obj):
+    return obj if type(obj) is not dict else create_BC(obj)
+
+
+def deparametrize_args_kwargs(params):
+    args = params["args"]
+    kwargs = params["kwargs"]
+    return [_create_BC_if_necessary(arg) for arg in args], {
+        key: create_BC(kwarg) for key, kwarg in kwargs.items()
+    }
+
+
+def recursive_get_attr(parent, attr):
+    if len(attr) == 0:
+        return parent
+    return recursive_get_attr(getattr(parent, attr[0]), attr[1:])
+
+
+def generate_uuid():
+    import uuid
+
+    return str(uuid.uuid4())
+
+
+class _RemoteRuntime:
+    def __init__(self) -> None:
+        self.reset_remote()
+
+    def get_remote_obj(self, remote_id):
+        if self.runner is None:
+            raise NotImplementedError(
+                "Not yet implemented remote object retrieval locally"
+            )
+        else:
+            return self.runner.get_remote_obj(remote_id)
+
+    def add_instruction(self, instruction):
+
+        self.instruction_buffer.append(instruction)
+
+        if instruction["type"] == "CALL":
+            x = self.send_instructions_blocking()
+            if x is not None and isinstance(x, str):
+                return json.loads(x)
+
+    def get_iterable(self, remote_id):
+        self.instruction_buffer.append({"type": "ITER", "REMOTE_ID": remote_id})
+        return [
+            BaseRemoteSymbol.from_rid(rid) for rid in self.send_instructions_blocking()
+        ]
+
+    def get_obj_repr(self, remote_id):
+        self.instruction_buffer.append({"type": "REPR", "REMOTE_ID": remote_id})
+        return self.send_instructions_blocking()
+
+    def send_instructions_blocking(self):
+        import requests
+
+        response = None
+
+        oas_connector.authenticate()
+        _, execution = oas_connector.logging_db.collection("executions").add(
+            {
+                "uid": oas_connector.uid,
+                "sid": self.session_id,
+                "instructions": json.dumps(self.instruction_buffer),
+                "status": "Queued",
+            }
+        )
+
+        eid = execution.id
+
+        response = requests.post(
+            f"{self.remote_url}/firestore/run_remote/",
+            params={
+                "eid": eid,
+                "uid": oas_connector.authenticate(),
+            },
+            headers={
+                "accept": "application/json",
+                "content-type": "application/x-www-form-urlencoded",
+            },
+        )
+
+        if response.status_code != 200:
+            self.instruction_buffer = []
+            raise ValueError(f"Error code on post to {self.remote_url} {response.status_code} - {response.text}. Length of instruction buffer: {len(json.dumps(self.instruction_buffer))}")
+
+        response = response.json()
+
+        execution = None
+
+        def on_snapshot(doc_snapshot, changes, read_time):
+            nonlocal execution
+            for doc in doc_snapshot:
+                execution = doc.to_dict()
+
+        listener = (
+            oas_connector.logging_db.collection("executions")
+            .document(eid)
+            .on_snapshot(on_snapshot)
+        )
+
+        while execution is None:
+            pass
+
+        while execution["status"] not in ("Finished", "Error"):
+            pass
+
+        listener.close()
+
+        if execution["status"] == "Error":
+            print(f"Remote Call: {self.instruction_buffer}\nResulted in Traceback:")
+            print(execution["traceback"])
+            oas_connector.logging_db.collection("executions").document(eid).delete()
+            self.instruction_buffer = []
+            raise ValueError("Remote call resulted in error")
+
+        elif execution["status"] == "Finished":
+            if len(execution["stdout"].strip()) > 0:
+                print(
+                    "\n".join([f"REMOTE: {x}" for x in execution["stdout"].split("\n")])
+                )
+            oas_connector.logging_db.collection("executions").document(eid).delete()
+            self.instruction_buffer = []
+            return execution["return"]
+        else:
+            raise NotImplementedError(f"Unknown status {execution['status']}")
+
+    @property
+    def is_local(self):
+        return self.runtime == "local"
+
+    def start_remote(self, session_id, debug, remote_url):
+        self.runtime = "remote"
+        self.instruction_buffer = []
+        self.session_id = session_id
+        self.runner = None
+        self.debug = debug
+        self.remote_url = remote_url
+
+    def reset_remote(self):
+        self.runtime = "local"
+        self.instruction_buffer = []
+        self.session_id = None
+        self.runner = None
+        self.debug = False
+        self.remote_url = None
+
+
+_runtime = _RemoteRuntime()  # internal runtime object
+
+
+def set_runner(runner):
+    _runtime.runner = runner
+
+
+def get_runtime():
+    return _runtime
+
+
+class Remote(object):
+    def __init__(self, remote_url, session_id=None, keep_alive=False, debug=False):
+        self.keep_alive = keep_alive
+        self.session_id = session_id
+        self.debug = debug
+        self.remote_url = remote_url
+
+    def __enter__(self):
+        self.session_id = generate_uuid()
+        _runtime.start_remote(self.session_id, self.debug, self.remote_url)
+        return self.session_id
+
+    def __exit__(self, type, value, traceback):
+        _runtime.reset_remote()
+        if not self.keep_alive:
+            import google
+
+            oas_connector.authenticate()
+
+            with contextlib.suppress(google.api_core.exceptions.PermissionDenied):
+                doc = (
+                    oas_connector.logging_db.collection("sessions")
+                    .document(self.session_id)
+                    .get()
+                    .to_dict()
+                )
+                if doc is not None:
+                    for rid in doc["objects"]:
+                        try:
+                            oas_connector.storage.delete(
+                                f"{oas_connector.uid}/sessions/{self.session_id}/{rid}.oce",
+                                oas_connector.uid_token,
+                            )
+                        except:
+                            import logging
+
+                            logging.warning(
+                                f"Unable to delete remote object {rid} from session {self.session_id}"
+                            )
+                    oas_connector.logging_db.collection("sessions").document(
+                        self.session_id
+                    ).delete()
+
+
+def pretty_args_kwargs(args, kwargs):
+    parameterized_args = [parameterize(arg) for arg in args]
+    parameterized_kwargs = {k: parameterize(v) for k, v in kwargs.items()}
+    return f"(args={parameterized_args}, kwargs={parameterized_kwargs})"
+
+
+def _paremetrize_if_necessary(obj):
+    if isinstance(obj, BaseRemoteSymbol):
+        return parameterize(obj)
+    elif isinstance(obj, pd.Series):
+        return list(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    with contextlib.suppress(Exception):
+        json.dumps(obj)
+        return obj
+
+
+def parametrize_args_kwargs(args, kwargs):
+    parameterized_args = [_paremetrize_if_necessary(arg) for arg in args]
+    parameterized_kwargs = {k: _paremetrize_if_necessary(v) for k, v in kwargs.items()}
+    return {"args": parameterized_args, "kwargs": parameterized_kwargs}
+
+
+def _truncate_json(json_obj, max_length=40):
+    if isinstance(json_obj, str):
+        if len(json_obj) > max_length:
+            return f"{json_obj[:max_length]}..."
+        else:
+            return json_obj
+    elif isinstance(json_obj, list):
+        return [_truncate_json(x, max_length) for x in json_obj]
+    elif isinstance(json_obj, dict):
+        return {
+            _truncate_json(k): _truncate_json(v, max_length)
+            for k, v in json_obj.items()
+        }
+    else:
+        return json_obj
+
+
+class BaseRemoteSymbol:
+    @log_arguments
+    def __init__(
+        self, REMOTE_SYMBOL_NAME, REMOTE_PARENT, args=None, kwargs=None
+    ) -> None:
+        if _runtime.is_local:
+            raise RuntimeError("Cannot instantiate RemoteObject in local _runtime")
+
+        if args is not None:
+            self.args = args
+        if kwargs is not None:
+            self.kwargs = kwargs
+
+        self.REMOTE_SYMBOL_NAME = REMOTE_SYMBOL_NAME
+        self.REMOTE_PARENT = REMOTE_PARENT
+        self.REMOTE_CHILDREN = {}
+
+        if REMOTE_SYMBOL_NAME == "CALL":
+            out = _runtime.add_instruction(
+                {
+                    "type": "CALL",
+                    "PARENT_REMOTE_ID": self.REMOTE_PARENT.REMOTE_ID,
+                    "REMOTE_ID": self.REMOTE_ID,
+                    "ARGUMENTS": parametrize_args_kwargs(args, kwargs),
+                }
+            )
+
+            if self.REMOTE_PARENT.REMOTE_SYMBOL_NAME == "render_ipynb":
+                from IPython.display import IFrame, display
+
+                display(IFrame(out, width=800, height=600))
+
+        else:
+            if not hasattr(self.REMOTE_PARENT, "REMOTE_ID"):
+                self.REMOTE_PARENT._upload_remote()
+
+            _runtime.add_instruction(
+                {
+                    "type": "SYMBOL",
+                    "PARENT_REMOTE_ID": self.REMOTE_PARENT.REMOTE_ID,
+                    "SYMBOL_NAME": self.REMOTE_SYMBOL_NAME,
+                    "REMOTE_ID": self.REMOTE_ID,
+                }
+            )
+
+    @classmethod
+    def from_rid(cls, rid):
+        x = object.__new__(cls)
+        x.REMOTE_ID = rid
+        return x
+
+    def __iter__(self):
+        return iter(_runtime.get_iterable(self.REMOTE_ID))
+
+    def __repr__(self):
+        return _runtime.get_obj_repr(self.REMOTE_ID)
+
+    def _upload_remote(self):
+        try:
+            oas_connector.authenticate()
+            REMOTE_ID = generate_uuid()
+
+            r = _runtime.runtime
+            _runtime.runtime = "local"
+
+            with tempfile.NamedTemporaryFile() as tmp:
+                pickle.dump(saves(self), tmp)
+                tmp.flush()
+                response = oas_connector.storage.child(
+                    f"{oas_connector.uid}/sessions/{_runtime.session_id}"
+                    + f"/{REMOTE_ID}.oce"
+                ).put(tmp.name, oas_connector.uid_token)
+        except:
+            import sys
+            import traceback
+
+            traceback.print_exc()
+            sys.exit(0)
+
+        _runtime.runtime = r
+        object.__setattr__(self, "REMOTE_ID", REMOTE_ID)
+
+        return REMOTE_ID
+
+    def __getattribute__(self, key):
+        if _runtime.is_local or key == "_upload_remote":
+            return object.__getattribute__(self, key)
+        if "ipython_canary_method_should_not_exist" in key:
+            return {}
+        if (
+            key.startswith("REMOTE")
+            or key.startswith("__")
+            or key in ["args", "kwargs"]
+        ):
+            return object.__getattribute__(self, key)
+
+        if not hasattr(self, "REMOTE_CHILDREN") or key not in self.REMOTE_CHILDREN:
+            if not hasattr(self, "REMOTE_CHILDREN"):
+                self.REMOTE_CHILDREN = {}
+            self.REMOTE_CHILDREN[key] = BaseRemoteSymbol(key, self)
+        return self.REMOTE_CHILDREN[key]
+
+    def __call__(self, *args, **kwargs):
+        # TODO: move logic from init to here, then return baseremotesymbol only if return type is none
+
+        remote_id = generate_uuid()
+        out = _runtime.add_instruction(
+            {
+                "type": "CALL",
+                "PARENT_REMOTE_ID": self.REMOTE_ID,
+                "REMOTE_ID": remote_id,
+                "ARGUMENTS": parametrize_args_kwargs(args, kwargs),
+            }
+        )
+
+        if self.REMOTE_SYMBOL_NAME == "render_ipynb":
+            from IPython.display import IFrame, display
+
+            display(IFrame(out, width=800, height=600))
+
+        return out if out is not None else RemoteObj(remote_id)
+
+
+class BaseClass(BaseRemoteSymbol):
     """BaseClass is the base class for all models.
 
     All classes in Oloren ChemEngine should inherit from BaseClass to enable for universal saving and loading of both
@@ -231,6 +641,9 @@ class BaseClass(ABC):
         _load: loads an instance of a BaseClass from a dictionary (abstract method to be implmented by subclasses)
     """
 
+    def __init__(self):
+        pass
+
     @classmethod
     def Registry(cls):
         """Registry is a recursive method to create a dictionary of all subclasses of BaseClass, with the key being the name of the subclass and the value being the subclass itself."""
@@ -248,12 +661,13 @@ def AllInstances(cls):
 
         Standard instances means that all required parameters for instantiation of the
         subclasses are set with canonical values."""
-        if not hasattr(cls, "__abstractmethods__") or len(cls.__abstractmethods__) == 0:
-            try:
-                return [cls()] + [o for sc in cls.__subclasses__() for o in sc.AllInstances()]
-            except:
-                return [o for sc in cls.__subclasses__() for o in sc.AllInstances()]
-        else:
+        if hasattr(cls, "__abstractmethods__") and len(cls.__abstractmethods__) != 0:
+            return [o for sc in cls.__subclasses__() for o in sc.AllInstances()]
+        try:
+            return [cls()] + [
+                o for sc in cls.__subclasses__() for o in sc.AllInstances()
+            ]
+        except Exception:
             return [o for sc in cls.__subclasses__() for o in sc.AllInstances()]
 
     @classmethod
@@ -284,6 +698,13 @@ def copy(self):
         return obj_copy
 
 
+class RemoteObj(BaseRemoteSymbol):
+    """Dummy object to represent remote objects."""
+
+    def __init__(self, remote_id):
+        self.REMOTE_ID = remote_id
+
+
 def parameterize(object: Union[BaseClass, list, int, float, str, None]) -> dict:
     """parameterize is a recursive method which creates a dictionary of all arguments necessary to instantiate a BaseClass object.
 
@@ -298,7 +719,11 @@ def parameterize(object: Union[BaseClass, list, int, float, str, None]) -> dict:
     Returns:
         dict: dictionary of parameters necessary to instantiate the object.
     """
-    if issubclass(type(object), BaseClass):
+    if issubclass(type(object), BaseClass) or (
+        issubclass(type(object), BaseRemoteSymbol) and hasattr(object, "REMOTE_ID")
+    ):
+        if hasattr(object, "REMOTE_ID"):
+            return {"REMOTE_ID": object.REMOTE_ID}
         return {
             **{"BC_class_name": type(object).__name__},
             **{"args": [parameterize(arg) for arg in object.args]},
@@ -314,8 +739,7 @@ def parameterize(object: Union[BaseClass, list, int, float, str, None]) -> dict:
     elif issubclass(type(object), list):
         return [parameterize(x) for x in object]
     else:
-        print(object)
-        raise ValueError
+        raise ValueError(f"Invalid object {object}")
 
 
 def model_name_from_params(param_dict: dict) -> str:
@@ -330,7 +754,9 @@ def model_name_from_params(param_dict: dict) -> str:
     return (
         param_dict["BC_class_name"]
         + " "
-        + base64.urlsafe_b64encode(hashlib.md5(str(param_dict).encode("utf-8")).digest())[:8].decode("utf-8")
+        + base64.urlsafe_b64encode(
+            hashlib.md5(str(param_dict).encode("utf-8")).digest()
+        )[:8].decode("utf-8")
     )
 
 
@@ -390,22 +816,41 @@ def create_BC(d: dict) -> BaseClass:
         BaseClass: the object created from the parameters
     """
     if isinstance(d, str):
-        d = d.replace("'", '"').replace("True", "true").replace("False", "false").replace("None", "null")
+        d = (
+            d.replace("'", '"')
+            .replace("True", "true")
+            .replace("False", "false")
+            .replace("None", "null")
+        )
         d = json.loads(d)
+
+    if "REMOTE_ID" in d.keys():
+        return _runtime.get_remote_obj(d["REMOTE_ID"])
+
     args = []
     if "args" in d.keys():
         for arg in d["args"]:
-            if isinstance(arg, dict) and "BC_class_name" in arg.keys():
+            if isinstance(arg, dict) and (
+                "BC_class_name" in arg.keys() or "REMOTE_ID" in arg.keys()
+            ):
                 args.append(create_BC(arg))
             else:
                 if isinstance(arg, list):
-                    arg = [create_BC(x) if isinstance(x, dict) and "BC_class_name" in x.keys() else x for x in arg]
+                    arg = [
+                        create_BC(x)
+                        if isinstance(x, dict)
+                        and ("BC_class_name" in x.keys() or "REMOTE_ID" in x.keys())
+                        else x
+                        for x in arg
+                    ]
                 args.append(arg)
 
     kwargs = {}
     if "kwargs" in d.keys():
         for k, v in d["kwargs"].items():
-            if isinstance(v, dict) and "BC_class_name" in v.keys():
+            if isinstance(v, dict) and (
+                "BC_class_name" in v.keys() or "REMOTE_ID" in v.keys()
+            ):
                 kwargs[k] = create_BC(v)
             else:
                 kwargs[k] = v
@@ -423,7 +868,7 @@ def loads(d: dict) -> BaseClass:
         BaseClass: the recreated object
     """
 
-    import olorenchemengine as olorenautoml # for backwards compatibility
+    import olorenchemengine as olorenautoml  # for backwards compatibility
 
     args = []
     for arg in d["args"]:
@@ -431,10 +876,18 @@ def loads(d: dict) -> BaseClass:
             args.append(loads(arg))
         else:
             if isinstance(arg, list):
-                arg = [loads(x) if isinstance(x, dict) and "BC_class_name" in x.keys() else x for x in arg]
+                arg = [
+                    loads(x)
+                    if isinstance(x, dict) and "BC_class_name" in x.keys()
+                    else x
+                    for x in arg
+                ]
             args.append(arg)
 
-    kwargs = {k: loads(v) if isinstance(v, dict) and "BC_class_name" in v.keys() else v for k, v in d["kwargs"].items()}
+    kwargs = {
+        k: loads(v) if isinstance(v, dict) and "BC_class_name" in v.keys() else v
+        for k, v in d["kwargs"].items()
+    }
     bc = BaseClass.Registry()[d["BC_class_name"]](*args, **kwargs)
     bc._load(d["instance_save"])
     return bc
@@ -447,9 +900,17 @@ def save(model: BaseClass, fname: str):
         model (BaseClass): the object to be saved
         fname (str): the file name to save the model to
     """
-    save_dict = saves(model)
-    with open(fname, "wb+") as f:
-        pickle.dump(save_dict, f)
+
+    if hasattr(model, "REMOTE_ID"):
+
+        REMOTE_ID = model.REMOTE_ID
+        oas_connector.storage.child(
+            f"{oas_connector.uid}/sessions/{_runtime.session_id}/{REMOTE_ID}.oce"
+        ).download("", fname, token=oas_connector.authenticate())
+    else:
+        save_dict = saves(model)
+        with open(fname, "wb+") as f:
+            pickle.dump(save_dict, f)
 
 
 def load(fname: str) -> BaseClass:
@@ -461,33 +922,51 @@ def load(fname: str) -> BaseClass:
     Returns:
         BaseClass: the BaseClass object which as been recreated from the file
     """
-    with open(fname, "rb") as f:
-        d = pickle.load(f)
-    return loads(d)
+    if _runtime.is_local:
+        with open(fname, "rb") as f:
+            d = pickle.load(f)
+        return loads(d)
+    else:
+        oas_connector.authenticate()
+        REMOTE_ID = generate_uuid()
+        response = oas_connector.storage.child(
+            f"{oas_connector.uid}/sessions/{_runtime.session_id}" + f"/{REMOTE_ID}.oce.delete"
+        ).put(fname, oas_connector.uid_token)
+        return RemoteObj(REMOTE_ID)
 
 
 def pretty_params(base: Union[BaseClass, dict]) -> dict:
-    """ Returns a dictionary of the parameters of the passed BaseClass object, formatted such that they are in a
+    """Returns a dictionary of the parameters of the passed BaseClass object, formatted such that they are in a
     human readable format, with the names of the arguments included."""
     if isinstance(base, dict):
         base = loads(base)
 
     if issubclass(type(base), BaseClass):
-        args = [arg for arg in base.args]
-        kwargs = {k: v for k, v in base.kwargs.items()}
-        base_object_parameters = list(inspect.signature(base.__init__).parameters.keys())
+        args = list(base.args)
+        kwargs = dict(base.kwargs.items())
+        base_object_parameters = list(
+            inspect.signature(base.__init__).parameters.keys()
+        )
 
         for kwarg in kwargs:
             if kwarg in base_object_parameters:
                 base_object_parameters.remove(kwarg)
 
-        labeled_args = {k: v for k, v in zip(base_object_parameters, args)}
-
+        labeled_args = dict(zip(base_object_parameters, args))
         fully_labeled_args = {k: pretty_params(v) for k, v in labeled_args.items()}
         fully_labeled_kwargs = {k: pretty_params(v) for k, v in kwargs.items()}
 
-        return {**{"BC_Class_name": type(base).__name__}, **fully_labeled_args, **fully_labeled_kwargs}
-    elif issubclass(type(base), int) or issubclass(type(base), float) or issubclass(type(base), str) or base is None:
+        return {
+            **{"BC_Class_name": type(base).__name__},
+            **fully_labeled_args,
+            **fully_labeled_kwargs,
+        }
+    elif (
+        issubclass(type(base), int)
+        or issubclass(type(base), float)
+        or issubclass(type(base), str)
+        or base is None
+    ):
         return base
     elif issubclass(type(base), list):
         return [pretty_params(x) for x in base]
@@ -499,13 +978,13 @@ def pretty_params(base: Union[BaseClass, dict]) -> dict:
 
 
 def pretty_params_str(base: Union[BaseClass, dict]) -> str:
-    """ Returns a string of the parameters of the passed BaseClass object, formatted such that they are in a human
+    """Returns a string of the parameters of the passed BaseClass object, formatted such that they are in a human
     readable format"""
     return json.dumps(pretty_params(base), indent=4)
 
 
 def json_params_str(base: Union[BaseClass, dict]) -> str:
-    """ Returns a json string of the parameters of the passed BaseClass object so that the model parameter dictionary can
+    """Returns a json string of the parameters of the passed BaseClass object so that the model parameter dictionary can
     be reconstructed with json.load(params_str)"""
     return (
         json.dumps(pretty_params(base))
@@ -515,14 +994,20 @@ def json_params_str(base: Union[BaseClass, dict]) -> str:
         .replace("None", "null")
     )
 
+
 def install_with_permission(package_name: str):
-    inp = input(f"The required package {package_name} is not installed. Do you want to install it? [y/N]? ")
+    inp = input(
+        f"The required package {package_name} is not installed. Do you want to install it? [y/N]? "
+    )
     if inp.lower() == "y":
         subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
     else:
-        print(f"Stopping program. You can install the package manually with: \n >> pip install {package_name}")
+        print(
+            f"Stopping program. You can install the package manually with: \n >> pip install {package_name}"
+        )
         os._exit(1)
 
+
 def import_or_install(package_name: str, statement: str = None, scope: dict = None):
     if scope is None:
         scope = globals()
@@ -533,4 +1018,4 @@ def import_or_install(package_name: str, statement: str = None, scope: dict = No
     except ImportError:
         install_with_permission(package_name)
     finally:
-        exec(statement, scope)
\ No newline at end of file
+        exec(statement, scope)
diff --git a/src/olorenchemengine/interpret.py b/src/olorenchemengine/interpret.py
index 04c5d107..ad265b3e 100644
--- a/src/olorenchemengine/interpret.py
+++ b/src/olorenchemengine/interpret.py
@@ -1,26 +1,20 @@
 import copy
 import io
-
 from abc import abstractmethod
-from tqdm import tqdm
 
 import numpy as np
+import PIL
+import selfies as sf
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from rdkit.DataStructs.cDataStructs import BulkTanimotoSimilarity, TanimotoSimilarity
+from tqdm import tqdm
 
 from .base_class import *
 from .dataset import *
-from .representations import *
 from .external.stoned import *
+from .representations import *
 
-import PIL
-
-try:
-    import selfies as sf
-except ImportError:
-    mock_imports(globals(), "sf")
-
-from rdkit.DataStructs.cDataStructs import BulkTanimotoSimilarity, TanimotoSimilarity
-from rdkit import Chem
-from rdkit.Chem import AllChem
 
 class PerturbationEngine(BaseClass):
 
@@ -47,6 +41,7 @@ def get_compound(self, smiles, n=1, **kwargs) -> str:
     def get_compound_list(self, smiles, **kwargs) -> list:
         pass
 
+
 class SwapMutations(PerturbationEngine):
     """SwapMutations replaces substructures with radius r with another substructure
     with radius < r. The substructure is chosen such that it has the same outgoing
@@ -60,72 +55,137 @@ class SwapMutations(PerturbationEngine):
             entire class of modifications"""
 
     @log_arguments
-    def __init__(self, radius = 0, log = True):
+    def __init__(self, radius=0, log=True):
         self.radius = radius
 
         if not radius in (0, 1, 2):
             raise ValueError("radius must be 0, 1, or 2")
 
-
         self.trans = dict()
-        for radius in range(radius+1):
-            transformation_path = download_public_file(f"swap-mutations/trans_{radius}.json")
+        for radius in range(radius + 1):
+            transformation_path = download_public_file(
+                f"swap-mutations/trans_{radius}.json"
+            )
             with open(transformation_path, "r") as f:
                 self.trans.update(json.load(f))
 
-    def get_substitution(self, m, idx, r = 1):
+    def get_substitution(self, m, idx, r=1):
 
         # finding atoms within radius r of atom idx
         substructure_atoms = set({idx})
         for i in range(r):
-            bonds = [b for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds()]
-            substructure_atoms_ = {b.GetBeginAtomIdx() for b in bonds}.union({b.GetEndAtomIdx() for b in bonds})
+            bonds = [
+                b for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds()
+            ]
+            substructure_atoms_ = {b.GetBeginAtomIdx() for b in bonds}.union(
+                {b.GetEndAtomIdx() for b in bonds}
+            )
             substructure_atoms = substructure_atoms.union(substructure_atoms_)
 
         # finding atoms connecting to substructure
-        surface_atoms = {b.GetEndAtomIdx() for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds() if not b.GetEndAtomIdx() in substructure_atoms}
-        surface_atoms.update({b.GetBeginAtomIdx() for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds() if not b.GetBeginAtomIdx() in substructure_atoms})
+        surface_atoms = {
+            b.GetEndAtomIdx()
+            for a_ in substructure_atoms
+            for b in m.GetAtomWithIdx(a_).GetBonds()
+            if not b.GetEndAtomIdx() in substructure_atoms
+        }
+        surface_atoms.update(
+            {
+                b.GetBeginAtomIdx()
+                for a_ in substructure_atoms
+                for b in m.GetAtomWithIdx(a_).GetBonds()
+                if not b.GetBeginAtomIdx() in substructure_atoms
+            }
+        )
         surface_atoms = list(surface_atoms)
         surface_l = list()
         for surface_atom in surface_atoms:
-            bonds = [b.GetIdx() for b in m.GetAtomWithIdx(surface_atom).GetBonds() if b.GetEndAtomIdx() in substructure_atoms or b.GetBeginAtomIdx() in substructure_atoms]
-            surface_l.append((m.GetAtomWithIdx(surface_atom).GetAtomicNum(),
-                              [m.GetBondWithIdx(b).GetBondTypeAsDouble() for b in bonds]))
-
-        tmp = [(x1, x2, x3) for x1, x2, x3 in zip([sum(surface_b) for surface_a, surface_b in surface_l], surface_atoms, surface_l)]
-        tmp.sort(key = lambda x: x[0])
+            bonds = [
+                b.GetIdx()
+                for b in m.GetAtomWithIdx(surface_atom).GetBonds()
+                if b.GetEndAtomIdx() in substructure_atoms
+                or b.GetBeginAtomIdx() in substructure_atoms
+            ]
+            surface_l.append(
+                (
+                    m.GetAtomWithIdx(surface_atom).GetAtomicNum(),
+                    [m.GetBondWithIdx(b).GetBondTypeAsDouble() for b in bonds],
+                )
+            )
+
+        tmp = [
+            (x1, x2, x3)
+            for x1, x2, x3 in zip(
+                [sum(surface_b) for surface_a, surface_b in surface_l],
+                surface_atoms,
+                surface_l,
+            )
+        ]
+        tmp.sort(key=lambda x: x[0])
         surface_l = [x for _, _, x in tmp]
         surface_atoms = [x for _, x, _ in tmp]
 
         return surface_l, surface_atoms, substructure_atoms
 
-    def get_entry(self, m, idx, r = 1):
+    def get_entry(self, m, idx, r=1):
         m = copy.deepcopy(m)
 
         # finding atoms within radius r of atom idx
         substructure_atoms = set({idx})
         for i in range(r):
-            bonds = [b for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds()]
-            substructure_atoms_ = {b.GetBeginAtomIdx() for b in bonds}.union({b.GetEndAtomIdx() for b in bonds})
+            bonds = [
+                b for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds()
+            ]
+            substructure_atoms_ = {b.GetBeginAtomIdx() for b in bonds}.union(
+                {b.GetEndAtomIdx() for b in bonds}
+            )
             substructure_atoms = substructure_atoms.union(substructure_atoms_)
 
         # finding atoms connecting to substructure
-        surface_atoms = {b.GetEndAtomIdx() for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds() if not b.GetEndAtomIdx() in substructure_atoms}
-        surface_atoms.update({b.GetBeginAtomIdx() for a_ in substructure_atoms for b in m.GetAtomWithIdx(a_).GetBonds() if not b.GetBeginAtomIdx() in substructure_atoms})
+        surface_atoms = {
+            b.GetEndAtomIdx()
+            for a_ in substructure_atoms
+            for b in m.GetAtomWithIdx(a_).GetBonds()
+            if not b.GetEndAtomIdx() in substructure_atoms
+        }
+        surface_atoms.update(
+            {
+                b.GetBeginAtomIdx()
+                for a_ in substructure_atoms
+                for b in m.GetAtomWithIdx(a_).GetBonds()
+                if not b.GetBeginAtomIdx() in substructure_atoms
+            }
+        )
         surface_atoms = list(surface_atoms)
         surface_l = list()
         for surface_atom in surface_atoms:
-            bonds = [b.GetIdx() for b in m.GetAtomWithIdx(surface_atom).GetBonds() if b.GetEndAtomIdx() in substructure_atoms or b.GetBeginAtomIdx() in substructure_atoms]
-            surface_l.append((m.GetAtomWithIdx(surface_atom).GetAtomicNum(),
-                              [m.GetBondWithIdx(b).GetBondTypeAsDouble() for b in bonds]))
-
-        tmp = [(x1, x2, x3) for x1, x2, x3 in zip([sum(surface_b) for surface_a, surface_b in surface_l], surface_atoms, surface_l)]
-        tmp.sort(key = lambda x: x[0])
+            bonds = [
+                b.GetIdx()
+                for b in m.GetAtomWithIdx(surface_atom).GetBonds()
+                if b.GetEndAtomIdx() in substructure_atoms
+                or b.GetBeginAtomIdx() in substructure_atoms
+            ]
+            surface_l.append(
+                (
+                    m.GetAtomWithIdx(surface_atom).GetAtomicNum(),
+                    [m.GetBondWithIdx(b).GetBondTypeAsDouble() for b in bonds],
+                )
+            )
+
+        tmp = [
+            (x1, x2, x3)
+            for x1, x2, x3 in zip(
+                [sum(surface_b) for surface_a, surface_b in surface_l],
+                surface_atoms,
+                surface_l,
+            )
+        ]
+        tmp.sort(key=lambda x: x[0])
         surface_l = [x for _, _, x in tmp]
         surface_atoms = [x for _, x, _ in tmp]
 
         for i, j in enumerate(surface_atoms):
-            m.GetAtomWithIdx(j).SetAtomMapNum(i+1)
+            m.GetAtomWithIdx(j).SetAtomMapNum(i + 1)
 
         substruct = Chem.RWMol(m)
 
@@ -164,7 +224,6 @@ def stitch(self, m):
 
                 removal.append(v_)
 
-
         for remove in removal:
             m.RemoveAtom(remove.GetIdx())
         return m
@@ -172,7 +231,7 @@ def stitch(self, m):
     def get_compound_at_idx(self, mol, idx, **kwargs):
         ref_m = mol
 
-        l, a, sub_a = self.get_substitution(ref_m, idx, r = self.radius)
+        l, a, sub_a = self.get_substitution(ref_m, idx, r=self.radius)
         # if the current substructure to replace has no replacements in the database
         if str(l) not in self.trans:
             return None
@@ -181,7 +240,7 @@ def get_compound_at_idx(self, mol, idx, **kwargs):
             ref_m.GetAtomWithIdx(x).SetAtomMapNum(10000)
 
         for k, j in enumerate(a):
-            ref_m.GetAtomWithIdx(j).SetAtomMapNum(k+1)
+            ref_m.GetAtomWithIdx(j).SetAtomMapNum(k + 1)
 
         sub = np.random.choice(self.trans[str(l)])
         m = Chem.CombineMols(ref_m, Chem.MolFromSmiles(sub, sanitize=False))
@@ -196,7 +255,7 @@ def get_compound_at_idx(self, mol, idx, **kwargs):
         for remove in removal:
             m.RemoveAtom(remove.GetIdx())
         if m is None:
-            errors+=1
+            errors += 1
         try:
             s1 = Chem.MolToSmiles(m)
             m = self.stitch(m)
@@ -211,7 +270,7 @@ def get_compound(self, smiles, **kwargs):
             smiles = Chem.MolToSmiles(mol)
         else:
             mol = Chem.MolFromSmiles(smiles)
-        idx = np.random.choice(mol.GetNumAtoms(), replace = False)
+        idx = np.random.choice(mol.GetNumAtoms(), replace=False)
         return self.get_compound_at_idx(mol, idx, **kwargs)
 
     def get_compound_list(self, smiles, idx: int = None, **kwargs) -> list:
@@ -223,13 +282,13 @@ def get_compound_list(self, smiles, idx: int = None, **kwargs) -> list:
             space = [idx]
         for i in tqdm(space):
             ref_m = Chem.MolFromSmiles(smiles)
-            l, a, sub_a = self.get_substitution(ref_m, i, r= self.radius)
+            l, a, sub_a = self.get_substitution(ref_m, i, r=self.radius)
 
             for x in sub_a:
                 ref_m.GetAtomWithIdx(x).SetAtomMapNum(10000)
 
             for k, j in enumerate(a):
-                ref_m.GetAtomWithIdx(j).SetAtomMapNum(k+1)
+                ref_m.GetAtomWithIdx(j).SetAtomMapNum(k + 1)
 
             surface_atoms = [ref_m.GetAtomWithIdx(j) for j in a]
             if str(l) not in self.trans.keys():
@@ -263,6 +322,7 @@ def _save(self) -> dict:
     def _load(self, d: dict):
         return super()._load(d)
 
+
 class STONEDMutations(PerturbationEngine):
     """Implements STONED-SELFIES algorithm for generating modified compounds.
 
@@ -277,7 +337,7 @@ class STONEDMutations(PerturbationEngine):
     """
 
     @log_arguments
-    def __init__(self, mutations: int = 1, log = True):
+    def __init__(self, mutations: int = 1, log=True):
         self.mutations = mutations
 
     def get_compound_at_idx(self, mol: Chem.Mol, idx: int, **kwargs) -> str:
@@ -285,7 +345,7 @@ def get_compound_at_idx(self, mol: Chem.Mol, idx: int, **kwargs) -> str:
         selfie = sf.encoder(smiles)
         selfie_chars = get_selfie_chars(selfie)
         max_molecules_len = len(selfie_chars) + self.mutations
-        selfie_mutated, _ = mutate_selfie(selfie, max_molecules_len, index = idx)
+        selfie_mutated, _ = mutate_selfie(selfie, max_molecules_len, index=idx)
         smiles_back = sf.decoder(selfie_mutated)
         return smiles_back
 
@@ -293,15 +353,19 @@ def get_compound(self, smiles: str, **kwargs) -> str:
         ref_m = Chem.MolFromSmiles(smiles)
         randomized_smiles_ordering = randomize_smiles(ref_m)
         selfie = sf.encoder(randomized_smiles_ordering)
-        selfie_mutated = get_mutated_SELFIES([selfie], num_mutations = self.mutations)[0]
+        selfie_mutated = get_mutated_SELFIES([selfie], num_mutations=self.mutations)[0]
         smiles_back = sf.decoder(selfie_mutated)
         return smiles_back
 
     def get_compound_list(self, smiles: str, num_samples: int = 1000, **kwargs) -> list:
         ref_m = Chem.MolFromSmiles(smiles)
-        randomized_smile_orderings = [randomize_smiles(ref_m) for _ in range(num_samples)]
+        randomized_smile_orderings = [
+            randomize_smiles(ref_m) for _ in range(num_samples)
+        ]
         selfies_ls = [sf.encoder(x) for x in randomized_smile_orderings]
-        selfies_mut = get_mutated_SELFIES(selfies_ls.copy(), num_mutations = self.mutations)
+        selfies_mut = get_mutated_SELFIES(
+            selfies_ls.copy(), num_mutations=self.mutations
+        )
         smiles_back = [sf.decoder(x) for x in selfies_mut]
         return list(set(smiles_back))
 
@@ -311,6 +375,7 @@ def _save(self) -> dict:
     def _load(self, d: dict):
         return super()._load(d)
 
+
 class CounterfactualEngine(BaseClass):
     """Generates counterfactual compounds based on:
 
@@ -319,7 +384,9 @@ class CounterfactualEngine(BaseClass):
     """
 
     @log_arguments
-    def __init__(self, model: BaseModel, perturbation_engine: PerturbationEngine = "default"):
+    def __init__(
+        self, model: BaseModel, perturbation_engine: PerturbationEngine = "default"
+    ):
         self.model = model
         if perturbation_engine == "default":
             self.perturbation_engine = SwapMutations(radius=1)
@@ -356,18 +423,19 @@ def generate_samples(self, smiles: str) -> None:
 
         from sklearn.cluster import DBSCAN
         from sklearn.decomposition import PCA
-        coords = PCA(n_components = 2).fit_transform(dmat)
+
+        coords = PCA(n_components=2).fit_transform(dmat)
         clustering = DBSCAN(eps=0.15, min_samples=5).fit(coords)
 
         samples = {
-            'SMILES': filtered_smiles,
-            'Similarity': scores,
-            'Value': values,
-            'Coordinate': list(coords),
-            'Cluster': clustering.labels_,
-            'Original': [True] + [False for _ in filtered_smiles[1:]]
+            "SMILES": filtered_smiles,
+            "Similarity": scores,
+            "Value": values,
+            "Coordinate": list(coords),
+            "Cluster": clustering.labels_,
+            "Original": [True] + [False for _ in filtered_smiles[1:]],
         }
-        self.samples = pd.DataFrame(samples).to_dict('records')
+        self.samples = pd.DataFrame(samples).to_dict("records")
 
     def get_samples(self) -> pd.DataFrame:
         """Returns candidate counterfactuals as a pandas dataframe.
@@ -377,23 +445,25 @@ def get_samples(self) -> pd.DataFrame:
         """
         return pd.DataFrame(self.samples)
 
-    def generate_cfs(self, delta: Union[int, float, Tuple] = (-1, 1), n: int = 4) -> None:
+    def generate_cfs(
+        self, delta: Union[int, float, Tuple] = (-1, 1), n: int = 4
+    ) -> None:
         """Generates counterfactuals and stores them in `self.cfs` as a list of dictionaries.
 
         Args:
             delta: margin defining counterfactuals for regression models
             n: number of counterfactuals
         """
-        base = self.samples[0]['Value']
+        base = self.samples[0]["Value"]
 
         if self.model.setting == "classification":
-            cfs = self._select_cfs(lambda s: s['Value'] != base, n)
+            cfs = self._select_cfs(lambda s: s["Value"] != base, n)
             self.cfs = self.samples[:1] + cfs
         else:
             if type(delta) in (float, int):
                 delta = (-delta, delta)
-            lcfs = self._select_cfs(lambda s: s['Value'] - delta[0] < base, n // 2)
-            hcfs = self._select_cfs(lambda s: s['Value'] - delta[1] > base, n // 2)
+            lcfs = self._select_cfs(lambda s: s["Value"] - delta[0] < base, n // 2)
+            hcfs = self._select_cfs(lambda s: s["Value"] - delta[1] > base, n // 2)
             self.cfs = self.samples[:1] + lcfs + hcfs
 
     def get_cfs(self) -> pd.DataFrame:
@@ -425,15 +495,15 @@ def cluster_score(s, i):
             Returns:
                 score: score for clustering
             """
-            return (s['Cluster'] == i) * fn(s) * s['Similarity']
+            return (s["Cluster"] == i) * fn(s) * s["Similarity"]
 
-        clusters = {s['Cluster'] for s in self.samples[1:]}
+        clusters = {s["Cluster"] for s in self.samples[1:]}
         for i in clusters:
             candidate = max(self.samples[1:], key=lambda s: cluster_score(s, i))
             if cluster_score(candidate, i):
                 cfs.append(candidate)
 
-        cfs = sorted(cfs, key=lambda s: s['Similarity'], reverse=True)[:n]
+        cfs = sorted(cfs, key=lambda s: s["Similarity"], reverse=True)[:n]
         return cfs
 
     def _save(self) -> dict:
@@ -442,7 +512,13 @@ def _save(self) -> dict:
     def _load(self, d: dict):
         return super()._load(d)
 
-def model_molecule_sensitivity(model: BaseModel, smiles: str, perturbation_engine: PerturbationEngine = "default", n: int = 30) -> Chem.Mol:
+
+def model_molecule_sensitivity(
+    model: BaseModel,
+    smiles: str,
+    perturbation_engine: PerturbationEngine = "default",
+    n: int = 30,
+) -> Chem.Mol:
 
     """
     Calculates the sensitivity of a model to perturbations in on each of a molecule's atoms,
@@ -484,4 +560,4 @@ def model_molecule_sensitivity(model: BaseModel, smiles: str, perturbation_engin
         v = np.std(predictions)
         a.SetProp("sensitivity", str(v))
         a.SetProp("atomNote", str(np.format_float_scientific(v, precision=2)))
-    return mol
\ No newline at end of file
+    return mol
diff --git a/src/olorenchemengine/manager.py b/src/olorenchemengine/manager.py
index 3b13f11d..c0d8b926 100644
--- a/src/olorenchemengine/manager.py
+++ b/src/olorenchemengine/manager.py
@@ -1,12 +1,13 @@
 from timeit import default_timer as timer
 from typing import List
 
+from tqdm import tqdm
+
 import olorenchemengine as oce
+
 from .base_class import *
 from .dataset import *
 
-from tqdm import tqdm
-
 
 def TOP_MODELS_ADMET() -> List[BaseModel]:
     """
@@ -58,7 +59,9 @@ def __init__(
         self.best_model = None
         self.best_primary_metric = float("inf")
 
-        self.model_database = pd.DataFrame(columns=["Model Name", "Model Parameters", "Fitting Time"] + self.metrics)
+        self.model_database = pd.DataFrame(
+            columns=["Model Name", "Model Parameters", "Fitting Time"] + self.metrics
+        )
 
         self.file_path = file_path
         if not self.file_path is None and os.path.exists(self.file_path):
@@ -74,7 +77,9 @@ def get_dataset(self):
     def primary_metric(self):
         return self.primary_metric
 
-    def run(self, models: Union[BaseModel, List[BaseModel]], return_models: bool = False):
+    def run(
+        self, models: Union[BaseModel, List[BaseModel]], return_models: bool = False
+    ):
         """Runs the model on the dataset and saves the results to the model_database.
 
         Parameters:
@@ -108,7 +113,9 @@ def run(self, models: Union[BaseModel, List[BaseModel]], return_models: bool = F
             y_pred = model.predict(eval_set[0])
             y = eval_set[1]
 
-            stats = {metric: metric_functions[metric](y, y_pred) for metric in self.metrics}
+            stats = {
+                metric: metric_functions[metric](y, y_pred) for metric in self.metrics
+            }
 
             if not self.file_path is None and os.path.exists(self.file_path):
                 with open(self.file_path, "rb") as f:
@@ -147,9 +154,13 @@ def run(self, models: Union[BaseModel, List[BaseModel]], return_models: bool = F
 
             # Using metric direction,
             if metric_direction[self.primary_metric] == "higher":
-                primary_metric_outputs.append(-1 * metric_functions[self.primary_metric](y, y_pred))
+                primary_metric_outputs.append(
+                    -1 * metric_functions[self.primary_metric](y, y_pred)
+                )
             else:
-                primary_metric_outputs.append(metric_functions[self.primary_metric](y, y_pred))
+                primary_metric_outputs.append(
+                    metric_functions[self.primary_metric](y, y_pred)
+                )
             if primary_metric_outputs[-1] < self.best_primary_metric:
                 self.best_model = model
                 self.best_primary_metric = primary_metric_outputs[-1]
@@ -199,7 +210,12 @@ def __init__(
     ):
         self.name = oce.model_name_from_model(self)
         super().__init__(
-            dataset, metrics, file_path=file_path, primary_metric=primary_metric, verbose=verbose, log=False
+            dataset,
+            metrics,
+            file_path=file_path,
+            primary_metric=primary_metric,
+            verbose=verbose,
+            log=False,
         )
 
     def _save(self) -> dict:
@@ -248,9 +264,13 @@ def __init__(
         import gspread
 
         if "GOOGLE_CREDENTIALS_FILENAME" in oce.CONFIG:
-            self.gc = gspread.oauth(credentials_filename=oce.CONFIG["GOOGLE_CREDENTIALS_FILENAME"],)
+            self.gc = gspread.oauth(
+                credentials_filename=oce.CONFIG["GOOGLE_CREDENTIALS_FILENAME"],
+            )
         elif "GOOGLE_SERVICE_ACC_FILENAME" in oce.CONFIG:
-            self.gc = gspread.service_account(filename=oce.CONFIG["GOOGLE_SERVICE_ACC_FILENAME"])
+            self.gc = gspread.service_account(
+                filename=oce.CONFIG["GOOGLE_SERVICE_ACC_FILENAME"]
+            )
         else:
             raise Exception(
                 "GOOGLE_CREDENTIALS_FILENAME and GOOGLE_SERVICE_ACC_FILENAME not found in CONFIG, please set one of them.\
@@ -269,11 +289,20 @@ def __init__(
 
         self.email = email
 
-        super().__init__(dataset, metrics, file_path=file_path, primary_metric=primary_metric, log=False)
+        super().__init__(
+            dataset,
+            metrics,
+            file_path=file_path,
+            primary_metric=primary_metric,
+            log=False,
+        )
 
     def _save(self) -> dict:
         self.wb.share(self.email, perm_type="anyone", role="writer")
-        self.ws.update([self.model_database.columns.values.tolist()] + self.model_database.values.tolist())
+        self.ws.update(
+            [self.model_database.columns.values.tolist()]
+            + self.model_database.values.tolist()
+        )
         self.url = f"https://docs.google.com/spreadsheets/d/{self.wb._properties['id']}/edit#gid=0"
         print(f"Saved to Google Sheets. Please visit {self.url} to view the results.")
 
@@ -321,17 +350,26 @@ def __init__(
         log=True,
     ):
         import firebase_admin
-        from firebase_admin import credentials
-        from firebase_admin import firestore
+        from firebase_admin import credentials, firestore
 
         if "FIREBASE_CREDENTIALS_CERTIFICATE" in oce.CONFIG:
-            cred = credentials.Certificate(oce.CONFIG["FIREBASE_CREDENTIALS_CERTIFICATE"])
+            cred = credentials.Certificate(
+                oce.CONFIG["FIREBASE_CREDENTIALS_CERTIFICATE"]
+            )
             firebase_admin.initialize_app(cred)
             self.db = firestore.client()
         else:
-            raise Exception("Firebase credentials in CONFIG were not found or not properly set. ")
+            raise Exception(
+                "Firebase credentials in CONFIG were not found or not properly set. "
+            )
 
-        super().__init__(dataset, metrics, file_path=file_path, primary_metric=primary_metric, log=False)
+        super().__init__(
+            dataset,
+            metrics,
+            file_path=file_path,
+            primary_metric=primary_metric,
+            log=False,
+        )
         self.uid = uid
 
     def run(self, models: Union[BaseModel, List[BaseModel]]):
@@ -361,14 +399,19 @@ def run(self, models: Union[BaseModel, List[BaseModel]]):
         # create a new document in datasets
         dataset_ref = self.db.collection("datasets").document()
         dataset_ref.set(
-            {"uid": self.uid, "data": data_map,}
+            {
+                "uid": self.uid,
+                "data": data_map,
+            }
         )
 
         # check if userID is already in users, if not, create a new document in users
         user_ref = self.db.collection("users").document(self.uid)
         if not user_ref.get().exists:
             user_ref.set(
-                {"uid": self.uid,}
+                {
+                    "uid": self.uid,
+                }
             )
 
         for model in model_database.to_dict("index").values():
@@ -392,7 +435,11 @@ def run(self, models: Union[BaseModel, List[BaseModel]]):
             # create a new document in models
             model_ref = self.db.collection("models").document()
             model_ref.set(
-                {"uid": self.uid, "did": dataset_ref.id, "data": model_map,}
+                {
+                    "uid": self.uid,
+                    "did": dataset_ref.id,
+                    "data": model_map,
+                }
             )
 
     def _save(self) -> dict:
diff --git a/src/olorenchemengine/online.py b/src/olorenchemengine/online.py
new file mode 100644
index 00000000..a1bdbd55
--- /dev/null
+++ b/src/olorenchemengine/online.py
@@ -0,0 +1,164 @@
+def _is_object(key):
+    return key.split(".")[-1][0].isupper()
+
+
+class RemoteSession:
+    def __init__(self):
+        pass
+
+    def construct(self, object):
+        if issubclass(type(object), RemoteObject):
+            if hasattr(object, "REMOTE_ID"):  # object already constructed
+                return object
+            if _is_object(object.REMOTE_BC_KEY):
+                object.REMOTE_ID = generate_random_uuid()
+                print(f"CONSTRUCT {object.REMOTE_ID} => {parameterize_remote(object)}")
+                return object
+            else:
+                return self.get_attribute(
+                    object.REMOTE_PARENT, object.REMOTE_BC_KEY.split(".")[-1]
+                )
+        elif (
+            object is None
+            or issubclass(type(object), int)
+            or issubclass(type(object), float)
+            or issubclass(type(object), str)
+        ):
+            return object
+        elif issubclass(type(object), list):
+            return [self.construct(x) for x in object]
+
+    def run_method(self, object, method_name, args, kwargs):
+        object = self.construct(object)
+        args = [self.construct(arg) for arg in args]
+        kwargs = {k: self.construct(v) for k, v in kwargs.items()}
+        print(
+            f"RUN {object.REMOTE_ID}.{method_name} with args = {args}, and kwargs = {kwargs}"
+        )
+
+    def get_attribute(self, object, attribute_name):
+        object = self.construct(object)
+        object.attribute_name.REMOTE_ID = generate_random_uuid()
+        print(
+            f"GET_ATTR {object.REMOTE_ID}.{attribute_name} => {object.attribute_name.REMOTE_ID}"
+        )
+        return object.attribute_name.REMOTE_ID
+
+    def __getattribute__(self, key):
+        EXCLUDED_WORDS = ["run_method", "construct", "get_attribute"]
+        if key in EXCLUDED_WORDS:
+            return object.__getattribute__(self, key)
+        return RemoteObject(key, None, self)
+
+
+def generate_random_uuid():
+    import uuid
+
+    return str(uuid.uuid4())
+
+
+def parameterize_remote(object):
+    if issubclass(type(object), RemoteObject):
+        if not hasattr(object, "args") or not hasattr(object, "kwargs"):
+            print("Warning: RemoteObject has no args or kwargs")
+            print(f"Object is {object.REMOTE_BC_KEY}")
+        return {
+            **(
+                {"BC_remote_id": object.REMOTE_ID}
+                if hasattr(object, "REMOTE_ID")
+                else {}
+            ),
+            **{"BC_class_name": object.REMOTE_BC_KEY.split(".")[-1]},
+            **{"args": [parameterize_remote(arg) for arg in object.args]},
+            **{"kwargs": {k: parameterize_remote(v) for k, v in object.kwargs.items()}},
+        }
+    elif (
+        object is None
+        or issubclass(type(object), int)
+        or issubclass(type(object), float)
+        or issubclass(type(object), str)
+    ):
+        return object
+    elif issubclass(type(object), list):
+        return [parameterize_remote(x) for x in object]
+    else:
+        raise ValueError(f"Cannot parameterize {object} of type {type(object)}")
+
+
+class RemoteObject:
+    def __init__(self, key, parent) -> None:
+        self.REMOTE_BC_KEY = key
+        self.REMOTE_PARENT = parent
+        self.REMOTE_CHILDREN = {}
+
+    def __getattribute__(self, key):
+        EXCLUDED_WORDS = [
+            "REMOTE_BC_KEY",
+            "REMOTE_SESSION",
+            "REMOTE_PARENT",
+            "REMOTE_ID",
+            "REMOTE_CHILDREN",
+            "args",
+            "kwargs",
+            "shabalaba",
+        ]
+
+        try:
+            return object.__getattribute__(self, key)
+        except AttributeError:
+            if key in self.REMOTE_CHILDREN:
+                return self.REMOTE_CHILDREN[key]
+            else:
+                self.REMOTE_CHILDREN[key] = RemoteObject(
+                    self.REMOTE_BC_KEY + "." + key, self, self.REMOTE_SESSION
+                )
+                return self.REMOTE_CHILDREN[key]
+
+    def __repr__(self) -> str:
+        if hasattr(self, "args") or hasattr(self, "kwargs"):
+            return parameterize_remote(self)
+        return "RemoteObject(" + self.REMOTE_BC_KEY + ")"
+
+    def __call__(self, *args, **kwargs):
+        if self.REMOTE_BC_KEY.split(".")[-1][0].isupper():  # instantiating an object
+            self.args = args
+            self.kwargs = kwargs
+            return self
+        else:
+            return self.REMOTE_SESSION.run_method(
+                self.REMOTE_PARENT, self.REMOTE_BC_KEY.split(".")[-1], args, kwargs
+            )
+            # parameterized_kwargs = {parameterize_remote(k): parameterize_remote(v) for k, v in kwargs.items()}
+            # print(f"Calling function {self.REMOTE_BC_KEY.split('.')[-1]} with: args: {[parameterize_remote(a) for a in args]}, kwargs: {parameterized_kwargs} on {parameterize_remote(self.REMOTE_PARENT)}" )
+            # return self # TODO: handle calling a function (like train)
+
+    # def __iter__(self):
+    #     return iter([0, 1])
+
+
+# TODO: track operators on RemoteObjects
+
+# We are basically making a Python "Compiler" that compiles down to a JSON dictionary - but we can use python
+# in the compiler lol
+
+
+if __name__ == "__main__":
+    # oce = RemoteSession()
+    # dataset = oce.ExampleDataset()
+    # dataset.shabalaba
+
+    import olorenchemengine as oce
+
+    print(dir(oce))
+
+    # sub_model1 = oce.RandomForestModel(oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=1000)
+    # model = oce.ensemble.BaseBoosting([
+    #         sub_model1,
+    #         oce.RandomForestModel(oce.OlorenCheckpoint("default"), n_estimators=1000)])
+    # model.fit(dataset.train_dataset)
+    # oce.save(model, "model.oce")
+
+    # CONSTRUCT model FROM JSON_PARAMS
+    # CONSTRUCT dataset FROM JSON_PARAMS
+    # CONSTRUCT train_dataset FROM dataset.train_dataset
+    # RUN fit on model with train_dataset
diff --git a/src/olorenchemengine/pyg/gcn.py b/src/olorenchemengine/pyg/gcn.py
index fbd546f5..1d5cae74 100644
--- a/src/olorenchemengine/pyg/gcn.py
+++ b/src/olorenchemengine/pyg/gcn.py
@@ -34,7 +34,10 @@ def __init__(self, emb_dim):
 
     def forward(self, x, edge_index, edge_attr):
         edge_embedding = self.bond_encoder(edge_attr)
-        out = self.mlp((1 + self.eps) * x + self.propagate(edge_index, x=x, edge_attr=edge_embedding))
+        out = self.mlp(
+            (1 + self.eps) * x
+            + self.propagate(edge_index, x=x, edge_attr=edge_embedding)
+        )
 
         return out
 
@@ -67,9 +70,9 @@ def forward(self, x, edge_index, edge_attr):
 
         norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
 
-        return self.propagate(edge_index, x=x, edge_attr=edge_embedding, norm=norm) + F.relu(
-            x + self.root_emb.weight
-        ) * 1.0 / deg.view(-1, 1)
+        return self.propagate(
+            edge_index, x=x, edge_attr=edge_embedding, norm=norm
+        ) + F.relu(x + self.root_emb.weight) * 1.0 / deg.view(-1, 1)
 
     def message(self, x_j, edge_attr, norm):
         return norm.view(-1, 1) * F.relu(x_j + edge_attr)
@@ -86,7 +89,13 @@ class GNN_node(torch.nn.Module):
     """
 
     def __init__(
-        self, num_layer, emb_dim, drop_ratio=0.5, JK="last", residual=False, gnn_type="gin",
+        self,
+        num_layer,
+        emb_dim,
+        drop_ratio=0.5,
+        JK="last",
+        residual=False,
+        gnn_type="gin",
     ):
         """
         emb_dim (int): node embedding dimensionality
@@ -165,7 +174,13 @@ class GNN_node_Virtualnode(torch.nn.Module):
     """
 
     def __init__(
-        self, num_layer, emb_dim, drop_ratio=0.5, JK="last", residual=False, gnn_type="gin",
+        self,
+        num_layer,
+        emb_dim,
+        drop_ratio=0.5,
+        JK="last",
+        residual=False,
+        gnn_type="gin",
     ):
         """
         emb_dim (int): node embedding dimensionality
@@ -254,7 +269,9 @@ def forward(self, batched_data):
             ### update the virtual nodes
             if layer < self.num_layer - 1:
                 ### add message from graph nodes to virtual nodes
-                virtualnode_embedding_temp = global_add_pool(h_list[layer], batch) + virtualnode_embedding
+                virtualnode_embedding_temp = (
+                    global_add_pool(h_list[layer], batch) + virtualnode_embedding
+                )
                 ### transform virtual nodes using MLP
 
                 if self.residual:
@@ -314,11 +331,21 @@ def __init__(
         ### GNN to generate node embeddings
         if virtual_node:
             self.gnn_node = GNN_node_Virtualnode(
-                num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type,
+                num_layer,
+                emb_dim,
+                JK=JK,
+                drop_ratio=drop_ratio,
+                residual=residual,
+                gnn_type=gnn_type,
             )
         else:
             self.gnn_node = GNN_node(
-                num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type,
+                num_layer,
+                emb_dim,
+                JK=JK,
+                drop_ratio=drop_ratio,
+                residual=residual,
+                gnn_type=gnn_type,
             )
 
         ### Pooling function to generate whole-graph embeddings
diff --git a/src/olorenchemengine/pyg/util.py b/src/olorenchemengine/pyg/util.py
index dbc4c69a..93f1c68d 100644
--- a/src/olorenchemengine/pyg/util.py
+++ b/src/olorenchemengine/pyg/util.py
@@ -1,19 +1,11 @@
 import numpy as np
-from olorenchemengine.internal import mock_imports
-
-try:
-    import torch
-    import torch_geometric.data
-except:
-    mock_imports(globals(), "torch", "torch_geometric")
-
-try:
-    from ogb.utils.features import atom_to_feature_vector, bond_to_feature_vector
-except ImportError:
-    mock_imports(globals(), "atom_to_feature_vector", "bond_to_feature_vector")
-
+import torch
+import torch_geometric.data
+from ogb.utils.features import atom_to_feature_vector, bond_to_feature_vector
 from rdkit import Chem
 
+from olorenchemengine.internal import mock_imports
+
 
 def molecule_to_graph(mol, include_mol=False):
     # Convert to RDKit molecule
diff --git a/src/olorenchemengine/remote.py b/src/olorenchemengine/remote.py
new file mode 100644
index 00000000..cbb01758
--- /dev/null
+++ b/src/olorenchemengine/remote.py
@@ -0,0 +1,1 @@
+"""Code to help execute OCE code on a remote machine"""
diff --git a/src/olorenchemengine/representations.py b/src/olorenchemengine/representations.py
index bc020478..aa4810ec 100644
--- a/src/olorenchemengine/representations.py
+++ b/src/olorenchemengine/representations.py
@@ -2,14 +2,12 @@
 """
 
 from abc import abstractmethod, abstractproperty
-from ctypes.wintypes import COLORREF
+from typing import Any, List, Union
 
 import numpy as np
-from tqdm import tqdm
 from pandas.api.types import is_numeric_dtype
 from rdkit import Chem
-
-from typing import List, Union, Any
+from tqdm import tqdm
 
 import olorenchemengine as oce
 
@@ -23,7 +21,7 @@ def get_all_reps():
 
 class BaseRepresentation(BaseClass):
 
-    """ BaseClass for all molecular representations (PyTorch Geometric graphs, descriptors, fingerprints, etc.)
+    """BaseClass for all molecular representations (PyTorch Geometric graphs, descriptors, fingerprints, etc.)
 
     Parameters:
         log (boolean): Whether to log the representation or not
@@ -39,7 +37,7 @@ def __init__(self, log=True):
 
     @abstractmethod
     def _convert(self, smiles: str, y: Union[int, float, np.number] = None) -> Any:
-        """ Converts a single structure (represented by a SMILES string) to a representation
+        """Converts a single structure (represented by a SMILES string) to a representation
 
         Parameters:
             smiles (str): SMILES string of the structure
@@ -50,8 +48,10 @@ def _convert(self, smiles: str, y: Union[int, float, np.number] = None) -> Any:
         """
         pass
 
-    def _convert_list(self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None) -> List[Any]:
-        """ Converts a list of structures (represented by a SMILES string) to a list of representations
+    def _convert_list(
+        self, smiles_list: List[str], ys: List[Union[int, float, np.number]] = None
+    ) -> List[Any]:
+        """Converts a list of structures (represented by a SMILES string) to a list of representations
 
         Parameters:
             smiles_list (List[str]): list of SMILES strings of the structures
@@ -67,8 +67,10 @@ def _convert_list(self, smiles_list: List[str], ys: List[Union[int, float, np.nu
         else:
             return [self._convert(s, y=y) for s, y in tqdm(zip(smiles_list, ys))]
 
-    def _convert_cache(self, smiles: str, y: Union[int, float, np.number] = None) -> Any:
-        """ Converts a single structure (represented by a SMILES string) to a representation
+    def _convert_cache(
+        self, smiles: str, y: Union[int, float, np.number] = None
+    ) -> Any:
+        """Converts a single structure (represented by a SMILES string) to a representation
 
         Parameters:
             smiles (str): SMILES string of the structure
@@ -82,9 +84,12 @@ def _convert_cache(self, smiles: str, y: Union[int, float, np.number] = None) ->
             return self._convert(smiles, y=y)
 
     def convert(
-        self, Xs: Union[list, pd.DataFrame, dict, str], ys: Union[list, pd.Series, np.ndarray] = None, **kwargs
+        self,
+        Xs: Union[list, pd.DataFrame, dict, str],
+        ys: Union[list, pd.Series, np.ndarray] = None,
+        **kwargs,
     ) -> List[Any]:
-        """ Converts input data to a list of representations
+        """Converts input data to a list of representations
 
         Parameters:
             Xs (Union[list, pd.DataFrame, dict, str]): input data
@@ -93,7 +98,9 @@ def convert(
         Returns:
             List[Any]: list of representations of the input data
         """
-        if isinstance(Xs, list) and (isinstance(Xs[0], list) or isinstance(Xs[0], tuple)):
+        if isinstance(Xs, list) and (
+            isinstance(Xs[0], list) or isinstance(Xs[0], tuple)
+        ):
             smiles = [X[0] for X in Xs]
         elif isinstance(Xs, pd.DataFrame) or isinstance(Xs, dict):
             if isinstance(Xs, pd.DataFrame):
@@ -123,7 +130,7 @@ def _load(self, d):
 
 
 class SMILESRepresentation(BaseRepresentation):
-    """ Extracts the SMILES strings from inputted data
+    """Extracts the SMILES strings from inputted data
 
     Methods:
         convert(Xs: Union[list, pd.DataFrame, dict, str], ys: Union[list, pd.Series, np.ndarray]=None) -> List[Any]: converts input data to a list of SMILES strings
@@ -138,7 +145,9 @@ def _convert(self, smiles, y=None):
         pass
 
     def convert(self, Xs, ys=None, **kwargs):
-        if isinstance(Xs, list) and (isinstance(Xs[0], list) or isinstance(Xs[0], tuple)):
+        if isinstance(Xs, list) and (
+            isinstance(Xs[0], list) or isinstance(Xs[0], tuple)
+        ):
             smiles = [X[0] for X in Xs]
         elif isinstance(Xs, pd.DataFrame) or isinstance(Xs, dict):
             if isinstance(Xs, pd.DataFrame):
@@ -168,7 +177,7 @@ def convert(self, Xs, ys=None, **kwargs):
 
 
 class AtomFeaturizer(BaseClass):
-    """ Abstract class for atom featurizers, which create a vector representation for a single atom.
+    """Abstract class for atom featurizers, which create a vector representation for a single atom.
 
     Methods:
         length(self) -> int: returns the length of the atom vector representation, to be implemented by subclasses
@@ -191,7 +200,7 @@ def _load(self, d: dict):
 
 
 class BondFeaturizer(BaseClass):
-    """ Abstract class for bond featurizers, which create a vector representation for a single bond.
+    """Abstract class for bond featurizers, which create a vector representation for a single bond.
 
     Methods:
         length(self) -> int: returns the length of the bond vector representation, to be implemented by subclasses
@@ -214,7 +223,7 @@ def _load(self, d: dict):
 
 
 class ConcatenatedAtomFeaturizers(AtomFeaturizer):
-    """ Concatenates multiple atom featurizers into a single vector representation.
+    """Concatenates multiple atom featurizers into a single vector representation.
 
     Methods:
         length(self) -> int: returns the length of the atom vector representation, to be implemented by subclasses
@@ -240,7 +249,7 @@ def _load(self, d: dict):
 
 
 class ConcatenatedBondFeaturizers(BondFeaturizer):
-    """ Concatenates multiple bond featurizers into a single vector representation.
+    """Concatenates multiple bond featurizers into a single vector representation.
 
     Methods:
         length(self) -> int: returns the length of the bond vector representation, to be implemented by subclasses
@@ -266,7 +275,7 @@ def _load(self, d: dict):
 
 
 class OGBAtomFeaturizer(AtomFeaturizer):
-    """ Creates a vector representation for a single atom using the Open Graph Benchmark's atom_to_feature_vector function."""
+    """Creates a vector representation for a single atom using the Open Graph Benchmark's atom_to_feature_vector function."""
 
     @log_arguments
     def __init__(self):
@@ -283,7 +292,7 @@ def convert(self, atom: Chem.Atom):
 
 
 class OGBBondFeaturizer(BondFeaturizer):
-    """ Creates a vector representation for a single bond using the Open Graph Benchmark's bond_to_feature_vector function."""
+    """Creates a vector representation for a single bond using the Open Graph Benchmark's bond_to_feature_vector function."""
 
     @log_arguments
     def __init__(self):
@@ -300,7 +309,7 @@ def convert(self, bond: Chem.Bond):
 
 
 class TorchGeometricGraph(BaseRepresentation):
-    """ Representation which returns torch_geometric.data.Data objects.
+    """Representation which returns torch_geometric.data.Data objects.
 
     Parameters:
         atom_featurizer (AtomFeaturizer): featurizer for atoms
@@ -329,8 +338,8 @@ def dimensions(self):
         return (self.atom_featurizer.length, self.bond_featurizer.length)
 
     def _convert(self, smiles, y=None, addHs=False, **kwargs):
+        from torch import Tensor, from_numpy
         from torch_geometric.data import Data
-        from torch import from_numpy, Tensor
 
         data = Data()
 
@@ -390,13 +399,19 @@ def _convert(self, smiles, y=None, addHs=False, **kwargs):
         return data
 
     def convert(
-        self, Xs: Union[list, pd.DataFrame, dict, str], ys: Union[list, pd.Series, np.ndarray] = None, **kwargs
+        self,
+        Xs: Union[list, pd.DataFrame, dict, str],
+        ys: Union[list, pd.Series, np.ndarray] = None,
+        **kwargs,
     ) -> List[Any]:
         Xs = SMILESRepresentation().convert(Xs)
-        return super().convert(Xs, ys = ys, **kwargs)
+        return super().convert(Xs, ys=ys, **kwargs)
 
     def _save(self):
-        return {"atom_featurizer": self.atom_featurizer._save(), "bond_featurizer": self.bond_featurizer._save()}
+        return {
+            "atom_featurizer": self.atom_featurizer._save(),
+            "bond_featurizer": self.bond_featurizer._save(),
+        }
 
     def _load(self, d):
         self.atom_featurizer._load(d["atom_featurizer"])
@@ -404,10 +419,18 @@ def _load(self, d):
 
 
 class BaseVecRepresentation(BaseRepresentation):
-    """ Representation where given input data, returns a vector representation for each compound."""
+    """Representation where given input data, returns a vector representation for each compound."""
 
     @log_arguments
-    def __init__(self, *args, collinear_thresh=1.01, scale=StandardScaler(), names = None, log=True, **kwargs):
+    def __init__(
+        self,
+        *args,
+        collinear_thresh=1.01,
+        scale=StandardScaler(),
+        names=None,
+        log=True,
+        **kwargs,
+    ):
         self.collinear_thresh = collinear_thresh
         self.to_drop = None
         if not scale is None:
@@ -417,12 +440,21 @@ def __init__(self, *args, collinear_thresh=1.01, scale=StandardScaler(), names =
         self._names = names
 
         from os import path
+
         if not path.exists(path.join(path.expanduser("~"), f".oce/cache/")):
             os.mkdir(path.join(path.expanduser("~"), f".oce/cache/"))
         if not path.exists(path.join(path.expanduser("~"), f".oce/cache/vecrep/")):
             os.mkdir(path.join(path.expanduser("~"), f".oce/cache/vecrep/"))
-        if not path.exists(path.join(path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}")):
-            os.mkdir(path.join(path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}"))
+        if not path.exists(
+            path.join(
+                path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}"
+            )
+        ):
+            os.mkdir(
+                path.join(
+                    path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}"
+                )
+            )
 
         super().__init__(*args, log=False, **kwargs)
 
@@ -431,7 +463,9 @@ def names(self):
         if not self._names is None:
             return self._names
         else:
-            raise ValueError(f"Names not set for representation {self.__class__.__name__}")
+            raise ValueError(
+                f"Names not set for representation {self.__class__.__name__}"
+            )
 
     def convert(
         self,
@@ -440,7 +474,7 @@ def convert(
         fit=False,
         **kwargs,
     ) -> List[np.ndarray]:
-        """ BaseVecRepresentation's convert returns a list of numpy arrays.
+        """BaseVecRepresentation's convert returns a list of numpy arrays.
 
         Args:
             Xs (Union[list, pd.DataFrame, dict, str]): input data
@@ -451,24 +485,46 @@ def convert(
         """
         import joblib
 
-        input_hash = (joblib.hash(Xs) +
-            joblib.hash(ys) +
-            joblib.hash(self._save()) +
-            joblib.hash(oce.parameterize(self)))
+        input_hash = (
+            joblib.hash(Xs)
+            + joblib.hash(ys)
+            + joblib.hash(self._save())
+            + joblib.hash(oce.parameterize(self))
+        )
 
         from os import path
-        if path.exists(path.join(path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy")):
-            return np.load(path.join(path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy"), allow_pickle = True)
+
+        if path.exists(
+            path.join(
+                path.expanduser("~"),
+                f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy",
+            )
+        ):
+            return np.load(
+                path.join(
+                    path.expanduser("~"),
+                    f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy",
+                ),
+                allow_pickle=True,
+            )
 
         feats = super().convert(Xs, ys)
         import pandas as pd
 
-        feats = pd.DataFrame.from_records(feats, columns=[f"col{i}" for i in range(len(feats[0]))])
+        feats = pd.DataFrame.from_records(
+            feats, columns=[f"col{i}" for i in range(len(feats[0]))]
+        )
         if fit and len(Xs) > 2:
             # collinear
             corr_matrix = feats.corr().abs()
-            upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(np.bool))
-            self.to_drop = [column for column in upper.columns if any(upper[column] > self.collinear_thresh)]
+            upper = corr_matrix.where(
+                np.triu(np.ones(corr_matrix.shape), k=1).astype(np.bool)
+            )
+            self.to_drop = [
+                column
+                for column in upper.columns
+                if any(upper[column] > self.collinear_thresh)
+            ]
             feats = feats.drop(columns=self.to_drop)
             # scale
             if not self.scale is None:
@@ -484,34 +540,46 @@ def convert(
                 feats.values[:] = np.nan_to_num(x.reshape(feats.values.shape))
 
         output = np.nan_to_num(np.array(feats.to_records(index=False).tolist()))
-        np.save(path.join(path.expanduser("~"), f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy"), output, allow_pickle=True)
+        np.save(
+            path.join(
+                path.expanduser("~"),
+                f".oce/cache/vecrep/{self.__class__.__name__}/{input_hash}.npy",
+            ),
+            output,
+            allow_pickle=True,
+        )
         return output
 
-    def calculate_distance(self, x1: Union[str, List[str]], x2: Union[str, List[str]],
-        metric: str = "cosine", **kwargs) -> np.ndarray:
-        """ Calculates the distance between two molecules or list of molecules.
-        
-        Returns a 2D array of distances between each pair of molecules of shape 
+    def calculate_distance(
+        self,
+        x1: Union[str, List[str]],
+        x2: Union[str, List[str]],
+        metric: str = "cosine",
+        **kwargs,
+    ) -> np.ndarray:
+        """Calculates the distance between two molecules or list of molecules.
+
+        Returns a 2D array of distances between each pair of molecules of shape
         len(x1) by len(x2).
-        
-        This uses pairwise_distances from sklearn.metrics to calculate distances 
+
+        This uses pairwise_distances from sklearn.metrics to calculate distances
         between the vector representations of the molecules. Options for distances
         are Valid values for metric are:
 
-            From scikit-learn: [‘cityblock’, ‘cosine’, ‘euclidean’, ‘l1’, ‘l2’, 
-                ‘manhattan’]. These metrics support sparse matrix inputs. 
+            From scikit-learn: [‘cityblock’, ‘cosine’, ‘euclidean’, ‘l1’, ‘l2’,
+                ‘manhattan’]. These metrics support sparse matrix inputs.
                 [‘nan_euclidean’] but it does not yet support sparse matrices.
             From scipy.spatial.distance: [‘braycurtis’, ‘canberra’, ‘chebyshev’,
-                ‘correlation’, ‘dice’, ‘hamming’, ‘jaccard’, ‘kulsinski’, 
-                ‘mahalanobis’, ‘minkowski’, ‘rogerstanimoto’, ‘russellrao’, 
-                ‘seuclidean’, ‘sokalmichener’, ‘sokalsneath’, ‘sqeuclidean’, 
+                ‘correlation’, ‘dice’, ‘hamming’, ‘jaccard’, ‘kulsinski’,
+                ‘mahalanobis’, ‘minkowski’, ‘rogerstanimoto’, ‘russellrao’,
+                ‘seuclidean’, ‘sokalmichener’, ‘sokalsneath’, ‘sqeuclidean’,
                 ‘yule’].
-            
+
             See the documentation for scipy.spatial.distance for details on these metrics.
         """
-        
+
         from sklearn.metrics import pairwise_distances
-        
+
         if isinstance(x1, str):
             x1 = [x1]
         if isinstance(x2, str):
@@ -521,7 +589,7 @@ def calculate_distance(self, x1: Union[str, List[str]], x2: Union[str, List[str]
         return pairwise_distances(x1, x2, metric=metric, **kwargs)
 
     def __add__(self, other):
-        """ Adds two representations together
+        """Adds two representations together
 
         Parameters:
             other (BaseVecRepresentation): representation to add to the current representation
@@ -548,8 +616,9 @@ def _load(self, d):
             self.scale._load(d["scale"])
             self.scale_fitted = d["scale_fitted"]
 
+
 class ConcatenatedVecRepresentation(BaseVecRepresentation):
-    """ Creates a structure vector representation by concatenating multiple representations.
+    """Creates a structure vector representation by concatenating multiple representations.
 
     Parameters:
         rep1 (BaseVecRepresentation): first representation to concatenate
@@ -570,7 +639,13 @@ class ConcatenatedVecRepresentation(BaseVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, rep1: BaseVecRepresentation, rep2: BaseVecRepresentation, log=True, **kwargs):
+    def __init__(
+        self,
+        rep1: BaseVecRepresentation,
+        rep2: BaseVecRepresentation,
+        log=True,
+        **kwargs,
+    ):
         self.rep1 = rep1
         self.rep2 = rep2
         try:
@@ -579,23 +654,24 @@ def __init__(self, rep1: BaseVecRepresentation, rep2: BaseVecRepresentation, log
             self._names = None
         super().__init__(names=self._names, log=False, **kwargs)
 
-    def _convert(self, smiles, y=None, fit = False):
-        converted_1 = self.rep1._convert(smiles, y=y, fit = fit)
-        converted_2 = self.rep2._convert(smiles, y=y, fit = fit)
+    def _convert(self, smiles, y=None, fit=False):
+        converted_1 = self.rep1._convert(smiles, y=y, fit=fit)
+        converted_2 = self.rep2._convert(smiles, y=y, fit=fit)
         return np.concatenate((converted_1, converted_2))
 
-    def _convert_list(self, smiles_list, ys=None, fit = False):
-        converted_1 = self.rep1._convert_list(smiles_list, ys=ys, fit = fit)
-        converted_2 = self.rep2._convert_list(smiles_list, ys=ys, fit = fit)
+    def _convert_list(self, smiles_list, ys=None, fit=False):
+        converted_1 = self.rep1._convert_list(smiles_list, ys=ys, fit=fit)
+        converted_2 = self.rep2._convert_list(smiles_list, ys=ys, fit=fit)
         return np.concatenate((converted_1, converted_2), axis=1)
 
-    def convert(self, smiles_list, ys = None, fit = False, **kwargs):
-        converted_1 = self.rep1.convert(smiles_list, ys=ys, fit = fit)
-        converted_2 = self.rep2.convert(smiles_list, ys=ys, fit = fit)
+    def convert(self, smiles_list, ys=None, fit=False, **kwargs):
+        converted_1 = self.rep1.convert(smiles_list, ys=ys, fit=fit)
+        converted_2 = self.rep2.convert(smiles_list, ys=ys, fit=fit)
         return np.concatenate((converted_1, converted_2), axis=1)
 
+
 class NoisyVec(BaseVecRepresentation):
-    """ Adds noise to a given BaseVecRepresentation
+    """Adds noise to a given BaseVecRepresentation
 
     Parameters:
         rep (BaseVecRepresentation): BaseVecRepresentation to add noise to
@@ -614,14 +690,16 @@ class NoisyVec(BaseVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, rep: BaseVecRepresentation, *args, a_std=0.1, m_std=0.1, **kwargs):
+    def __init__(
+        self, rep: BaseVecRepresentation, *args, a_std=0.1, m_std=0.1, **kwargs
+    ):
         self.a_std = a_std
         self.m_std = m_std
         self.rep = rep
         super().__init__(*args, **kwargs)
 
     def _convert(self, smiles: str, y=None) -> np.ndarray:
-        """ Returns specified vector representation of inputted compound with added noise.
+        """Returns specified vector representation of inputted compound with added noise.
 
         Args:
             smiles (str): SMILES representation of compound
@@ -631,11 +709,14 @@ def _convert(self, smiles: str, y=None) -> np.ndarray:
             np.ndarrary: vector representation of inputted compound with added noise
         """
         x = self.rep._convert(smiles)
-        x = (x + np.random.normal(0, self.a_std, x.shape)) * np.random.normal(1, self.m_std, x.shape)
+        x = (x + np.random.normal(0, self.a_std, x.shape)) * np.random.normal(
+            1, self.m_std, x.shape
+        )
         return x
 
+
 class DatasetFeatures(BaseVecRepresentation):
-    """ Selects features from the input dataset as the vector representation """
+    """Selects features from the input dataset as the vector representation"""
 
     def _convert(self, smiles, y=None):
         pass
@@ -643,13 +724,15 @@ def _convert(self, smiles, y=None):
     def convert(self, X, **kwargs):
         assert isinstance(X, pd.DataFrame)
         self.numeric_cols = [c for c in X.columns.tolist() if is_numeric_dtype(X[c])]
-        assert not len(self.numeric_cols) > 0, "No numeric feature columns found in dataset"
+        assert (
+            not len(self.numeric_cols) > 0
+        ), "No numeric feature columns found in dataset"
         X = X[self.numeric_cols].to_numpy()
         return X
 
 
 class BaseCompoundVecRepresentation(BaseVecRepresentation):
-    """ Computes a vector representation from each structure.
+    """Computes a vector representation from each structure.
 
     Parameters:
         normalize (bool): whether to normalize the vector representation or not
@@ -674,13 +757,13 @@ def convert(
         fit=False,
         **kwargs,
     ) -> np.ndarray:
-        """ Computes a vector representation from each structure in Xs."""
+        """Computes a vector representation from each structure in Xs."""
         feats = np.array(super().convert(Xs, ys, fit=fit))
         out = np.nan_to_num(feats.astype(np.float32))
         return out
 
     def inverse(self, Xs):
-        """ Inverts the vector representation to the original feature values
+        """Inverts the vector representation to the original feature values
 
         Parameters:
             Xs (np.ndarray): vector representation of the structures
@@ -690,8 +773,9 @@ def inverse(self, Xs):
 
         pass
 
+
 class ConcatenatedStructVecRepresentation(BaseCompoundVecRepresentation):
-    """ Creates a structure vector representation by concatenating multiple representations.
+    """Creates a structure vector representation by concatenating multiple representations.
 
     DEPRECEATED, use ConcatenatedVecRepresentation instead.
 
@@ -702,7 +786,13 @@ class ConcatenatedStructVecRepresentation(BaseCompoundVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, rep1: BaseCompoundVecRepresentation, rep2: BaseCompoundVecRepresentation, log=True, **kwargs):
+    def __init__(
+        self,
+        rep1: BaseCompoundVecRepresentation,
+        rep2: BaseCompoundVecRepresentation,
+        log=True,
+        **kwargs,
+    ):
         self.rep1 = rep1
         self.rep2 = rep2
         try:
@@ -723,7 +813,7 @@ def _convert_list(self, smiles_list, ys=None):
 
 
 class DescriptastorusDescriptor(BaseCompoundVecRepresentation):
-    """ Wrapper for DescriptaStorus descriptors (https://github.com/bp-kelley/descriptastorus)
+    """Wrapper for DescriptaStorus descriptors (https://github.com/bp-kelley/descriptastorus)
 
     Parameters:
         name (str): name of the descriptor. Either "atompaircounts", "morgan3counts",
@@ -768,7 +858,7 @@ def AllInstances(cls):
 
 class LipinskiDescriptor(BaseCompoundVecRepresentation):
 
-    """ Wrapper for Lipinski descriptors (https://www.rdkit.org/docs/RDKit_Book.html#Lipinski_Descriptors)
+    """Wrapper for Lipinski descriptors (https://www.rdkit.org/docs/RDKit_Book.html#Lipinski_Descriptors)
 
     Parameters:
         log (bool): whether to log the representations or not"""
@@ -825,7 +915,7 @@ def _convert(self, smiles, y=None):
 
 
 class FragmentIndicator(BaseCompoundVecRepresentation):
-    """ Indicator variables for all fragments in rdkit.Chem.Fragments
+    """Indicator variables for all fragments in rdkit.Chem.Fragments
 
     http://rdkit.org/docs/source/rdkit.Chem.Fragments.html
     """
@@ -1021,6 +1111,7 @@ def _convert(self, smiles, y=None):
 
 
 from rdkit.Chem import Descriptors, rdMolDescriptors
+
 from .external import calc_pI
 
 
@@ -1028,14 +1119,23 @@ class PeptideDescriptors1(BaseCompoundVecRepresentation):
     @log_arguments
     def __init__(self, log=True, **kwargs):
         self.calc_pI = calc_pI(ph=7.4)
-        super().__init__(names=["NumAtoms", "MW", *self.calc_pI.names, "NumAmideBonds"], log=False, **kwargs)
+        super().__init__(
+            names=["NumAtoms", "MW", *self.calc_pI.names, "NumAmideBonds"],
+            log=False,
+            **kwargs,
+        )
 
     def _convert(self, smiles, y=None):
         m = Chem.MolFromSmiles(smiles)
         calc_pI_props = self.calc_pI._convert(smiles).tolist()
         x = np.nan_to_num(
             np.array(
-                [m.GetNumAtoms(), Descriptors.MolWt(m), *calc_pI_props, rdMolDescriptors.CalcNumAmideBonds(m),]
+                [
+                    m.GetNumAtoms(),
+                    Descriptors.MolWt(m),
+                    *calc_pI_props,
+                    rdMolDescriptors.CalcNumAmideBonds(m),
+                ]
             ).astype(float)
         )
         return x
@@ -1061,9 +1161,12 @@ def _convert(self, smiles):
     def info(self, smiles):
         info = {}
         m = Chem.MolFromSmiles(smiles)
-        fp = AllChem.GetMorganFingerprintAsBitVect(m, self.radius, nBits=self.nbits, bitInfo=info)
+        fp = AllChem.GetMorganFingerprintAsBitVect(
+            m, self.radius, nBits=self.nbits, bitInfo=info
+        )
         return info
 
+
 class MACCSKeys(BaseCompoundVecRepresentation):
     """Calculate MACCS (Molecular ACCess System) Keys fingerprint.
 
@@ -1077,6 +1180,7 @@ def __init__(self):
     def _convert(self, s: str) -> np.ndarray:
         return AllChem.GetMACCSKeysFingerprint(Chem.MolFromSmiles(s))
 
+
 class PubChemFingerprint(BaseCompoundVecRepresentation):
     """PubChem Fingerprint.
     The PubChem fingerprint is a 881 bit structural key,
@@ -1093,6 +1197,7 @@ class PubChemFingerprint(BaseCompoundVecRepresentation):
     Certain SMILES strings encountered in datasets are not recognized by PubChem.
     In these cases, the fingerprint is returned with all bits equal to zero.
     """
+
     @log_arguments
     def __init__(self):
         super().__init__(log=False)
@@ -1101,17 +1206,19 @@ def _convert(self, s: str) -> np.ndarray:
         oce.import_or_install("pubchempy")
 
         import pubchempy as pcp
-        #Check if retrieval of compound and subsequent descriptor calculation succeed without error
+
+        # Check if retrieval of compound and subsequent descriptor calculation succeed without error
         try:
-            pubchem_compound = pcp.get_compounds(s, 'smiles')[0]
+            pubchem_compound = pcp.get_compounds(s, "smiles")[0]
             fp = [int(bit) for bit in pubchem_compound.cactvs_fingerprint]
         except:
             fp = [0] * 881
         return np.array(fp)
 
+
 class MordredDescriptor(BaseCompoundVecRepresentation):
 
-    """ Wrapper for Mordred descriptors (https://github.com/mordred-descriptor/mordred)
+    """Wrapper for Mordred descriptors (https://github.com/mordred-descriptor/mordred)
 
     Parameters:
         log (bool): whether to log the representations or not
@@ -1120,7 +1227,13 @@ class MordredDescriptor(BaseCompoundVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, descriptor_set: Union[str, list] = "all", log: bool = True, normalize: bool = False, **kwargs):
+    def __init__(
+        self,
+        descriptor_set: Union[str, list] = "all",
+        log: bool = True,
+        normalize: bool = False,
+        **kwargs,
+    ):
         oce.import_or_install("mordred")
 
         from mordred import Calculator, descriptors
@@ -1136,7 +1249,7 @@ def _convert(self, smiles, y=None):
         pass
 
     def convert(self, Xs, ys=None, **kwargs):
-        """ Convert list of SMILES to descriptors in the form of a numpy array.
+        """Convert list of SMILES to descriptors in the form of a numpy array.
 
         Parameters:
             Xs (list): List of SMILES strings.
@@ -1163,7 +1276,7 @@ def convert(self, Xs, ys=None, **kwargs):
 
 class GobbiPharma2D(BaseCompoundVecRepresentation):
 
-    """ 2D Gobbi pharmacophore descriptor (implemented in RDKit, from https://doi.org/10.1002/(SICI)1097-0290(199824)61:1<47::AID-BIT9>3.0.CO;2-Z) """
+    """2D Gobbi pharmacophore descriptor (implemented in RDKit, from https://doi.org/10.1002/(SICI)1097-0290(199824)61:1<47::AID-BIT9>3.0.CO;2-Z)"""
 
     def _convert(self, smiles, y=None):
         m = Chem.MolFromSmiles(smiles)
@@ -1178,7 +1291,7 @@ def _convert(self, smiles, y=None):
 
 class GobbiPharma3D(BaseCompoundVecRepresentation):
 
-    """ 3D Gobbi pharmacophore descriptor (implemented in RDKit, from https://doi.org/10.1002/(SICI)1097-0290(199824)61:1<47::AID-BIT9>3.0.CO;2-Z) """
+    """3D Gobbi pharmacophore descriptor (implemented in RDKit, from https://doi.org/10.1002/(SICI)1097-0290(199824)61:1<47::AID-BIT9>3.0.CO;2-Z)"""
 
     def _convert(self, smiles, y=None):
         m = Chem.MolFromSmiles(smiles)
@@ -1195,22 +1308,14 @@ def _convert(self, smiles, y=None):
 
 from collections import OrderedDict
 
-try:
-    import torch
-except ImportError:
-    oce.mock_imports(globals(), "torch")
-
-try:
-    import torch_geometric.data
-    from torch_geometric.data import DataLoader as PyGDataLoader
-except:
-    oce.mock_imports(globals(), "torch_geometric", "PyGDataLoader")
-
+import torch
+import torch_geometric.data
 from rdkit import Chem
+from torch_geometric.data import DataLoader as PyGDataLoader
 
 
 class OlorenCheckpoint(BaseCompoundVecRepresentation):
-    """ Use OlorenVec from checkpoint as a molecular representation
+    """Use OlorenVec from checkpoint as a molecular representation
 
     Parameters:
         model_path (str): path to checkpoint file for OlorenVec. Use "default" if unsure
@@ -1219,7 +1324,9 @@ class OlorenCheckpoint(BaseCompoundVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, model_path: str, num_tasks: int = 2048, log: bool = True, **kwargs):
+    def __init__(
+        self, model_path: str, num_tasks: int = 2048, log: bool = True, **kwargs
+    ):
         if model_path == "default":
             import os
             from os.path import expanduser
@@ -1240,14 +1347,21 @@ def __init__(self, model_path: str, num_tasks: int = 2048, log: bool = True, **k
         state_dict = OrderedDict(
             [
                 (k.replace("model.", ""), v)
-                for k, v in torch.load(path, map_location=oce.CONFIG["DEVICE"])["state_dict"].items()
+                for k, v in torch.load(path, map_location=oce.CONFIG["DEVICE"])[
+                    "state_dict"
+                ].items()
             ]
         )
 
         from olorenchemengine.pyg.gcn import GNN
 
         self.model = GNN(
-            gnn_type="gcn", num_tasks=num_tasks, num_layer=5, emb_dim=300, drop_ratio=0.5, virtual_node=False
+            gnn_type="gcn",
+            num_tasks=num_tasks,
+            num_layer=5,
+            emb_dim=300,
+            drop_ratio=0.5,
+            virtual_node=False,
         )
 
         self.model.to(oce.CONFIG["DEVICE"])
@@ -1275,7 +1389,11 @@ def _convert(self, smiles, y=None):
     def _convert_list(self, smiles_list, ys=None):
         xs = [self.smiles2pyg(s, None) for s in smiles_list]
 
-        kwargs = {"num_workers": oce.CONFIG["NUM_WORKERS"], "pin_memory": True} if oce.CONFIG["USE_CUDA"] else {}
+        kwargs = (
+            {"num_workers": oce.CONFIG["NUM_WORKERS"], "pin_memory": True}
+            if oce.CONFIG["USE_CUDA"]
+            else {}
+        )
 
         dataloader = PyGDataLoader(xs, batch_size=64, **kwargs)
 
@@ -1293,7 +1411,7 @@ def AllInstances(cls):
         return [cls("default", num_tasks=2048)]
 
     def molecule2graph(self, mol, include_mol=False):
-        """ Convert a molecule to a PyG graph with features and labels
+        """Convert a molecule to a PyG graph with features and labels
 
         Parameters:
             mol (rdkit.Chem.rdchem.Mol): molecule to convert
@@ -1353,7 +1471,7 @@ def molecule2graph(self, mol, include_mol=False):
         return graph
 
     def smiles2pyg(self, smiles_str, y, morgan_params={"radius": 2, "nBits": 1024}):
-        """ Convert a SMILES string to a PyG graph with features and labels
+        """Convert a SMILES string to a PyG graph with features and labels
 
         Parameters:
             smiles_str (str): SMILES string to convert
@@ -1391,7 +1509,7 @@ def smiles2pyg(self, smiles_str, y, morgan_params={"radius": 2, "nBits": 1024}):
 
 
 class MCSClusterRep(BaseCompoundVecRepresentation):
-    """ Clusters a train set of compounds and then finds the maximum common
+    """Clusters a train set of compounds and then finds the maximum common
     substructure (MCS) within each set. The presence of each cluster's MCS is
     used as a feature
     """
@@ -1409,8 +1527,8 @@ def __init__(
         **kwargs,
     ):
         if not self.kwargs["cached"]:
-            from rdkit.Chem import AllChem
             from rdkit import DataStructs
+            from rdkit.Chem import AllChem
 
             if eval_set == "train":
                 eval_set = dataset.data[dataset.data["split"] == "train"]
@@ -1421,17 +1539,25 @@ def __init__(
             else:
                 eval_set = dataset.data
 
-            fingerprinter = lambda s: AllChem.GetMorganFingerprint(Chem.MolFromSmiles(s), 4)
+            fingerprinter = lambda s: AllChem.GetMorganFingerprint(
+                Chem.MolFromSmiles(s), 4
+            )
             eval_set["fps"] = eval_set[dataset.structure_col].apply(fingerprinter)
             corr = np.array(
-                [DataStructs.BulkTanimotoSimilarity(fp, eval_set["fps"].tolist()) for fp in eval_set["fps"]]
+                [
+                    DataStructs.BulkTanimotoSimilarity(fp, eval_set["fps"].tolist())
+                    for fp in eval_set["fps"]
+                ]
             )
 
-            from sklearn.cluster import AgglomerativeClustering
             from rdkit.Chem.rdFMCS import FindMCS
+            from sklearn.cluster import AgglomerativeClustering
 
             clustering = AgglomerativeClustering(
-                n_clusters=None, distance_threshold=0.4, affinity="precomputed", linkage="single"
+                n_clusters=None,
+                distance_threshold=0.4,
+                affinity="precomputed",
+                linkage="single",
             ).fit(1 - corr)
             eval_set["cluster_id"] = clustering.labels_
 
@@ -1441,14 +1567,21 @@ def __init__(
                 if len(data_) < 2:
                     continue
                 data_["mols"] = data_[dataset.structure_col].apply(Chem.MolFromSmiles)
-                mcs = FindMCS(data_["mols"].tolist(), threshold=threshold, timeout=timeout)
+                mcs = FindMCS(
+                    data_["mols"].tolist(), threshold=threshold, timeout=timeout
+                )
                 self.smarts.append(mcs.smartsString)
             self.args = [""]
             self.kwargs["cached"] = True
         super().__init__(*args, **kwargs, log=False)
 
     def _convert(self, smiles):
-        return np.array([Chem.MolFromSmiles(smiles).HasSubstructMatch(Chem.MolFromSmarts(s)) for s in self.smarts])
+        return np.array(
+            [
+                Chem.MolFromSmiles(smiles).HasSubstructMatch(Chem.MolFromSmarts(s))
+                for s in self.smarts
+            ]
+        )
 
     def _save(self):
         d = super()._save()
@@ -1461,7 +1594,7 @@ def _load(self, d):
 
 
 class ModelAsRep(BaseCompoundVecRepresentation):
-    """ Uses a trained model itself as a representation.
+    """Uses a trained model itself as a representation.
 
     If we are trying to predict property A, and there is a highly related property B
     that has a lot of data we could train a model on property B and use that model
@@ -1478,8 +1611,14 @@ class ModelAsRep(BaseCompoundVecRepresentation):
     """
 
     @log_arguments
-    def __init__(self, model: Union[BaseModel, str], name="ModelAsRep", 
-            download_public_file = False, log=True, **kwargs):
+    def __init__(
+        self,
+        model: Union[BaseModel, str],
+        name="ModelAsRep",
+        download_public_file=False,
+        log=True,
+        **kwargs,
+    ):
         if isinstance(model, str):
             if download_public_file:
                 self.model = oce.load(oce.download_public_file(model))
diff --git a/src/olorenchemengine/splitters.py b/src/olorenchemengine/splitters.py
index 4ddcecde..13bddf67 100644
--- a/src/olorenchemengine/splitters.py
+++ b/src/olorenchemengine/splitters.py
@@ -1,14 +1,14 @@
 """ For creating splits on the data"""
 
+import itertools
+import random
 from abc import abstractmethod
 
 import numpy as np
-from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import QuantileTransformer
 from rdkit import Chem
 from rdkit.Chem.Scaffolds import MurckoScaffold
-import itertools
-import random
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import QuantileTransformer
 
 from .base_class import *
 from .dataset import *
@@ -45,8 +45,13 @@ def split(self, data, *args, **kwargs):
         pass
 
     def transform(self, dataset: BaseDataset, *args, **kwargs) -> BaseDataset:
-        train, val, test = self.split(dataset.data, *args,
-            structure_col=dataset.structure_col, date_col=dataset.date_col, **kwargs)
+        train, val, test = self.split(
+            dataset.data,
+            *args,
+            structure_col=dataset.structure_col,
+            date_col=dataset.date_col,
+            **kwargs
+        )
         train["split"] = "train"
         val["split"] = "valid"
         test["split"] = "test"
@@ -183,14 +188,17 @@ def split(self, data, date_col, *args, **kwargs):
         data[date_col] = pd.to_datetime(data[date_col])
         data = data.sort_values(by=date_col)
 
-        splits = [0] + [int(sum(self.split_proportions[:i+1])*len(data)) for i in range(len(self.split_proportions))]
+        splits = [0] + [
+            int(sum(self.split_proportions[: i + 1]) * len(data))
+            for i in range(len(self.split_proportions))
+        ]
 
         out = []
-        for i in range(len(splits)-1):
-            if splits[i] == splits[i+1]:
+        for i in range(len(splits) - 1):
+            if splits[i] == splits[i + 1]:
                 out.append(pd.DataFrame())
             else:
-                out.append(data.iloc[splits[i]:splits[i+1]])
+                out.append(data.iloc[splits[i] : splits[i + 1]])
         return out
 
 
@@ -223,7 +231,13 @@ class ScaffoldSplit(BaseSplitter):
     """
 
     @log_arguments
-    def __init__(self, scaffold_filter_threshold: int = 0, split_type="murcko", log=True, **kwargs):
+    def __init__(
+        self,
+        scaffold_filter_threshold: int = 0,
+        split_type="murcko",
+        log=True,
+        **kwargs
+    ):
         self.scaffold_filter_threshold = scaffold_filter_threshold
         self.split_type = split_type
         super().__init__(log=False, **kwargs)
@@ -280,6 +294,7 @@ def _kmeans_murcko_split(self, data):
         mtc = DescriptastorusDescriptor("morgan3counts").convert(filtered_scaffolds)
 
         from sklearn.cluster import KMeans
+
         cluster = KMeans(n_clusters=cluster_count).fit_predict(mtc)
 
         """This loops through the array of cluster labels. Checks if the train/val/test split is full, if not then adds the scaffold corresponding to cluster label.
@@ -289,9 +304,17 @@ def _kmeans_murcko_split(self, data):
             cluster = [2 if i == 1 else 0 for i in cluster]
         for i in range(len(cluster)):
             cluster_id = cluster[i]
-            if len(out[cluster_id].index) / data_size <= self.split_proportions[cluster_id]:
+            if (
+                len(out[cluster_id].index) / data_size
+                <= self.split_proportions[cluster_id]
+            ):
                 out[cluster_id] = pd.concat(
-                    [out[cluster_id], data.loc[data["murcko"] == filtered_scaffolds[i]]], ignore_index=True, sort=False
+                    [
+                        out[cluster_id],
+                        data.loc[data["murcko"] == filtered_scaffolds[i]],
+                    ],
+                    ignore_index=True,
+                    sort=False,
                 )
             else:
                 if cluster_count == 2 and cluster_id == 2:
@@ -325,18 +348,25 @@ def _murcko_split(self, data):
         split_reached = [False] * 3
 
         if scaff_size == 0:
-            raise Exception("No scaffolds with enough molecules to meet filter threshold. Please lower filter threshold.")
+            raise Exception(
+                "No scaffolds with enough molecules to meet filter threshold. Please lower filter threshold."
+            )
 
-        '''
+        """
         Cycle through train/val/test sets.
         If split proportion is not met, append data with the next most common scaffold.
-        If all split proportions are met or scaffolds are fully cycled return split.'''
+        If all split proportions are met or scaffolds are fully cycled return split."""
         for i in itertools.cycle([0, 1, 2]):
             if self.split_proportions[i] == 0:
                 split_reached[i] = True
             elif len(out[i].index) / data_size <= self.split_proportions[i]:
                 out[i] = pd.concat(
-                    [out[i], data.loc[data["murcko"] == filtered_scaffolds[scaff_index]]], ignore_index=True, sort=False
+                    [
+                        out[i],
+                        data.loc[data["murcko"] == filtered_scaffolds[scaff_index]],
+                    ],
+                    ignore_index=True,
+                    sort=False,
                 )
                 scaff_index += 1
             else:
@@ -355,6 +385,7 @@ def split(self, data: pd.DataFrame, *args, structure_col: str = "Smiles", **kwar
         else:
             return "Invalid split_type initializer."
 
+
 class dc_ScaffoldSplit(BaseSplitter):
     """
     Split data into train/val/test sets by scaffold using DeepChem implementation. https://deepchem.readthedocs.io/en/latest/api_reference/splitters.html#scaffoldsplitter
@@ -388,12 +419,22 @@ def split(self, data: pd.DataFrame, *args, structure_col: str = "Smiles", **kwar
         try:
             import deepchem as dc
         except:
-            print("Error, deepchem not installed. Please use ScaffoldSplit instead or install deepchem")
-        dc_dataset = dc.data.DiskDataset.from_numpy(X=data.index, ids=data[structure_col])
+            print(
+                "Error, deepchem not installed. Please use ScaffoldSplit instead or install deepchem"
+            )
+        dc_dataset = dc.data.DiskDataset.from_numpy(
+            X=data.index, ids=data[structure_col]
+        )
         scaffoldsplitter = dc.splits.ScaffoldSplitter()
-        train,valid,test = scaffoldsplitter.train_valid_test_split(dc_dataset, frac_train=self.split_proportions[0], frac_valid=self.split_proportions[1], frac_test=self.split_proportions[2])
+        train, valid, test = scaffoldsplitter.train_valid_test_split(
+            dc_dataset,
+            frac_train=self.split_proportions[0],
+            frac_valid=self.split_proportions[1],
+            frac_test=self.split_proportions[2],
+        )
         return data.loc[train.X], data.loc[valid.X], data.loc[test.X]
 
+
 class gg_ScaffoldSplit(BaseSplitter):
     """
     Split data into train/val/test sets by scaffold using implementation from https://www.nature.com/articles/s42256-021-00438-4, https://github.com/PaddlePaddle/PaddleHelix
@@ -424,18 +465,22 @@ def __init__(self, log=True, **kwargs):
         super().__init__(log=False, **kwargs)
 
     def split(self, data: pd.DataFrame, *args, structure_col: str = "smiles", **kwargs):
-        return self.gg_split(data,
+        return self.gg_split(
+            data,
             frac_train=self.split_proportions[0],
             frac_valid=self.split_proportions[1],
             frac_test=self.split_proportions[2],
-            structure_col = structure_col)
-
-    def gg_split(self,
-            dataset,
-            frac_train=None,
-            frac_valid=None,
-            frac_test=None,
-            structure_col = "smiles"):
+            structure_col=structure_col,
+        )
+
+    def gg_split(
+        self,
+        dataset,
+        frac_train=None,
+        frac_valid=None,
+        frac_test=None,
+        structure_col="smiles",
+    ):
         """
         Args:
             dataset(InMemoryDataset): the dataset to split. Make sure each element in
@@ -451,7 +496,9 @@ def gg_split(self,
         # create dict of the form {scaffold_i: [idx1, idx....]}
         all_scaffolds = {}
         for i in range(N):
-            scaffold = self.generate_scaffold(dataset.iloc[i][structure_col], include_chirality=True)
+            scaffold = self.generate_scaffold(
+                dataset.iloc[i][structure_col], include_chirality=True
+            )
             if scaffold not in all_scaffolds:
                 all_scaffolds[scaffold] = [i]
             else:
@@ -460,8 +507,10 @@ def gg_split(self,
         # sort from largest to smallest sets
         all_scaffolds = {key: sorted(value) for key, value in all_scaffolds.items()}
         all_scaffold_sets = [
-            scaffold_set for (scaffold, scaffold_set) in sorted(
-                all_scaffolds.items(), key=lambda x: (len(x[1]), x[1][0]), reverse=True)
+            scaffold_set
+            for (scaffold, scaffold_set) in sorted(
+                all_scaffolds.items(), key=lambda x: (len(x[1]), x[1][0]), reverse=True
+            )
         ]
 
         # get train, valid test indices
@@ -513,9 +562,11 @@ def generate_scaffold(self, smiles, include_chirality=False):
             the scaffold of the given smiles.
         """
         scaffold = MurckoScaffold.MurckoScaffoldSmiles(
-            smiles=smiles, includeChirality=include_chirality)
+            smiles=smiles, includeChirality=include_chirality
+        )
         return scaffold
 
+
 class PropertySplit(BaseSplitter):
     """Split molecules into train/val/test based on user-defined property.
 
@@ -544,7 +595,15 @@ class PropertySplit(BaseSplitter):
     """
 
     @log_arguments
-    def __init__(self, property_col, threshold=None, noise=0.1, categorical=False, log=True, **kwargs):
+    def __init__(
+        self,
+        property_col,
+        threshold=None,
+        noise=0.1,
+        categorical=False,
+        log=True,
+        **kwargs
+    ):
         self.property_col = property_col
         self.threshold = threshold
         self.noise = noise
@@ -562,9 +621,14 @@ def _quantile_uniform(self, data, noise):
         """
         property = data[self.property_col]
 
-        data["prop_norm"] = QuantileTransformer(output_distribution="uniform").fit_transform(property)
+        data["prop_norm"] = QuantileTransformer(
+            output_distribution="uniform"
+        ).fit_transform(property)
         data["prop_norm"] = (
-            (data["prop_norm"] + [random.uniform(-noise, noise) for i in range(len(data))])
+            (
+                data["prop_norm"]
+                + [random.uniform(-noise, noise) for i in range(len(data))]
+            )
             .clip(lower=0, upper=1)
             .sample(frac=1)
         )
@@ -574,20 +638,33 @@ def _quantile_uniform(self, data, noise):
         out.append(
             data.loc[
                 (data["prop_norm"] > self.split_proportions[0])
-                & (data["prop_norm"] <= self.split_proportions[0] + self.split_proportions[1])
+                & (
+                    data["prop_norm"]
+                    <= self.split_proportions[0] + self.split_proportions[1]
+                )
+            ]
+        )
+        out.append(
+            data.loc[
+                data["prop_norm"]
+                > (self.split_proportions[0] + self.split_proportions[1])
             ]
         )
-        out.append(data.loc[data["prop_norm"] > (self.split_proportions[0] + self.split_proportions[1])])
 
         if self.categorical:
-            thresh1 = np.percentile(data[self.property_col], self.split_proportions[0] * 10)
+            thresh1 = np.percentile(
+                data[self.property_col], self.split_proportions[0] * 10
+            )
             thresh2 = np.percentile(
-                data[self.property_col], (self.split_proportions[0] + self.split_proportions[1]) * 10
+                data[self.property_col],
+                (self.split_proportions[0] + self.split_proportions[1]) * 10,
             )
             for df in out:
                 df.loc[df[self.property_col] <= thresh1, self.property_col] = int(0)
                 df.loc[
-                    (df[self.property_col] > thresh1) & (data[self.property_col] <= thresh2), self.property_col
+                    (df[self.property_col] > thresh1)
+                    & (data[self.property_col] <= thresh2),
+                    self.property_col,
                 ] = int(1)
                 df.loc[df[self.property_col] > thresh2, self.property_col] = int(2)
 
@@ -601,12 +678,18 @@ def _quantile_uniform_thresh(self, data, threshold, noise):
         data["prop_norm"] = transformer.transform(property.to_numpy().reshape(-1, 1))
 
         threshold_scaled = (
-            (transformer.transform(np.array([threshold]).reshape(1, -1)) - data["prop_norm"].min())
+            (
+                transformer.transform(np.array([threshold]).reshape(1, -1))
+                - data["prop_norm"].min()
+            )
             / (data["prop_norm"].max() - data["prop_norm"].min())
         )[0][0]
 
         data["prop_norm"] = (
-            (data["prop_norm"] + [random.uniform(-noise, noise) for i in range(len(data))])
+            (
+                data["prop_norm"]
+                + [random.uniform(-noise, noise) for i in range(len(data))]
+            )
             .clip(lower=0, upper=1)
             .sample(frac=1)
         )
@@ -617,7 +700,9 @@ def _quantile_uniform_thresh(self, data, threshold, noise):
 
         if self.categorical:
             for df in out:
-                df[self.property_col] = np.where(df[self.property_col] <= self.threshold, int(0), int(1))
+                df[self.property_col] = np.where(
+                    df[self.property_col] <= self.threshold, int(0), int(1)
+                )
 
         data = data.drop("prop_norm", 1)
         return out
diff --git a/src/olorenchemengine/toolkit.py b/src/olorenchemengine/toolkit.py
index 0cb85713..6440c134 100644
--- a/src/olorenchemengine/toolkit.py
+++ b/src/olorenchemengine/toolkit.py
@@ -2,4 +2,4 @@
 
 Although most manipulations are done with RDKit, the aim of this file is for
 custom implementations.
-"""
\ No newline at end of file
+"""
diff --git a/src/olorenchemengine/uncertainty.py b/src/olorenchemengine/uncertainty.py
index 61b36309..dab2d21f 100644
--- a/src/olorenchemengine/uncertainty.py
+++ b/src/olorenchemengine/uncertainty.py
@@ -1,27 +1,28 @@
 from cmath import exp
 from tkinter import N
-import olorenchemengine as oce
-
-from .base_class import *
-from .representations import *
-from .dataset import *
-from .basics import *
 
-import pandas as pd
 import numpy as np
-
+import pandas as pd
+from rdkit.Chem import AllChem
+from rdkit.DataStructs.cDataStructs import BulkTanimotoSimilarity
 from sklearn.cross_decomposition import PLSRegression
 from sklearn.decomposition import PCA
 from sklearn.neighbors import NearestNeighbors
 from sklearn.preprocessing import OneHotEncoder
-from rdkit.DataStructs.cDataStructs import BulkTanimotoSimilarity
-from rdkit.Chem import AllChem
+
+import olorenchemengine as oce
+
+from .base_class import *
+from .basics import *
+from .dataset import *
+from .representations import *
 
 
 class BaseFingerprintModel(BaseErrorModel):
-    """ BaseFingerprintModel is the base class for error models that require the
-        computation of Morgan Fingerprints.
+    """BaseFingerprintModel is the base class for error models that require the
+    computation of Morgan Fingerprints.
     """
+
     def build(
         self,
         model: BaseModel,
@@ -46,8 +47,9 @@ def _load(self, d) -> None:
         if "train_fps" in d.keys():
             self.train_fps = d["train_fps"]
 
+
 class SDC(BaseFingerprintModel):
-    """ SDC is an error model that predicts error bars based on the Sum of
+    """SDC is an error model that predicts error bars based on the Sum of
         Distance-weighted Contributions: `Molecular Similarity-Based Domain
         Applicability Metric Efficiently Identifies Out-of-Domain Compounds
         <http://dx.doi.org/10.1021/acs.jcim.8b00597>`_
@@ -80,7 +82,7 @@ def sdc(smi):
 
 
 class TargetDistDC(SDC):
-    """ TargetDistDC is an error model that calculates the root-mean-square
+    """TargetDistDC is an error model that calculates the root-mean-square
         difference between the predicted activity of the target molecule and
         the observed activities of all training molecules, weighted by the DC.
 
@@ -96,6 +98,7 @@ class TargetDistDC(SDC):
     model.predict(test["Drug"], return_ci = True)
     ------------------------------
     """
+
     def calculate(self, X, y_pred):
         def dist(smi, pred):
             mol = Chem.MolFromSmiles(smi)
@@ -107,8 +110,9 @@ def dist(smi, pred):
         X = np.array(X).flatten()
         return np.array([dist(smi, y_pred[i]) for i, smi in tqdm(np.ndenumerate(X))])
 
+
 class TrainDistDC(SDC):
-    """ TrainDistDC is an error model that calculates the root-mean-square
+    """TrainDistDC is an error model that calculates the root-mean-square
         difference between the predicted and observed activities of all
         training molecules, weighted by the DC.
 
@@ -124,19 +128,22 @@ class TrainDistDC(SDC):
     model.predict(test["Drug"], return_ci = True)
     ------------------------------
     """
+
     def calculate(self, X, y_pred):
         residuals = np.abs(self.y_train - self.y_pred_train)
+
         def dist(smi):
             mol = Chem.MolFromSmiles(smi)
             ref_fp = AllChem.GetMorganFingerprint(mol, 2)
             TD = 1 - np.array(BulkTanimotoSimilarity(ref_fp, self.train_fps))
             DC = np.exp(-self.a * TD / (1 - TD))
-            return np.sqrt(np.dot(DC, residuals) / np.sum(DC) )
+            return np.sqrt(np.dot(DC, residuals) / np.sum(DC))
 
         return np.array([dist(smi) for _, smi in tqdm(np.ndenumerate(X))])
 
+
 class KNNSimilarity(BaseFingerprintModel):
-    """ NNSimilarity is an error model that calculates mean Tanimoto similarity
+    """NNSimilarity is an error model that calculates mean Tanimoto similarity
         between the target molecule and the k most similar training molecules
         using a Morgan Fingerprint with a radius of 2 bonds.
 
@@ -152,6 +159,7 @@ class KNNSimilarity(BaseFingerprintModel):
     model.predict(test["Drug"], return_ci = True)
     ------------------------------
     """
+
     @log_arguments
     def __init__(self, k=5):
         self.k = k
@@ -161,12 +169,13 @@ def mean_sim(smi):
             mol = Chem.MolFromSmiles(smi)
             ref_fp = AllChem.GetMorganFingerprint(mol, 2)
             similarity = np.array(BulkTanimotoSimilarity(ref_fp, self.train_fps))
-            return np.mean(sorted(similarity)[-self.k:])
+            return np.mean(sorted(similarity)[-self.k :])
 
         return np.array([mean_sim(smi) for _, smi in tqdm(np.ndenumerate(X))])
 
+
 class TargetDistKNN(KNNSimilarity):
-    """ TargetDistKNN is an error model that calculates the root-mean-square
+    """TargetDistKNN is an error model that calculates the root-mean-square
         difference between the predicted activity of the target molecule and
         the observed activities of the k most similar training molecules,
         weighted by their similarity.
@@ -183,20 +192,24 @@ class TargetDistKNN(KNNSimilarity):
     model.predict(test["Drug"], return_ci = True)
     ------------------------------
     """
+
     def calculate(self, X, y_pred):
         def dist(smi, pred):
             mol = Chem.MolFromSmiles(smi)
             ref_fp = AllChem.GetMorganFingerprint(mol, 2)
             similarity = np.array(BulkTanimotoSimilarity(ref_fp, self.train_fps))
             mae = np.abs(self.y_train - pred)
-            idxs = np.argsort(similarity)[-self.k:]
-            return np.sqrt(np.dot(similarity[idxs], mae[idxs]) / np.sum(similarity[idxs]))
+            idxs = np.argsort(similarity)[-self.k :]
+            return np.sqrt(
+                np.dot(similarity[idxs], mae[idxs]) / np.sum(similarity[idxs])
+            )
 
         X = np.array(X).flatten()
         return np.array([dist(smi, y_pred[i]) for i, smi in tqdm(np.ndenumerate(X))])
 
+
 class TrainDistKNN(KNNSimilarity):
-    """ TrainDistKNN is an error model that calculates the root-mean-square
+    """TrainDistKNN is an error model that calculates the root-mean-square
         difference between the predicted and observed activities of the k most
         similar training molecules to the target molecule, weighted by their
         similarity.
@@ -213,19 +226,24 @@ class TrainDistKNN(KNNSimilarity):
     model.predict(test["Drug"], return_ci = True)
     ------------------------------
     """
+
     def calculate(self, X, y_pred):
         residuals = np.abs(self.y_train - self.y_pred_train)
+
         def dist(smi):
             mol = Chem.MolFromSmiles(smi)
             ref_fp = AllChem.GetMorganFingerprint(mol, 2)
             similarity = np.array(BulkTanimotoSimilarity(ref_fp, self.train_fps))
-            idxs = np.argsort(similarity)[-self.k:]
-            return np.sqrt(np.dot(similarity[idxs], residuals[idxs]) / np.sum(similarity[idxs]))
+            idxs = np.argsort(similarity)[-self.k :]
+            return np.sqrt(
+                np.dot(similarity[idxs], residuals[idxs]) / np.sum(similarity[idxs])
+            )
 
         return np.array([dist(smi) for _, smi in tqdm(np.ndenumerate(X))])
 
+
 class Predicted(BaseErrorModel):
-    """ Predicted is an error model that predicts error bars based on only the
+    """Predicted is an error model that predicts error bars based on only the
         predicted value of a molecule. It is best used as part of an aggregate
         error model rather than by itself.
 
@@ -246,8 +264,9 @@ def __init__(self):
     def calculate(self, X, y_pred):
         return y_pred
 
+
 class ADAN(BaseErrorModel):
-    """ ADAN is an error model that predicts error bars based on one or
+    """ADAN is an error model that predicts error bars based on one or
         multiple ADAN categories: `Applicability Domain Analysis (ADAN): A
         Robust Method for Assessing the Reliability of Drug Property Predictions
         <https://doi.org/10.1021/ci500172z>`_
@@ -277,7 +296,7 @@ def build(
         rep: BaseCompoundVecRepresentation = None,
         dim_reduction: str = "pls",
         explvar: float = 0.8,
-        threshold: float = 0.95
+        threshold: float = 0.95,
     ):
         super().build(model, X, y)
         self.rep = rep
@@ -287,12 +306,12 @@ def build(
         max_components = min(100, min(self.X_train.shape))
 
         if dim_reduction == "pls":
-            self.reduction = PLSRegression(n_components = max_components)
+            self.reduction = PLSRegression(n_components=max_components)
             self.reduction.fit(self.X_train, self.y_train)
             x_var = np.var(self.reduction.x_scores_, axis=0)
             x_var /= np.sum(x_var)
         elif dim_reduction == "pca":
-            self.reduction = PCA(n_components = max_components)
+            self.reduction = PCA(n_components=max_components)
             self.reduction.fit(self.X_train, self.y_train)
             x_var = self.reduction.explained_variance_ratio_
         else:
@@ -305,11 +324,11 @@ def build(
         else:
             self.n_components = max_components
 
-        self.Xp_train = self.reduction.transform(self.X_train)[:, :self.n_components]
+        self.Xp_train = self.reduction.transform(self.X_train)[:, : self.n_components]
         self.Xp_mean = np.mean(self.Xp_train, axis=0)
         self.y_mean = np.mean(self.y_train)
 
-        nbrs = NearestNeighbors(n_neighbors = 2).fit(self.Xp_train)
+        nbrs = NearestNeighbors(n_neighbors=2).fit(self.Xp_train)
         distances, indices = nbrs.kneighbors(self.Xp_train)
 
         centroid_dist = np.linalg.norm(self.Xp_train - self.Xp_mean, axis=1)
@@ -317,7 +336,7 @@ def build(
         model_dist = self.DModX(self.X_train, self.Xp_train)
         y_mean_dist = np.abs(self.y_train - self.y_mean)
         y_nei_dist = np.abs(self.y_train - self.y_train[indices[:, 1]])
-        SDEP_dist = self.SDEP(self.Xp_train, n_drop = 1)
+        SDEP_dist = self.SDEP(self.Xp_train, n_drop=1)
 
         self.training_ = {
             "A_raw": centroid_dist,
@@ -338,10 +357,10 @@ def build(
         }
 
     def calculate_full(self, X):
-        criteria = ["A","B","C","D","E","F"]
+        criteria = ["A", "B", "C", "D", "E", "F"]
         y_pred = np.array(self.model.predict(X)).flatten()
         X = self.preprocess(X)
-        Xp = self.reduction.transform(X)[:, :self.n_components]
+        Xp = self.reduction.transform(X)[:, : self.n_components]
 
         self.results = {}
         for c in criteria:
@@ -354,18 +373,17 @@ def calculate_full(self, X):
         self.results = pd.DataFrame(self.results)
 
     def calculate(self, X, y_pred):
-        """Calcualtes confidence scores.
-        """
+        """Calcualtes confidence scores."""
         X = self.preprocess(X)
-        Xp = self.reduction.transform(X)[:, :self.n_components]
+        Xp = self.reduction.transform(X)[:, : self.n_components]
 
         return self._calculate(X, Xp, y_pred, self.criterion)
-    
+
     def _calculate(self, X, Xp, y_pred, criterion: str, standardize: bool = True):
         if criterion in ("A", "A_raw"):
             dist = np.linalg.norm(Xp - self.Xp_mean, axis=1)
         elif criterion in ("B", "B_raw"):
-            nbrs = NearestNeighbors(n_neighbors = 1).fit(self.Xp_train)
+            nbrs = NearestNeighbors(n_neighbors=1).fit(self.Xp_train)
             distances = nbrs.kneighbors(Xp)[0]
             dist = distances.flatten()
         elif criterion in ("C", "C_raw"):
@@ -375,7 +393,7 @@ def _calculate(self, X, Xp, y_pred, criterion: str, standardize: bool = True):
         elif criterion in ("E", "E_raw"):
             dist = np.abs(y_pred - self.y_mean)
         elif criterion in ("F", "F_raw"):
-            nbrs = NearestNeighbors(n_neighbors = 1).fit(self.Xp_train)
+            nbrs = NearestNeighbors(n_neighbors=1).fit(self.Xp_train)
             indices = nbrs.kneighbors(Xp)[1]
             dist = np.abs(y_pred - self.y_train[indices.flatten()])
         else:
@@ -383,21 +401,24 @@ def _calculate(self, X, Xp, y_pred, criterion: str, standardize: bool = True):
 
         if "raw" in criterion:
             if standardize:
-                return (dist - np.mean(self.training_[criterion])) / np.std(self.training_[criterion])
+                return (dist - np.mean(self.training_[criterion])) / np.std(
+                    self.training_[criterion]
+                )
             else:
                 return dist
         else:
             return np.where(dist > self.thresholds_[criterion], 1, 0)
 
-    def preprocess(self, X, y = None):
-        """Preprocesses data into the appropriate representation.
-        """
+    def preprocess(self, X, y=None):
+        """Preprocesses data into the appropriate representation."""
         if self.rep is None:
             X = self.model.preprocess(X, y)
         else:
             X = np.array(self.rep.convert(X))
 
-        assert isinstance(X, np.ndarray), "The preprocess for the model must return a np.ndarray, e.g. the model must use a BaseCompoundVecRepresentation or a representation must be passed"
+        assert isinstance(
+            X, np.ndarray
+        ), "The preprocess for the model must return a np.ndarray, e.g. the model must use a BaseCompoundVecRepresentation or a representation must be passed"
 
         return X
 
@@ -413,12 +434,24 @@ def DModX(self, X: np.ndarray, Xp: np.ndarray) -> np.ndarray:
             Xp (np.ndarray): queries transformed into latent space
         """
         if self.dim_reduction == "pls":
-            X_reconstructed = Xp @ self.reduction.x_loadings_[:, :self.n_components].T * self.reduction._x_std + self.reduction._x_mean
+            X_reconstructed = (
+                Xp
+                @ self.reduction.x_loadings_[:, : self.n_components].T
+                * self.reduction._x_std
+                + self.reduction._x_mean
+            )
         elif self.dim_reduction == "pca":
-            X_reconstructed = np.dot(Xp, self.reduction.components_[:self.n_components, :]) + self.reduction.mean_
-        return np.linalg.norm(X - X_reconstructed, axis=1) / np.sqrt(self.X_train.shape[1] - self.n_components)
-
-    def SDEP(self, Xp: np.ndarray, n_drop: int = 0, neighbor_thresh: float = 0.05) -> np.ndarray:
+            X_reconstructed = (
+                np.dot(Xp, self.reduction.components_[: self.n_components, :])
+                + self.reduction.mean_
+            )
+        return np.linalg.norm(X - X_reconstructed, axis=1) / np.sqrt(
+            self.X_train.shape[1] - self.n_components
+        )
+
+    def SDEP(
+        self, Xp: np.ndarray, n_drop: int = 0, neighbor_thresh: float = 0.05
+    ) -> np.ndarray:
         """Computes the standard deviation error of predictions (SDEP).
 
         Computes the standard deviation training error of the `neighbor_thresh`
@@ -430,7 +463,7 @@ def SDEP(self, Xp: np.ndarray, n_drop: int = 0, neighbor_thresh: float = 0.05) -
             neighbor_thresh (float): fraction of closest training queries to consider
         """
         n_neighbors = int(self.X_train.shape[0] * neighbor_thresh) + n_drop
-        nbrs = NearestNeighbors(n_neighbors = n_neighbors).fit(self.Xp_train)
+        nbrs = NearestNeighbors(n_neighbors=n_neighbors).fit(self.Xp_train)
         distances, indices = nbrs.kneighbors(Xp)
 
         y_sqerr = np.array((self.y_train - self.y_pred_train) ** 2)
@@ -438,8 +471,9 @@ def SDEP(self, Xp: np.ndarray, n_drop: int = 0, neighbor_thresh: float = 0.05) -
 
         return np.sqrt(y_mse).flatten()
 
+
 class RandomForestErrorModel(BaseAggregateErrorModel):
-    """ RandomForestErrorModel is an aggregate error model that predicts error
+    """RandomForestErrorModel is an aggregate error model that predicts error
         bars based on an aggregate score determined by a random forest model
         trained on several BaseErrorModels.
 
@@ -467,7 +501,7 @@ def _predict(self, scores, **kwargs):
 
 
 class LinearRegressionErrorModel(BaseAggregateErrorModel):
-    """ LinearRegressionErrorModel is an aggregate error model that predicts error
+    """LinearRegressionErrorModel is an aggregate error model that predicts error
         bars based on a linear combination score from several BaseErrorModels.
 
         Parameters:
@@ -494,8 +528,8 @@ def _predict(self, scores, **kwargs):
 
 
 class BaseDomainApplicability(BaseClass):
-    """ Depricated. See BaseErrorModel.
-    
+    """Depricated. See BaseErrorModel.
+
     Base class for analyzing domain applicability of models and estimating
     prediction uncertainty.
 
@@ -559,7 +593,12 @@ def _test(self, X, y_pred, **kwargs):
         """
         pass
 
-    def test(self, dataset: Union[BaseDataset, list, str], use_entire_dataset: bool = False, **kwargs):
+    def test(
+        self,
+        dataset: Union[BaseDataset, list, str],
+        use_entire_dataset: bool = False,
+        **kwargs
+    ):
         """Evaluates model uncertainty on a test dataset. Calls the _test method on the provided dataset.
 
         Parameters:
@@ -580,7 +619,7 @@ def test(self, dataset: Union[BaseDataset, list, str], use_entire_dataset: bool
 
         return self._test(X, y_pred, **kwargs)
 
-    def preprocess(self, X, y = None):
+    def preprocess(self, X, y=None):
         """Preprocesses data into the appropriate representation.
 
         The preprocess for the model must return a np.ndarray, e.g. the model must use a
@@ -598,7 +637,9 @@ def preprocess(self, X, y = None):
         else:
             X = np.array(self.rep.convert(X))
 
-        assert isinstance(X, np.ndarray), "The preprocess for the model must return a np.ndarray, e.g. the model must use a BaseCompoundVecRepresentation or a representation must be passed"
+        assert isinstance(
+            X, np.ndarray
+        ), "The preprocess for the model must return a np.ndarray, e.g. the model must use a BaseCompoundVecRepresentation or a representation must be passed"
 
         return X
 
@@ -606,4 +647,4 @@ def _save(self) -> dict:
         return super()._save()
 
     def _load(self, d: dict):
-        return super()._load(d)
\ No newline at end of file
+        return super()._load(d)
diff --git a/src/olorenchemengine/visualizations/__init__.py b/src/olorenchemengine/visualizations/__init__.py
index b0584a2a..ed52e98a 100644
--- a/src/olorenchemengine/visualizations/__init__.py
+++ b/src/olorenchemengine/visualizations/__init__.py
@@ -3,4 +3,4 @@
 from .matched_pairs import *
 from .exploratory_analysis import *
 from .visualize_interpret import *
-from .util import sample_colorscale
\ No newline at end of file
+from .util import sample_colorscale
diff --git a/src/olorenchemengine/visualizations/attribute_section.py b/src/olorenchemengine/visualizations/attribute_section.py
index f7a4df05..160c3ff0 100644
--- a/src/olorenchemengine/visualizations/attribute_section.py
+++ b/src/olorenchemengine/visualizations/attribute_section.py
@@ -7,11 +7,12 @@
 from abc import ABC, abstractmethod
 
 from olorenchemengine.base_class import *
-from olorenchemengine.representations import *
 from olorenchemengine.dataset import *
-from olorenchemengine.uncertainty import *
-from olorenchemengine.interpret import *
 from olorenchemengine.internal import *
+from olorenchemengine.interpret import *
+from olorenchemengine.representations import *
+from olorenchemengine.uncertainty import *
+
 
 class BaseAttribute(ABC):
     """Base class for all attributes for a visualization
@@ -19,6 +20,7 @@ class BaseAttribute(ABC):
     Parameters:
         name (str): The name of the attribute
     """
+
     def __init__(self, name: str):
         self.name = name
 
@@ -46,12 +48,14 @@ class SimpleAttribute(BaseAttribute):
     def to_json(self) -> list:
         return [self.name, self.attribute_name]
 
+
 class ModelSelector(SimpleAttribute):
     """Class to create a model selector attribute
 
     Parameters:
         name (str): The name of the attribute
     """
+
     attribute_name = "modelSelector"
 
 
@@ -61,32 +65,40 @@ class DatasetSelector(SimpleAttribute):
     Parameters:
         name (str): The name of the attribute
     """
+
     attribute_name = "datasetSelector"
 
+
 class InputNumber(SimpleAttribute):
     """Class to create an input number attribute
 
     Parameters:
         name (str): The name of the attribute
     """
+
     attribute_name = "inputNumber"
 
+
 class InputString(SimpleAttribute):
     """Class to create an input string attribute
 
     Parameters:
         name (str): The name of the attribute
     """
+
     attribute_name = "inputString"
 
+
 class ColorPicker(SimpleAttribute):
     """Class to create a color picker attribute
 
     Parameters:
         name (str): The name of the attribute
     """
+
     attribute_name = "colorPicker"
 
+
 class InputThreshold(BaseAttribute):
     """Class to create an input threshold attribute
 
@@ -95,15 +107,25 @@ class InputThreshold(BaseAttribute):
         min (int): The minimum value of the threshold
         max (int): The maximum value of the threshold
     """
+
     @log_arguments
-    def __init__(self, name:  str, min: int = 0, max: int = 100,  ):
+    def __init__(
+        self,
+        name: str,
+        min: int = 0,
+        max: int = 100,
+    ):
         super().__init__(name)
         self.min = min
         self.max = max
 
-
     def to_json(self):
-        inputThreshold = [self.name, "threshold", "min: " + str(self.min), "max: " + str(self.max)]
+        inputThreshold = [
+            self.name,
+            "threshold",
+            "min: " + str(self.min),
+            "max: " + str(self.max),
+        ]
         return inputThreshold
 
 
@@ -114,23 +136,19 @@ class AttributeSection(BaseAttribute):
         attribute_name (str): The name of the attribute section
         attributes (list): The list of attributes of type BaseAttribute
     """
+
     @log_arguments
-    def __init__(self, attribute_name:  str, attributes: List[BaseAttribute],  ):
+    def __init__(
+        self,
+        attribute_name: str,
+        attributes: List[BaseAttribute],
+    ):
         super().__init__(attribute_name)
         self.attributes = attributes
 
-
     def to_json(self):
-        attribute_section = {self.name: [attribute.to_json() for attribute in self.attributes]}
+        attribute_section = {
+            self.name: [attribute.to_json() for attribute in self.attributes]
+        }
 
         return attribute_section
-
-
-
-
-
-
-
-
-
-
diff --git a/src/olorenchemengine/visualizations/exploratory_analysis.py b/src/olorenchemengine/visualizations/exploratory_analysis.py
index 02c9eab2..d670e8e0 100644
--- a/src/olorenchemengine/visualizations/exploratory_analysis.py
+++ b/src/olorenchemengine/visualizations/exploratory_analysis.py
@@ -3,14 +3,15 @@
 """
 
 from olorenchemengine.base_class import *
-from olorenchemengine.representations import *
 from olorenchemengine.dataset import *
-from olorenchemengine.uncertainty import *
-from olorenchemengine.interpret import *
 from olorenchemengine.internal import *
+from olorenchemengine.interpret import *
 from olorenchemengine.manager import *
+from olorenchemengine.representations import *
+from olorenchemengine.uncertainty import *
 from olorenchemengine.visualizations import *
 
+
 class VisualizeFeatures(BaseVisualization):
     """
     Visualize how the selected features are correlated with the dataset.
@@ -30,39 +31,53 @@ def __init__(self, dataset: BaseDataset, features: list = [], log=True, **kwargs
         self.packages = ["plotly", "olorenrenderer"]
         self.data = {
             "SMILES": self.dataset.data[self.dataset.structure_col].tolist(),
-            self.dataset.property_col: self.dataset.data[self.dataset.property_col].tolist(),
+            self.dataset.property_col: self.dataset.data[
+                self.dataset.property_col
+            ].tolist(),
         }
         self.d = {}
         self.feature_cols = []
         self.spearman_corr = []
         from scipy.stats import spearmanr
+
         for feature in self.features:
             if isinstance(feature, str):
                 self.feature_cols.append(features)
                 self.d[feature] = self.dataset.data[feature].tolist()
-                self.spearman_corr.append(spearmanr(self.d[feature], self.dataset.data[self.dataset.property_col].tolist()))
+                self.spearman_corr.append(
+                    spearmanr(
+                        self.d[feature],
+                        self.dataset.data[self.dataset.property_col].tolist(),
+                    )
+                )
             elif issubclass(type(feature), BaseRepresentation):
                 rep_features = feature.convert(self.dataset.entire_dataset[0])
                 rep_features = np.array(rep_features).transpose()
                 try:
                     names = feature.names
                 except:
-                    names = [f"{feature.__class__.__name__}-{i}" for i in range(rep_features.shape[0])]
+                    names = [
+                        f"{feature.__class__.__name__}-{i}"
+                        for i in range(rep_features.shape[0])
+                    ]
                 for i, name in enumerate(names):
                     self.feature_cols.append(name)
                     self.d[name] = rep_features[i].tolist()
-                    self.spearman_corr.append(spearmanr(self.d[name], self.dataset.data[self.dataset.property_col]))
+                    self.spearman_corr.append(
+                        spearmanr(
+                            self.d[name], self.dataset.data[self.dataset.property_col]
+                        )
+                    )
 
     def get_data(self):
-        d = {"datacols": self.d.copy(),
+        d = {
+            "datacols": self.d.copy(),
             "SMILES": self.data["SMILES"],
             "PROPERTY": self.dataset.property_col,
             "PROPERTY_VALUES": self.data[self.dataset.property_col],
             "FEATURE_COLS": self.feature_cols,
             "SPEARMAN_COEF": [np.around(x[0], decimals=2) for x in self.spearman_corr],
             "SPEARMAN_PVAL": [np.around(x[1], decimals=2) for x in self.spearman_corr],
-            "hoverSize": 0.8}
+            "hoverSize": 0.8,
+        }
         return d
-
-
-
diff --git a/src/olorenchemengine/visualizations/matched_pairs.py b/src/olorenchemengine/visualizations/matched_pairs.py
index cc369dc3..d0d27de7 100644
--- a/src/olorenchemengine/visualizations/matched_pairs.py
+++ b/src/olorenchemengine/visualizations/matched_pairs.py
@@ -7,7 +7,12 @@
 
 class MatchedPairsTable(BaseVisualization):
     @log_arguments
-    def __init__(self, dataset: BaseDataset, mode: str = "features", annotations: Union[str, List[str]] = []):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        mode: str = "features",
+        annotations: Union[str, List[str]] = [],
+    ):
 
         self.packages = ["d3"]
 
@@ -31,21 +36,35 @@ def __init__(self, dataset: BaseDataset, mode: str = "features", annotations: Un
                         "initial": diff[1],
                         "final": diff[2],
                         **{
-                            f"Delta (Final - Initial) {property_col}": r2[property_col + " mean"]
+                            f"Delta (Final - Initial) {property_col}": r2[
+                                property_col + " mean"
+                            ]
                             - r1[property_col + " mean"]
                             for property_col in dataset.property_col
                         },
                         **{
-                            f"Fold (Final/Initial) {property_col}": r2[property_col + " mean"]
+                            f"Fold (Final/Initial) {property_col}": r2[
+                                property_col + " mean"
+                            ]
                             / r1[property_col + " mean"]
                             for property_col in dataset.property_col
                         },
-                        **{f"Initial {annotation}": r1[annotation] for annotation in self.annotations},
-                        **{f"Final {annotation}": r2[annotation] for annotation in self.annotations},
+                        **{
+                            f"Initial {annotation}": r1[annotation]
+                            for annotation in self.annotations
+                        },
+                        **{
+                            f"Final {annotation}": r2[annotation]
+                            for annotation in self.annotations
+                        },
                     },
                     ignore_index=True,
                 )
-            self.mp_diffs = self.mp_diffs.sort_values(by="col", ascending=True).reset_index().drop(columns=["index"])
+            self.mp_diffs = (
+                self.mp_diffs.sort_values(by="col", ascending=True)
+                .reset_index()
+                .drop(columns=["index"])
+            )
         elif mode == "property":
             self.mp_diffs = pd.DataFrame()
             for mp in self.mps:
@@ -62,13 +81,21 @@ def __init__(self, dataset: BaseDataset, mode: str = "features", annotations: Un
                                 "final": diff[2],
                                 f"Initial {property}": r1[property],
                                 f"Final {property}": r2[property],
-                                **{f"Initial {annotation}": r1[annotation] for annotation in annotations},
-                                **{f"Final {annotation}": r2[annotation] for annotation in annotations},
+                                **{
+                                    f"Initial {annotation}": r1[annotation]
+                                    for annotation in annotations
+                                },
+                                **{
+                                    f"Final {annotation}": r2[annotation]
+                                    for annotation in annotations
+                                },
                             },
                             ignore_index=True,
                         )
             self.mp_diffs = (
-                self.mp_diffs.sort_values(by=["col", "initial", "final"], ascending=True)
+                self.mp_diffs.sort_values(
+                    by=["col", "initial", "final"], ascending=True
+                )
                 .reset_index()
                 .drop(columns=["index"])
             )
@@ -95,7 +122,15 @@ def matched_pairs(self, df, cols, dist=1):
     def get_data(self):
         data = []
         for i, col in enumerate(self.mp_diffs["col"].unique()):
-            data.append({"col": col, "i": i, "col_data": self.mp_diffs.loc[self.mp_diffs["col"] == col].to_dict("r")})
+            data.append(
+                {
+                    "col": col,
+                    "i": i,
+                    "col_data": self.mp_diffs.loc[self.mp_diffs["col"] == col].to_dict(
+                        "r"
+                    ),
+                }
+            )
         return data
 
     @property
@@ -114,15 +149,18 @@ def get_data(self):
             property_col = f"Delta (Final - Initial) {property_col}"
             for i, col in enumerate(self.mp_diffs["col"].unique()):
                 col_data = self.mp_diffs[self.mp_diffs["col"] == col]
-                mapping = {str(x): i for i, x in enumerate(col_data["initial"].unique())}
+                mapping = {
+                    str(x): i for i, x in enumerate(col_data["initial"].unique())
+                }
                 col_heatmap = np.zeros((len(mapping), len(mapping)))
                 for i, col_row in col_data.iterrows():
                     try:
                         x = float(col_row[property_col])
                         if not np.isnan(x):
-                            col_heatmap[mapping[str(col_row["final"])], mapping[str(col_row["initial"])]] += col_row[
-                                property_col
-                            ]
+                            col_heatmap[
+                                mapping[str(col_row["final"])],
+                                mapping[str(col_row["initial"])],
+                            ] += col_row[property_col]
                     except:
                         pass
 
diff --git a/src/olorenchemengine/visualizations/model_comparisons.py b/src/olorenchemengine/visualizations/model_comparisons.py
index 4b02ea21..6420fdd4 100644
--- a/src/olorenchemengine/visualizations/model_comparisons.py
+++ b/src/olorenchemengine/visualizations/model_comparisons.py
@@ -5,16 +5,17 @@
 """
 
 from olorenchemengine.base_class import *
-from olorenchemengine.representations import *
 from olorenchemengine.dataset import *
-from olorenchemengine.uncertainty import *
-from olorenchemengine.interpret import *
 from olorenchemengine.internal import *
+from olorenchemengine.interpret import *
 from olorenchemengine.manager import *
+from olorenchemengine.representations import *
+from olorenchemengine.uncertainty import *
 from olorenchemengine.visualizations import *
 
+
 class VisualizationModelManager_Bar(BaseVisualization):
-    """ Visualize the database of a model manager in a bar plot, with the model
+    """Visualize the database of a model manager in a bar plot, with the model
     parameters and name displayed on hover
 
     Parameters:
@@ -23,7 +24,7 @@ class VisualizationModelManager_Bar(BaseVisualization):
     """
 
     @log_arguments
-    def __init__(self, mm: BaseModelManager, top = 5, log=True, **kwargs):
+    def __init__(self, mm: BaseModelManager, top=5, log=True, **kwargs):
         self.df = mm.get_model_database()
         self.metrics = mm.metrics
         self.mm = mm
@@ -35,7 +36,7 @@ def __init__(self, mm: BaseModelManager, top = 5, log=True, **kwargs):
     def JS_NAME(self) -> str:
         return "ModelManagerBar"
 
-    def get_data(self, metric = None) -> dict:
+    def get_data(self, metric=None) -> dict:
         """
         Get the data for the visualization
 
@@ -45,11 +46,16 @@ def get_data(self, metric = None) -> dict:
         Returns:
             dict: The data for the visualization
         """
-        df2 = self.df.sort_values(by=metric, ascending=(True if self.mm.metric_direction[metric] == "lower" else False))
-        df2 = df2.iloc[:self.top]
-        d =  df2.to_dict("l")
+        df2 = self.df.sort_values(
+            by=metric,
+            ascending=(True if self.mm.metric_direction[metric] == "lower" else False),
+        )
+        df2 = df2.iloc[: self.top]
+        d = df2.to_dict("l")
         if not metric is None:
-            assert metric in self.metrics, f"Designated metric, {metric}, not in model manager metrics: {self.metrics}"
+            assert (
+                metric in self.metrics
+            ), f"Designated metric, {metric}, not in model manager metrics: {self.metrics}"
             d["metric"] = metric
         else:
             d["metric"] = self.metrics[0]
@@ -70,10 +76,13 @@ def render_ipynb(self, metric=None, *args, print_html_js=False, **kwargs) -> str
             str: The html code for the visualization
         """
         d = self.get_data(metric)
-        return self.render_ipynb_from_dict(d, *args, print_html_js=print_html_js, **kwargs)
+        return self.render_ipynb_from_dict(
+            d, *args, print_html_js=print_html_js, **kwargs
+        )
+
 
 class ModelOverTime(CompoundScatterPlot):
-    """ Visualizes the performance of a given model architecture over a length of time.
+    """Visualizes the performance of a given model architecture over a length of time.
 
     This visualization will use a date split and for each date i, will train a model
     on the data on dates 1, ..., i-1, evaluating on date i. Note, any existing splits
@@ -85,8 +94,17 @@ class ModelOverTime(CompoundScatterPlot):
     """
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True,
-        title="Model Performance over Time", yaxis_title="Error", xaxis_title="Date", **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        *args,
+        log=True,
+        title="Model Performance over Time",
+        yaxis_title="Error",
+        xaxis_title="Date",
+        **kwargs,
+    ):
         self.dataset = dataset
         self.model = model
         self.title = title
@@ -96,16 +114,25 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True,
         self.dataset.data = self.dataset.data.sort_values(by=[self.dataset.date_col])
         self.dates = self.dataset.data[self.dataset.date_col].unique()
         results = pd.DataFrame()
-        for i in tqdm(range(len(self.dates)-1)):
+        for i in tqdm(range(len(self.dates) - 1)):
             # try:
             train_date = self.dates[i]
-            test_date = self.dates[i+1]
-            train = self.dataset.data[self.dataset.data[self.dataset.date_col] <= train_date]
+            test_date = self.dates[i + 1]
+            train = self.dataset.data[
+                self.dataset.data[self.dataset.date_col] <= train_date
+            ]
             train["split"] = "train"
-            test = self.dataset.data[self.dataset.data[self.dataset.date_col] == test_date]
+            test = self.dataset.data[
+                self.dataset.data[self.dataset.date_col] == test_date
+            ]
             test["split"] = "test"
             DATA = pd.concat([train, test])
-            dataset = (oce.BaseDataset(data = DATA.to_csv(), structure_col="SMILES", property_col = "log AUC") + oce.CleanStructures())
+            dataset = (
+                oce.BaseDataset(
+                    data=DATA.to_csv(), structure_col="SMILES", property_col="log AUC"
+                )
+                + oce.CleanStructures()
+            )
             print(f"Train size is {len(train)} and test size is {len(test)}")
             self.model.fit(*dataset.train_dataset)
 
@@ -115,10 +142,15 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True,
             results = pd.concat([results, test_set])
 
             if self.model.setting == "regression":
-                results["Error (|Predicted - True|)"] = np.abs(results["pred"] - results[self.dataset.property_col])
+                results["Error (|Predicted - True|)"] = np.abs(
+                    results["pred"] - results[self.dataset.property_col]
+                )
                 self.yaxis_title = "Error (|Predicted - True|)"
             elif self.model.setting == "classification":
-                results["Error (p or 1-p)"] = np.abs(results[self.dataset.property_col]+(-2*results[self.dataset.property_col]+1)*results["pred"])
+                results["Error (p or 1-p)"] = np.abs(
+                    results[self.dataset.property_col]
+                    + (-2 * results[self.dataset.property_col] + 1) * results["pred"]
+                )
                 self.yaxis_title = "Error (p or 1-p)"
 
         results["X"] = results[self.dataset.date_col]
@@ -129,23 +161,38 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True,
             results["Y"] = results["Error (p or 1-p)"]
             self.yaxis_title = yaxis_title
         results["SMILES"] = results[self.dataset.structure_col]
-        super().__init__(*args,df = results, xaxis_type="date",
-            title=title, xaxis_title=self.xaxis_title, yaxis_title=self.yaxis_title, log=False, **kwargs)
+        super().__init__(
+            *args,
+            df=results,
+            xaxis_type="date",
+            title=title,
+            xaxis_title=self.xaxis_title,
+            yaxis_title=self.yaxis_title,
+            log=False,
+            **kwargs,
+        )
 
     def get_data(self):
         d = super().get_data()
         window = 10
         quantile = 0.80
-        results = list(zip(pd.to_datetime(self.df["X"]).apply(lambda x: x.timestamp()), self.df["Y"]))
+        results = list(
+            zip(
+                pd.to_datetime(self.df["X"]).apply(lambda x: x.timestamp()),
+                self.df["Y"],
+            )
+        )
         scores, residuals = zip(*sorted(results))
         X = pd.Series(scores).rolling(window).mean()
         y = pd.Series(residuals).rolling(window).quantile(quantile)
         X = np.array(X[~np.isnan(X)])
         y = np.array(y[~np.isnan(y)])
 
-        d["trace_update"] = {'x':self.df["X"].tolist(),
-            'y': y.tolist(),
-            'mode': 'markers',
-            'type': 'scatter',
-            'marker': {'color': 'red'}}
-        return d
\ No newline at end of file
+        d["trace_update"] = {
+            "x": self.df["X"].tolist(),
+            "y": y.tolist(),
+            "mode": "markers",
+            "type": "scatter",
+            "marker": {"color": "red"},
+        }
+        return d
diff --git a/src/olorenchemengine/visualizations/util.py b/src/olorenchemengine/visualizations/util.py
index ce5f4062..257ca265 100644
--- a/src/olorenchemengine/visualizations/util.py
+++ b/src/olorenchemengine/visualizations/util.py
@@ -587,9 +587,7 @@ def make_colorscale(colors, scale=None):
 
     # validate minimum colors length of 2
     if len(colors) < 2:
-        raise Exception(
-            "You must input a list of colors that has at least two colors."
-        )
+        raise Exception("You must input a list of colors that has at least two colors.")
 
     if scale is None:
         scale_incr = 1.0 / (len(colors) - 1)
@@ -597,9 +595,7 @@ def make_colorscale(colors, scale=None):
 
     else:
         if len(colors) != len(scale):
-            raise Exception(
-                "The length of colors and scale must be the same."
-            )
+            raise Exception("The length of colors and scale must be the same.")
 
         validate_scale_values(scale)
 
@@ -881,4 +877,4 @@ def sample_colorscale(colorscale, samplepoints, low=0.0, high=1.0, colortype="rg
         interpolant = (point - scale[low]) / (scale[high] - scale[low])
         sampled_color = find_intermediate_color(colors[low], colors[high], interpolant)
         sampled_colors.append(sampled_color)
-    return validate_colors(sampled_colors, colortype=colortype)
\ No newline at end of file
+    return validate_colors(sampled_colors, colortype=colortype)
diff --git a/src/olorenchemengine/visualizations/visualization.py b/src/olorenchemengine/visualizations/visualization.py
index 125c5f09..a653001c 100644
--- a/src/olorenchemengine/visualizations/visualization.py
+++ b/src/olorenchemengine/visualizations/visualization.py
@@ -1,29 +1,28 @@
+import dataclasses
+import os
+import urllib.parse
 from cmath import e
 from dataclasses import dataclass
-import dataclasses
 from lib2to3.pgen2.literals import simple_escapes
-import os
 from re import escape
-import urllib.parse
 
 from IPython.display import IFrame
 from pandas.core.indexes.accessors import NoNewAttributesMixin
-
+from rdkit import Chem
+from rdkit.Chem import Draw
 from sklearn.model_selection import PredefinedSplit
 
 from olorenchemengine.base_class import *
-from olorenchemengine.representations import *
 from olorenchemengine.dataset import *
-from olorenchemengine.uncertainty import *
-from olorenchemengine.interpret import *
-from olorenchemengine.internal import *
 from olorenchemengine.ensemble import *
+from olorenchemengine.internal import *
+from olorenchemengine.interpret import *
+from olorenchemengine.representations import *
+from olorenchemengine.uncertainty import *
 
-from rdkit import Chem
-from rdkit.Chem import Draw
 
 def get_oas_compatible_visualizations():
-    """ Returns a list of all available visualizations that are compatible with OAS,
+    """Returns a list of all available visualizations that are compatible with OAS,
     meaning that they can be created and specified in OAS. All visualizations
     can be shared via OAS."""
 
@@ -35,7 +34,10 @@ def get_oas_compatible_visualizations():
         vis.__name__
         for vis in all_visualizations
         if hasattr(vis, "get_attributes")
-        and (not hasattr(vis.__bases__[0], "get_attributes") or vis.get_attributes != vis.__bases__[0].get_attributes)
+        and (
+            not hasattr(vis.__bases__[0], "get_attributes")
+            or vis.get_attributes != vis.__bases__[0].get_attributes
+        )
     ]
 
 
@@ -65,8 +67,7 @@ def __init__(self, log=True, **kwargs):
         self.packages = []
 
     def get_data() -> dict:
-        """Get data for visualization in JSON-like dictionary.
-        """
+        """Get data for visualization in JSON-like dictionary."""
 
         return {}
 
@@ -76,7 +77,9 @@ def get_attributes() -> list:
 
         Returns:
             list: List of attributes for OAS to display."""
-        raise NotImplementedError("get_attributes must be implemented for this visualization")
+        raise NotImplementedError(
+            "get_attributes must be implemented for this visualization"
+        )
 
     @classmethod
     def from_attributes(cls, attributes: dict) -> "BaseVisualization":
@@ -90,7 +93,9 @@ def from_attributes(cls, attributes: dict) -> "BaseVisualization":
             vis = cls(**attributes)
             return vis
         except Exception as e:
-            raise ValueError(f"Error creating visualization {cls.__name__} using {attributes}: {e}")
+            raise ValueError(
+                f"Error creating visualization {cls.__name__} using {attributes}: {e}"
+            )
 
     @property
     def JS_NAME(self) -> str:
@@ -103,7 +108,7 @@ def JS_NAME(self) -> str:
         return self.__class__.__name__
 
     def get_html(self, data_str: str, packages: str, **kwargs):
-        """ Returns the HTML code for the visualization.
+        """Returns the HTML code for the visualization.
 
         Parameters:
             data_str (str): JSON string of data to be used in visualization.
@@ -142,7 +147,10 @@ def get_js(self, *args, **kwargs) -> str:
             str: JavaScript code for visualization.
         """
 
-        with open(os.path.join(os.path.dirname(__file__), "scripts/", self.JS_NAME + ".js"), "r") as f:
+        with open(
+            os.path.join(os.path.dirname(__file__), "scripts/", self.JS_NAME + ".js"),
+            "r",
+        ) as f:
             js = f.read()
 
         resize = (
@@ -166,7 +174,7 @@ def render(self, data: dict = None, print_html_js=False, **kwargs) -> str:
         # Formats the data to be properly JSON
         if data is None:
             data_str = (
-                json.dumps(self.get_data(**kwargs), separators=(',', ':'))
+                json.dumps(self.get_data(**kwargs), separators=(",", ":"))
                 .replace('"', "&quot;")
                 .replace("None", "null")
                 .replace("NaN", "null")
@@ -175,7 +183,7 @@ def render(self, data: dict = None, print_html_js=False, **kwargs) -> str:
             )
         else:
             data_str = (
-                json.dumps(data, separators=(',', ':'))
+                json.dumps(data, separators=(",", ":"))
                 .replace('"', "&quot;")
                 .replace("None", "null")
                 .replace("NaN", "null")
@@ -184,7 +192,12 @@ def render(self, data: dict = None, print_html_js=False, **kwargs) -> str:
             )
 
         # Compiles package imports into HTML
-        packages = "".join([f'<script src = "{self.package_urls[package]}"></script>' for package in self.packages])
+        packages = "".join(
+            [
+                f'<script src = "{self.package_urls[package]}"></script>'
+                for package in self.packages
+            ]
+        )
         # Get base HTML code for visualization
         html = self.get_html(data_str, packages)
 
@@ -211,7 +224,9 @@ def render_ipynb(self, data: dict = None, print_html_js=False, **kwargs) -> str:
             print_html_js (bool): Whether or not to print the JavaScript code for the
                 visualization.
         """
-        return IFrame(self.render_data_url(data, print_html_js, **kwargs), width=800, height=600)
+        return IFrame(
+            self.render_data_url(data, print_html_js, **kwargs), width=800, height=600
+        )
 
     def render_data_url(self, data: dict = None, print_html_js=False, **kwargs) -> str:
         """Render visualization to data_url string.
@@ -222,7 +237,9 @@ def render_data_url(self, data: dict = None, print_html_js=False, **kwargs) -> s
             print_html_js (bool): Whether or not to print the JavaScript code for the
                 visualization.
         """
-        data_url = "data:text/html," + urllib.parse.quote(self.render(data, print_html_js, **kwargs), safe="")
+        data_url = "data:text/html," + urllib.parse.quote(
+            self.render(data, print_html_js, **kwargs), safe=""
+        )
         return data_url
 
     def upload_oas(self):
@@ -242,6 +259,7 @@ def _save(self) -> dict:
     def _load(self, d: dict):
         pass
 
+
 class VisualizeError(BaseVisualization):
     """Visualizes property values and errors for a compound
 
@@ -256,6 +274,7 @@ class VisualizeError(BaseVisualization):
             if dataset is a BaseDataset.
         yaxis_title (str): Title for y-axis.
     """
+
     @log_arguments
     def __init__(
         self,
@@ -270,7 +289,7 @@ def __init__(
         yaxis_title=None,
         width=800,
         height=600,
-        **kwargs
+        **kwargs,
     ):
         assert error >= 0, "error must be nonnegative"
 
@@ -312,13 +331,13 @@ def __init__(
             self.ci = int(ci * 100)
 
         self.title = title
-        self.width=width
-        self.height=height
+        self.width = width
+        self.height = height
         super().__init__(**kwargs)
         self.packages = ["plotly"]
 
     def get_data(self):
-        d =  {
+        d = {
             "reference": self.reference,
             "value": self.value,
             "error": self.error,
@@ -337,6 +356,7 @@ def get_data(self):
 
         return d
 
+
 class VisualizeCompounds(BaseVisualization):
     """Visualizes the compounds in a set.
 
@@ -356,12 +376,27 @@ class VisualizeCompounds(BaseVisualization):
         shuffle = True (bool): Whether or not to shuffle the compounds."""
 
     @log_arguments
-    def __init__(self, dataset: Union[BaseDataset, list, pd.Series, str, Chem.Mol], table_width: int = 1, table_height: int = 5,
-        compound_width: int = 500, compound_height: int = 500, annotations = None, kekulize = True,
-        box=False, shuffle=False,log=True, **kwargs):
+    def __init__(
+        self,
+        dataset: Union[BaseDataset, list, pd.Series, str, Chem.Mol],
+        table_width: int = 1,
+        table_height: int = 5,
+        compound_width: int = 500,
+        compound_height: int = 500,
+        annotations=None,
+        kekulize=True,
+        box=False,
+        shuffle=False,
+        log=True,
+        **kwargs,
+    ):
         self.dataset = dataset
         if issubclass(type(dataset), BaseDataset):
-            self.compounds = dataset.data[dataset.structure_col].head(table_width * table_height).tolist()
+            self.compounds = (
+                dataset.data[dataset.structure_col]
+                .head(table_width * table_height)
+                .tolist()
+            )
         elif isinstance(dataset, pd.Series):
             self.compounds = dataset.tolist()
         elif isinstance(dataset, list):
@@ -375,6 +410,7 @@ def __init__(self, dataset: Union[BaseDataset, list, pd.Series, str, Chem.Mol],
             self.compounds = [Chem.MolToSmiles(dataset)]
         if shuffle:
             import random
+
             random.shuffle(self.compounds)
 
         if not issubclass(type(dataset), BaseDataset):
@@ -396,23 +432,39 @@ def __init__(self, dataset: Union[BaseDataset, list, pd.Series, str, Chem.Mol],
 
     def get_data(self):
         if self.kekulize:
-            self.compounds = [Chem.MolToSmiles(Chem.MolFromSmiles(s, sanitize=False), kekuleSmiles = True) for s in self.compounds]
-            d =  {"smiles": self.compounds,
+            self.compounds = [
+                Chem.MolToSmiles(
+                    Chem.MolFromSmiles(s, sanitize=False), kekuleSmiles=True
+                )
+                for s in self.compounds
+            ]
+            d = {
+                "smiles": self.compounds,
                 "table_width": self.table_width,
                 "table_height": self.table_height,
                 "compound_width": self.compound_width,
                 "compound_height": self.compound_height,
-                "box": self.box}
+                "box": self.box,
+            }
 
         if not self.annotations is None:
-            d.update({"annotations": {col: self.dataset.data[col].tolist() for col in self.annotations}})
+            d.update(
+                {
+                    "annotations": {
+                        col: self.dataset.data[col].tolist() for col in self.annotations
+                    }
+                }
+            )
 
         return d
 
+
 class VisualizeDatasetCompounds(VisualizeCompounds):
     """Alias for VisualizeCompounds"""
+
     pass
 
+
 class ScatterPlot(BaseVisualization):
     """Scatter plot visualization.
 
@@ -435,7 +487,7 @@ def get_data(self) -> dict:
 
 
 class CompoundScatterPlot(BaseVisualization):
-    """ Scatter plot visualization, where molecules are displayed on hover on
+    """Scatter plot visualization, where molecules are displayed on hover on
     an x-y plane.
 
     Parameters:
@@ -487,8 +539,8 @@ def __init__(
         title: str = "Compound Scatter Plot",
         xaxis_title: str = "X axis",
         yaxis_title: str = "y axis",
-        x_col: str =    None,
-        y_col: str =    None,
+        x_col: str = None,
+        y_col: str = None,
         smiles_col: str = None,
         kekulize: bool = True,
         color_col: str = None,
@@ -589,7 +641,15 @@ def get_data(self, include_data=True) -> dict:
 
         if include_data:
             d = self.df.to_dict("l")
-            d["SMILES"] = self.df["SMILES"].apply(lambda x: Chem.MolToSmiles(Chem.MolFromSmiles(x), kekuleSmiles=self.kekulize)).tolist()
+            d["SMILES"] = (
+                self.df["SMILES"]
+                .apply(
+                    lambda x: Chem.MolToSmiles(
+                        Chem.MolFromSmiles(x), kekuleSmiles=self.kekulize
+                    )
+                )
+                .tolist()
+            )
         else:
             d = dict()
         d["title"] = self.title
@@ -629,7 +689,15 @@ def get_attributes():
         Returns:
             list: List of attributes for OAS to display.
         """
-        return [{"Labels": [["title", "inputString"], ["xaxis_title", "inputString"], ["yaxis_title", "inputString"]]}]
+        return [
+            {
+                "Labels": [
+                    ["title", "inputString"],
+                    ["xaxis_title", "inputString"],
+                    ["yaxis_title", "inputString"],
+                ]
+            }
+        ]
 
 
 from sklearn.decomposition import PCA
@@ -637,7 +705,7 @@ def get_attributes():
 
 
 class ChemicalSpacePlot(CompoundScatterPlot):
-    """ Visualize chemical space by calculating a vector representation and
+    """Visualize chemical space by calculating a vector representation and
     performing dimensionality reduction to 2 dimensions.
 
     Parameters:
@@ -653,18 +721,28 @@ class ChemicalSpacePlot(CompoundScatterPlot):
             chemical space plot
         colorscale (str, optional): Color scale to use for coloring the chemical space plot.
             Other options can be found
-        """
+    """
 
     @log_arguments
-    def __init__(self, dataset: Union[BaseDataset, list, pd.Series, pd.DataFrame], rep: BaseCompoundVecRepresentation,
-            *args, dim_reduction="tsne",
-            smiles_col=None, title ="Chemical Space Plot", log=True, **kwargs):
+    def __init__(
+        self,
+        dataset: Union[BaseDataset, list, pd.Series, pd.DataFrame],
+        rep: BaseCompoundVecRepresentation,
+        *args,
+        dim_reduction="tsne",
+        smiles_col=None,
+        title="Chemical Space Plot",
+        log=True,
+        **kwargs,
+    ):
 
         # Sets visualization instance variables
         if issubclass(type(dataset), BaseDataset):
             self.structures = dataset.data[dataset.structure_col]
         elif isinstance(dataset, pd.DataFrame):
-            assert smiles_col is not None, "smiles_col must be defined if `dataset` parameter is pd.DataFrame"
+            assert (
+                smiles_col is not None
+            ), "smiles_col must be defined if `dataset` parameter is pd.DataFrame"
             self.structures = dataset[smiles_col]
         else:
             self.structures = dataset
@@ -695,9 +773,16 @@ def __init__(self, dataset: Union[BaseDataset, list, pd.Series, pd.DataFrame], r
 
         self.title = title
 
-        super().__init__(self.df, *args, title=self.title,
-            xaxis_title="Component 1", yaxis_title="Component 2", smiles_col = smiles_col,
-            log=False, **kwargs)
+        super().__init__(
+            self.df,
+            *args,
+            title=self.title,
+            xaxis_title="Component 1",
+            yaxis_title="Component 2",
+            smiles_col=smiles_col,
+            log=False,
+            **kwargs,
+        )
 
     def tsne_df(self, chem_rep_list):
         """
@@ -729,23 +814,30 @@ def pca_df(self, chem_rep_list):
         pca_arr = pca.fit_transform(chem_rep_list)
         return pd.DataFrame(pca_arr, columns=["Component 1", "Component 2"])
 
-    def get_data(self, color_col: str = None, size_col: str = None, SMILES: str = None) -> dict:
+    def get_data(
+        self, color_col: str = None, size_col: str = None, SMILES: str = None
+    ) -> dict:
         d = super().get_data(include_data=True)
 
         self.color_col = color_col
         if not self.color_col is None:
-            assert self.color_col in self.df.columns, f"specified color column, {self.color_col}, not in columns"
+            assert (
+                self.color_col in self.df.columns
+            ), f"specified color column, {self.color_col}, not in columns"
             d["color"] = self.df[self.color_col].tolist()
 
         self.size_col = size_col
         if not self.size_col is None:
-            assert self.size_col in self.df.columns, f"specified size column, {self.size_col}, not in columns"
+            assert (
+                self.size_col in self.df.columns
+            ), f"specified size column, {self.size_col}, not in columns"
             d["size"] = self.df[self.size_col].tolist()
 
         return d
 
+
 class VisualizeMoleculePerturbations(ChemicalSpacePlot):
-    """ Visualize perturbations of a single molecule given from a PerturbationEngine
+    """Visualize perturbations of a single molecule given from a PerturbationEngine
     in a ChemicalSpacePlot.
 
     Parameters:
@@ -756,14 +848,17 @@ class VisualizeMoleculePerturbations(ChemicalSpacePlot):
             dimensionality reduction"""
 
     @log_arguments
-    def __init__(self, smiles: str,
-            perturbation_engine: PerturbationEngine = None,
-            rep: BaseVecRepresentation = None,
-            idx: int = None,
-            n: int = None):
+    def __init__(
+        self,
+        smiles: str,
+        perturbation_engine: PerturbationEngine = None,
+        rep: BaseVecRepresentation = None,
+        idx: int = None,
+        n: int = None,
+    ):
         self.smiles = smiles
         if perturbation_engine is None:
-            self.perturbation_engine = SwapMutations(radius = 0)
+            self.perturbation_engine = SwapMutations(radius=0)
         else:
             self.perturbation_engine = perturbation_engine
         if rep is None:
@@ -783,20 +878,33 @@ def __init__(self, smiles: str,
         if idx is None:
             df["SMILES"] = self.perturbation_engine.get_compound_list(smiles) + [smiles]
         else:
-            df["SMILES"] = self.perturbation_engine.get_compound_list(smiles, idx = idx)+ [smiles]
+            df["SMILES"] = self.perturbation_engine.get_compound_list(
+                smiles, idx=idx
+            ) + [smiles]
         df["mols"] = [Chem.MolFromSmiles(s) for s in df["SMILES"]]
-        df = df.dropna(subset = ["mols"])
+        df = df.dropna(subset=["mols"])
 
-        fps = [AllChem.GetMorganFingerprintAsBitVect(m, 2, nBits=2048, useChirality=False) for m in df["mols"]]
+        fps = [
+            AllChem.GetMorganFingerprintAsBitVect(m, 2, nBits=2048, useChirality=False)
+            for m in df["mols"]
+        ]
+
+        df["sim"] = DataStructs.BulkTanimotoSimilarity(
+            AllChem.GetMorganFingerprintAsBitVect(
+                Chem.MolFromSmiles(self.smiles), 2, nBits=2048, useChirality=False
+            ),
+            fps,
+        )
 
-        df["sim"] = DataStructs.BulkTanimotoSimilarity(AllChem.GetMorganFingerprintAsBitVect(Chem.MolFromSmiles(self.smiles), 2, nBits=2048, useChirality=False) , fps)
+        super().__init__(
+            df,
+            self.rep,
+            title="Chemical Space Plot of Molecular Perturbations<br><sub>Points colored by tanimoto similarity to the reference compund</sub>",
+            color_col="sim",
+            smiles_col="SMILES",
+            colorscale="YlOrRd",
+        )
 
-        super().__init__(df,
-                self.rep,
-                title = "Chemical Space Plot of Molecular Perturbations<br><sub>Points colored by tanimoto similarity to the reference compund</sub>",
-                color_col = "sim",
-                smiles_col = "SMILES",
-                colorscale = "YlOrRd")
 
 class VisualizeDatasetSplit(ChemicalSpacePlot):
     """Visualize a dataset by seeing where train/test compounds are in a dimensionality
@@ -843,7 +951,7 @@ def __init__(
                 *args,
                 title="Dataset Visualization<br><sub>Red outline is train; Blue outline is test</sub>",
                 dim_reduction="tsne",
-                opacity = self.opacity,
+                opacity=self.opacity,
                 log=False,
                 **kwargs,
             )
@@ -854,7 +962,7 @@ def __init__(
                 *args,
                 title="Model Residual Visualization<br><sub>Red outline is train; Blue outline is test; Marker size is residual magnitude</sub>",
                 dim_reduction="tsne",
-                opacity = self.opacity,
+                opacity=self.opacity,
                 log=False,
                 **kwargs,
             )
@@ -863,11 +971,20 @@ def get_data(self) -> dict:
         d = super().get_data(SMILES=self.dataset.structure_col)
 
         if self.model is None:
-            d["color"] = ["red" if x == "train" else "blue" for x in self.dataset.entire_dataset_split]
+            d["color"] = [
+                "red" if x == "train" else "blue"
+                for x in self.dataset.entire_dataset_split
+            ]
         else:
-            d["outline"] = ["red" if x == "train" else "blue" for x in self.dataset.entire_dataset_split]
-
-            residuals = (self.dataset.entire_dataset[1] - self.model.predict(self.dataset.entire_dataset[0])) ** 2
+            d["outline"] = [
+                "red" if x == "train" else "blue"
+                for x in self.dataset.entire_dataset_split
+            ]
+
+            residuals = (
+                self.dataset.entire_dataset[1]
+                - self.model.predict(self.dataset.entire_dataset[0])
+            ) ** 2
             if self.res_lim is None:
                 self.res_lim = np.percentile(residuals, 90)
             residuals = residuals / self.res_lim
@@ -882,8 +999,16 @@ def get_data(self) -> dict:
     @staticmethod
     def get_attributes():
         return [
-            {"Dataset Selection": [["dataset", "datasetSelector"],]},
-            {"Configure Chemical Space": [["rep", "selector", ["oce", "rdkit2d-normed", "morgan"]],]},
+            {
+                "Dataset Selection": [
+                    ["dataset", "datasetSelector"],
+                ]
+            },
+            {
+                "Configure Chemical Space": [
+                    ["rep", "selector", ["oce", "rdkit2d-normed", "morgan"]],
+                ]
+            },
         ]
 
     @classmethod
@@ -900,8 +1025,10 @@ def from_attributes(cls, attributes: dict) -> "BaseVisualization":
 
         return super(ChemicalSpacePlot, cls).from_attributes(attributes)
 
+
 from rdkit.Chem import Draw
 
+
 class MorganContributions(BaseVisualization):
     """Visualize a morgan fingerprint bit contributions of a molecule.
 
@@ -922,26 +1049,26 @@ class MorganContributions(BaseVisualization):
             Default is 2.
         nbits(int, optional): Parameter for the Morgan representation.
             Default is 1024.
-        """
+    """
 
     @log_arguments
     def __init__(
-            self,
-            smiles: str,
-            dataset: BaseDataset,
-            compound_id: str = "The compound",
-            radius=2,
-            nbits=1024,
-            *args,
-            **kwargs
-        ):
+        self,
+        smiles: str,
+        dataset: BaseDataset,
+        compound_id: str = "The compound",
+        radius=2,
+        nbits=1024,
+        *args,
+        **kwargs,
+    ):
         self.radius = radius
         self.nbits = nbits
         self.dataset = dataset
         self.property = dataset.property_col
 
         self.smiles = smiles
-        self.compound_id = compound_id,
+        self.compound_id = (compound_id,)
         self.model = self._train_model()
         self.original_prediction, self.predictions = self._make_predictions(self.smiles)
         self.args = args
@@ -949,13 +1076,16 @@ def __init__(
         self.packages = ["plotly", "rdkit", "olorenrenderer"]
 
     def _train_model(self):
-        """ Train random forest model based on the morgan vec representation"""
-        model = oce.RandomForestModel(MorganVecRepresentation(scale=None, collinear_thresh=1.01), n_estimators=1000)
+        """Train random forest model based on the morgan vec representation"""
+        model = oce.RandomForestModel(
+            MorganVecRepresentation(scale=None, collinear_thresh=1.01),
+            n_estimators=1000,
+        )
         model.fit(*self.dataset.entire_dataset)
         return model
 
     def _make_predictions(self, smiles):
-        """ Make predictions on the dataset with systematic experiments to test
+        """Make predictions on the dataset with systematic experiments to test
         the effect of bit flips."""
         morgan_rep = self._to_morgan(smiles)
         original_prediction = self.model.predict([morgan_rep], skip_preprocess=True)
@@ -965,9 +1095,8 @@ def _make_predictions(self, smiles):
 
         return original_prediction, predictions
 
-
     def _to_morgan(self, smiles):
-        """ Return morgan representation of given compound (represented by
+        """Return morgan representation of given compound (represented by
         smiles)."""
         return MorganVecRepresentation._convert(self, smiles)
 
@@ -1043,8 +1172,12 @@ def visualize_morganfp(self):
         self.predictions = self._calculate_effect()
         bottom_3_bits, top_3_bits = self._top_3_predictions()
 
-        mfp_svgs_min = {str(i): Draw.DrawMorganBit(mol, i, info, useSVG=True) for i in bottom_3_bits}
-        mfp_svgs_max = {str(i): Draw.DrawMorganBit(mol, i, info, useSVG=True) for i in top_3_bits}
+        mfp_svgs_min = {
+            str(i): Draw.DrawMorganBit(mol, i, info, useSVG=True) for i in bottom_3_bits
+        }
+        mfp_svgs_max = {
+            str(i): Draw.DrawMorganBit(mol, i, info, useSVG=True) for i in top_3_bits
+        }
 
         return mfp_svgs_min, mfp_svgs_max
 
@@ -1070,7 +1203,10 @@ def get_data(self):
 
         bottom_3_min, bottom_3_max = self._filter(bottom_3_min, bottom_3_max)
         print(bottom_3_min, bottom_3_max)
-        bottom_bits, top_bits = bottom_3_bits[0:len(bottom_3_min)], top_3_bits[0:len(bottom_3_max)]
+        bottom_bits, top_bits = (
+            bottom_3_bits[0 : len(bottom_3_min)],
+            top_3_bits[0 : len(bottom_3_max)],
+        )
         substructures = self.get_substructures(bottom_bits, top_bits)
         print(substructures)
 
@@ -1123,13 +1259,17 @@ def get_substructures(self, bottom_3_bits, top_3_bits):
             atom_tuples = info[i]
             atomID = atom_tuples[0][0]
             radius = atom_tuples[0][1]
-            sub_structures[str(i)] = str(self.getSubstructSmi(mol, int(atomID), int(radius))[0])
+            sub_structures[str(i)] = str(
+                self.getSubstructSmi(mol, int(atomID), int(radius))[0]
+            )
 
         for i in top_3_bits:
             atom_tuples = info[i]
             atomID = atom_tuples[0][0]
             radius = atom_tuples[0][1]
-            sub_structures[str(i)] = str(self.getSubstructSmi(mol, int(atomID), int(radius))[0])
+            sub_structures[str(i)] = str(
+                self.getSubstructSmi(mol, int(atomID), int(radius))[0]
+            )
         return sub_structures
 
     # Functions for providing detailed descriptions of MFP bits from Nadine Schneider
@@ -1139,41 +1279,48 @@ def get_substructures(self, bottom_3_bits, top_3_bits):
     # The link to the source: http://rdkit.blogspot.com/2016/02/morgan-fingerprint-bit-statistics.html
 
     def _includeRingMembership(self, s, n):
-        r=';R]'
-        d="]"
-        return r.join([d.join(s.split(d)[:n]),d.join(s.split(d)[n:])])
+        r = ";R]"
+        d = "]"
+        return r.join([d.join(s.split(d)[:n]), d.join(s.split(d)[n:])])
 
     def _includeDegree(self, s, n, d):
-        r=';D'+str(d)+']'
-        d="]"
-        return r.join([d.join(s.split(d)[:n]),d.join(s.split(d)[n:])])
+        r = ";D" + str(d) + "]"
+        d = "]"
+        return r.join([d.join(s.split(d)[:n]), d.join(s.split(d)[n:])])
 
-    def _writePropsToSmiles(self, mol,smi,order):
-        #finalsmi = copy.deepcopy(smi)
+    def _writePropsToSmiles(self, mol, smi, order):
+        # finalsmi = copy.deepcopy(smi)
         finalsmi = smi
-        for i,a in enumerate(order):
+        for i, a in enumerate(order):
             atom = mol.GetAtomWithIdx(a)
             if atom.IsInRing():
-                finalsmi = self._includeRingMembership(finalsmi,i+1)
-            finalsmi = self._includeDegree(finalsmi, i+1, atom.GetDegree())
+                finalsmi = self._includeRingMembership(finalsmi, i + 1)
+            finalsmi = self._includeDegree(finalsmi, i + 1, atom.GetDegree())
         return finalsmi
 
-    def getSubstructSmi(self, mol,atomID,radius):
-        if radius>0:
-            env = Chem.FindAtomEnvironmentOfRadiusN(mol,radius,atomID)
-            atomsToUse=[]
+    def getSubstructSmi(self, mol, atomID, radius):
+        if radius > 0:
+            env = Chem.FindAtomEnvironmentOfRadiusN(mol, radius, atomID)
+            atomsToUse = []
             for b in env:
                 atomsToUse.append(mol.GetBondWithIdx(b).GetBeginAtomIdx())
                 atomsToUse.append(mol.GetBondWithIdx(b).GetEndAtomIdx())
             atomsToUse = list(set(atomsToUse))
         else:
             atomsToUse = [atomID]
-            env=None
-
-        smi = Chem.MolFragmentToSmiles(mol,atomsToUse,bondsToUse=env,allHsExplicit=True, allBondsExplicit=True, rootedAtAtom=atomID)
+            env = None
+
+        smi = Chem.MolFragmentToSmiles(
+            mol,
+            atomsToUse,
+            bondsToUse=env,
+            allHsExplicit=True,
+            allBondsExplicit=True,
+            rootedAtAtom=atomID,
+        )
         order = eval(mol.GetProp("_smilesAtomOutputOrder"))
-        smi2 = self._writePropsToSmiles(mol,smi,order)
-        return smi,smi2
+        smi2 = self._writePropsToSmiles(mol, smi, order)
+        return smi, smi2
 
     @staticmethod
     def get_attributes():
@@ -1195,6 +1342,7 @@ def JS_NAME(self) -> str:
         """
         return self.__class__.__name__
 
+
 class VisualizeADAN(CompoundScatterPlot):
     """Visualize a model trained on a dataset, by seeing predicted vs true
     colored by ADAN criteria.
@@ -1229,13 +1377,13 @@ def __init__(
         self.model = model
         self.ADAN = oce.ADAN(None)
         self.ADAN.build(
-            self.model, 
-            X=dataset.train_dataset[0], 
-            y=np.array(dataset.train_dataset[1]), 
-            rep=rep, 
+            self.model,
+            X=dataset.train_dataset[0],
+            y=np.array(dataset.train_dataset[1]),
+            rep=rep,
             dim_reduction=dim_reduction,
             explvar=explvar,
-            threshold=threshold
+            threshold=threshold,
         )
         self.ADAN.calculate_full(dataset.test_dataset[0])
 
@@ -1270,12 +1418,15 @@ def get_data(self, criterion="B") -> dict:
                 f"ADAN Criterion {criterion} correlation with squared error is {r2_score(self.adan.results[criterion], self.df['Z'])}"
             )
             import matplotlib.pyplot as plt
+
             plt.scatter(self.adan.results[criterion], self.df["Z"])
             plt.show()
         else:
             print(f"ADAN Criterion {criterion}; RMSE by class")
             for i in sorted(self.adan.results[criterion].unique()):
-                print(f"Class {i}: RMSE {np.sqrt(np.mean(self.df[self.adan.results[criterion] == i]['Z']))}")
+                print(
+                    f"Class {i}: RMSE {np.sqrt(np.mean(self.df[self.adan.results[criterion] == i]['Z']))}"
+                )
         return super().get_data()
 
     @staticmethod
@@ -1292,7 +1443,7 @@ def get_attributes():
 
 
 class VisualizeModelSim(CompoundScatterPlot):
-    """ Visualize a model's predicted vs true plot on given dataset, where each
+    """Visualize a model's predicted vs true plot on given dataset, where each
     point is colored by a compounds similarity to the train set
 
     Parameters:
@@ -1303,7 +1454,15 @@ class VisualizeModelSim(CompoundScatterPlot):
             or 'valid'."""
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", *args, log=True, **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        eval_set="test",
+        *args,
+        log=True,
+        **kwargs,
+    ):
         if not type(dataset) is BaseDataset:
             raise TypeError(
                 "dataset must be a BaseDataset. If you're using a Pandas Dataframe consider using df.to_csv() in the"
@@ -1312,11 +1471,14 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", *a
         self.dataset = dataset
         self.model = model
 
-        from rdkit.Chem import AllChem
         from rdkit import DataStructs
+        from rdkit.Chem import AllChem
 
         fingerprinter = lambda s: AllChem.GetMorganFingerprint(Chem.MolFromSmiles(s), 4)
-        train_fps = [fingerprinter(s) for s in self.dataset.train_dataset[0][self.dataset.structure_col]]
+        train_fps = [
+            fingerprinter(s)
+            for s in self.dataset.train_dataset[0][self.dataset.structure_col]
+        ]
 
         if eval_set == "test":
             self.eval_set = self.dataset.test_dataset
@@ -1325,8 +1487,12 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", *a
         elif eval_set == "valid":
             self.eval_set = self.dataset.valid_dataset
 
-        eval_fps = [fingerprinter(s) for s in self.eval_set[0][self.dataset.structure_col]]
-        self.sim = np.array([max(DataStructs.BulkTanimotoSimilarity(fp, train_fps)) for fp in eval_fps])
+        eval_fps = [
+            fingerprinter(s) for s in self.eval_set[0][self.dataset.structure_col]
+        ]
+        self.sim = np.array(
+            [max(DataStructs.BulkTanimotoSimilarity(fp, train_fps)) for fp in eval_fps]
+        )
 
         self.df = pd.DataFrame(
             {
@@ -1350,15 +1516,30 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", *a
     def get_data(self) -> dict:
         d = super().get_data()
         d["color"] = self.sim.tolist()
-        d["trace_update"] = {'x':[0.9*min(self.df["X"].min(), self.df["Y"].min()), 1.1*max(self.df["X"].max(), self.df["Y"].max())],
-            'y':[0.9*min(self.df["X"].min(), self.df["Y"].min()), 1.1*max(self.df["X"].max(), self.df["Y"].max())],
-            'mode': 'line',
-            'type': 'scatter'}
+        d["trace_update"] = {
+            "x": [
+                0.9 * min(self.df["X"].min(), self.df["Y"].min()),
+                1.1 * max(self.df["X"].max(), self.df["Y"].max()),
+            ],
+            "y": [
+                0.9 * min(self.df["X"].min(), self.df["Y"].min()),
+                1.1 * max(self.df["X"].max(), self.df["Y"].max()),
+            ],
+            "mode": "line",
+            "type": "scatter",
+        }
         return d
 
     @staticmethod
     def get_attributes():
-        return [{"Select a Model/Dataset": [["model", "modelSelector"], ["dataset", "datasetSelector"],]}]
+        return [
+            {
+                "Select a Model/Dataset": [
+                    ["model", "modelSelector"],
+                    ["dataset", "datasetSelector"],
+                ]
+            }
+        ]
 
 
 class VisualizeCounterfactual(CompoundScatterPlot):
@@ -1395,17 +1576,24 @@ def __init__(
         if type(delta) in (float, int):
             delta = (-delta, delta)
         self.delta = delta
-        self.cf_engine = CounterfactualEngine(model=self.model, perturbation_engine=self.perturbation_engine)
+        self.cf_engine = CounterfactualEngine(
+            model=self.model, perturbation_engine=self.perturbation_engine
+        )
         self.cf_engine.generate_samples(self.smiles)
 
         self.cf_engine.generate_cfs(delta=self.delta, n=n // 2)
         self.base = self.cf_engine.samples[0]["Value"]
 
         if self.model.setting == "classification":
-            self.factuals = self.cf_engine._select_cfs(lambda s: s["Value"] == self.base, n // 2)
+            self.factuals = self.cf_engine._select_cfs(
+                lambda s: s["Value"] == self.base, n // 2
+            )
         else:
             self.factuals = self.cf_engine._select_cfs(
-                lambda s: s["Value"] - self.delta[0] > self.base > s["Value"] - self.delta[1], n // 2
+                lambda s: s["Value"] - self.delta[0]
+                > self.base
+                > s["Value"] - self.delta[1],
+                n // 2,
             )
 
         samples = self.cf_engine.cfs + self.factuals
@@ -1446,8 +1634,16 @@ def get_data(self) -> dict:
             get_cf_type = lambda s: "Counterfactual"
         else:
             get_color = lambda s: "#EF553B" if s["Value"] > self.base else "#636EFA"
-            get_cf_type = lambda s: "High Counterfactual" if s["Value"] > self.base else "Low Counterfactual"
-        self.df["color"] = ["#00CC96"] + [get_color(s) for s in self.cf_engine.cfs[1:]] + ["#FECB52"] * len(self.factuals)
+            get_cf_type = (
+                lambda s: "High Counterfactual"
+                if s["Value"] > self.base
+                else "Low Counterfactual"
+            )
+        self.df["color"] = (
+            ["#00CC96"]
+            + [get_color(s) for s in self.cf_engine.cfs[1:]]
+            + ["#FECB52"] * len(self.factuals)
+        )
 
         format_label = lambda s: "Similarity = {}<br>f(x) = {}".format(
             np.round(s["Similarity"], 2), np.round(s["Value"], 3)
@@ -1457,12 +1653,17 @@ def get_data(self) -> dict:
         f_template = [format_label(s) for s in self.factuals]
         self.df["hovertemplate"] = base_template + cf_template + f_template
 
-        self.df["group"] = ["Base"] + [get_cf_type(s) for s in self.cf_engine.cfs[1:]] + ["Factual"] * len(self.factuals)
+        self.df["group"] = (
+            ["Base"]
+            + [get_cf_type(s) for s in self.cf_engine.cfs[1:]]
+            + ["Factual"] * len(self.factuals)
+        )
         self.df["size"] = 12
         return super().get_data()
 
+
 class VisualizeModelSim2(CompoundScatterPlot):
-    """ Visualize the connection between a model's error on a given compound and
+    """Visualize the connection between a model's error on a given compound and
     specific variables.
 
     Parameters:
@@ -1476,7 +1677,16 @@ class VisualizeModelSim2(CompoundScatterPlot):
             is 'sim'."""
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", var_name = "sim", *args, log=True, **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        eval_set="test",
+        var_name="sim",
+        *args,
+        log=True,
+        **kwargs,
+    ):
 
         if not type(dataset) is BaseDataset:
             raise TypeError(
@@ -1486,11 +1696,14 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", va
         self.dataset = dataset
         self.model = model
 
-        from rdkit.Chem import AllChem
         from rdkit import DataStructs
+        from rdkit.Chem import AllChem
 
         fingerprinter = lambda s: AllChem.GetMorganFingerprint(Chem.MolFromSmiles(s), 4)
-        train_fps = [fingerprinter(s) for s in self.dataset.train_dataset[0][self.dataset.structure_col]]
+        train_fps = [
+            fingerprinter(s)
+            for s in self.dataset.train_dataset[0][self.dataset.structure_col]
+        ]
 
         if eval_set == "test":
             self.eval_set = self.dataset.test_dataset
@@ -1499,8 +1712,12 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, eval_set = "test", va
         elif eval_set == "valid":
             self.eval_set = self.dataset.valid_dataset
 
-        eval_fps = [fingerprinter(s) for s in self.eval_set[0][self.dataset.structure_col]]
-        sim = np.array([max(DataStructs.BulkTanimotoSimilarity(fp, train_fps)) for fp in eval_fps])
+        eval_fps = [
+            fingerprinter(s) for s in self.eval_set[0][self.dataset.structure_col]
+        ]
+        sim = np.array(
+            [max(DataStructs.BulkTanimotoSimilarity(fp, train_fps)) for fp in eval_fps]
+        )
 
         preds = self.model.predict(self.eval_set[0])
         error = preds - self.eval_set[1]
@@ -1538,10 +1755,18 @@ def get_data(self) -> dict:
 
     @staticmethod
     def get_attributes():
-        return [{"Select a Model/Dataset": [["model", "modelSelector"], ["dataset", "datasetSelector"],]}]
+        return [
+            {
+                "Select a Model/Dataset": [
+                    ["model", "modelSelector"],
+                    ["dataset", "datasetSelector"],
+                ]
+            }
+        ]
+
 
 class ModelROC(BaseVisualization):
-    """ Visualize the ROC curve for a model.
+    """Visualize the ROC curve for a model.
 
     Parameters:
         dataset (Dataset): Dataset to evaluate model on
@@ -1551,8 +1776,16 @@ class ModelROC(BaseVisualization):
             or 'valid'. Default is 'test'."""
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test",
-        log=True, model_name=None, **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        *args,
+        eval_set="test",
+        log=True,
+        model_name=None,
+        **kwargs,
+    ):
 
         if not type(dataset) is BaseDataset:
             raise TypeError(
@@ -1581,7 +1814,12 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test
 
         fpr, tpr, thresholds = metrics.roc_curve(y_true, preds)
         ix = np.sort(np.random.choice(np.arange(len(fpr)), size=1000, replace=True))
-        self.df = pd.DataFrame({"X": fpr[ix], "Y": tpr[ix],})
+        self.df = pd.DataFrame(
+            {
+                "X": fpr[ix],
+                "Y": tpr[ix],
+            }
+        )
 
         self.score = metrics.roc_auc_score(y_true, preds)
 
@@ -1594,7 +1832,14 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test
 
     @staticmethod
     def get_attributes():
-        return [{"Select a Model/Dataset": [["model", "modelSelector"], ["dataset", "datasetSelector"],]}]
+        return [
+            {
+                "Select a Model/Dataset": [
+                    ["model", "modelSelector"],
+                    ["dataset", "datasetSelector"],
+                ]
+            }
+        ]
 
     @property
     def JS_NAME(self) -> str:
@@ -1606,8 +1851,9 @@ def get_data(self) -> dict:
         d["score"] = np.around(self.score, decimals=2)
         return d
 
+
 class ModelROCThreshold(BaseVisualization):
-    """ Visualize the ROC curve for a model.
+    """Visualize the ROC curve for a model.
 
     Parameters:
         dataset (Dataset): Dataset to evaluate model on
@@ -1617,8 +1863,16 @@ class ModelROCThreshold(BaseVisualization):
             or 'valid'. Default is 'test'."""
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test",
-        log=True, model_name=None, **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        *args,
+        eval_set="test",
+        log=True,
+        model_name=None,
+        **kwargs,
+    ):
 
         if not type(dataset) is BaseDataset:
             raise TypeError(
@@ -1647,8 +1901,9 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test
 
         fpr, tpr, thresholds = metrics.roc_curve(y_true, preds)
         ix = np.sort(np.random.choice(np.arange(len(fpr)), size=1000, replace=True))
-        self.df = pd.DataFrame({"X": fpr[ix], "Y": tpr[ix],
-            "thresholds": thresholds[ix]},)
+        self.df = pd.DataFrame(
+            {"X": fpr[ix], "Y": tpr[ix], "thresholds": thresholds[ix]},
+        )
 
         self.score = metrics.roc_auc_score(y_true, preds)
 
@@ -1661,22 +1916,32 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, eval_set="test
 
     @staticmethod
     def get_attributes():
-        return [{"Select a Model/Dataset": [["model", "modelSelector"], ["dataset", "datasetSelector"],]}]
+        return [
+            {
+                "Select a Model/Dataset": [
+                    ["model", "modelSelector"],
+                    ["dataset", "datasetSelector"],
+                ]
+            }
+        ]
 
     @property
     def JS_NAME(self) -> str:
-            return "ModelROCThreshold"
+        return "ModelROCThreshold"
 
     def get_data(self) -> dict:
         d = self.df.to_dict("l")
         d["model_name"] = self.model_name
         d["score"] = np.around(self.score, decimals=2)
-        d["P"] = (sum(self.dataset.test_dataset[1]))/len(self.dataset.test_dataset[1])
-        d["N"] = (len(self.dataset.test_dataset[1]) - sum(self.dataset.test_dataset[1]))/len(self.dataset.test_dataset[1])
+        d["P"] = (sum(self.dataset.test_dataset[1])) / len(self.dataset.test_dataset[1])
+        d["N"] = (
+            len(self.dataset.test_dataset[1]) - sum(self.dataset.test_dataset[1])
+        ) / len(self.dataset.test_dataset[1])
         return d
 
+
 class ModelPR(BaseVisualization):
-    """ Visualize the Precision-Recall curve for a model.
+    """Visualize the Precision-Recall curve for a model.
 
     Parameters:
         dataset (Dataset): Dataset to evaluate model on
@@ -1686,7 +1951,15 @@ class ModelPR(BaseVisualization):
             or 'valid'. Default is 'test'."""
 
     @log_arguments
-    def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True, model_name=None, **kwargs):
+    def __init__(
+        self,
+        dataset: BaseDataset,
+        model: BaseModel,
+        *args,
+        log=True,
+        model_name=None,
+        **kwargs,
+    ):
         self.dataset = dataset
         self.model = model
 
@@ -1698,11 +1971,20 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True, mode
         if isinstance(y_true, pd.Series):
             y_true = y_true.tolist()
         precision, recall, thresholds = metrics.precision_recall_curve(y_true, preds)
-        ix = np.sort(np.random.choice(np.arange(len(precision)), size=1000, replace=True))
-        self.df = pd.DataFrame({"X": recall[ix], "Y": precision[ix],})
+        ix = np.sort(
+            np.random.choice(np.arange(len(precision)), size=1000, replace=True)
+        )
+        self.df = pd.DataFrame(
+            {
+                "X": recall[ix],
+                "Y": precision[ix],
+            }
+        )
 
         self.score = metrics.auc(recall, precision)
-        self.baseline = sum(self.dataset.test_dataset[1])/len(self.dataset.test_dataset[1])
+        self.baseline = sum(self.dataset.test_dataset[1]) / len(
+            self.dataset.test_dataset[1]
+        )
 
         if model_name is None:
             self.model_name = model_name_from_params(oce.parameterize(model))
@@ -1713,7 +1995,14 @@ def __init__(self, dataset: BaseDataset, model: BaseModel, *args, log=True, mode
 
     @staticmethod
     def get_attributes():
-        return [{"Select a Model/Dataset": [["model", "modelSelector"], ["dataset", "datasetSelector"],]}]
+        return [
+            {
+                "Select a Model/Dataset": [
+                    ["model", "modelSelector"],
+                    ["dataset", "datasetSelector"],
+                ]
+            }
+        ]
 
     @property
     def JS_NAME(self) -> str:
@@ -1726,25 +2015,27 @@ def get_data(self) -> dict:
         d["baseline"] = self.baseline
         return d
 
+
 class BaseErrorWaterfall(BaseVisualization):
-    """ Visualize the error waterfall for a base boosting model.
+    """Visualize the error waterfall for a base boosting model.
 
-   Args:
-        model (BaseBoosting): Model to evaluate on. must be base boosting model.
-        x_data (Union[pd.DataFrame, np.ndarray]): Data to predict on using the model
-        y_data (Union[pd.Series, list, np.ndarray], optional): True values to compare to. Defaults to None. If None, then the waterfall plot will be for residuals.
-        normalization (bool, optional): If the data is normalized. Defaults to False.
+    Args:
+         model (BaseBoosting): Model to evaluate on. must be base boosting model.
+         x_data (Union[pd.DataFrame, np.ndarray]): Data to predict on using the model
+         y_data (Union[pd.Series, list, np.ndarray], optional): True values to compare to. Defaults to None. If None, then the waterfall plot will be for residuals.
+         normalization (bool, optional): If the data is normalized. Defaults to False.
     """
 
     @log_arguments
     def __init__(
-        self, 
-        model: BaseBoosting, 
-        x_data: Union[pd.DataFrame, np.ndarray], 
-        *y_data: Union[pd.Series, list, np.ndarray], 
-        normalization = False, 
-        log=True, 
-        **kwargs):
+        self,
+        model: BaseBoosting,
+        x_data: Union[pd.DataFrame, np.ndarray],
+        *y_data: Union[pd.Series, list, np.ndarray],
+        normalization=False,
+        log=True,
+        **kwargs,
+    ):
 
         self.model = model
         self.x_data = x_data
@@ -1755,7 +2046,7 @@ def __init__(
         self.normalization = normalization
         super().__init__(log=False)
         self.packages += ["plotly"]
-        
+
         self.df = pd.DataFrame({})
 
     def get_data(self) -> dict:
@@ -1763,9 +2054,11 @@ def get_data(self) -> dict:
 
         Returns:
             dict: Data for visualization."""
-        predictions = self.model.predict(self.x_data, waterfall = True, normalize = self.normalization)
-        data = self.model._waterfall(y_data = self.y_data, normalize = self.normalization)        
- 
+        predictions = self.model.predict(
+            self.x_data, waterfall=True, normalize=self.normalization
+        )
+        data = self.model._waterfall(y_data=self.y_data, normalize=self.normalization)
+
         diffs = [j - i for i, j in zip(data, data[1:])]
         diffs.insert(0, data[0])
         diffs.append(0)
@@ -1774,19 +2067,19 @@ def get_data(self) -> dict:
             model_names = ["Model " + str(i) for i in range(1, len(data))]
             model_names.insert(0, "Dataset Baseline")
         else:
-            model_names = ["Model " + str(i+1) for i in range(1, len(data))]
+            model_names = ["Model " + str(i + 1) for i in range(1, len(data))]
             model_names.insert(0, "Model 1 Baseline")
         model_names.append("Boosted Model")
 
-        self.df['diffs'] = diffs
-        self.df['model_names'] = model_names
+        self.df["diffs"] = diffs
+        self.df["model_names"] = model_names
 
-        text = [f'{val:.2f}' for val in diffs]
-        text[-1] = f'{data[-1]:.2f}'
-        self.df['text'] = text
+        text = [f"{val:.2f}" for val in diffs]
+        text[-1] = f"{data[-1]:.2f}"
+        self.df["text"] = text
 
-        w_type = ['relative' for val in diffs]
-        w_type[-1] = 'total'
-        self.df['w_type'] = w_type
+        w_type = ["relative" for val in diffs]
+        w_type[-1] = "total"
+        self.df["w_type"] = w_type
 
-        return self.df.to_dict("l")
\ No newline at end of file
+        return self.df.to_dict("l")
diff --git a/src/olorenchemengine/visualizations/visualize_interpret.py b/src/olorenchemengine/visualizations/visualize_interpret.py
index baeefe22..e7fdb0ea 100644
--- a/src/olorenchemengine/visualizations/visualize_interpret.py
+++ b/src/olorenchemengine/visualizations/visualize_interpret.py
@@ -3,17 +3,18 @@
 certain prediction. 
 """
 
+from rdkit import Chem
+from tqdm import tqdm
+
 from olorenchemengine.base_class import *
-from olorenchemengine.representations import *
 from olorenchemengine.dataset import *
-from olorenchemengine.uncertainty import *
-from olorenchemengine.interpret import *
 from olorenchemengine.internal import *
+from olorenchemengine.interpret import *
 from olorenchemengine.manager import *
+from olorenchemengine.representations import *
+from olorenchemengine.uncertainty import *
 from olorenchemengine.visualizations import *
 
-from tqdm import tqdm
-from rdkit import Chem
 
 class VisualizePredictionSensitivity(BaseVisualization):
     """
@@ -28,14 +29,24 @@ class VisualizePredictionSensitivity(BaseVisualization):
     """
 
     @log_arguments
-    def __init__(self, model: BaseModel, query_compound: str, radius: int = 2,
-        n: int = 200, colorscale = "viridis", bottom_quantile = 0.75,
-        top_quantile = 0.95, nbins = 3, log=True, **kwargs):
+    def __init__(
+        self,
+        model: BaseModel,
+        query_compound: str,
+        radius: int = 2,
+        n: int = 200,
+        colorscale="viridis",
+        bottom_quantile=0.75,
+        top_quantile=0.95,
+        nbins=3,
+        log=True,
+        **kwargs
+    ):
 
         super().__init__(log=False, **kwargs)
         self.packages = ["olorenrenderer"]
 
-        self.mutator = SwapMutations(radius = radius)
+        self.mutator = SwapMutations(radius=radius)
         self.model = model
         self.colorscale = colorscale
         self.bottom_quantile = bottom_quantile
@@ -49,7 +60,9 @@ def __init__(self, model: BaseModel, query_compound: str, radius: int = 2,
         for i, a in tqdm(enumerate(self.mol.GetAtoms())):
             smiles_list = []
             for i in range(n):
-                smiles = self.mutator.get_compound_at_idx(Chem.Mol(self.mol), a.GetIdx())
+                smiles = self.mutator.get_compound_at_idx(
+                    Chem.Mol(self.mol), a.GetIdx()
+                )
                 if smiles is None or Chem.MolFromSmiles(smiles) is None:
                     continue
                 else:
@@ -61,7 +74,7 @@ def __init__(self, model: BaseModel, query_compound: str, radius: int = 2,
                 vals.append(stdev)
             else:
                 print("Not enough perturbations for atom", a.GetIdx())
-        
+
         bottom_threshold = np.quantile(vals, self.bottom_quantile)
         top_threshold = np.quantile(vals, self.top_quantile)
 
@@ -74,25 +87,27 @@ def __init__(self, model: BaseModel, query_compound: str, radius: int = 2,
                     a.SetDoubleProp("bin", 1)
                     a.SetAtomMapNum(self.nbins)
                 else:
-                    normalized_val = (val - bottom_threshold) / (top_threshold - bottom_threshold)
+                    normalized_val = (val - bottom_threshold) / (
+                        top_threshold - bottom_threshold
+                    )
                     bin_number = np.around(normalized_val * self.nbins)
                     if int(bin_number) > 0:
                         a.SetAtomMapNum(int(bin_number))
 
     def get_data(self):
-        
+
         from . import sample_colorscale
 
-        def rgb_to_hex(x, colorscale = self.colorscale):
+        def rgb_to_hex(x, colorscale=self.colorscale):
             if not isinstance(x, list):
                 x = [x]
             x = sample_colorscale(colorscale, x, colortype="hex")
-            x = np.rint(np.array(x)*255).astype(int)
-            return ['#%02x%02x%02x' % (x_[0], x_[1], x_[2]) for x_ in x]
+            x = np.rint(np.array(x) * 255).astype(int)
+            return ["#%02x%02x%02x" % (x_[0], x_[1], x_[2]) for x_ in x]
 
         return {
             "SMILES": Chem.MolToSmiles(self.mol),
             "highlights": [
-                [i+1, rgb_to_hex((i+1)/self.nbins)[0]]  for i in range(self.nbins)
-            ]
+                [i + 1, rgb_to_hex((i + 1) / self.nbins)[0]] for i in range(self.nbins)
+            ],
         }


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index ad5e53ea..5867fd54 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -8,3 +8,30 @@
 """
 
 # import pytest
+
+import pytest
+
+import olorenchemengine as oce
+
+# def pytest_configure(config):
+#     config.addinivalue_line(
+#         "markers",
+#         "no_profiling: mark test to not use sql profiling"
+#     )
+
+# @pytest.fixture(scope="session", autouse=True)
+# def cleanup(request):
+#     """Cleanup a testing directory once we are finished."""
+#     remote = oce.Remote("http://api.oloren.ai:5000")
+#     print("Opening remote")
+#     remote.__enter__()
+#     def close_remote():
+#         print("Closing remote")
+#         remote.__exit__(None, None, None)
+#     request.addfinalizer(close_remote)
+
+
+@pytest.fixture(autouse=True)
+def run_around_tests():
+    with oce.Remote("http://api.oloren.ai:5000") as remote:
+        yield
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
index 00000000..d631def2
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,3 @@
+pytest-cov
+pytest
+pytest-timeout
\ No newline at end of file
diff --git a/tests/test_backward_compatibility.py b/tests/test_backward_compatibility.py
index 1808052c..22d36f8e 100644
--- a/tests/test_backward_compatibility.py
+++ b/tests/test_backward_compatibility.py
@@ -2,53 +2,82 @@
 from os.path import expanduser
 
 import numpy as np
+
 import olorenchemengine as oce
 
 
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 def get_model_for_testing(name):
     path = oce.download_public_file(f"test_models/{name}.oam")
     return oce.load(path)
 
+
 test_molecules = ["O=C(C)Oc1ccccc1C(=O)O", "CC(CC(c1ccccc1)C2=C(c3c(OC2=O)cccc3)O)=O"]
 
+
 def test_reg_0():
     model = get_model_for_testing("reg model 0")
-    assert np.allclose(model.predict(test_molecules), [ 0.5017606, -0.8976576], atol=1e-3)
+    assert np.allclose(
+        model.predict(test_molecules), [0.5017606, -0.8976576], atol=1e-3
+    )
+
 
 def test_reg_1():
     model = get_model_for_testing("reg model 1")
-    assert np.allclose(model.predict(test_molecules), [ 0.25767391, -0.15310819], atol=1e-3)
+    assert np.allclose(
+        model.predict(test_molecules), [0.25767391, -0.15310819], atol=1e-3
+    )
+
 
 def test_reg_2():
     model = get_model_for_testing("reg model 2")
-    assert np.allclose(model.predict(test_molecules), [ 0.14640912, -0.19515301], atol=1e-3)
+    assert np.allclose(
+        model.predict(test_molecules), [0.14640912, -0.19515301], atol=1e-3
+    )
+
 
 def test_reg_3():
     model = get_model_for_testing("reg model 3")
     print(model.predict(test_molecules))
-    assert np.allclose(model.predict(test_molecules), [0.5101875,  0.27005324], atol=1e-3)
+    assert np.allclose(
+        model.predict(test_molecules), [0.5101875, 0.27005324], atol=1e-3
+    )
+
 
 def test_reg_4():
     model = get_model_for_testing("reg model 4")
-    assert np.allclose(model.predict(test_molecules), [0.69562391, 0.1398656 ], atol=1e-3)
+    assert np.allclose(
+        model.predict(test_molecules), [0.69562391, 0.1398656], atol=1e-3
+    )
+
 
 def test_class_0():
     model = get_model_for_testing("class model 0")
     assert np.allclose(model.predict(test_molecules), [0.9907742, 0.9907742])
 
+
 def test_class_1():
     model = get_model_for_testing("class model 1")
-    assert np.allclose(model.predict(test_molecules), [1., 1.])
+    assert np.allclose(model.predict(test_molecules), [1.0, 1.0])
+
 
 def test_class_2():
     model = get_model_for_testing("class model 2")
-    assert np.allclose(model.predict(test_molecules), [1., 1.])
+    assert np.allclose(model.predict(test_molecules), [1.0, 1.0])
+
 
 def test_class_3():
     model = get_model_for_testing("class model 3")
-    assert np.allclose(model.predict(test_molecules), [0.580161,   0.56691015])
+    assert np.allclose(model.predict(test_molecules), [0.580161, 0.56691015])
+
 
 def test_class_4():
     model = get_model_for_testing("class model 4")
-    assert np.allclose(model.predict(test_molecules), [1., 1.])
-
+    assert np.allclose(model.predict(test_molecules), [1.0, 1.0])
diff --git a/tests/test_data_manage.py b/tests/test_data_manage.py
index 1db0f305..2039e46d 100644
--- a/tests/test_data_manage.py
+++ b/tests/test_data_manage.py
@@ -1,26 +1,33 @@
-import pytest
-import olorenchemengine as oce
-import numpy as np
-from olorenchemengine.interpret import *
+import os
 
+import numpy as np
 import pandas as pd
-from sklearn.model_selection import train_test_split
-from olorenchemengine.internal import download_public_file
-
+import pytest
 from rdkit import Chem
+from sklearn.model_selection import train_test_split
 
-import os
+import olorenchemengine as oce
+from olorenchemengine.internal import download_public_file
+from olorenchemengine.interpret import *
 
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
 
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 def test_datasetfromcsv():
     dataset = oce.DatasetFromCSV(
         file_path=download_public_file("sample-csvs/sample_data1.csv"),
         structure_col="Smiles",
         property_col="pChEMBL Value",
-    ) + oce.RandomSplit(split_proportions=[0.8,0.0,2])
+    ) + oce.RandomSplit(split_proportions=[0.8, 0.0, 2])
     x_train, y_train = dataset.train_dataset
     x_test, y_test = dataset.test_dataset
     x_val, y_val = dataset.valid_dataset
@@ -34,12 +41,20 @@ def test_manage():
         file_path=download_public_file("sample-csvs/sample_data1.csv"),
         structure_col="Smiles",
         property_col="pChEMBL Value",
-    ) + oce.RandomSplit(split_proportions=[0.8,0.1,0.1])
+    ) + oce.RandomSplit(split_proportions=[0.8, 0.1, 0.1])
     if os.path.exists("tmp.oce"):
         os.remove("tmp.oce")
-    manager = oce.ModelManager(dataset, ["Root Mean Squared Error"], file_path="tmp.oce")
-    manager.run(oce.RandomForestModel(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10))
-    manager.run(oce.MLP(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10))
+    manager = oce.ModelManager(
+        dataset, ["Root Mean Squared Error"], file_path="tmp.oce"
+    )
+    manager.run(
+        oce.RandomForestModel(
+            oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10
+        )
+    )
+    manager.run(
+        oce.MLP(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10)
+    )
     manager2 = oce.load("tmp.oce")
 
     error1 = manager.get_model_database()["Root Mean Squared Error"].to_numpy()
@@ -55,8 +70,15 @@ def test_manage():
 
 
 def test_sheets_manager():
-    if "GOOGLE_CREDENTIALS_FILENAME" not in oce.CONFIG and "GOOGLE_SERVICE_ACC_FILENAME" not in oce.CONFIG:
-        warnings.warn(UserWarning("Skipping sheets manager test due to lack of supplied credentials."))
+    if (
+        "GOOGLE_CREDENTIALS_FILENAME" not in oce.CONFIG
+        and "GOOGLE_SERVICE_ACC_FILENAME" not in oce.CONFIG
+    ):
+        warnings.warn(
+            UserWarning(
+                "Skipping sheets manager test due to lack of supplied credentials."
+            )
+        )
         return
 
     dataset = oce.DatasetFromCSV(
@@ -66,9 +88,17 @@ def test_sheets_manager():
     )
     if os.path.exists("tmp.oce"):
         os.remove("tmp.oce")
-    manager = oce.SheetsModelManager(dataset, ["Root Mean Squared Error"], file_path="tmp.oce")
-    manager.run(oce.RandomForestModel(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10))
-    manager.run(oce.MLP(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10))
+    manager = oce.SheetsModelManager(
+        dataset, ["Root Mean Squared Error"], file_path="tmp.oce"
+    )
+    manager.run(
+        oce.RandomForestModel(
+            oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10
+        )
+    )
+    manager.run(
+        oce.MLP(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=10)
+    )
     manager2 = oce.load("tmp.oce")
 
     error1 = manager.get_model_database()["Root Mean Squared Error"].to_numpy()
diff --git a/tests/test_framework.py b/tests/test_framework.py
index 7880e3fd..afdd3030 100644
--- a/tests/test_framework.py
+++ b/tests/test_framework.py
@@ -3,13 +3,23 @@
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
+
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 def test_main():
     oce.test_oce()
 
+
 def test_config():
     assert isinstance(oce.CONFIG, dict)
     oce.set_config_param("TESTING PARAM", "1")
     assert oce.CONFIG["TESTING PARAM"] == "1"
     oce.test_oce()
     oce.remove_config_param("TESTING PARAM")
-    assert "TESTING PARAM" not in oce.CONFIG
\ No newline at end of file
+    assert "TESTING PARAM" not in oce.CONFIG
diff --git a/tests/test_interpret.py b/tests/test_interpret.py
index daa3529b..a2e11964 100644
--- a/tests/test_interpret.py
+++ b/tests/test_interpret.py
@@ -1,16 +1,24 @@
-import pytest
-import olorenchemengine as oce
 import numpy as np
-from olorenchemengine.interpret import *
-
 import pandas as pd
+import pytest
+from rdkit import Chem
 from sklearn.model_selection import train_test_split
 
-from rdkit import Chem
+import olorenchemengine as oce
+from olorenchemengine.interpret import *
 
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
+
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 def test_swap_mutations():
     s = "CC1([C@@H](N2[C@H](S1)[C@@H](C2=O)NC(=O)Cc3ccccc3)C(=O)O)C"
     sm = SwapMutations(radius=1)
@@ -18,4 +26,4 @@ def test_swap_mutations():
 
     assert not s_out == s
     assert not s_out is None
-    assert isinstance(s_out, str)
\ No newline at end of file
+    assert isinstance(s_out, str)
diff --git a/tests/test_known_errors.py b/tests/test_known_errors.py
index 594ede7f..484da5d1 100644
--- a/tests/test_known_errors.py
+++ b/tests/test_known_errors.py
@@ -1,41 +1,66 @@
+import numpy as np
+import pandas as pd
 import pytest
+
 import olorenchemengine as oce
-import pandas as pd
-import numpy as np
 from olorenchemengine.internal import download_public_file
 
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
+
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 @pytest.fixture
 def example_data1():
     file_path = download_public_file("sample-csvs/sample_data1.csv")
     df = pd.read_csv(file_path)
     return df
 
-#RuntimeError: Not compiled with CUDA support
+
+# RuntimeError: Not compiled with CUDA support
+
+
 def test_TorchGeometric(example_data1):
-    data = oce.RandomSplit(split_proportions=[0.8 ,0.0, 0.2]).split(example_data1)
+    data = oce.RandomSplit(split_proportions=[0.8, 0.0, 0.2]).split(example_data1)
     train, valid, test = data[0], data[1], data[2]
     s = "{'BC_class_name': 'BaseTorchGeometricModel', 'args': [{'BC_class_name': 'GINModel', 'args': [], 'kwargs': {'batch_size': 100, 'conv_radius': 3, 'conv_type': 'gin+', 'dataset': 'molpcba', 'dropout': 0.5, 'hidden': 100, 'layers': 3, 'lr': 0.001, 'optim': 'adam', 'task_type': 'classification', 'virtual_node': False}}], 'kwargs': {'auto_lr_find': True, 'batch_size': 16, 'epochs': 5, 'log': True, 'lr': 0.0001, 'pos_weight': 'balanced', 'preinitialized': False, 'representation': {'BC_class_name': 'TorchGeometricGraph', 'args': [], 'kwargs': {'atom_featurizer': {'BC_class_name': 'OGBAtomFeaturizer', 'args': [], 'kwargs': {}}, 'bond_featurizer': {'BC_class_name': 'OGBBondFeaturizer', 'args': [], 'kwargs': {}}}}}}"
-    s = s.replace("'", "\"").replace("True", "true").replace("False", "false").replace("None", "null")
-    model=(oce.create_BC(s))
+    s = (
+        s.replace("'", '"')
+        .replace("True", "true")
+        .replace("False", "false")
+        .replace("None", "null")
+    )
+    model = oce.create_BC(s)
+
+    model.fit(train["Smiles"], train["pChEMBL Value"])
+    preds = model.predict(test["Smiles"])
+
+
+# TypeError: unsupported operand type(s) for +: 'Mol2Vec' and 'DescriptastorusDescriptor'
 
-    model.fit(train['Smiles'], train['pChEMBL Value'])
-    preds=model.predict(test['Smiles'])
 
-#TypeError: unsupported operand type(s) for +: 'Mol2Vec' and 'DescriptastorusDescriptor'
 def test_ConcatMol2Vec():
-    oce.DescriptastorusDescriptor('morgan3counts') + oce.Mol2Vec()
-    oce.Mol2Vec() + oce.DescriptastorusDescriptor('morgan3counts')
+    oce.DescriptastorusDescriptor("morgan3counts") + oce.Mol2Vec()
+    oce.Mol2Vec() + oce.DescriptastorusDescriptor("morgan3counts")
+
+
+# ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)
+
 
-#ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)
 def test_ConcatMordred(example_data1):
-    rep = oce.DescriptastorusDescriptor('morgan3counts') + oce.MordredDescriptor()
-    rep.convert('C')
+    rep = oce.DescriptastorusDescriptor("morgan3counts") + oce.MordredDescriptor()
+    rep.convert("C")
 
     model = oce.RandomForestModel(representation=rep, n_estimators=10)
-    model.fit(example_data1['Smiles'], example_data1['pChEMBL Value'])
+    model.fit(example_data1["Smiles"], example_data1["pChEMBL Value"])
+
 
 '''
 Alternative code I found working taken from DeepChem
@@ -64,4 +89,4 @@ def _convert(self, s: str) -> np.ndarray:
             for val in feature
         ]
         return np.array(feature)
-'''
\ No newline at end of file
+'''
diff --git a/tests/test_models.py b/tests/test_models.py
index 4b254d17..25c70130 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,38 +1,54 @@
+import numpy as np
+import pandas as pd
 import pytest
+from sklearn.model_selection import train_test_split
+
 import olorenchemengine as oce
-import numpy as np
 from olorenchemengine.base_class import *
 from olorenchemengine.basics import *
-from olorenchemengine.representations import *
-from olorenchemengine.ensemble import *
 from olorenchemengine.benchmarks import *
-import pandas as pd
+from olorenchemengine.ensemble import *
 from olorenchemengine.internal import download_public_file
-from sklearn.model_selection import train_test_split
+from olorenchemengine.representations import *
 
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
 
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 @pytest.fixture
 def example_data():
     file_path = download_public_file("sample-csvs/sample_data1.csv")
     df = pd.read_csv(file_path)[:10]
-    return train_test_split(df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42)
+    return train_test_split(
+        df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42
+    )
 
 
 @pytest.fixture
 def example_data2():
     file_path = download_public_file("sample-csvs/sample_data2.csv")
     df = pd.read_csv(file_path)[:10]
-    return train_test_split(df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42)
+    return train_test_split(
+        df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42
+    )
 
 
 @pytest.fixture
 def example_data3():
     file_path = download_public_file("sample-csvs/sample_data3.csv")
     df = pd.read_csv(file_path)[:10]
-    return train_test_split(df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42)
+    return train_test_split(
+        df["Smiles"], df["pChEMBL Value"], test_size=0.33, random_state=42
+    )
+
 
 def bf_model(model):
     s1 = parameterize(model)
@@ -63,7 +79,9 @@ def train_predict_slf(model, example_data):
 
 
 def test_create_rf_model():
-    model = RandomForestModel(DescriptastorusDescriptor("morgan3counts"), n_estimators=10)
+    model = RandomForestModel(
+        DescriptastorusDescriptor("morgan3counts"), n_estimators=10
+    )
     assert model is not None
 
 
@@ -73,82 +91,146 @@ def test_rf_back_and_forth():
 
 
 def test_basic_train(example_data):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     X_train, X_test, y_train, y_test = example_data
     model.fit(X_train, y_train)
     p1 = model.predict(X_test)
     assert len(p1) == len(y_test)
 
+
 def test_basic_test(example_data):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     X_train, X_test, y_train, y_test = example_data
     model.fit(X_train, y_train)
     model.test(X_test, y_test)
 
+
 def test_basic_sl(example_data):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     train_predict_sl(model, example_data)
 
 
 def test_basic_slf(example_data):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data)
 
 
 def test_basic_slf2(example_data):
     model = RandomForestModel(
-        OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10, normalization=QuantileTransformer(),
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+        normalization=QuantileTransformer(),
     )
     train_predict_slf(model, example_data)
 
 
 def test_basic_slf3(example_data3):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data3)
 
 
 def test_basic_slf4(example_data3):
     model = RandomForestModel(
-        OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10, normalization=QuantileTransformer(),
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+        normalization=QuantileTransformer(),
     )
     train_predict_slf(model, example_data3)
 
 
-@pytest.mark.timeout(20)
 def test_torchgeometric_save_load(example_data):
     from olorenchemengine.external import GINModel
 
-    model = oce.BaseTorchGeometricModel(GINModel(hidden=1, layers=1, batch_size=10), epochs=1)
+    model = oce.BaseTorchGeometricModel(
+        GINModel(hidden=1, layers=1, batch_size=10), epochs=1
+    )
     train_predict_sl(model, example_data)
 
 
 def test_mlp_saurus(example_data):
-    model = oce.SklearnMLP(oce.DescriptastorusDescriptor("morgan3counts"), hidden_layer_sizes=[16, 4], activation="relu",)
+    model = oce.SklearnMLP(
+        oce.DescriptastorusDescriptor("morgan3counts"),
+        hidden_layer_sizes=[16, 4],
+        activation="relu",
+    )
     train_predict_slf(model, example_data)
 
+
 def test_mlp_torch(example_data):
-    model = oce.TorchMLP(oce.DescriptastorusDescriptor("morgan3counts"), hidden_layer_sizes=[16, 4], activation="relu",)
+    model = oce.TorchMLP(
+        oce.DescriptastorusDescriptor("morgan3counts"),
+        hidden_layer_sizes=[16, 4],
+        activation="relu",
+    )
     train_predict_slf(model, example_data)
 
+
 def test_boost(example_data):
     model = oce.BaseBoosting(
         [
-            oce.RandomForestModel(oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=1),
+            oce.RandomForestModel(
+                oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=1
+            ),
             oce.RandomForestModel(oce.OlorenCheckpoint("default"), n_estimators=1),
         ],
         n=1,
     )
     train_predict_sl(model, example_data)
 
+
 def test_basic_save_load_file(example_data):
-    model = RandomForestModel(OlorenCheckpoint("default", num_tasks=2048,), n_estimators=10,)
+    model = RandomForestModel(
+        OlorenCheckpoint(
+            "default",
+            num_tasks=2048,
+        ),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data)
 
 
 def test_boost_file(example_data):
     model = oce.BaseBoosting(
         [
-            oce.RandomForestModel(oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=10),
+            oce.RandomForestModel(
+                oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=10
+            ),
             oce.RandomForestModel(oce.OlorenCheckpoint("default"), n_estimators=10),
         ],
         n=1,
@@ -165,7 +247,9 @@ def test_attentive_fp(example_data):
 def test_rfstacker(example_data):
     model = oce.RFStacker(
         [
-            oce.RandomForestModel(oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=10),
+            oce.RandomForestModel(
+                oce.DescriptastorusDescriptor("rdkit2dnormalized"), n_estimators=10
+            ),
             oce.RandomForestModel(oce.OlorenCheckpoint("default"), n_estimators=10),
         ],
         n=1,
@@ -173,7 +257,6 @@ def test_rfstacker(example_data):
     train_predict_slf(model, example_data)
 
 
-
 import warnings
 
 
@@ -210,42 +293,63 @@ def test_chemprop_2(example_data3):
 
 
 def test_zwk_xgboost_model1(example_data):
-    model = oce.basics.ZWK_XGBoostModel(oce.DescriptastorusDescriptor("morgan3counts"), n_iter=2, cv=2)
+    model = oce.basics.ZWK_XGBoostModel(
+        oce.DescriptastorusDescriptor("morgan3counts"), n_iter=2, cv=2
+    )
     train_predict_slf(model, example_data)
 
 
 def test_zwk_xgboost_model2(example_data3):
-    model = oce.basics.ZWK_XGBoostModel(oce.DescriptastorusDescriptor("morgan3counts"), n_iter=2, cv=2)
+    model = oce.basics.ZWK_XGBoostModel(
+        oce.DescriptastorusDescriptor("morgan3counts"), n_iter=2, cv=2
+    )
     train_predict_slf(model, example_data3)
 
+
 def test_add_reps(example_data):
-    model = oce.RandomForestModel(oce.OlorenCheckpoint("default") + oce.DescriptastorusDescriptor("morgan3counts"))
+    model = oce.RandomForestModel(
+        oce.OlorenCheckpoint("default") + oce.DescriptastorusDescriptor("morgan3counts")
+    )
     train_predict_slf(model, example_data)
 
 
 def test_pep_desc(example_data):
-    model = oce.RandomForestModel(oce.PeptideDescriptors1(), n_estimators=10,)
+    model = oce.RandomForestModel(
+        oce.PeptideDescriptors1(),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data)
 
 
 def test_mol2vec(example_data):
     from olorenchemengine.external import Mol2Vec
 
-    model = oce.RandomForestModel(Mol2Vec(), n_estimators=10,)
+    model = oce.RandomForestModel(
+        Mol2Vec(),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data)
 
+
 def test_smiles_transformer(example_data):
     from olorenchemengine.external import HondaSTRep
-    model = RandomForestModel(HondaSTRep(), n_estimators=10,)
+
+    model = RandomForestModel(
+        HondaSTRep(),
+        n_estimators=10,
+    )
     train_predict_slf(model, example_data)
 
 
-@pytest.mark.parametrize("rep", [(rep) for rep in oce.BaseCompoundVecRepresentation.AllInstances()])
-def test_all_base_compound_vecs(example_data, rep):
-    if isinstance(rep, oce.MordredDescriptor):
-        return  # Mordred does not play well with pytest
-    model = oce.RandomForestModel(rep, n_estimators=10)
-    train_predict_slf(model, example_data)
+# @pytest.mark.parametrize("rep", [(rep) for rep in oce.BaseCompoundVecRepresentation.AllInstances()])
+#
+#
+# def test_all_base_compound_vecs(example_data, rep):
+#     if isinstance(rep, oce.MordredDescriptor):
+#         return  # Mordred does not play well with pytest
+#     model = oce.RandomForestModel(rep, n_estimators=10)
+#     train_predict_slf(model, example_data)
+
 
 def test_molnet_single_task():
     mn_datasets = [
@@ -261,12 +365,18 @@ def test_molnet_single_task():
         "delaney" or "ESOL",
         "freesolv",
         "lipo",
-        "bace_regression"
+        "bace_regression",
     ]
 
-    MTC = oce.RandomForestModel(oce.DescriptastorusDescriptor("morgan3counts"), max_features="log2", n_estimators=1)
-    RDK = oce.RandomForestModel(oce.FragmentIndicator(), max_features="log2", n_estimators=1)
+    MTC = oce.RandomForestModel(
+        oce.DescriptastorusDescriptor("morgan3counts"),
+        max_features="log2",
+        n_estimators=1,
+    )
+    RDK = oce.RandomForestModel(
+        oce.FragmentIndicator(), max_features="log2", n_estimators=1
+    )
     models = [MTC, RDK]
 
     mn_bench = BenchmarkMolNet(datasets=[mn_datasets[0]])
-    mn_bench.run(models)
\ No newline at end of file
+    mn_bench.run(models)
diff --git a/tests/test_splitters.py b/tests/test_splitters.py
index 0ec09d3d..8166c4ef 100644
--- a/tests/test_splitters.py
+++ b/tests/test_splitters.py
@@ -1,15 +1,25 @@
 from random import random
+
+import pandas as pd
 import pytest
+
 import olorenchemengine as oce
-import pandas as pd
 from olorenchemengine.internal import download_public_file
 
-
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
-'''There are currently no tests for DateSplitter as sample data does not have date values.
-'''
+"""There are currently no tests for DateSplitter as sample data does not have date values.
+"""
+
+
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
 
 @pytest.fixture
 def example_data1():
@@ -17,77 +27,107 @@ def example_data1():
     df = pd.read_csv(file_path)
     return df
 
+
 @pytest.fixture
 def example_data2():
     file_path = download_public_file("sample-csvs/sample_data2.csv")
     df = pd.read_csv(file_path)
     return df
 
+
 @pytest.fixture
 def example_data3():
     file_path = download_public_file("sample-csvs/sample_data3.csv")
     df = pd.read_csv(file_path)
     return df
 
+
 def test_run_random(example_data3):
     splitter = oce.RandomSplit(split_proportions=[0.8, 0.1, 0.1])
     for i in splitter.split(example_data3):
-        assert(isinstance(i, pd.DataFrame))
-        assert(len(i) > 0)
+        assert isinstance(i, pd.DataFrame)
+        assert len(i) > 0
+
 
-'''
+"""
 Fail Info:
 Stratify by categorical column of 0, 1.
 Test passes with split_proportions = [0.8, 0.0, 0.2].
 Test fails with split_proportions = [0.8, 0.1, 0.1]: "ValueError: The test_size = 1 should be greater or equal to the number of classes = 2"
 Test fails with split_proportions = [0.8, 0.2, 0.0]: "ValueError: train_size=1.0 should be either positive and smaller than the number of samples 10 or a float in the (0, 1) range"
-'''
+"""
+
+
 def test_run_stratified(example_data3):
-    splitter = oce.StratifiedSplitter(split_proportions=[0.8, 0.1, 0.1], value_col='pChEMBL Value')
+    splitter = oce.StratifiedSplitter(
+        split_proportions=[0.8, 0.1, 0.1], value_col="pChEMBL Value"
+    )
     for i in splitter.split(example_data3):
-        assert(isinstance(i, pd.DataFrame))
-        assert(len(i) > 0)
+        assert isinstance(i, pd.DataFrame)
+        assert len(i) > 0
+
 
 def test_run_scaffold(example_data3):
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='murcko', split_proportions=[0.8, 0.1, 0.1])
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="murcko",
+        split_proportions=[0.8, 0.1, 0.1],
+    )
     for i in splitter.split(example_data3):
-        assert(isinstance(i, pd.DataFrame))
-        assert(len(i) > 0)
-
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='kmeans_murcko', split_proportions=[0.8, 0.1, 0.1])
+        assert isinstance(i, pd.DataFrame)
+        assert len(i) > 0
+
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="kmeans_murcko",
+        split_proportions=[0.8, 0.1, 0.1],
+    )
     for i in splitter.split(example_data3):
-        assert(isinstance(i, pd.DataFrame))
-        assert(len(i) > 0)
+        assert isinstance(i, pd.DataFrame)
+        assert len(i) > 0
+
 
 def test_run_property(example_data1):
-    splitter = oce.PropertySplit(property_col='pChEMBL Value', threshold=None, noise=0.1, categorical=False, split_proportions=[0.8, 0.1, 0.1])
+    splitter = oce.PropertySplit(
+        property_col="pChEMBL Value",
+        threshold=None,
+        noise=0.1,
+        categorical=False,
+        split_proportions=[0.8, 0.1, 0.1],
+    )
     for i in splitter.split(example_data1):
-        assert(isinstance(i, pd.DataFrame))
-        assert(len(i) > 0)
+        assert isinstance(i, pd.DataFrame)
+        assert len(i) > 0
 
-    splitter = oce.PropertySplit(property_col='pChEMBL Value', threshold=7, noise=0.1, categorical=False)
+    splitter = oce.PropertySplit(
+        property_col="pChEMBL Value", threshold=7, noise=0.1, categorical=False
+    )
     split = splitter.split(example_data1)
     for i in range(len(split)):
-        assert(isinstance(split[i], pd.DataFrame))
+        assert isinstance(split[i], pd.DataFrame)
         if i == 1:
-            assert(len(split[i]) == 0)
+            assert len(split[i]) == 0
         else:
-            assert(len(split[i]) > 0)
+            assert len(split[i]) > 0
+
 
-'''Split Proportion Tests:
+"""Split Proportion Tests:
  - Check if true split proportions are within 0.05 or 0.1 of user-defined split proportions
  - Test case with empty validation split (TT/train test split)
-'''
+"""
+
+
 def test_split_props_random(example_data3):
     split_proportions = [0.8, 0.1, 0.1]
 
     splitter = oce.RandomSplit(split_proportions=split_proportions)
+    print(example_data3)
     split = splitter.split(example_data3)
     total_samples = len(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
-        assert(abs(actual_split_prop - split_proportions[i]) < 0.05)
+        assert abs(actual_split_prop - split_proportions[i]) < 0.05
 
     split_proportions_TT = [0.8, 0.0, 0.2]
 
@@ -97,135 +137,198 @@ def test_split_props_random(example_data3):
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
         if i == 1:
-            assert(actual_split_prop == 0)
+            assert actual_split_prop == 0
         else:
-            assert(abs(actual_split_prop - split_proportions_TT[i]) < 0.05)
+            assert abs(actual_split_prop - split_proportions_TT[i]) < 0.05
 
-'''For the [0.8, 0.0, 0.2] case, it is found that the validation split has samples in it despite 0 split setting.
+
+"""For the [0.8, 0.0, 0.2] case, it is found that the validation split has samples in it despite 0 split setting.
 Test won't show this as there is still the run error from the first test case [0.8, 0.1, 0.1], described in test_run_stratified.
-'''
+"""
+
+
 def test_split_props_stratified(example_data3):
     split_proportions = [0.8, 0.1, 0.1]
 
-    splitter = oce.StratifiedSplitter(split_proportions=split_proportions, value_col='pChEMBL Value')
+    splitter = oce.StratifiedSplitter(
+        split_proportions=split_proportions, value_col="pChEMBL Value"
+    )
     split = splitter.split(example_data3)
     total_samples = len(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
-        assert(abs(actual_split_prop - split_proportions[i]) < 0.05)
+        assert abs(actual_split_prop - split_proportions[i]) < 0.05
 
     split_proportions_TT = [0.8, 0.0, 0.2]
 
-    splitter = oce.StratifiedSplitter(split_proportions=split_proportions_TT, value_col='pChEMBL Value')
+    splitter = oce.StratifiedSplitter(
+        split_proportions=split_proportions_TT, value_col="pChEMBL Value"
+    )
     split = splitter.split(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
         if i == 1:
-            assert(actual_split_prop == 0)
+            assert actual_split_prop == 0
         else:
-            assert(abs(actual_split_prop - split_proportions_TT[i]) < 0.05)
+            assert abs(actual_split_prop - split_proportions_TT[i]) < 0.05
+
 
 def test_split_props_scaffold_murcko(example_data3):
     split_proportions = [0.8, 0.1, 0.1]
 
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='murcko', split_proportions=split_proportions)
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="murcko",
+        split_proportions=split_proportions,
+    )
     split = splitter.split(example_data3)
     total_samples = len(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
-        assert(abs(actual_split_prop - split_proportions[i]) < 0.05)
+        assert abs(actual_split_prop - split_proportions[i]) < 0.05
 
     split_proportions_TT = [0.8, 0.0, 0.2]
 
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='murcko', split_proportions=split_proportions_TT)
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="murcko",
+        split_proportions=split_proportions_TT,
+    )
     split = splitter.split(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
         if i == 1:
-            assert(actual_split_prop == 0)
+            assert actual_split_prop == 0
         else:
-            assert(abs(actual_split_prop - split_proportions_TT[i]) < 0.05)
+            assert abs(actual_split_prop - split_proportions_TT[i]) < 0.05
+
 
 def test_split_props_scaffold_kmeans_murcko(example_data3):
     split_proportions = [0.33, 0.33, 0.33]
 
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='kmeans_murcko', split_proportions=split_proportions)
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="kmeans_murcko",
+        split_proportions=split_proportions,
+    )
     split = splitter.split(example_data3)
     total_samples = len(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
-        assert(abs(actual_split_prop - split_proportions[i]) < 0.1)
+        assert abs(actual_split_prop - split_proportions[i]) < 0.1
 
     split_proportions_TT = [0.7, 0.0, 0.3]
 
-    splitter = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='kmeans_murcko', split_proportions=split_proportions_TT)
+    splitter = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="kmeans_murcko",
+        split_proportions=split_proportions_TT,
+    )
     split = splitter.split(example_data3)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
         if i == 1:
-            assert(actual_split_prop == 0)
+            assert actual_split_prop == 0
         else:
-            assert(abs(actual_split_prop - split_proportions_TT[i]) < 0.1)
+            assert abs(actual_split_prop - split_proportions_TT[i]) < 0.1
+
 
 def test_split_props_property(example_data1):
     split_proportions = [0.8, 0.1, 0.1]
 
-    splitter = oce.PropertySplit(property_col='pChEMBL Value', threshold=None, noise=0.1, categorical=False, split_proportions=split_proportions)
+    splitter = oce.PropertySplit(
+        property_col="pChEMBL Value",
+        threshold=None,
+        noise=0.1,
+        categorical=False,
+        split_proportions=split_proportions,
+    )
     split = splitter.split(example_data1)
     total_samples = len(example_data1)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
-        assert(abs(actual_split_prop - split_proportions[i]) < 0.1)
+        assert abs(actual_split_prop - split_proportions[i]) < 0.1
 
     split_proportions_TT = [0.8, 0.0, 0.2]
 
-    splitter = oce.PropertySplit(property_col='pChEMBL Value', threshold=None, noise=0.1, categorical=False, split_proportions=split_proportions_TT)
+    splitter = oce.PropertySplit(
+        property_col="pChEMBL Value",
+        threshold=None,
+        noise=0.1,
+        categorical=False,
+        split_proportions=split_proportions_TT,
+    )
     split = splitter.split(example_data1)
 
     for i in range(3):
         actual_split_prop = len(split[i]) / total_samples
         if i == 1:
-            assert(actual_split_prop == 0)
+            assert actual_split_prop == 0
         else:
-            assert(abs(actual_split_prop - split_proportions_TT[i]) < 0.1)
+            assert abs(actual_split_prop - split_proportions_TT[i]) < 0.1
 
 
 def test_split_props_property_thresh(example_data1):
-    splitter = oce.PropertySplit(property_col='pChEMBL Value', threshold=7, noise=0.1, categorical=False)
+    splitter = oce.PropertySplit(
+        property_col="pChEMBL Value", threshold=7, noise=0.1, categorical=False
+    )
     split = splitter.split(example_data1)
 
     for i in range(3):
         if i == 1:
-            assert(len(split[i]) == 0)
+            assert len(split[i]) == 0
         else:
-            assert(len(split[i]) > 0)
+            assert len(split[i]) > 0
+
+
+"""Check if saved and loaded splitter parameters equal the true parameters before save/load.
+"""
+
 
-'''Check if saved and loaded splitter parameters equal the true parameters before save/load.
-'''
 def test_save_load(example_data1, example_data3, tmp_path):
     d = tmp_path / "sub"
     d.mkdir()
 
     random = oce.RandomSplit(split_proportions=[0.8, 0.1, 0.1])
-    stratified = oce.StratifiedSplitter(split_proportions=[0.8, 0.0, 0.2], value_col='pChEMBL Value')
-    scaffold = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='murcko', split_proportions=[0.8, 0.1, 0.1])
-    scaffold_kmeans = oce.ScaffoldSplit(scaffold_filter_threshold=1, split_type='kmeans_murcko', split_proportions=[0.8, 0.1, 0.1])
-    property = oce.PropertySplit(property_col='pChEMBL Value', threshold=None, noise=0.1, categorical=False, split_proportions=[0.8, 0.1, 0.1])
-
-    splitters = {'Random':random, 'Stratified':stratified, 'Scaffold':scaffold, 'Scaffold_KMeans':scaffold_kmeans, 'Property':property}
+    stratified = oce.StratifiedSplitter(
+        split_proportions=[0.8, 0.0, 0.2], value_col="pChEMBL Value"
+    )
+    scaffold = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="murcko",
+        split_proportions=[0.8, 0.1, 0.1],
+    )
+    scaffold_kmeans = oce.ScaffoldSplit(
+        scaffold_filter_threshold=1,
+        split_type="kmeans_murcko",
+        split_proportions=[0.8, 0.1, 0.1],
+    )
+    property = oce.PropertySplit(
+        property_col="pChEMBL Value",
+        threshold=None,
+        noise=0.1,
+        categorical=False,
+        split_proportions=[0.8, 0.1, 0.1],
+    )
+
+    splitters = {
+        "Random": random,
+        "Stratified": stratified,
+        "Scaffold": scaffold,
+        "Scaffold_KMeans": scaffold_kmeans,
+        "Property": property,
+    }
     for name in splitters.keys():
         splitter = splitters[name]
         params = oce.base_class.parameterize(splitter)
-        oce.base_class.save(splitter, d / f'{name}')
-
-        loaded_splitter = oce.base_class.load(d / f'{name}')
-        assert(oce.base_class.parameterize(loaded_splitter) == params)
-
+        oce.base_class.save(splitter, d / f"{name}")
 
+        loaded_splitter = oce.base_class.load(d / f"{name}")
+        assert oce.base_class.parameterize(loaded_splitter) == params
diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py
index 5a399059..a731ee0f 100644
--- a/tests/test_tutorials.py
+++ b/tests/test_tutorials.py
@@ -1,22 +1,40 @@
 from random import random
+
+import pandas as pd
 import pytest
+
 import olorenchemengine as oce
-import pandas as pd
 
 __author__ = "Oloren AI"
 __copyright__ = "Oloren AI"
 
+
+def remote(func):
+    def wrapper(*args, **kwargs):
+        with oce.Remote("http://api.oloren.ai:5000") as remote:
+            func(*args, **kwargs)
+
+    return wrapper
+
+
 @pytest.mark.timeout(300)
 def test_0():
-    import olorenchemengine as oce
     import pandas as pd
 
-    df = pd.read_csv("https://storage.googleapis.com/oloren-public-data/CHEMBL%20Datasets/997_2298%20-%20VEGFR1%20(CHEMBL1868).csv")
-    dataset = (oce.BaseDataset(data = df.to_csv(),
-        structure_col = "Smiles", property_col = "pChEMBL Value") +
-               oce.CleanStructures() +
-               oce.RandomSplit()
+    import olorenchemengine as oce
+
+    df = pd.read_csv(
+        "https://storage.googleapis.com/oloren-public-data/CHEMBL%20Datasets/997_2298%20-%20VEGFR1%20(CHEMBL1868).csv"
+    )
+    dataset = (
+        oce.BaseDataset(
+            data=df.to_csv(), structure_col="Smiles", property_col="pChEMBL Value"
+        )
+        + oce.CleanStructures()
+        + oce.RandomSplit()
+    )
+    model = oce.RandomForestModel(
+        oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=1000
     )
-    model = oce.RandomForestModel(oce.DescriptastorusDescriptor("morgan3counts"), n_estimators=1000)
-    _ = model.fit_cv(*dataset.train_dataset, error_model = oce.kNNwRMSD1())
+    _ = model.fit_cv(*dataset.train_dataset, error_model=oce.kNNwRMSD1())
     oce.save(model, "vis-model.oce")