standardization of bd93c86

Author: rabah-khalek

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

@@ -1,5 +1,6 @@
 import logging
 import pickle
+import yaml
 from abc import ABC, abstractmethod
 from inspect import isfunction, signature
 from pathlib import Path
@@ -29,17 +30,17 @@ class WrapperModel(BaseModel, ABC):
 
     @configured_validate_arguments
     def __init__(
-        self,
-        model: Any,
-        model_type: ModelType,
-        data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
-        model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
-        name: Optional[str] = None,
-        feature_names: Optional[Iterable] = None,
-        classification_threshold: Optional[float] = 0.5,
-        classification_labels: Optional[Iterable] = None,
-        batch_size: Optional[int] = None,
-        **kwargs,
+            self,
+            model: Any,
+            model_type: ModelType,
+            data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
+            model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
+            name: Optional[str] = None,
+            feature_names: Optional[Iterable] = None,
+            classification_threshold: Optional[float] = 0.5,
+            classification_labels: Optional[Iterable] = None,
+            batch_size: Optional[int] = None,
+            **kwargs,
     ) -> None:
         """
         Parameters
@@ -188,13 +189,11 @@ def model_predict(self, data):
 
     def save(self, local_path: Union[str, Path]) -> None:
         super().save(local_path)
-
+        self.save_wrapper_meta(local_path)
         if self.data_preprocessing_function:
             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:
@@ -215,15 +214,21 @@ 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))
+    def save_wrapper_meta(self, local_path):
+        with open(Path(local_path) / "giskard-model-wrapper-meta.yaml", "w") as f:
+            yaml.dump(
+                {
+                    "batch_size": self.batch_size,
+                },
+                f,
+                default_flow_style=False,
+            )
 
     @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)
+        kwargs.update(cls.load_wrapper_meta(local_dir))
         model_id, meta = cls.read_meta_from_local_dir(local_dir)
         constructor_params = meta.__dict__
         constructor_params["id"] = model_id
@@ -262,10 +267,15 @@ def load_model_postprocessing_function(cls, local_path: Union[str, Path]):
                 return cloudpickle.load(f)
         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
+    @classmethod
+    def load_wrapper_meta(cls, local_dir):
+        wrapper_meta_file = Path(local_dir) / "giskard-model-wrapper-meta.yaml"
+        if wrapper_meta_file.exists():
+            with open(wrapper_meta_file) as f:
+                wrapper_meta = yaml.load(f, Loader=yaml.Loader)
+                wrapper_meta["batch_size"] = int(wrapper_meta["batch_size"]) if wrapper_meta["batch_size"] else None
+                return wrapper_meta
+        else:
+            raise ValueError(
+                f"Cannot load model ({cls.__module__}.{cls.__name__}), " f"{wrapper_meta_file} file not found"
+            )

python-client/giskard/models/pytorch.py

@@ -60,20 +60,20 @@ def __getitem__(self, idx):
 
 class PyTorchModel(MLFlowSerializableModel):
     def __init__(
-        self,
-        model,
-        model_type: ModelType,
-        torch_dtype: TorchDType = "float32",
-        device="cpu",
-        name: Optional[str] = None,
-        data_preprocessing_function=None,
-        model_postprocessing_function=None,
-        feature_names=None,
-        classification_threshold=0.5,
-        classification_labels=None,
-        iterate_dataset: bool = True,
-        batch_size: Optional[int] = None,
-        **kwargs,
+            self,
+            model,
+            model_type: ModelType,
+            torch_dtype: TorchDType = "float32",
+            device="cpu",
+            name: Optional[str] = None,
+            data_preprocessing_function=None,
+            model_postprocessing_function=None,
+            feature_names=None,
+            classification_threshold=0.5,
+            classification_labels=None,
+            iterate_dataset: bool = True,
+            batch_size: Optional[int] = None,
+            **kwargs,
     ) -> None:
         """Automatically wraps a PyTorch model.
 
@@ -214,14 +214,19 @@ def save(self, local_path: Union[str, Path]) -> None:
 
     @classmethod
     def load(cls, local_dir, **kwargs):
+        kwargs.update(cls.load_pytorch_meta(local_dir))
+        return super().load(local_dir, **kwargs)
+
+    @classmethod
+    def load_pytorch_meta(cls, local_dir):
         pytorch_meta_file = Path(local_dir) / "giskard-model-pytorch-meta.yaml"
         if pytorch_meta_file.exists():
             with open(pytorch_meta_file) as f:
                 pytorch_meta = yaml.load(f, Loader=yaml.Loader)
-                kwargs["device"] = pytorch_meta["device"]
-                kwargs["torch_dtype"] = pytorch_meta["torch_dtype"]
-                kwargs["iterate_dataset"] = pytorch_meta.get("iterate_dataset")
-                return super().load(local_dir, **kwargs)
+                pytorch_meta["device"] = pytorch_meta.get("device")
+                pytorch_meta["torch_dtype"] = pytorch_meta.get("torch_dtype")
+                pytorch_meta["iterate_dataset"] = pytorch_meta.get("iterate_dataset")
+                return pytorch_meta
         else:
             raise ValueError(
                 f"Cannot load model ({cls.__module__}.{cls.__name__}), " f"{pytorch_meta_file} file not found"