feature: extracted analysis config generation for explainability

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -1307,6 +1307,36 @@ def run_explainability(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
+        analysis_config = _AnalysisConfigGenerator.explainability(
+            data_config,
+            model_config,
+            model_scores,
+            explainability_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")
+        return self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
+
+
+class _AnalysisConfigGenerator:
+    @staticmethod
+    def explainability(
+        data_config,
+        model_config,
+        model_scores,
+        explainability_config
+    ):
         analysis_config = data_config.get_config()
         predictor_config = model_config.get_predictor_config()
         if isinstance(model_scores, ModelPredictedLabelConfig):
@@ -1329,34 +1359,21 @@ def run_explainability(
             if not len(explainability_methods.keys()) == len(explainability_config):
                 raise ValueError("Duplicate explainability configs are provided")
             if (
-                "shap" not in explainability_methods
-                and explainability_methods["pdp"].get("features", None) is None
+                    "shap" not in explainability_methods
+                    and explainability_methods["pdp"].get("features", None) is None
             ):
                 raise ValueError("PDP features must be provided when ShapConfig is not provided")
         else:
             if (
-                isinstance(explainability_config, PDPConfig)
-                and explainability_config.get_explainability_config()["pdp"].get("features", None)
-                is None
+                    isinstance(explainability_config, PDPConfig)
+                    and explainability_config.get_explainability_config()["pdp"].get("features", None)
+                    is None
             ):
                 raise ValueError("PDP features must be provided when ShapConfig is not provided")
             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")
-        return self._run(
-            data_config,
-            analysis_config,
-            wait,
-            logs,
-            job_name,
-            kms_key,
-            experiment_config,
-        )
+        return analysis_config
 
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):

tests/unit/test_clarify.py

@@ -29,6 +29,7 @@
     SHAPConfig,
     TextConfig,
     ImageConfig,
+    _AnalysisConfigGenerator,
 )
 
 JOB_NAME_PREFIX = "my-prefix"
@@ -1277,3 +1278,28 @@ def test_shap_with_image_config(
         expected_predictor_config,
         expected_image_config=expected_image_config,
     )
+
+
+def test_analysis_config_generator_for_explainability(data_config, model_config):
+    model_scores = ModelPredictedLabelConfig(
+        probability="pr",
+        label_headers=["success"],
+    )
+    actual = _AnalysisConfigGenerator.explainability(
+        data_config,
+        model_config,
+        model_scores,
+        SHAPConfig(),
+    )
+    expected = {'dataset_type': 'text/csv',
+                'headers': ['Label', 'F1', 'F2', 'F3', 'F4'],
+                'joinsource_name_or_index': 'F4',
+                'label': 'Label',
+                'methods': {'shap': {'save_local_shap_values': True, 'use_logit': False}},
+                'predictor': {'initial_instance_count': 1,
+                              'instance_type': 'ml.c5.xlarge',
+                              'label_headers': ['success'],
+                              'model_name': 'xgboost-model',
+                              'probability': 'pr'}}
+    assert actual == expected
+