wehub-resource-sync
292 lines
12 KiB
Python
292 lines
12 KiB
Python
import functools
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
from packaging.version import Version
|
|
from sklearn.pipeline import Pipeline as sk_Pipeline
|
|
|
|
import mlflow
|
|
from mlflow import MlflowException
|
|
from mlflow.models.evaluation.base import EvaluationMetric, EvaluationResult, _ModelType
|
|
from mlflow.models.evaluation.default_evaluator import (
|
|
BuiltInEvaluator,
|
|
_extract_predict_fn,
|
|
_extract_raw_model,
|
|
_get_dataframe_with_renamed_columns,
|
|
)
|
|
from mlflow.models.evaluation.evaluators.classifier import (
|
|
_is_continuous,
|
|
_suppress_class_imbalance_errors,
|
|
)
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.pyfunc import _ServedPyFuncModel
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
_SUPPORTED_SHAP_ALGORITHMS = ("exact", "permutation", "partition", "kernel")
|
|
_DEFAULT_SAMPLE_ROWS_FOR_SHAP = 2000
|
|
|
|
|
|
def _shap_predict_fn(x, predict_fn, feature_names):
|
|
return predict_fn(_get_dataframe_with_renamed_columns(x, feature_names))
|
|
|
|
|
|
class ShapEvaluator(BuiltInEvaluator):
|
|
"""
|
|
A built-in evaluator to get SHAP explainability insights for classifier and regressor models.
|
|
|
|
This evaluator often run with the main evaluator for the model like ClassifierEvaluator.
|
|
"""
|
|
|
|
name = "shap"
|
|
|
|
@classmethod
|
|
def can_evaluate(cls, *, model_type, evaluator_config, **kwargs):
|
|
return model_type in (_ModelType.CLASSIFIER, _ModelType.REGRESSOR) and evaluator_config.get(
|
|
"log_model_explainability", True
|
|
)
|
|
|
|
def _evaluate(
|
|
self,
|
|
model: Optional["mlflow.pyfunc.PyFuncModel"],
|
|
extra_metrics: list[EvaluationMetric],
|
|
custom_artifacts=None,
|
|
**kwargs,
|
|
) -> EvaluationResult | None:
|
|
if isinstance(model, _ServedPyFuncModel):
|
|
_logger.warning(
|
|
"Skipping model explainability because a model server is used for environment "
|
|
"restoration."
|
|
)
|
|
return
|
|
|
|
model_loader_module, raw_model = _extract_raw_model(model)
|
|
if model_loader_module == "mlflow.spark":
|
|
# TODO: Shap explainer need to manipulate on each feature values,
|
|
# but spark model input dataframe contains Vector type feature column
|
|
# which shap explainer does not support.
|
|
# To support this, we need expand the Vector type feature column into
|
|
# multiple scalar feature columns and pass it to shap explainer.
|
|
_logger.warning(
|
|
"Logging model explainability insights is not currently supported for PySpark "
|
|
"models."
|
|
)
|
|
return
|
|
|
|
self.y_true = self.dataset.labels_data
|
|
self.label_list = self.evaluator_config.get("label_list")
|
|
self.pos_label = self.evaluator_config.get("pos_label")
|
|
|
|
if not (np.issubdtype(self.y_true.dtype, np.number) or self.y_true.dtype == np.bool_):
|
|
# Note: python bool type inherits number type but np.bool_ does not inherit np.number.
|
|
_logger.warning(
|
|
"Skip logging model explainability insights because it requires all label "
|
|
"values to be numeric or boolean."
|
|
)
|
|
return
|
|
|
|
algorithm = self.evaluator_config.get("explainability_algorithm", None)
|
|
if algorithm is not None and algorithm not in _SUPPORTED_SHAP_ALGORITHMS:
|
|
raise MlflowException(
|
|
message=f"Specified explainer algorithm {algorithm} is unsupported. Currently only "
|
|
f"support {','.join(_SUPPORTED_SHAP_ALGORITHMS)} algorithms.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if algorithm != "kernel":
|
|
feature_dtypes = list(self.X.get_original().dtypes)
|
|
for feature_dtype in feature_dtypes:
|
|
if not np.issubdtype(feature_dtype, np.number):
|
|
_logger.warning(
|
|
"Skip logging model explainability insights because the shap explainer "
|
|
f"{algorithm} requires all feature values to be numeric, and each feature "
|
|
"column must only contain scalar values."
|
|
)
|
|
return
|
|
|
|
try:
|
|
import shap
|
|
from matplotlib import pyplot
|
|
except ImportError:
|
|
_logger.warning(
|
|
"SHAP or matplotlib package is not installed, so model explainability insights "
|
|
"will not be logged."
|
|
)
|
|
return
|
|
|
|
if Version(shap.__version__) < Version("0.40"):
|
|
_logger.warning(
|
|
"Shap package version is lower than 0.40, Skip log model explainability."
|
|
)
|
|
return
|
|
|
|
sample_rows = self.evaluator_config.get(
|
|
"explainability_nsamples", _DEFAULT_SAMPLE_ROWS_FOR_SHAP
|
|
)
|
|
|
|
X_df = self.X.copy_to_avoid_mutation()
|
|
|
|
sampled_X = shap.sample(X_df, sample_rows, random_state=0)
|
|
|
|
mode_or_mean_dict = _compute_df_mode_or_mean(X_df)
|
|
sampled_X = sampled_X.fillna(mode_or_mean_dict)
|
|
|
|
# shap explainer might call provided `predict_fn` with a `numpy.ndarray` type
|
|
# argument, this might break some model inference, so convert the argument into
|
|
# a pandas dataframe.
|
|
# The `shap_predict_fn` calls model's predict function, we need to restore the input
|
|
# dataframe with original column names, because some model prediction routine uses
|
|
# the column name.
|
|
|
|
predict_fn = _extract_predict_fn(model)
|
|
shap_predict_fn = functools.partial(
|
|
_shap_predict_fn, predict_fn=predict_fn, feature_names=self.dataset.feature_names
|
|
)
|
|
|
|
if self.label_list is None:
|
|
# If label list is not specified, infer label list from model output.
|
|
# We need to copy the input data as the model might mutate the input data.
|
|
y_pred = predict_fn(X_df.copy()) if predict_fn else self.dataset.predictions_data
|
|
self.label_list = np.unique(np.concatenate([self.y_true, y_pred]))
|
|
|
|
try:
|
|
if algorithm:
|
|
if algorithm == "kernel":
|
|
# We need to lazily import shap, so lazily import `_PatchedKernelExplainer`
|
|
from mlflow.models.evaluation._shap_patch import _PatchedKernelExplainer
|
|
|
|
kernel_link = self.evaluator_config.get(
|
|
"explainability_kernel_link", "identity"
|
|
)
|
|
if kernel_link not in ["identity", "logit"]:
|
|
raise ValueError(
|
|
"explainability_kernel_link config can only be set to 'identity' or "
|
|
f"'logit', but got '{kernel_link}'."
|
|
)
|
|
background_X = shap.sample(X_df, sample_rows, random_state=3)
|
|
background_X = background_X.fillna(mode_or_mean_dict)
|
|
|
|
explainer = _PatchedKernelExplainer(
|
|
shap_predict_fn, background_X, link=kernel_link
|
|
)
|
|
else:
|
|
explainer = shap.Explainer(
|
|
shap_predict_fn,
|
|
sampled_X,
|
|
feature_names=self.dataset.feature_names,
|
|
algorithm=algorithm,
|
|
)
|
|
else:
|
|
if (
|
|
raw_model
|
|
and not len(self.label_list) > 2
|
|
and not isinstance(raw_model, sk_Pipeline)
|
|
):
|
|
# For mulitnomial classifier, shap.Explainer may choose Tree/Linear explainer
|
|
# for raw model, this case shap plot doesn't support it well, so exclude the
|
|
# multinomial_classifier case here.
|
|
explainer = shap.Explainer(
|
|
raw_model, sampled_X, feature_names=self.dataset.feature_names
|
|
)
|
|
else:
|
|
# fallback to default explainer
|
|
explainer = shap.Explainer(
|
|
shap_predict_fn, sampled_X, feature_names=self.dataset.feature_names
|
|
)
|
|
|
|
_logger.info(f"Shap explainer {explainer.__class__.__name__} is used.")
|
|
|
|
shap_values = explainer(sampled_X)
|
|
except Exception as e:
|
|
# Shap evaluation might fail on some edge cases, e.g., unsupported input data values
|
|
# or unsupported model on specific shap explainer. Catch exception to prevent it
|
|
# breaking the whole `evaluate` function.
|
|
|
|
if not self.evaluator_config.get("ignore_exceptions", True):
|
|
raise e
|
|
|
|
_logger.warning(
|
|
f"Shap evaluation failed. Reason: {e!r}. "
|
|
"Set logging level to DEBUG to see the full traceback."
|
|
)
|
|
_logger.debug("", exc_info=True)
|
|
return
|
|
|
|
if self.evaluator_config.get("log_explainer", False):
|
|
try:
|
|
mlflow.shap.log_explainer(explainer, name="explainer")
|
|
except Exception as e:
|
|
# TODO: The explainer saver is buggy, if `get_underlying_model_flavor` return
|
|
# "unknown", then fallback to shap explainer saver, and shap explainer will call
|
|
# `model.save` for sklearn model, there is no `.save` method, so error will
|
|
# happen.
|
|
_logger.warning(
|
|
f"Logging explainer failed. Reason: {e!r}. "
|
|
"Set logging level to DEBUG to see the full traceback."
|
|
)
|
|
_logger.debug("", exc_info=True)
|
|
|
|
def _adjust_color_bar():
|
|
pyplot.gcf().axes[-1].set_aspect("auto")
|
|
pyplot.gcf().axes[-1].set_box_aspect(50)
|
|
|
|
def _adjust_axis_tick():
|
|
pyplot.xticks(fontsize=10)
|
|
pyplot.yticks(fontsize=10)
|
|
|
|
def plot_beeswarm():
|
|
shap.plots.beeswarm(shap_values, show=False, color_bar=True)
|
|
_adjust_color_bar()
|
|
_adjust_axis_tick()
|
|
|
|
with _suppress_class_imbalance_errors(ValueError, log_warning=False):
|
|
self._log_image_artifact(
|
|
plot_beeswarm,
|
|
"shap_beeswarm_plot",
|
|
)
|
|
|
|
def plot_summary():
|
|
shap.summary_plot(shap_values, show=False, color_bar=True)
|
|
_adjust_color_bar()
|
|
_adjust_axis_tick()
|
|
|
|
with _suppress_class_imbalance_errors(TypeError, log_warning=False):
|
|
self._log_image_artifact(
|
|
plot_summary,
|
|
"shap_summary_plot",
|
|
)
|
|
|
|
def plot_feature_importance():
|
|
shap.plots.bar(shap_values, show=False)
|
|
_adjust_axis_tick()
|
|
|
|
with _suppress_class_imbalance_errors(IndexError, log_warning=False):
|
|
self._log_image_artifact(
|
|
plot_feature_importance,
|
|
"shap_feature_importance_plot",
|
|
)
|
|
|
|
return EvaluationResult(
|
|
metrics=self.aggregate_metrics,
|
|
artifacts=self.artifacts,
|
|
run_id=self.run_id,
|
|
)
|
|
|
|
|
|
def _compute_df_mode_or_mean(df):
|
|
"""
|
|
Compute mean (for continuous columns) and compute mode (for other columns) for the
|
|
input dataframe, return a dict, key is column name, value is the corresponding mode or
|
|
mean value, this function calls `_is_continuous` to determine whether the
|
|
column is continuous column.
|
|
"""
|
|
continuous_cols = [c for c in df.columns if _is_continuous(df[c])]
|
|
df_cont = df[continuous_cols]
|
|
df_non_cont = df.drop(continuous_cols, axis=1)
|
|
|
|
means = {} if df_cont.empty else df_cont.mean().to_dict()
|
|
modes = {} if df_non_cont.empty else df_non_cont.mode().loc[0].to_dict()
|
|
return means | modes
|