Add support for white-box explainers to alibi-explain runtime

runtimes/alibi-explain/mlserver_alibi_explain/alibi_dependency_reference.py

@@ -17,6 +17,9 @@ class ExplainerDependencyReference:
 _ANCHOR_TABULAR_TAG = "anchor_tabular"
 _KERNEL_SHAP_TAG = "kernel_shap"
 _INTEGRATED_GRADIENTS_TAG = "integrated_gradients"
+_TREE_SHAP_TAG = "tree_shap"
+_TREE_PARTIAL_DEPENDENCE_TAG = "tree_partial_dependence"
+_TREE_PARTIAL_DEPENDENCE_VARIANCE_TAG = "tree_partial_dependence_variance"
 
 
 # NOTE: to add new explainers populate the below dict with a new
@@ -30,6 +33,7 @@ class ExplainerDependencyReference:
 
 _BLACKBOX_MODULE = "mlserver_alibi_explain.explainers.black_box_runtime"
 _INTEGRATED_GRADIENTS_MODULE = "mlserver_alibi_explain.explainers.integrated_gradients"
+_WHITEBOX_SKLEARN_MODULE = "mlserver_alibi_explain.explainers.sklearn_api_runtime"
 
 _TAG_TO_RT_IMPL: Dict[str, ExplainerDependencyReference] = {
     _ANCHOR_IMAGE_TAG: ExplainerDependencyReference(
@@ -57,6 +61,21 @@ class ExplainerDependencyReference:
         runtime_class=f"{_INTEGRATED_GRADIENTS_MODULE}.IntegratedGradientsWrapper",
         alibi_class="alibi.explainers.IntegratedGradients",
     ),
+    _TREE_SHAP_TAG: ExplainerDependencyReference(
+        explainer_name=_TREE_SHAP_TAG,
+        runtime_class=f"{_WHITEBOX_SKLEARN_MODULE}.SKLearnRuntime",
+        alibi_class="alibi.explainers.TreeShap",
+    ),
+    _TREE_PARTIAL_DEPENDENCE_TAG: ExplainerDependencyReference(
+        explainer_name=_TREE_PARTIAL_DEPENDENCE_TAG,
+        runtime_class=f"{_WHITEBOX_SKLEARN_MODULE}.SKLearnRuntime",
+        alibi_class="alibi.explainers.TreePartialDependence",
+    ),
+    _TREE_PARTIAL_DEPENDENCE_VARIANCE_TAG: ExplainerDependencyReference(
+        explainer_name=_TREE_PARTIAL_DEPENDENCE_VARIANCE_TAG,
+        runtime_class=f"{_WHITEBOX_SKLEARN_MODULE}.SKLearnRuntime",
+        alibi_class="alibi.explainers.PartialDependenceVariance",
+    ),
 }
 
 
@@ -66,6 +85,9 @@ class ExplainerEnum(str, Enum):
     anchor_tabular = _ANCHOR_TABULAR_TAG
     kernel_shap = _KERNEL_SHAP_TAG
     integrated_gradients = _INTEGRATED_GRADIENTS_TAG
+    tree_shap = _TREE_SHAP_TAG
+    tree_partial_dependence = _TREE_PARTIAL_DEPENDENCE_TAG
+    tree_partial_dependence_variance = _TREE_PARTIAL_DEPENDENCE_VARIANCE_TAG
 
 
 def get_mlmodel_class_as_str(tag: Union[ExplainerEnum, str]) -> str:

runtimes/alibi-explain/mlserver_alibi_explain/explainers/black_box_runtime.py

@@ -44,8 +44,9 @@ async def load(self) -> bool:
         # TODO: use init explainer field instead?
         if self.alibi_explain_settings.init_parameters is not None:
             init_parameters = self.alibi_explain_settings.init_parameters
-            init_parameters["predictor"] = self._infer_impl
-            self._model = self._explainer_class(**init_parameters)  # type: ignore
+            self._model = self._explainer_class(
+                self._infer_impl, **init_parameters  # type: ignore
+            )
         else:
             self._model = await self._load_from_uri(self._infer_impl)
 

