feat: validation for asymmetric shapley value config baseline

doc: fix baseline doc to be sphinx-compliant

Author: rvasahu-amazon

src/sagemaker/clarify.py

@@ -332,8 +332,24 @@
                 SchemaOptional("baseline"): Or(
                     str,
                     {
-                        SchemaOptional("target_time_series", default="zero"): str,
-                        SchemaOptional("related_time_series"): str,
+                        SchemaOptional("target_time_series", default="zero"): And(
+                            str,
+                            Use(str.lower),
+                            lambda s: s
+                            in (
+                                "zero",
+                                "mean",
+                            ),
+                        ),
+                        SchemaOptional("related_time_series"): And(
+                            str,
+                            Use(str.lower),
+                            lambda s: s
+                            in (
+                                "zero",
+                                "mean",
+                            ),
+                        ),
                         SchemaOptional("static_covariates"): {Or(str, int): [Or(str, int, float)]},
                     },
                 ),
@@ -1769,13 +1785,14 @@ def __init__(
                 (static covariates), a baseline value for each covariate should be provided for
                 each possible item_id. An example config follows, where ``item1`` and ``item2``
                 are item ids.::
+
                     {
+                        "target_time_series": "zero",
                         "related_time_series": "zero",
                         "static_covariates": {
                             "item1": [1, 1],
                             "item2": [0, 1]
-                        },
-                        "target_time_series": "zero"
+                        }
                     }
 
         Raises:
@@ -1803,13 +1820,27 @@ def __init__(
             ), f"{direction} and {granularity} granularity are not supported together."
         elif num_samples:  # validate num_samples is not provided when unnecessary
             raise ValueError("``num_samples`` is only used for fine-grained explanations.")
+        # validate baseline if provided as a dictionary
+        if isinstance(baseline, dict):
+            temporal_baselines = ["zero", "mean"]  # possible baseline options for temporal fields
+            if "target_time_series" in baseline:
+                target_baseline = baseline.get("target_time_series")
+                assert target_baseline in temporal_baselines, (
+                    f"Provided value {target_baseline} for ``target_time_series`` is "
+                    f"invalid. Please select one of {temporal_baselines}."
+                )
+            if "related_time_series" in baseline:
+                related_baseline = baseline.get("related_time_series")
+                assert related_baseline in temporal_baselines, (
+                    f"Provided value {related_baseline} for ``related_time_series`` is "
+                    f"invalid. Please select one of {temporal_baselines}."
+                )
         # set explanation type and (if provided) num_samples in internal config dictionary
         _set(direction, "direction", self.asymmetric_shapley_value_config)
         _set(granularity, "granularity", self.asymmetric_shapley_value_config)
         _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):

tests/unit/test_clarify.py

@@ -17,7 +17,7 @@
 
 import pytest
 from mock import ANY, MagicMock, Mock, patch
-from typing import List, NamedTuple, Optional, Union
+from typing import Any, Dict, List, NamedTuple, Optional, Union
 
 from sagemaker import Processor, image_uris
 from sagemaker.clarify import (
@@ -1283,9 +1283,10 @@ def test_shap_config_no_parameters():
 class AsymmetricShapleyValueConfigCase(NamedTuple):
     direction: str
     granularity: str
-    num_samples: Optional[int]
-    error: Exception
-    error_message: str
+    num_samples: Optional[int] = None
+    baseline: Optional[Union[str, Dict[str, Any]]] = None
+    error: Exception = None
+    error_message: str = None
 
 
 class TestAsymmetricShapleyValueConfig:
@@ -1296,22 +1297,28 @@ class TestAsymmetricShapleyValueConfig:
                 direction=direction,
                 granularity="timewise",
                 num_samples=None,
-                error=None,
-                error_message=None,
             )
             for direction in ASYM_SHAP_VAL_EXPLANATION_DIRECTIONS
         ]
         + [
-            AsymmetricShapleyValueConfigCase(  # cases for fine_grained granularity
+            AsymmetricShapleyValueConfigCase(  # case for fine_grained granularity
                 direction="chronological",
                 granularity="fine_grained",
                 num_samples=1,
-                error=None,
-                error_message=None,
-            )
+            ),
+            AsymmetricShapleyValueConfigCase(  # case for target time series baseline
+                direction="chronological",
+                granularity="timewise",
+                baseline={"target_time_series": "mean"},
+            ),
+            AsymmetricShapleyValueConfigCase(  # case for related time series baseline
+                direction="chronological",
+                granularity="timewise",
+                baseline={"related_time_series": "zero"},
+            ),
         ],
     )
-    def test_asymmetric_shapley_value_config(self, test_case):
+    def test_asymmetric_shapley_value_config(self, test_case: AsymmetricShapleyValueConfigCase):
         """
         GIVEN valid arguments for an AsymmetricShapleyValueConfig object
         WHEN AsymmetricShapleyValueConfig object is instantiated with those arguments
@@ -1325,11 +1332,14 @@ def test_asymmetric_shapley_value_config(self, test_case):
         }
         if test_case.granularity == "fine_grained":
             expected_config["num_samples"] = test_case.num_samples
+        if test_case.baseline:
+            expected_config["baseline"] = test_case.baseline
         # WHEN
         asym_shap_val_config = AsymmetricShapleyValueConfig(
             direction=test_case.direction,
             granularity=test_case.granularity,
             num_samples=test_case.num_samples,
+            baseline=test_case.baseline,
         )
         # THEN
         assert asym_shap_val_config.asymmetric_shapley_value_config == expected_config
@@ -1380,6 +1390,20 @@ def test_asymmetric_shapley_value_config(self, test_case):
                 error=AssertionError,
                 error_message="not supported together.",
             ),
+            AsymmetricShapleyValueConfigCase(  # case for unsupported target time series baseline value
+                direction="chronological",
+                granularity="timewise",
+                baseline={"target_time_series": "median"},
+                error=AssertionError,
+                error_message="for ``target_time_series`` is invalid.",
+            ),
+            AsymmetricShapleyValueConfigCase(  # case for unsupported related time series baseline value
+                direction="chronological",
+                granularity="timewise",
+                baseline={"related_time_series": "mode"},
+                error=AssertionError,
+                error_message="for ``related_time_series`` is invalid.",
+            ),
         ],
     )
     def test_asymmetric_shapley_value_config_invalid(self, test_case):
@@ -1394,6 +1418,7 @@ def test_asymmetric_shapley_value_config_invalid(self, test_case):
                 direction=test_case.direction,
                 granularity=test_case.granularity,
                 num_samples=test_case.num_samples,
+                baseline=test_case.baseline,
             )