refactored _AnalysisConfigGenerator

to re-use for generating a config for both bias and explainability at once

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -1368,68 +1368,70 @@ def run_explainability(
             experiment_config,
         )
 
+    def run_bias_and_explainability(self):
+        """
+        TODO:
+        - add doc string
+        - add logic
+        - add tests
+        """
+        raise NotImplementedError(
+            "Please choose a method of run_pre_training_bias, run_post_training_bias or run_explainability."
+        )
+
 
 class _AnalysisConfigGenerator:
     """
     Creates analysis_config objects for different type of runs.
     """
 
     @classmethod
-    def explainability(
+    def bias_and_explainability(
         cls,
         data_config: DataConfig,
         model_config: ModelConfig,
-        model_scores: ModelPredictedLabelConfig,
-        explainability_config: ExplainabilityConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]],
+        bias_config: BiasConfig,
+        pre_training_methods: Union[str, List[str]] = "all",
+        post_training_methods: Union[str, List[str]] = "all",
     ):
-        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()
-            _set(probability_threshold, "probability_threshold", analysis_config)
-            predictor_config.update(predicted_label_config)
-        else:
-            _set(model_scores, "label", predictor_config)
+        analysis_config = {**data_config.get_config(), **bias_config.get_config()}
+        analysis_config = cls._add_methods(
+            analysis_config,
+            pre_training_methods=pre_training_methods,
+            post_training_methods=post_training_methods,
+            explainability_config=explainability_config,
+        )
+        analysis_config = cls._add_predictor(
+            analysis_config, model_config, model_predicted_label_config
+        )
+        return analysis_config
 
-        explainability_methods = {}
-        if isinstance(explainability_config, list):
-            if len(explainability_config) == 0:
-                raise ValueError("Please provide at least one explainability config.")
-            for config in explainability_config:
-                explain_config = config.get_explainability_config()
-                explainability_methods.update(explain_config)
-            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
-            ):
-                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
-            ):
-                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
-        return cls._common(analysis_config)
+    @classmethod
+    def explainability(
+        cls,
+        data_config: DataConfig,
+        model_config: ModelConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]],
+    ):
+        analysis_config = data_config.analysis_config
+        analysis_config = cls._add_predictor(
+            analysis_config, model_config, model_predicted_label_config
+        )
+        analysis_config = cls._add_methods(
+            analysis_config, explainability_config=explainability_config
+        )
+        return analysis_config
 
     @classmethod
     def bias_pre_training(
         cls, data_config: DataConfig, bias_config: BiasConfig, methods: Union[str, List[str]]
     ):