runtimes/alibi-explain/mlserver_alibi_explain/explainers/sklearn_api_runtime.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+import joblib
+from xgboost.core import XGBoostError
+
+from mlserver_xgboost.xgboost import _load_sklearn_interface as load_xgb_model
+from mlserver.errors import InvalidModelURI
+from mlserver_alibi_explain.explainers.white_box_runtime import (
+    AlibiExplainWhiteBoxRuntime,
+)
+
+
+class SKLearnRuntime(AlibiExplainWhiteBoxRuntime):
+    """
+    Runtime for white-box explainers that require access to a tree-based model matching
+    the SKLearn API, such as a sklearn, XGBoost, or LightGBM model. Example explainers
+    include TreeShap and TreePartialDependence.
+    """
+
+    async def _get_inference_model(self) -> Any:
+        inference_model_path = self.alibi_explain_settings.infer_uri
+        # Attempt to load model.
+        try:
+            # Try to load as joblib model first
+            model = joblib.load(inference_model_path)
+        except (IndexError, KeyError, IOError):
+            try:
+                # Try to load as XGBoost model
+                model = load_xgb_model(inference_model_path)
+            except XGBoostError:
+                raise InvalidModelURI(self.name, inference_model_path)
+        return model

runtimes/alibi-explain/mlserver_alibi_explain/explainers/white_box_runtime.py

@@ -1,7 +1,7 @@
 from abc import ABC
-from typing import Any, Type
+from typing import Any, Type, Dict
 
-from alibi.api.interfaces import Explainer
+from alibi.api.interfaces import Explainer, Explanation
 
 from mlserver import ModelSettings
 from mlserver_alibi_explain.common import AlibiExplainSettings
@@ -29,17 +29,25 @@ def __init__(self, settings: ModelSettings, explainer_class: Type[Explainer]):
         super().__init__(settings, explainer_settings)
 
     async def load(self) -> bool:
+        # white box explainers requires access to the full inference model
         self._inference_model = await self._get_inference_model()
 
         if self.alibi_explain_settings.init_parameters is not None:
+            # Instantiate explainer with init parameters (and give it inference model)
             init_parameters = self.alibi_explain_settings.init_parameters
-            # white box explainers requires access to the inference model
-            init_parameters["model"] = self._inference_model
-            self._model = self._explainer_class(**init_parameters)  # type: ignore
+            self._model = self._explainer_class(
+                self._inference_model, **init_parameters  # type: ignore
+            )
         else:
+            # Load explainer from URI (and give it full inference model)
             self._model = await self._load_from_uri(self._inference_model)
 
         return True
 
+    def _explain_impl(self, input_data: Any, explain_parameters: Dict) -> Explanation:
+        # TODO: how are we going to deal with that?
+        assert self._inference_model is not None, "Inference model is not set"
+        return self._model.explain(input_data, **explain_parameters)
+
     async def _get_inference_model(self) -> Any:
         raise NotImplementedError

runtimes/alibi-explain/mlserver_alibi_explain/runtime.py

@@ -123,9 +123,7 @@ async def _async_explain_impl(
         )
 
     async def _load_from_uri(self, predictor: Any) -> Explainer:
-        # load the model from disk
-        # full model is passed as `predictor`
-        # load the model from disk
+        """Load the explainer from disk, and pass the predictor"""
         model_parameters: Optional[ModelParameters] = self.settings.parameters
         if model_parameters is None:
             raise ModelParametersMissing(self.name)

runtimes/alibi-explain/pyproject.toml

@@ -10,11 +10,17 @@ packages = [{include = "mlserver_alibi_explain"}]
 [tool.poetry.dependencies]
 python = "^3.8.1,<3.12"
 mlserver = "*"
+mlserver_sklearn = "*"
+mlserver_xgboost = "*"
+mlserver_lightgbm = "*"
 orjson = "*"
 alibi = {extras = ["shap", "tensorflow"], version = "*"}
 
 [tool.poetry.group.dev.dependencies]
 mlserver = {path = "../..", develop = true}
+mlserver_sklearn = {path = "../../runtimes/sklearn", develop = true}
+mlserver_xgboost = {path = "../../runtimes/xgboost", develop = true}
+mlserver_lightgbm = {path = "../../runtimes/lightgbm", develop = true}
 tensorflow = "~2.12.0"
 requests-mock = "~1.10.0"
 types-requests = "~2.28.11.5"