Merge pull request #1490 from Giskard-AI/feature/gsk-1921-save-langchain-chain-without-ml_flow-to-reduce-compatibility

1921 & GSK-1851 & GSK-1901 - save langchain chain without ml flow to reduce compatibility

Author: kevinmessiaen

giskard/models/automodel.py

@@ -5,9 +5,9 @@
 
 import pandas as pd
 
-from ..core.core import ModelType, SupportedModelTypes
 from .base.serialization import CloudpickleSerializableModel
 from .function import PredictionFunctionModel
+from ..core.core import ModelType, SupportedModelTypes
 
 logger = logging.getLogger(__name__)
 
@@ -128,8 +128,6 @@ def __new__(
                     giskard_cls = CloudpickleSerializableModel
                 # if save_model and load_model are overriden, replace them, if not, these equalities will be identities.
                 possibly_overriden_cls = cls
-                possibly_overriden_cls.save_model = giskard_cls.save_model
-                possibly_overriden_cls.load_model = giskard_cls.load_model
                 possibly_overriden_cls.should_save_model_class = True
             elif giskard_cls:
                 input_type = "'prediction_function'" if giskard_cls == PredictionFunctionModel else "'model'"

giskard/models/base/serialization.py

@@ -7,6 +7,7 @@
 
 from .wrapper import WrapperModel
 
+
 # @TODO: decouple the serialization logic from models. These abstract classes
 # could be implemented as mixins and then used in the models that need them.
 # The logic of saving the model should be moved to the serialization classes.

giskard/models/base/wrapper.py

@@ -11,10 +11,10 @@
 import pandas as pd
 import yaml
 
+from .model import BaseModel
+from ..utils import warn_once
 from ...core.core import ModelType
 from ...core.validation import configured_validate_arguments
-from ..utils import warn_once
-from .model import BaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -242,6 +242,12 @@ def save_wrapper_meta(self, local_path):
 
     @classmethod
     def load(cls, local_dir, **kwargs):
+        constructor_params = cls.load_constructor_params(local_dir, **kwargs)
+
+        return cls(model=cls.load_model(local_dir), **constructor_params)
+
+    @classmethod
+    def load_constructor_params(cls, local_dir, **kwargs):
         params = cls.load_wrapper_meta(local_dir)
         params["data_preprocessing_function"] = cls.load_data_preprocessing_function(local_dir)
         params["model_postprocessing_function"] = cls.load_model_postprocessing_function(local_dir)
@@ -253,7 +259,7 @@ def load(cls, local_dir, **kwargs):
         constructor_params = constructor_params.copy()
         constructor_params.update(params)
 
-        return cls(model=cls.load_model(local_dir), **constructor_params)
+        return constructor_params
 
     @classmethod
     @abstractmethod

giskard/models/huggingface.py

@@ -103,7 +103,6 @@ class explicitly using :class:`giskard.models.huggingface.HuggingFaceModel`.
 from giskard.core.core import ModelType
 from giskard.core.validation import configured_validate_arguments
 from giskard.models.base import WrapperModel
-
 from ..client.python_utils import warning
 
 try:

giskard/models/langchain.py

@@ -1,29 +1,29 @@
-from typing import Any, Callable, Iterable, Optional
+from pathlib import Path
+from typing import Any, Callable, Iterable, Optional, Union, Dict
 
-import mlflow
 import pandas as pd
 
 from giskard.core.core import SupportedModelTypes
 from giskard.core.validation import configured_validate_arguments
-from giskard.models.base import MLFlowSerializableModel
+from giskard.models.base import WrapperModel
 
 
-class LangchainModel(MLFlowSerializableModel):
+class LangchainModel(WrapperModel):
     @configured_validate_arguments
     def __init__(
-            self,
-            model,
-            model_type: SupportedModelTypes,
-            name: Optional[str] = None,
-            data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
-            model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
-            feature_names: Optional[Iterable] = None,
-            classification_threshold: Optional[float] = 0.5,
-            classification_labels: Optional[Iterable] = None,
-            **kwargs,
+        self,
+        model,
+        model_type: SupportedModelTypes,
+        name: Optional[str] = None,
+        data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
+        model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
+        feature_names: Optional[Iterable] = None,
+        classification_threshold: Optional[float] = 0.5,
+        classification_labels: Optional[Iterable] = None,
+        **kwargs,
     ) -> None:
         assert (
-                model_type == SupportedModelTypes.TEXT_GENERATION
+            model_type == SupportedModelTypes.TEXT_GENERATION
         ), "LangchainModel only support text_generation ModelType"
 
         super().__init__(
@@ -38,15 +38,49 @@ def __init__(
             **kwargs,
         )
 
-    def save_model(self, local_path, mlflow_meta):
-        mlflow.langchain.save_model(self.model, path=local_path, mlflow_model=mlflow_meta)
+    def save(self, local_path: Union[str, Path]) -> None:
+        super().save(local_path)
+        self.save_model(local_path)
+        self.save_artifacts(Path(local_path) / "artifacts")
+
+    def save_model(self, local_path: Union[str, Path]) -> None:
+        path = Path(local_path)
+        self.model.save(path / "chain.json")
+
+    def save_artifacts(self, artifact_dir) -> None:
+        ...
+
+    @classmethod
+    def load(cls, local_dir, **kwargs):
+        constructor_params = cls.load_constructor_params(local_dir, **kwargs)
+
+        artifacts = cls.load_artifacts(Path(local_dir) / "artifacts") or dict()
+        constructor_params.update(artifacts)
+
+        return cls(model=cls.load_model(local_dir, **artifacts), **constructor_params)
+
+    @classmethod
+    def load_model(cls, local_dir, **kwargs):
+        from langchain.chains import load_chain
+
+        path = Path(local_dir)
+        return load_chain(path / "chain.json", **kwargs)
 
     @classmethod
-    def load_model(cls, local_dir):
-        return mlflow.langchain.load_model(local_dir)
+    def load_artifacts(cls, local_path: Union[str, Path]) -> Optional[Dict[str, Any]]:
+        ...
 
     def model_predict(self, df):
-        return [self.model.predict(**data) for data in df.to_dict("records")]
+        generations = [self.model(data) for data in df.to_dict("records")]
+        output_keys = self.model.output_keys
+
+        if len(output_keys) == 1:
+            return [generation[output_keys[0]] for generation in generations]
+        else:
+            return [
+                str({key: value for key, value in generation.items() if key in output_keys})
+                for generation in generations
+            ]
 
     def rewrite_prompt(self, template, input_variables=None, **kwargs):
         from langchain import LLMChain
@@ -70,6 +104,3 @@ def rewrite_prompt(self, template, input_variables=None, **kwargs):
         model_kwargs.update(kwargs)
 
         return self.__class__(chain, **model_kwargs)
-
-    def to_mlflow(self, artifact_path: str = "langchain-model-from-giskard", **kwargs):
-        return mlflow.langchain.log_model(self.model, artifact_path, **kwargs)

giskard/models/pytorch.py

@@ -10,10 +10,10 @@
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset as torch_dataset
 
-from ..client.python_utils import warning
-from ..core.core import ModelType
 from .base.serialization import MLFlowSerializableModel
 from .utils import map_to_tuples
+from ..client.python_utils import warning
+from ..core.core import ModelType
 
 TorchDType = Literal[
     "float32",
@@ -60,21 +60,21 @@ 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,
-            id: Optional[str] = None,
-            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,
+        id: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        **kwargs,
     ) -> None:
         """Automatically wraps a PyTorch model.
 

giskard/models/sklearn.py

@@ -3,9 +3,9 @@
 import mlflow
 import pandas as pd
 
+from .base.serialization import MLFlowSerializableModel
 from ..core.core import ModelType, SupportedModelTypes
 from ..core.validation import configured_validate_arguments
-from .base.serialization import MLFlowSerializableModel
 
 
 class SKLearnModel(MLFlowSerializableModel):
@@ -15,18 +15,18 @@ class SKLearnModel(MLFlowSerializableModel):
 
     @configured_validate_arguments
     def __init__(
-            self,
-            model,
-            model_type: ModelType,
-            name: Optional[str] = None,
-            data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
-            model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
-            feature_names: Optional[Iterable] = None,
-            classification_threshold: Optional[float] = 0.5,
-            classification_labels: Optional[Iterable] = None,
-            id: Optional[str] = None,
-            batch_size: Optional[int] = None,
-            **kwargs,
+        self,
+        model,
+        model_type: ModelType,
+        name: Optional[str] = None,
+        data_preprocessing_function: Optional[Callable[[pd.DataFrame], Any]] = None,
+        model_postprocessing_function: Optional[Callable[[Any], Any]] = None,
+        feature_names: Optional[Iterable] = None,
+        classification_threshold: Optional[float] = 0.5,
+        classification_labels: Optional[Iterable] = None,
+        id: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        **kwargs,
     ) -> None:
         model_type = SupportedModelTypes(model_type) if isinstance(model_type, str) else model_type
         if model_type == SupportedModelTypes.CLASSIFICATION:
@@ -75,9 +75,7 @@ def model_predict(self, df):
         else:
             return self.model.predict_proba(df)
 
-    def to_mlflow(self,
-                  artifact_path="sklearn-model-from-giskard",
-                  **kwargs):
-        return mlflow.sklearn.log_model(sk_model=self.model, artifact_path=artifact_path,
-                                        pyfunc_predict_fn=self._get_pyfunc_predict_fn(),
-                                        **kwargs)
+    def to_mlflow(self, artifact_path="sklearn-model-from-giskard", **kwargs):
+        return mlflow.sklearn.log_model(
+            sk_model=self.model, artifact_path=artifact_path, pyfunc_predict_fn=self._get_pyfunc_predict_fn(), **kwargs
+        )

giskard/models/tensorflow.py

@@ -4,27 +4,27 @@
 import mlflow
 import pandas as pd
 
+from .base import MLFlowSerializableModel
 from ..core.core import ModelType
 from ..core.validation import configured_validate_arguments
-from .base import MLFlowSerializableModel
 
 logger = logging.getLogger(__name__)
 
 
 class TensorFlowModel(MLFlowSerializableModel):
     @configured_validate_arguments
     def __init__(
-            self,
-            model,
-            model_type: ModelType,
-            name: Optional[str] = None,
-            data_preprocessing_function: Callable[[pd.DataFrame], Any] = None,
-            model_postprocessing_function: Callable[[Any], Any] = None,
-            feature_names: Optional[Iterable] = None,
-            classification_threshold: Optional[float] = 0.5,
-            classification_labels: Optional[Iterable] = None,
-            id: Optional[str] = None,
-            **kwargs,
+        self,
+        model,
+        model_type: ModelType,
+        name: Optional[str] = None,
+        data_preprocessing_function: Callable[[pd.DataFrame], Any] = None,
+        model_postprocessing_function: Callable[[Any], Any] = None,
+        feature_names: Optional[Iterable] = None,
+        classification_threshold: Optional[float] = 0.5,
+        classification_labels: Optional[Iterable] = None,
+        id: Optional[str] = None,
+        **kwargs,
     ):
         super().__init__(
             model=model,