Wrap binarize and disabled_outputs options into an additional proto message so that we can better add default values when no value is given.

PiperOrigin-RevId: 292483034

Author: mdreves

RELEASE.md

@@ -12,6 +12,10 @@
 
 ## Breaking changes
 
+*   `tfma.BinarizeOptions.class_ids`, `tfma.BinarizeOptions.k_list`,
+    `tfma.BinarizeOptions.top_k_list`, and `tfma.Options.disabled_outputs` are
+    now wrapped in an additional proto message.
+
 ## Deprecations
 
 # Release 0.21.0

g3doc/metrics.md

@@ -192,7 +192,7 @@ from google.protobuf import text_format
 
 metrics_specs = text_format.Parse("""
   metrics_specs {
-    binarize: { class_ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] }
+    binarize: { class_ids: { values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] } }
     // Metrics to binarize
     metrics { class_name: "AUC" }
     ...
@@ -209,7 +209,8 @@ metrics = [
     ...
 ]
 metrics_specs = tfma.metrics.specs_from_metrics(
-    metrics, binarize=tfma.BinarizationOptions(class_ids=[0,1,2,3,4,5,6,7,8,9]))
+    metrics, binarize=tfma.BinarizationOptions(
+        class_ids={'values': [0,1,2,3,4,5,6,7,8,9]}))
 ```
 
 ### Multi-class/Multi-label Aggregate Metrics
@@ -259,7 +260,7 @@ from google.protobuf import text_format
 
 metrics_specs = text_format.Parse("""
   metrics_specs {
-    binarize: { class_ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] }
+    binarize: { class_ids: { values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] } }
     aggregeate: { macro_average: true }
     // Metrics to both binarize and aggregate
     metrics { class_name: "AUC" }
