Save the batch size

Author: mattbit

python-client/giskard/models/base/wrapper.py

@@ -193,6 +193,8 @@ def save(self, local_path: Union[str, Path]) -> None:
             self.save_data_preprocessing_function(local_path)
         if self.model_postprocessing_function:
             self.save_model_postprocessing_function(local_path)
+        if self.batch_size:
+            self.save_batch_size(local_path)
 
     @abstractmethod
     def save_model(self, path: Union[str, Path]) -> None:
@@ -213,10 +215,15 @@ def save_model_postprocessing_function(self, local_path: Union[str, Path]):
         with open(Path(local_path) / "giskard-model-postprocessing-function.pkl", "wb") as f:
             cloudpickle.dump(self.model_postprocessing_function, f, protocol=pickle.DEFAULT_PROTOCOL)
 
+    def save_batch_size(self, local_path: Union[str, Path]):
+        with Path(local_path).joinpath("batch_size").open("w") as f:
+            f.write(str(self.batch_size))
+
     @classmethod
     def load(cls, local_dir, **kwargs):
         kwargs["data_preprocessing_function"] = cls.load_data_preprocessing_function(local_dir)
         kwargs["model_postprocessing_function"] = cls.load_model_postprocessing_function(local_dir)
+        kwargs["batch_size"] = cls._load_batch_size(local_dir)
         model_id, meta = cls.read_meta_from_local_dir(local_dir)
         constructor_params = meta.__dict__
         constructor_params["id"] = model_id
@@ -244,8 +251,7 @@ def load_data_preprocessing_function(cls, local_path: Union[str, Path]):
         if file_path.exists():
             with open(file_path, "rb") as f:
                 return cloudpickle.load(f)
-        else:
-            return None
+        return None
 
     @classmethod
     def load_model_postprocessing_function(cls, local_path: Union[str, Path]):
@@ -254,5 +260,12 @@ def load_model_postprocessing_function(cls, local_path: Union[str, Path]):
         if file_path.exists():
             with open(file_path, "rb") as f:
                 return cloudpickle.load(f)
-        else:
-            return None
+        return None
+
+    @staticmethod
+    def _load_batch_size(local_path: Union[str, Path]):
+        file_path = Path(local_path) / "batch_size"
+        if file_path.exists():
+            with file_path.open("r") as f:
+                return int(f.read())
+        return None

python-client/tests/models/test_wrapper_model.py

@@ -1,7 +1,10 @@
+import tempfile
+
 import numpy as np
 import pandas as pd
+from sklearn.linear_model import LogisticRegression
 
-from giskard import Dataset
+from giskard import Dataset, Model
 from giskard.models.base.wrapper import WrapperModel
 
 
@@ -36,3 +39,20 @@ def save_model(self, path):
     model.expected_batch_size = [20, 1]
 
     model.predict(dataset)
+
+
+def test_wrapper_model_saves_and_loads_batch_size():
+    base_model = LogisticRegression()
+    model = Model(
+        base_model.predict_proba,
+        model_type="classification",
+        feature_names=["one", "two"],
+        classification_threshold=0.5,
+        classification_labels=[0, 1],
+    )
+    with tempfile.TemporaryDirectory() as tmpdir:
+        model.batch_size = 127
+        model.save(tmpdir)
+        loaded_model = Model.load(tmpdir)
+
+        assert loaded_model.batch_size == 127