change: Update BiasConfig to accept multiple facet params (#2243)

Co-authored-by: Georgios Schinas <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: georschi

src/sagemaker/clarify.py

@@ -88,21 +88,34 @@ def __init__(
         Args:
             label_values_or_threshold (Any): List of label values or threshold to indicate positive
                 outcome used for bias metrics.
-            facet_name (str): Sensitive attribute in the input data for which we like to compare
-                metrics.
+            facet_name (str or [str]): String or List of strings of sensitive attribute(s) in the
+            input data for which we like to compare metrics.
             facet_values_or_threshold (list): Optional list of values to form a sensitive group or
                 threshold for a numeric facet column that defines the lower bound of a sensitive
                 group. Defaults to considering each possible value as sensitive group and
                 computing metrics vs all the other examples.
+                If facet_name is a list, this needs to be None or a List consisting of lists or None
+                with the same length as facet_name list.
             group_name (str): Optional column name or index to indicate a group column to be used
                 for the bias metric 'Conditional Demographic Disparity in Labels - CDDL' or
                 'Conditional Demographic Disparity in Predicted Labels - CDDPL'.
         """
-        facet = {"name_or_index": facet_name}
-        _set(facet_values_or_threshold, "value_or_threshold", facet)
+        if isinstance(facet_name, str):
+            facet = {"name_or_index": facet_name}
+            _set(facet_values_or_threshold, "value_or_threshold", facet)
+            facet_list = [facet]
+        elif facet_values_or_threshold is None or len(facet_name) == len(facet_values_or_threshold):
+            facet_list = []
+            for i, single_facet_name in enumerate(facet_name):
+                facet = {"name_or_index": single_facet_name}
+                if facet_values_or_threshold is not None:
+                    _set(facet_values_or_threshold[i], "value_or_threshold", facet)
+                facet_list.append(facet)
+        else:
+            raise ValueError("Wrong combination of argument values passed")
         self.analysis_config = {
             "label_values_or_threshold": label_values_or_threshold,
-            "facet": [facet],
+            "facet": facet_list,
         }
         _set(group_name, "group_variable", self.analysis_config)
 

tests/unit/test_clarify.py

@@ -89,6 +89,54 @@ def test_data_bias_config():
     assert expected_config == data_bias_config.get_config()
 
 
+def test_data_bias_config_multi_facet():
+    label_values = [1]
+    facet_name = ["Facet1", "Facet2"]
+    facet_threshold = [[0], [1, 2]]
+    group_name = "A151"
+
+    data_bias_config = BiasConfig(
+        label_values_or_threshold=label_values,
+        facet_name=facet_name,
+        facet_values_or_threshold=facet_threshold,
+        group_name=group_name,
+    )
+
+    expected_config = {
+        "label_values_or_threshold": label_values,
+        "facet": [
+            {"name_or_index": facet_name[0], "value_or_threshold": facet_threshold[0]},
+            {"name_or_index": facet_name[1], "value_or_threshold": facet_threshold[1]},
+        ],
+        "group_variable": group_name,
+    }
+    assert expected_config == data_bias_config.get_config()
+
+
+def test_data_bias_config_multi_facet_not_all_with_value():
+    label_values = [1]
+    facet_name = ["Facet1", "Facet2"]
+    facet_threshold = [[0], None]
+    group_name = "A151"
+
+    data_bias_config = BiasConfig(
+        label_values_or_threshold=label_values,
+        facet_name=facet_name,
+        facet_values_or_threshold=facet_threshold,
+        group_name=group_name,
+    )
+
+    expected_config = {
+        "label_values_or_threshold": label_values,
+        "facet": [
+            {"name_or_index": facet_name[0], "value_or_threshold": facet_threshold[0]},
+            {"name_or_index": facet_name[1]},
+        ],
+        "group_variable": group_name,
+    }
+    assert expected_config == data_bias_config.get_config()
+
+
 def test_model_config():
     model_name = "xgboost-model"
     instance_type = "ml.c5.xlarge"