@@ -278,7 +279,8 @@ metrics = [
 ]
 metrics_specs = tfma.metrics.specs_from_metrics(
     metrics,
-    binarize=tfma.BinarizationOptions(class_ids=[0,1,2,3,4,5,6,7,8,9]),
+    binarize=tfma.BinarizationOptions(
+        class_ids={'values': [0,1,2,3,4,5,6,7,8,9]}),
     aggregate=tfma.AggregationOptions(macro_average=True))
 ```
 
@@ -293,7 +295,7 @@ from google.protobuf import text_format
 metrics_specs = text_format.Parse("""
   metrics_specs {
     query_key: "doc_id"
-    binarize { top_k: [1, 2] }
+    binarize { top_k_list: { values: [1, 2] } }
     metrics { class_name: "NDCG" config: '"gain_key": "gain"' }
   }
   metrics_specs {
@@ -310,7 +312,8 @@ metrics = [
     tfma.metrics.NDCG(name='ndcg', gain_key='gain'),
 ]
 metrics_specs = tfma.metrics.specs_from_metrics(
-    metrics, query_key='doc_id', binarize=tfma.BinarizationOptions(top_k=[1,2]))
+    metrics, query_key='doc_id', binarize=tfma.BinarizationOptions(
+        top_k_list={'values': [1,2]}))
 
 metrics = [
     tfma.metrics.MinLabelPosition(name='min_label_position')

tensorflow_model_analysis/api/model_eval_lib.py

@@ -363,7 +363,6 @@ def default_evaluators(  # pylint: disable=invalid-name
     desired_batch_size: Optional[int] = None,
     serialize: bool = False,
     random_seed_for_testing: Optional[int] = None) -> List[evaluator.Evaluator]:
-
   """Returns the default evaluators for use in ExtractAndEvaluate.
 
   Args:
@@ -378,8 +377,8 @@ def default_evaluators(  # pylint: disable=invalid-name
     random_seed_for_testing: Provide for deterministic tests only.
   """
   disabled_outputs = []
-  if eval_config and eval_config.options:
-    disabled_outputs = eval_config.options.disabled_outputs
+  if eval_config:
+    disabled_outputs = eval_config.options.disabled_outputs.values
   if (constants.METRICS_KEY in disabled_outputs and
       constants.PLOTS_KEY in disabled_outputs):
     return []
@@ -702,7 +701,7 @@ def ExtractEvaluateAndWriteResults(  # pylint: disable=invalid-name
     options.compute_confidence_intervals.value = compute_confidence_intervals
     options.k_anonymization_count.value = k_anonymization_count
     if not write_config:
-      options.disabled_outputs.append(_EVAL_CONFIG_FILE)
+      options.disabled_outputs.values.append(_EVAL_CONFIG_FILE)
     eval_config = config.EvalConfig(
         model_specs=model_specs, slicing_specs=slicing_specs, options=options)
 
@@ -744,7 +743,7 @@ def ExtractEvaluateAndWriteResults(  # pylint: disable=invalid-name
           extractors=extractors, evaluators=evaluators)
       | 'WriteResults' >> WriteResults(writers=writers))
 
-  if _EVAL_CONFIG_FILE not in eval_config.options.disabled_outputs:
+  if _EVAL_CONFIG_FILE not in eval_config.options.disabled_outputs.values:
     data_location = '<user provided PCollection>'
     if display_only_data_location is not None:
       data_location = display_only_data_location
@@ -781,7 +780,8 @@ def run_model_analysis(
     compute_confidence_intervals: Optional[bool] = False,
     k_anonymization_count: int = 1,
     desired_batch_size: Optional[int] = None,
-    random_seed_for_testing: Optional[int] = None) -> Union[EvalResult, EvalResults]:
+    random_seed_for_testing: Optional[int] = None
+) -> Union[EvalResult, EvalResults]:
   """Runs TensorFlow model analysis.
 
   It runs a Beam pipeline to compute the slicing metrics exported in TensorFlow
@@ -856,7 +856,7 @@ def run_model_analysis(
     options.compute_confidence_intervals.value = compute_confidence_intervals
     options.k_anonymization_count.value = k_anonymization_count
     if not write_config:
-      options.disabled_outputs.append(_EVAL_CONFIG_FILE)
+      options.disabled_outputs.values.append(_EVAL_CONFIG_FILE)
     eval_config = config.EvalConfig(
         model_specs=model_specs, slicing_specs=slicing_specs, options=options)
 

tensorflow_model_analysis/api/model_eval_lib_test.py

@@ -402,7 +402,7 @@ def testRunModelAnalysisWithKerasModel(self):
           config.MetricConfig(
               class_name=cfg['class_name'], config=json.dumps(cfg['config'])))
     for class_id in (0, 5, 9):
-      metrics_spec.binarize.class_ids.append(class_id)
+      metrics_spec.binarize.class_ids.values.append(class_id)
     eval_config = config.EvalConfig(
         model_specs=[config.ModelSpec(label_key='label')],
         metrics_specs=[metrics_spec])
@@ -470,7 +470,7 @@ def testRunModelAnalysisWithQueryBasedMetrics(self):
         slicing_specs=slicing_specs,
         metrics_specs=metric_specs.specs_from_metrics(
             [ndcg.NDCG(gain_key='age', name='ndcg')],
-            binarize=config.BinarizationOptions(top_k_list=[1]),
+            binarize=config.BinarizationOptions(top_k_list={'values': [1]}),
             query_key='language'))
     eval_shared_model = model_eval_lib.default_eval_shared_model(
         eval_saved_model_path=model_location, tags=[tf.saved_model.SERVING])

tensorflow_model_analysis/evaluators/metrics_and_plots_evaluator_v2_test.py

@@ -598,7 +598,8 @@ def testEvaluateWithMultiClassModel(self):
         slicing_specs=[config.SlicingSpec()],
         metrics_specs=metric_specs.specs_from_metrics(
             [calibration.MeanLabel('mean_label')],
-            binarize=config.BinarizationOptions(class_ids=range(n_classes))))
+            binarize=config.BinarizationOptions(
+                class_ids={'values': range(n_classes)})))
     eval_shared_model = self.createTestEvalSharedModel(
         eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING])
 
@@ -898,7 +899,7 @@ def testEvaluateWithQueryBasedMetrics(self):
         ],
         metrics_specs=metric_specs.specs_from_metrics(
             [ndcg.NDCG(gain_key='fixed_float', name='ndcg')],
-            binarize=config.BinarizationOptions(top_k_list=[1, 2]),
+            binarize=config.BinarizationOptions(top_k_list={'values': [1, 2]}),
             query_key='fixed_string'))
     eval_shared_model = self.createTestEvalSharedModel(
         eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING])

tensorflow_model_analysis/metrics/metric_specs.py

@@ -305,13 +305,13 @@ def default_multi_class_classification_specs(
   metrics.append(
       multi_class_confusion_matrix_plot.MultiClassConfusionMatrixPlot())
   if binarize is not None:
-    for top_k in binarize.top_k_list:
+    for top_k in binarize.top_k_list.values:
       metrics.extend([
           tf.keras.metrics.Precision(name='precision', top_k=top_k),
           tf.keras.metrics.Recall(name='recall', top_k=top_k)
       ])
     binarize = config.BinarizationOptions().CopyFrom(binarize)
-    binarize.ClearField('top_k')
+    binarize.ClearField('top_k_list')
   multi_class_metrics = specs_from_metrics(
       metrics, model_names=model_names, output_names=output_names)
   if aggregate is None:
@@ -526,14 +526,14 @@ def _create_sub_keys(
   sub_keys = None
   if spec.HasField('binarize'):
     sub_keys = []
-    if spec.binarize.class_ids:
-      for v in spec.binarize.class_ids:
+    if spec.binarize.class_ids.values:
+      for v in spec.binarize.class_ids.values:
         sub_keys.append(metric_types.SubKey(class_id=v))
-    if spec.binarize.k_list:
-      for v in spec.binarize.k_list:
+    if spec.binarize.k_list.values:
+      for v in spec.binarize.k_list.values:
         sub_keys.append(metric_types.SubKey(k=v))
-    if spec.binarize.top_k_list:
-      for v in spec.binarize.top_k_list:
+    if spec.binarize.top_k_list.values:
+      for v in spec.binarize.top_k_list.values:
         sub_keys.append(metric_types.SubKey(top_k=v))
     if spec.aggregate.micro_average:
       # Micro averaging is performed by flattening the labels and predictions

tensorflow_model_analysis/metrics/metric_specs_test.py

@@ -42,7 +42,7 @@ def testSpecsFromMetrics(self):
             ]
         },
         model_names=['model_name1', 'model_name2'],
-        binarize=config.BinarizationOptions(class_ids=[0, 1]),
+        binarize=config.BinarizationOptions(class_ids={'values': [0, 1]}),
         aggregate=config.AggregationOptions(macro_average=True))
 
     self.assertLen(metrics_specs, 5)
@@ -80,7 +80,7 @@ def testSpecsFromMetrics(self):
             ],
             model_names=['model_name1', 'model_name2'],
             output_names=['output_name1'],
-            binarize=config.BinarizationOptions(class_ids=[0, 1]),
+            binarize=config.BinarizationOptions(class_ids={'values': [0, 1]}),
             aggregate=config.AggregationOptions(macro_average=True)))
     self.assertProtoEquals(
         metrics_specs[3],
@@ -109,7 +109,7 @@ def testSpecsFromMetrics(self):
             ],
             model_names=['model_name1', 'model_name2'],
             output_names=['output_name2'],
-            binarize=config.BinarizationOptions(class_ids=[0, 1]),
+            binarize=config.BinarizationOptions(class_ids={'values': [0, 1]}),
             aggregate=config.AggregationOptions(macro_average=True)))
 
   def testToComputations(self):
@@ -122,7 +122,7 @@ def testToComputations(self):
                 ]
             },
             model_names=['model_name'],
-            binarize=config.BinarizationOptions(class_ids=[0, 1]),
+            binarize=config.BinarizationOptions(class_ids={'values': [0, 1]}),
             aggregate=config.AggregationOptions(macro_average=True)),
         config.EvalConfig())
 

tensorflow_model_analysis/proto/config.proto

@@ -119,17 +119,19 @@ message AggregationOptions {
 message BinarizationOptions {
   // Creates binary classification metrics based on one-vs-rest for each
   // value of class_id provided.
-  repeated int32 class_ids = 1;
+  RepeatedInt32Value class_ids = 4;
   // Creates binary classification metrics based on the kth predicted value
   // for each value of k provided.
-  repeated int32 k_list = 2;
+  RepeatedInt32Value k_list = 5;
   // Creates binary classification metrics based on the top k predicted values
   // for each value of top_k provided. When used to create calibration plots
   // the histogram will contain a mix of all labels and predictions in the top
   // k predictions. Note that precision@k and recall@k can also be configured
   // directly as multi-class classification metrics by setting top_k on the
   // metric itself.
-  repeated int32 top_k_list = 3;
+  RepeatedInt32Value top_k_list = 6;
+
+  reserved 1, 2, 3;
 }
 
 // Metric configuration.
@@ -183,9 +185,9 @@ message Options {
   google.protobuf.Int32Value k_anonymization_count = 3;
   // List of outputs that should not be written (e.g.  'metrics', 'plots',
   // 'analysis', 'eval_config.json').
-  repeated string disabled_outputs = 6;
+  RepeatedStringValue disabled_outputs = 7;
 
-  reserved 4, 5;
+  reserved 4, 5, 6;
 }
 
 // Tensorflow model analaysis config settings.
@@ -219,6 +221,16 @@ message EvalConfig {
   reserved 1, 3, 7;
 }
 
+// Repeated string value. Used to allow a default if no values are given.
+message RepeatedStringValue {
+  repeated string values = 1;
+}
+
+// Repeated int32 value. Used to allow a default if no values are given.
+message RepeatedInt32Value {
+  repeated int32 values = 1;
+}
+
 // Config and version.
 message EvalConfigAndVersion {
   EvalConfig eval_config = 1;

tensorflow_model_analysis/writers/metrics_and_plots_writer_test.py

@@ -59,7 +59,8 @@ def testWriteMetricsAndPlots(self):
             None, temp_eval_export_dir))
     eval_config = config.EvalConfig(
         model_specs=[config.ModelSpec()],
-        options=config.Options(disabled_outputs=['eval_config.json']))
+        options=config.Options(
+            disabled_outputs={'values': ['eval_config.json']}))
     eval_shared_model = self.createTestEvalSharedModel(
         eval_saved_model_path=eval_export_dir,
         add_metrics_callbacks=[