-        analysis_config = {
-            **data_config.get_config(),
-            **bias_config.get_config(),
-            "methods": {"pre_training_bias": {"methods": methods}},
-        }
-        return cls._common(analysis_config)
+        analysis_config = {**data_config.get_config(), **bias_config.get_config()}
+        analysis_config = cls._add_methods(analysis_config, pre_training_methods=methods)
+        return analysis_config
 
     @classmethod
     def bias_post_training(
@@ -1440,21 +1442,12 @@ def bias_post_training(
         methods: Union[str, List[str]],
         model_config: ModelConfig,
     ):
-        analysis_config = {
-            **data_config.get_config(),
-            **bias_config.get_config(),
-            "predictor": {**model_config.get_predictor_config()},
-            "methods": {"post_training_bias": {"methods": methods}},
-        }
-        if model_predicted_label_config:
-            (
-                probability_threshold,
-                predictor_config,
-            ) = model_predicted_label_config.get_predictor_config()
-            if predictor_config:
-                analysis_config["predictor"].update(predictor_config)
-            _set(probability_threshold, "probability_threshold", analysis_config)
-        return cls._common(analysis_config)
+        analysis_config = {**data_config.get_config(), **bias_config.get_config()}
+        analysis_config = cls._add_methods(analysis_config, post_training_methods=methods)
+        analysis_config = cls._add_predictor(
+            analysis_config, model_config, model_predicted_label_config
+        )
+        return analysis_config
 
     @classmethod
     def bias(
@@ -1466,33 +1459,95 @@ def bias(
         pre_training_methods: Union[str, List[str]] = "all",
         post_training_methods: Union[str, List[str]] = "all",
     ):
-        analysis_config = {
-            **data_config.get_config(),
-            **bias_config.get_config(),
-            "predictor": model_config.get_predictor_config(),
-            "methods": {
-                "pre_training_bias": {"methods": pre_training_methods},
-                "post_training_bias": {"methods": post_training_methods},
-            },
-        }
-        if model_predicted_label_config:
+        analysis_config = {**data_config.get_config(), **bias_config.get_config()}
+        analysis_config = cls._add_methods(
+            analysis_config,
+            pre_training_methods=pre_training_methods,
+            post_training_methods=post_training_methods,
+        )
+        analysis_config = cls._add_predictor(
+            analysis_config, model_config, model_predicted_label_config
+        )
+        return analysis_config
+
+    @classmethod
+    def _add_predictor(cls, analysis_config, model_config, model_predicted_label_config):
+        analysis_config = {**analysis_config}
+        analysis_config["predictor"] = model_config.get_predictor_config()
+        if isinstance(model_predicted_label_config, ModelPredictedLabelConfig):
             (
                 probability_threshold,
                 predictor_config,
             ) = model_predicted_label_config.get_predictor_config()
             if predictor_config:
                 analysis_config["predictor"].update(predictor_config)
             _set(probability_threshold, "probability_threshold", analysis_config)
-        return cls._common(analysis_config)
+        else:
+            _set(model_predicted_label_config, "label", analysis_config["predictor"])
+        return analysis_config
 
-    @staticmethod
-    def _common(analysis_config):
-        analysis_config["methods"]["report"] = {
-            "name": "report",
-            "title": "Analysis Report",
-        }
+    @classmethod
+    def _add_methods(
+        cls,
+        analysis_config,
+        pre_training_methods=None,
+        post_training_methods=None,
+        explainability_config=None,
+        report=True,
+    ):
+        # validate
+        params = [pre_training_methods, post_training_methods, explainability_config]
+        if all([1 if p is None else 0 for p in params]):
+            raise AttributeError(
+                "analysis_config must have at least one working method: "
+                "One of the `pre_training_methods`, `post_training_methods`, `explainability_config`."
+            )
+
+        # main logic
+        analysis_config = {**analysis_config}
+        if "methods" not in analysis_config:
+            analysis_config["methods"] = {}
+
+        if report:
+            analysis_config["methods"]["report"] = {"name": "report", "title": "Analysis Report"}
+
+        if pre_training_methods:
+            analysis_config["methods"]["pre_training_bias"] = {"methods": pre_training_methods}
+
+        if post_training_methods:
+            analysis_config["methods"]["post_training_bias"] = {"methods": post_training_methods}
+
+        if explainability_config is not None:
+            explainability_methods = cls._merge_explainability_configs(explainability_config)
+            analysis_config["methods"] = {**analysis_config["methods"], **explainability_methods}
         return analysis_config
 
+    @classmethod
+    def _merge_explainability_configs(
+        cls, explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]]
+    ):
+        if isinstance(explainability_config, list):
+            explainability_methods = {}
+            if len(explainability_config) == 0:
+                raise ValueError("Please provide at least one explainability config.")
+            for config in explainability_config:
+                explain_config = config.get_explainability_config()
+                explainability_methods.update(explain_config)
+            if not len(explainability_methods) == len(explainability_config):
+                raise ValueError("Duplicate explainability configs are provided")
+            if (
+                "shap" not in explainability_methods
+                and "features" not in explainability_methods["pdp"]
+            ):
+                raise ValueError("PDP features must be provided when ShapConfig is not provided")
+            return explainability_methods
+        if (
+            isinstance(explainability_config, PDPConfig)
+            and "features" not in explainability_config.get_explainability_config()["pdp"]
+        ):
+            raise ValueError("PDP features must be provided when ShapConfig is not provided")
+        return explainability_config.get_explainability_config()
+
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):
     """Uploads the local ``analysis_config_file`` to the ``s3_output_path``.

tests/unit/test_clarify.py

@@ -1094,7 +1094,9 @@ def test_explainability_with_invalid_config(
         "initial_instance_count": 1,
     }
     with pytest.raises(
-        AttributeError, match="'NoneType' object has no attribute 'get_explainability_config'"
+        AttributeError,
+        match="analysis_config must have at least one working method: "
+        "One of the `pre_training_methods`, `post_training_methods`, `explainability_config`.",
     ):
         _run_test_explain(
             name_from_base,
@@ -1320,6 +1322,80 @@ def test_analysis_config_generator_for_explainability(data_config, model_config)
     assert actual == expected
 
 
+def test_analysis_config_generator_for_explainability_failing(data_config, model_config):
+    model_scores = ModelPredictedLabelConfig(
+        probability="pr",
+        label_headers=["success"],
+    )
+    with pytest.raises(
+        ValueError, match="PDP features must be provided when ShapConfig is not provided"
+    ):
+        _AnalysisConfigGenerator.explainability(
+            data_config,
+            model_config,
+            model_scores,
+            PDPConfig(),
+        )
+
+    with pytest.raises(ValueError, match="Duplicate explainability configs are provided"):
+        _AnalysisConfigGenerator.explainability(
+            data_config,
+            model_config,
+            model_scores,
+            [SHAPConfig(), SHAPConfig()],
+        )
+
+    with pytest.raises(ValueError, match="Please provide at least one explainability config."):
+        _AnalysisConfigGenerator.explainability(
+            data_config,
+            model_config,
+            model_scores,
+            [],
+        )
+
+
+def test_analysis_config_generator_for_bias_explainability(
+    data_config, data_bias_config, model_config
+):
+    model_predicted_label_config = ModelPredictedLabelConfig(
+        probability="pr",
+        label_headers=["success"],
+    )
+    actual = _AnalysisConfigGenerator.bias_and_explainability(
+        data_config,
+        model_config,
+        model_predicted_label_config,
+        [SHAPConfig(), PDPConfig()],
+        data_bias_config,
+        pre_training_methods="all",
+        post_training_methods="all",
+    )
+    expected = {
+        "dataset_type": "text/csv",
+        "facet": [{"name_or_index": "F1"}],
+        "group_variable": "F2",
+        "headers": ["Label", "F1", "F2", "F3", "F4"],
+        "joinsource_name_or_index": "F4",
+        "label": "Label",
+        "label_values_or_threshold": [1],
+        "methods": {
+            "pdp": {"grid_resolution": 15, "top_k_features": 10},
+            "post_training_bias": {"methods": "all"},
+            "pre_training_bias": {"methods": "all"},
+            "report": {"name": "report", "title": "Analysis Report"},
+            "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
+
+
 def test_analysis_config_generator_for_bias_pre_training(data_config, data_bias_config):
     actual = _AnalysisConfigGenerator.bias_pre_training(
         data_config, data_bias_config, methods="all"