typehinting and refactoring load function

Author: michaelraczycki

pymc_experimental/model_builder.py

@@ -28,8 +28,6 @@ class ModelBuilder:
     """
     ModelBuilder can be used to provide an easy-to-use API (similar to scikit-learn) for models
     and help with deployment.
-
-    Extends the pymc.Model class.
     """
 
     _model_type = "BaseClass"
@@ -65,7 +63,7 @@ def __init__(
         self.idata = None  # inference data object
         self.data = data
 
-    def build(self):
+    def build(self) -> None:
         """
         Builds the defined model.
         """
@@ -136,7 +134,7 @@ def create_sample_input():
 
         raise NotImplementedError
 
-    def save(self, fname):
+    def save(self, fname: str) -> None:
         """
         Saves inference data of the model.
 
@@ -160,8 +158,9 @@ def save(self, fname):
         self.idata.to_netcdf(file)
 
     @classmethod
-    def load(cls, fname):
+    def load(cls, fname: str):
         """
+        Creates a ModelBuilder instance from a file,
         Loads inference data for the model.
 
         Parameters
@@ -171,7 +170,7 @@ def load(cls, fname):
 
         Returns
         -------
-        Returns an instance of pm.Model, that is loaded from local data.
+        Returns an instance of ModelBuilder.
 
         Raises
         ------
@@ -188,21 +187,23 @@ def load(cls, fname):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
-        self = cls(
+        model_builder = cls(
             dict(zip(idata.attrs["model_config_keys"], idata.attrs["model_config_values"])),
             dict(zip(idata.attrs["sample_config_keys"], idata.attrs["sample_config_values"])),
             idata.fit_data.to_dataframe(),
         )
-        self.idata = idata
-        self.build()
-        if self.id != idata.attrs["id"]:
+        model_builder.idata = idata
+        model_builder.build()
+        if model_builder.id != idata.attrs["id"]:
             raise ValueError(
                 f"The file '{fname}' does not contain an inference data of the same model or configuration as '{self._model_type}'"
             )
 
-        return self.model
+        return model_builder
 
-    def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
+    def fit(
+        self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
+    ) -> az.InferenceData:
         """
         As the name suggests fit can be used to fit a model using the data that is passed as a parameter.
         Sets attrs to inference data of the model.
@@ -248,7 +249,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
     def predict(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
-    ):
+    ) -> dict:
         """
         Uses model to predict on unseen data and return point prediction of all the samples
 
@@ -288,7 +289,7 @@ def predict(
     def predict_posterior(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
-    ):
+    ) -> Dict[str, np.array]:
         """
         Uses model to predict samples on unseen data.
 

pymc_experimental/tests/test_model_builder.py

@@ -28,7 +28,7 @@ class test_ModelBuilder(ModelBuilder):
     _model_type = "LinearModel"
     version = "0.1"
 
-    def build_model(self, model_instance, model_config, data=None):
+    def build_model(self, model_instance: ModelBuilder, model_config: dict, data: dict = None):
         model_instance.model_config = model_config
         model_instance.data = data
         self.model_config = model_config
@@ -86,7 +86,7 @@ def create_sample_input(self):
         return data, model_config, sampler_config
 
     @staticmethod
-    def initial_build_and_fit(check_idata=True):
+    def initial_build_and_fit(check_idata=True) -> ModelBuilder:
         data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
         model_builder = test_ModelBuilder(model_config, sampler_config, data)
         model_builder.idata = model_builder.fit(data=data)
@@ -113,8 +113,7 @@ def test_save_load():
     test_builder = test_ModelBuilder.initial_build_and_fit()
     temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
     test_builder.save(temp.name)
-    test_builder2 = test_ModelBuilder.initial_build_and_fit()
-    test_builder2.model = test_ModelBuilder.load(temp.name)
+    test_builder2 = test_ModelBuilder.load(temp.name)
     assert test_builder.idata.groups() == test_builder2.idata.groups()
 
     x_pred = np.random.uniform(low=0, high=1, size=100)