Added data types and Formatted

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -25,7 +25,7 @@
 
 import tempfile
 from abc import ABC, abstractmethod
-from typing import List, Union
+from typing import List, Union, Dict
 
 from sagemaker import image_uris, s3, utils
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
@@ -173,7 +173,11 @@ def __init__(
         _set(joinsource, "joinsource_name_or_index", self.analysis_config)
         _set(facet_dataset_uri, "facet_dataset_uri", self.analysis_config)
         _set(facet_headers, "facet_headers", self.analysis_config)
-        _set(predicted_label_dataset_uri, "predicted_label_dataset_uri", self.analysis_config)
+        _set(
+            predicted_label_dataset_uri,
+            "predicted_label_dataset_uri",
+            self.analysis_config,
+        )
         _set(predicted_label_headers, "predicted_label_headers", self.analysis_config)
         _set(predicted_label, "predicted_label", self.analysis_config)
         _set(excluded_columns, "excluded_columns", self.analysis_config)
@@ -239,7 +243,8 @@ def __init__(
             assert len(facet_name) > 0, "Please provide at least one facet"
             if facet_values_or_threshold is None:
                 facet_list = [
-                    {"name_or_index": single_facet_name} for single_facet_name in facet_name
+                    {"name_or_index": single_facet_name}
+                    for single_facet_name in facet_name
                 ]
             elif len(facet_values_or_threshold) == len(facet_name):
                 facet_list = []
@@ -492,7 +497,10 @@ def __init__(self, features=None, grid_resolution=15, top_k_features=10):
             top_k_features (int): Sets the number of top SHAP attributes used to compute
                 partial dependence plots.
         """  # noqa E501
-        self.pdp_config = {"grid_resolution": grid_resolution, "top_k_features": top_k_features}
+        self.pdp_config = {
+            "grid_resolution": grid_resolution,
+            "top_k_features": top_k_features,
+        }
         if features is not None:
             self.pdp_config["features"] = features
 
@@ -825,9 +833,14 @@ def __init__(
             image_config (:class:`~sagemaker.clarify.ImageConfig`): Config for handling image
                 features. Default is None.
         """  # noqa E501  # pylint: disable=c0301
-        if agg_method is not None and agg_method not in ["mean_abs", "median", "mean_sq"]:
+        if agg_method is not None and agg_method not in [
+            "mean_abs",
+            "median",
+            "mean_sq",
+        ]:
             raise ValueError(
-                f"Invalid agg_method {agg_method}." f" Please choose mean_abs, median, or mean_sq."
+                f"Invalid agg_method {agg_method}."
+                f" Please choose mean_abs, median, or mean_sq."
             )
         if num_clusters is not None and baseline is not None:
             raise ValueError(
@@ -923,7 +936,9 @@ def __init__(
             job_name_prefix (str): Processing job name prefix.
             version (str): Clarify version to use.
         """  # noqa E501  # pylint: disable=c0301
-        container_uri = image_uris.retrieve("clarify", sagemaker_session.boto_region_name, version)
+        container_uri = image_uris.retrieve(
+            "clarify", sagemaker_session.boto_region_name, version
+        )
         self._last_analysis_config = None
         self.job_name_prefix = job_name_prefix
         super(SageMakerClarifyProcessor, self).__init__(
@@ -996,7 +1011,8 @@ def _run(
                 json.dump(analysis_config, f)
             s3_analysis_config_file = _upload_analysis_config(
                 analysis_config_file,
-                data_config.s3_analysis_config_output_path or data_config.s3_output_path,
+                data_config.s3_analysis_config_output_path
+                or data_config.s3_output_path,
                 self.sagemaker_session,
                 kms_key,
             )
@@ -1168,7 +1184,11 @@ def run_post_training_bias(
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
         analysis_config = _AnalysisConfigGenerator.bias_post_training(
-            data_config, data_bias_config, model_predicted_label_config, methods, model_config
+            data_config,
+            data_bias_config,
+            model_predicted_label_config,
+            methods,
+            model_config,
         )
         # when name is either not provided (is None) or an empty string ("")
         job_name = job_name or utils.name_from_base(
@@ -1267,7 +1287,9 @@ def run_bias(
             post_training_methods,
         )
         # when name is either not provided (is None) or an empty string ("")
-        job_name = job_name or utils.name_from_base(self.job_name_prefix or "Clarify-Bias")
+        job_name = job_name or utils.name_from_base(
+            self.job_name_prefix or "Clarify-Bias"
+        )
         return self._run(
             data_config,
             analysis_config,
@@ -1450,8 +1472,8 @@ def run_bias_and_explainability(
                 "`FT <https://docs.aws.amazon.com/sagemaker/latest/dg/clarify-post-training-bias-metric-ft.html>`_"].
                 Defaults to str "all" to run all metrics if left unspecified.
             model_predicted_label_config (
-                int or 
-                str or 
+                int or
+                str or
                 :class:`~sagemaker.clarify.ModelPredictedLabelConfig`
             ):
                 Index or JSONPath to locate the predicted scores in the model output. This is not
@@ -1552,11 +1574,16 @@ def explainability(
 
     @classmethod
     def bias_pre_training(
-        cls, data_config: DataConfig, bias_config: BiasConfig, methods: Union[str, List[str]]
+        cls,
+        data_config: DataConfig,
+        bias_config: BiasConfig,
+        methods: Union[str, List[str]],
     ):
         """Generates a config for Bias Pre Training"""
         analysis_config = {**data_config.get_config(), **bias_config.get_config()}
-        analysis_config = cls._add_methods(analysis_config, pre_training_methods=methods)
+        analysis_config = cls._add_methods(
+            analysis_config, pre_training_methods=methods
+        )
         return analysis_config
 
     @classmethod
@@ -1570,7 +1597,9 @@ def bias_post_training(
     ):
         """Generates a config for Bias Post Training"""
         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_methods(
+            analysis_config, post_training_methods=methods
+        )
         analysis_config = cls._add_predictor(
             analysis_config, model_config, model_predicted_label_config
         )
@@ -1599,7 +1628,12 @@ def bias(
         return analysis_config
 
     @classmethod
-    def _add_predictor(cls, analysis_config, model_config, model_predicted_label_config):
+    def _add_predictor(
+        cls,
+        analysis_config: Dict,
+        model_config: ModelConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+    ):
         """Extends analysis config with predictor."""
         analysis_config = {**analysis_config}
         analysis_config["predictor"] = model_config.get_predictor_config()
@@ -1618,10 +1652,12 @@ def _add_predictor(cls, analysis_config, model_config, model_predicted_label_con
     @classmethod
     def _add_methods(
         cls,
-        analysis_config,
-        pre_training_methods=None,
-        post_training_methods=None,
-        explainability_config=None,
+        analysis_config: Dict,
+        pre_training_methods: Union[str, List[str]] = "all",
+        post_training_methods: Union[str, List[str]] = "all",
+        explainability_config: Union[
+            ExplainabilityConfig, List[ExplainabilityConfig]
+        ] = None,
         report=True,
     ):
         """Extends analysis config with methods."""
@@ -1640,22 +1676,35 @@ def _add_methods(
             analysis_config["methods"] = {}
 
         if report:
-            analysis_config["methods"]["report"] = {"name": "report", "title": "Analysis Report"}
+            analysis_config["methods"]["report"] = {
+                "name": "report",
+                "title": "Analysis Report",
+            }
 
         if pre_training_methods:
-            analysis_config["methods"]["pre_training_bias"] = {"methods": 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}
+            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}
+            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]]
+        cls,
+        explainability_config: Union[ExplainabilityConfig, List[ExplainabilityConfig]],
     ):
         """Merges explainability configs, when more than one."""
         if isinstance(explainability_config, list):
@@ -1671,17 +1720,24 @@ def _merge_explainability_configs(
                 "shap" not in explainability_methods
                 and "features" not in explainability_methods["pdp"]
             ):
-                raise ValueError("PDP features must be provided when ShapConfig is not provided")
+                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"]
+            and "features"
+            not in explainability_config.get_explainability_config()["pdp"]
         ):
-            raise ValueError("PDP features must be provided when ShapConfig is not provided")
+            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):
+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``.
 
     Args: