Added data types

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
@@ -1450,8 +1450,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
@@ -1599,7 +1599,13 @@ 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 +1624,10 @@ 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."""