feature: Add support for Partial Dependence Plots(PDP) in SageMaker Clarify

Author: keerthanvasist

src/sagemaker/clarify.py

@@ -20,6 +20,8 @@
 import os
 import tempfile
 import re
+from typing import List
+
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 from sagemaker import image_uris, s3, utils
 
@@ -54,7 +56,11 @@ def __init__(
                 "ShardedByS3Key".
             s3_compression_type (str): Valid options are "None" or "Gzip".
         """
-        if dataset_type not in ["text/csv", "application/jsonlines", "application/x-parquet"]:
+        if dataset_type not in [
+            "text/csv",
+            "application/jsonlines",
+            "application/x-parquet",
+        ]:
             raise ValueError(
                 f"Invalid dataset_type '{dataset_type}'."
                 f" Please check the API documentation for the supported dataset types."
@@ -292,7 +298,30 @@ class ExplainabilityConfig(ABC):
     @abstractmethod
     def get_explainability_config(self):
         """Returns config."""
-        return None
+
+
+class PDPConfig(ExplainabilityConfig):
+    """Config class for Partial Dependence Plots (PDP)"""
+
+    def __init__(self, features=None, grid_resolution=None):
+        """Initializes config for PDP.
+
+        Args:
+            features (None or list): List of features names or indices for which partial dependence
+                plots must be computed and plotted.
+            grid_resolution (int): In case of numerical features, this number represents that
+                number of buckets that range of values must be divided into. This decides the
+                granularity of the grid in which the PDP are plotted.
+        """
+        self.pdp_config = {}
+        if features is not None:
+            self.pdp_config["features"] = features
+        if grid_resolution is not None:
+            self.pdp_config["grid_resolution"] = grid_resolution
+
+    def get_explainability_config(self):
+        """Returns config."""
+        return {"pdp": copy.deepcopy(self.pdp_config)}
 
 
 class SHAPConfig(ExplainabilityConfig):
@@ -466,7 +495,10 @@ def _run(
                 will be unassociated.
                 * `TrialComponentDisplayName` is used for display in Studio.
         """
-        analysis_config["methods"]["report"] = {"name": "report", "title": "Analysis Report"}
+        analysis_config["methods"]["report"] = {
+            "name": "report",
+            "title": "Analysis Report",
+        }
         with tempfile.TemporaryDirectory() as tmpdirname:
             analysis_config_file = os.path.join(tmpdirname, "analysis_config.json")
             with open(analysis_config_file, "w") as f:
@@ -568,7 +600,15 @@ def run_pre_training_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Pretraining-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
+        self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
 
     def run_post_training_bias(
         self,
@@ -646,7 +686,15 @@ def run_post_training_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Posttraining-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
+        self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
 
     def run_bias(
         self,
@@ -741,7 +789,15 @@ def run_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
+        self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
 
     def run_explainability(
         self,
@@ -772,7 +828,8 @@ def run_explainability(
             model_config (:class:`~sagemaker.clarify.ModelConfig`): Config of the model and its
                 endpoint to be created.
             explainability_config (:class:`~sagemaker.clarify.ExplainabilityConfig`): Config of the
-                specific explainability method. Currently, only SHAP is supported.
+                specific explainability method or a list of ExplainabilityConfig objects. Currently,
+                SHAP and PDP are the two methods supported.
             model_scores(str|int|ModelPredictedLabelConfig):  Index or JSONPath location in the
                 model output for the predicted scores to be explained. This is not required if the
                 model output is a single score. Alternatively, an instance of
@@ -781,7 +838,7 @@ def run_explainability(
             logs (bool): Whether to show the logs produced by the job.
                 Only meaningful when ``wait`` is True (default: True).
             job_name (str): Processing job name. When ``job_name`` is not specified, if
-                ``job_name_prefix`` in :class:`SageMakerClarifyProcessor` specified, the job name
+                `job_name_prefix` in :class:`SageMakerClarifyProcessor` specified, the job name
                 will be composed of ``job_name_prefix`` and current timestamp; otherwise use
                 "Clarify-Explainability" as prefix.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
@@ -801,19 +858,44 @@ def run_explainability(
         analysis_config = data_config.get_config()
         predictor_config = model_config.get_predictor_config()
         if isinstance(model_scores, ModelPredictedLabelConfig):
-            probability_threshold, predicted_label_config = model_scores.get_predictor_config()
+            (
+                probability_threshold,
+                predicted_label_config,
+            ) = model_scores.get_predictor_config()
             _set(probability_threshold, "probability_threshold", analysis_config)
             predictor_config.update(predicted_label_config)
         else:
             _set(model_scores, "label", predictor_config)
-        analysis_config["methods"] = explainability_config.get_explainability_config()
+
+        explainability_methods = {}
+        if isinstance(explainability_config, List):  # pylint: disable=W1116
+            for config in explainability_config:
+                if not isinstance(config, ExplainabilityConfig):
+                    raise ValueError(
+                        f"Invalid input: Excepted ExplainabilityConfig, got {type(config)} instead"
+                    )
+                explain_config = config.get_explainability_config()
+                explainability_methods[list(explain_config.keys())[0]] = explain_config[
+                    list(explain_config.keys())[0]
+                ]
+        elif isinstance(explainability_config, ExplainabilityConfig):
+            explainability_methods = explainability_config.get_explainability_config()
+        analysis_config["methods"] = explainability_methods
         analysis_config["predictor"] = predictor_config
         if job_name is None:
             if self.job_name_prefix:
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Explainability")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
+        self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
 
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):

tests/integ/test_clarify.py

@@ -13,7 +13,6 @@
 
 from __future__ import print_function, absolute_import
 
-
 import json
 import math
 import numpy as np
@@ -31,14 +30,14 @@
     ModelConfig,
     ModelPredictedLabelConfig,
     SHAPConfig,
+    PDPConfig,
 )
 
 from sagemaker.amazon.linear_learner import LinearLearner, LinearLearnerPredictor
 from sagemaker import utils
 from tests import integ
 from tests.integ import timeout
 
-
 CLARIFY_DEFAULT_TIMEOUT_MINUTES = 15
 
 
@@ -177,6 +176,11 @@ def shap_config():
     )
 
 
+@pytest.fixture(scope="module")
+def pdp_config():
+    return PDPConfig(features=["F1"], grid_resolution=10)
+
+
 def test_pre_training_bias(clarify_processor, data_config, data_bias_config, sagemaker_session):
     with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
         clarify_processor.run_pre_training_bias(
@@ -258,6 +262,57 @@ def test_shap(clarify_processor, data_config, model_config, shap_config, sagemak
         check_analysis_config(data_config, sagemaker_session, "shap")
 
 
+def test_pdp(clarify_processor, data_config, model_config, pdp_config, sagemaker_session):
+    with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
+        clarify_processor.run_explainability(
+            data_config,
+            model_config,
+            pdp_config,
+            model_scores="score",
+            job_name=utils.unique_name_from_base("clarify-explainability-pdp"),
+            wait=True,
+        )
+        analysis_result_json = s3.S3Downloader.read_file(
+            data_config.s3_output_path + "/analysis.json",
+            sagemaker_session,
+        )
+        analysis_result = json.loads(analysis_result_json)
+        print(analysis_result)
+        assert analysis_result["explanations"]["pdp"][0]["feature_name"] == "F1"
+
+        check_analysis_config(data_config, sagemaker_session, "pdp")
+
+
+def test_shap_and_pdp(
+    clarify_processor, data_config, model_config, shap_config, pdp_config, sagemaker_session
+):
+    with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
+        clarify_processor.run_explainability(
+            data_config,
+            model_config,
+            [shap_config, pdp_config],
+            model_scores="score",
+            job_name=utils.unique_name_from_base("clarify-explainability"),
+            wait=True,
+        )
+        analysis_result_json = s3.S3Downloader.read_file(
+            data_config.s3_output_path + "/analysis.json",
+            sagemaker_session,
+        )
+        analysis_result = json.loads(analysis_result_json)
+        print(analysis_result)
+        assert (
+            math.fabs(
+                analysis_result["explanations"]["kernel_shap"]["label0"]["global_shap_values"]["F2"]
+            )
+            <= 1
+        )
+        assert analysis_result["explanations"]["pdp"][0]["feature_name"] == "F1"
+
+        check_analysis_config(data_config, sagemaker_session, "pdp")
+        check_analysis_config(data_config, sagemaker_session, "shap")
+
+
 def check_analysis_config(data_config, sagemaker_session, method):
     analysis_config_json = s3.S3Downloader.read_file(
         data_config.s3_output_path + "/analysis_config.json",