diff --git a/python-package/xgboost/__init__.py b/python-package/xgboost/__init__.py index f271e073cfd3..dc1a0869b30c 100644 --- a/python-package/xgboost/__init__.py +++ b/python-package/xgboost/__init__.py @@ -3,8 +3,11 @@ Contributors: https://github.com/dmlc/xgboost/blob/master/CONTRIBUTORS.md """ -from . import tracker # noqa -from . import collective +from . import ( + collective, + interpret, + tracker, # noqa +) from ._c_api import _py_version from .core import ( Booster, @@ -62,4 +65,6 @@ "XGBRFRegressor", # collective "collective", + # interpretability + "interpret", ] diff --git a/python-package/xgboost/interpret.py b/python-package/xgboost/interpret.py new file mode 100644 index 000000000000..e2e247396678 --- /dev/null +++ b/python-package/xgboost/interpret.py @@ -0,0 +1,146 @@ +"""Interpretability functions for XGBoost models.""" + +from typing import Optional, Tuple, Union + +import numpy as np + +from ._typing import ArrayLike, IterationRange +from .core import Booster, DMatrix + + +def _as_booster(model: object) -> Booster: + if isinstance(model, Booster): + return model + get_booster = getattr(model, "get_booster", None) + if get_booster is None: + raise TypeError( + "`model` must be an xgboost.Booster or an object with get_booster()." + ) + booster = get_booster() + if not isinstance(booster, Booster): + raise TypeError("`model.get_booster()` must return an xgboost.Booster.") + return booster + + +def _get_iteration_range( + model: object, iteration_range: Optional[IterationRange] +) -> IterationRange: + get_iteration_range = getattr(model, "_get_iteration_range", None) + if get_iteration_range is not None: + return get_iteration_range(iteration_range) + if iteration_range is None: + return (0, 0) + return iteration_range + + +def _as_prediction_dmatrix(model: object, X: Union[DMatrix, ArrayLike]) -> DMatrix: + if isinstance(X, DMatrix): + return X + + return DMatrix( + X, + missing=getattr(model, "missing", None), + nthread=getattr(model, "n_jobs", None), + feature_types=getattr(model, "feature_types", None), + enable_categorical=getattr(model, "enable_categorical", False), + ) + + +def _predict_contribs( + booster: Booster, + data: DMatrix, + *, + device: Optional[str], + kwargs: dict, +) -> np.ndarray: + if device is None: + return booster.predict(data, **kwargs) + + config = booster.save_config() + try: + booster.set_param({"device": device}) + return booster.predict(data, **kwargs) + finally: + booster.load_config(config) + + +def shap_values( # pylint: disable=too-many-arguments + model: object, + X: Union[DMatrix, ArrayLike], + *, + X_background: Optional[Union[DMatrix, ArrayLike]] = None, + device: Optional[str] = None, + output_margin: bool = False, + iteration_range: Optional[IterationRange] = None, + approx: bool = False, + validate_features: bool = True, + return_bias: bool = False, +) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + """Return SHAP values for an XGBoost model. + + This function accepts either a :py:class:`xgboost.Booster` or an sklearn-style + XGBoost model and wraps :py:meth:`xgboost.Booster.predict` with + ``pred_contribs=True``. The final bias column returned by ``predict`` is + removed from the default return value. + + Parameters + ---------- + model : + XGBoost booster or sklearn-style XGBoost model. + X : + Input data. + X_background : + Background data for interventional SHAP values. This is reserved for a + future implementation and is currently unsupported. + device : + Optional prediction device override, such as ``"cpu"``, ``"cuda"``, or + ``"cuda:0"``. The model's original configuration is restored after + prediction. This option temporarily mutates the underlying Booster and + is not safe for concurrent use of the same model. + output_margin : + Accepted for API compatibility. SHAP contributions currently correspond + to the model margin, matching ``Booster.predict(pred_contribs=True)``. + iteration_range : + Specifies which layer of trees are used in prediction. + approx : + Use approximate SHAP contributions. + validate_features : + Validate feature names between the model and input data. + return_bias : + When True, return ``(values, bias)``. + + Returns + ------- + values : + Feature SHAP values, excluding the bias term. + values, bias : + Returned when ``return_bias`` is True. + """ + if X_background is not None: + raise NotImplementedError("`X_background` is not yet supported.") + # Existing SHAP prediction always returns margin contributions. Keep this + # argument in the initial API so callers can use the proposed signature. + _ = output_margin + + booster = _as_booster(model) + data = _as_prediction_dmatrix(model, X) + contribs = _predict_contribs( + booster, + data, + device=device, + kwargs={ + "pred_contribs": True, + "approx_contribs": approx, + "validate_features": validate_features, + "iteration_range": _get_iteration_range(model, iteration_range), + }, + ) + + values = contribs[..., :-1] + bias = contribs[..., -1] + if return_bias: + return values, bias + return values + + +__all__ = ["shap_values"] diff --git a/tests/python/test_interpret.py b/tests/python/test_interpret.py new file mode 100644 index 000000000000..28ac568afd0b --- /dev/null +++ b/tests/python/test_interpret.py @@ -0,0 +1,92 @@ +import numpy as np +import pytest +import xgboost as xgb +from xgboost import interpret + + +def test_shap_values_matches_predict() -> None: + rng = np.random.RandomState(1994) + X = rng.randn(16, 4) + y = rng.randn(16) + booster = xgb.train({"tree_method": "hist"}, xgb.DMatrix(X, label=y), 4) + + values, bias = interpret.shap_values(booster, X, return_bias=True) + contribs = booster.predict(xgb.DMatrix(X), pred_contribs=True) + + np.testing.assert_allclose(values, contribs[:, :-1]) + np.testing.assert_allclose(bias, contribs[:, -1]) + np.testing.assert_allclose(interpret.shap_values(booster, X), contribs[:, :-1]) + + +def test_shap_values_accepts_sklearn_model() -> None: + rng = np.random.RandomState(1995) + X = rng.randn(16, 4) + y = rng.randn(16) + reg = xgb.XGBRegressor(n_estimators=4, tree_method="hist") + reg.fit(X, y) + + values = interpret.shap_values(reg, X) + contribs = reg.get_booster().predict(xgb.DMatrix(X), pred_contribs=True) + + np.testing.assert_allclose(values, contribs[:, :-1]) + + +def test_shap_values_uses_sklearn_iteration_range() -> None: + rng = np.random.RandomState(1996) + X = rng.randn(64, 4) + y = rng.randn(64) + reg = xgb.XGBRegressor(n_estimators=8, tree_method="hist") + reg.fit(X, y) + reg.get_booster().set_attr(best_iteration="3") + + values = interpret.shap_values(reg, X, iteration_range=(0, 0)) + contribs = reg.get_booster().predict( + xgb.DMatrix(X), pred_contribs=True, iteration_range=(0, 4) + ) + + np.testing.assert_allclose(values, contribs[:, :-1]) + + +def test_shap_values_rejects_background_data() -> None: + rng = np.random.RandomState(1997) + X = rng.randn(16, 4) + y = rng.randn(16) + booster = xgb.train({"tree_method": "hist"}, xgb.DMatrix(X, label=y), 4) + + with pytest.raises(NotImplementedError, match="X_background"): + interpret.shap_values(booster, X, X_background=X) + + +def test_shap_values_device_override_restores_config() -> None: + rng = np.random.RandomState(1998) + X = rng.randn(16, 4) + y = rng.randn(16) + booster = xgb.train({"tree_method": "hist"}, xgb.DMatrix(X, label=y), 4) + config = booster.save_config() + + values = interpret.shap_values(booster, X, device="cpu") + contribs = booster.predict(xgb.DMatrix(X), pred_contribs=True) + + np.testing.assert_allclose(values, contribs[:, :-1]) + assert booster.save_config() == config + + +def test_shap_values_device_override_restores_config_on_error() -> None: + rng = np.random.RandomState(1999) + X = rng.randn(16, 4) + y = rng.randn(16) + booster = xgb.train( + {"tree_method": "hist"}, + xgb.DMatrix(X, label=y, feature_names=["a", "b", "c", "d"]), + 4, + ) + config = booster.save_config() + + with pytest.raises(ValueError, match="feature_names mismatch"): + interpret.shap_values( + booster, + xgb.DMatrix(X, feature_names=["q", "b", "c", "d"]), + device="cpu", + ) + + assert booster.save_config() == config