feat: add (early version of) baseline config to asym shap val config

Author: rvasahu-amazon

src/sagemaker/clarify.py

@@ -327,6 +327,14 @@
                     ),
                 ),
                 SchemaOptional("num_samples"): int,
+                SchemaOptional("baseline"): Or(
+                    str,
+                    {
+                        SchemaOptional("target_ts", default="zero"): str,
+                        SchemaOptional("related_ts"): str,
+                        SchemaOptional("static_covariates"): [Or(str, int, float)],
+                    },
+                ),
             },
         },
         SchemaOptional("predictor"): {
@@ -1661,6 +1669,7 @@ def __init__(
             "fine_grained",
         ] = ASYM_SHAP_VAL_DEFAULT_EXPLANATION_GRANULARITY,
         num_samples: Optional[int] = None,
+        baseline: Optional[Union[str, Dict[str, Any]]] = None,
     ):
         """Initialises config for time series explainability with Asymmetric Shapley Values.
 
@@ -1675,6 +1684,8 @@ def __init__(
             num_samples (None or int): Number of samples to be used in the Asymmetric Shapley
                 Value forecasting algorithm. Only applicable when using ``"fine_grained"``
                 explanations.
+            baseline (str or dict): Link to a baseline configuration or a dictionary for it.
+            # TODO: improve above.
 
         Raises:
             AssertionError: when ``direction`` or ``granularity`` are not valid,
@@ -1707,6 +1718,8 @@ def __init__(
         _set(
             num_samples, "num_samples", self.asymmetric_shapley_value_config
         )  # _set() does nothing if a given argument is None
+        # TODO: add sdk-side validation to baseline
+        _set(baseline, "baseline", self.asymmetric_shapley_value_config)
 
     def get_explainability_config(self):
         """Returns an asymmetric shap config dictionary."""