added data types in _AnalysisConfigGenerator methods

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -25,6 +25,8 @@
 
 import tempfile
 from abc import ABC, abstractmethod
+from typing import List, Union
+
 from sagemaker import image_uris, s3, utils
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 
@@ -1314,10 +1316,10 @@ class _AnalysisConfigGenerator:
     @classmethod
     def explainability(
         cls,
-        data_config,
-        model_config,
-        model_scores,
-        explainability_config
+        data_config: DataConfig,
+        model_config: ModelConfig,
+        model_scores: ModelPredictedLabelConfig,
+        explainability_config: ExplainabilityConfig,
     ):
         analysis_config = data_config.get_config()
         predictor_config = model_config.get_predictor_config()
@@ -1358,7 +1360,7 @@ def explainability(
         return cls._common(analysis_config)
 
     @classmethod
-    def bias_pre_training(cls, data_config, bias_config, methods):
+    def bias_pre_training(cls, data_config: DataConfig, bias_config: BiasConfig, methods: List[str]):
         analysis_config = {
             **data_config.get_config(),
             **bias_config.get_config(),
@@ -1369,11 +1371,11 @@ def bias_pre_training(cls, data_config, bias_config, methods):
     @classmethod
     def bias_post_training(
         cls,
-        data_config,
-        bias_config,
-        model_predicted_label_config,
-        methods,
-        model_config
+        data_config: DataConfig,
+        bias_config: BiasConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        methods: List[str],
+        model_config: ModelConfig,
     ):
         analysis_config = {
             **data_config.get_config(),
@@ -1391,12 +1393,12 @@ def bias_post_training(
     @classmethod
     def bias(
         cls,
-        data_config,
-        bias_config,
-        model_config,
-        model_predicted_label_config,
-        pre_training_methods="all",
-        post_training_methods="all",
+        data_config: DataConfig,
+        bias_config: BiasConfig,
+        model_config: ModelConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        pre_training_methods: Union[str, List[str]] = "all",
+        post_training_methods: Union[str, List[str]] = "all",
     ):
         analysis_config = {
             **data_config.get_config(),