code: def testModelAgnosticConstructFn(self):
      # End to end test for the entire flow going from tf.Examples -> metrics
      # with slicing.
      with beam.Pipeline() as pipeline:
          # Set up the inputs. All we need is are tf.Examples and an example parsing
          # spec with explicit mapping for key to (Features, Predictions, Labels).
          # TODO(b/119788402): Add a fairness data examples/callbacks as another
          # test.
          examples = [
              self._makeExample(
                  age=3.0,
                  language="english",
                  probabilities=1.0,
                  labels=1.0,
              ),
              self._makeExample(
                  age=3.0,
                  language="chinese",
                  probabilities=3.0,
                  labels=0.0,
              ),
              self._makeExample(
                  age=4.0,
                  language="english",
                  probabilities=2.0,
                  labels=1.0,
              ),
              self._makeExample(
                  age=5.0,
                  language="chinese",
                  probabilities=3.0,
                  labels=0.0,
              ),
              # Add some examples with no language.
              self._makeExample(
                  age=5.0, probabilities=2.0, labels=10.0
              ),
              self._makeExample(age=6.0, probabilities=1.0, labels=0.0),
          ]
          serialized_examples = [
              e.SerializeToString() for e in examples
          ]
  
          # Set up a config to bucket our example keys.
          feature_map = {
              "age": tf.io.FixedLenFeature([], tf.float32),
              "language": tf.io.VarLenFeature(tf.string),
              "probabilities": tf.io.FixedLenFeature([], tf.float32),
              "labels": tf.io.FixedLenFeature([], tf.float32),
          }
  
          model_agnostic_config = agnostic_predict.ModelAgnosticConfig(
              label_keys=["labels"],
              prediction_keys=["probabilities"],
              feature_spec=feature_map,
          )
  
          # Set up the Model Agnostic Extractor
          extractors = [
              model_agnostic_extractor.ModelAgnosticExtractor(
                  model_agnostic_config=model_agnostic_config,
                  desired_batch_size=3,
              ),
              slice_key_extractor.SliceKeyExtractor(
                  [
                      slicer.SingleSliceSpec(),
                      slicer.SingleSliceSpec(columns=["language"]),
                  ]
              ),
          ]
  
          # Set up the metrics we wish to calculate via a metric callback. In
          # particular, this metric calculates the mean and sum of all labels.
          eval_shared_model = types.EvalSharedModel(
              add_metrics_callbacks=[add_mean_callback],
              construct_fn=model_agnostic_evaluate_graph.make_construct_fn(
                  add_metrics_callbacks=[add_mean_callback],
                  config=model_agnostic_config,
              ),
          )
  
          # Run our pipeline doing Extract -> Slice -> Fanout -> Calculate Metrics.
          (metrics, _), _ = (
              pipeline
              | "Create Examples" >> beam.Create(serialized_examples)
              | "InputsToExtracts" >> model_eval_lib.InputsToExtracts()
              | "Extract"
              >> tfma_unit.Extract(
                  extractors=extractors
              )  # pylint: disable=no-value-for-parameter
              | "ComputeMetricsAndPlots"
              >> legacy_metrics_and_plots_evaluator._ComputeMetricsAndPlots(  # pylint: disable=protected-access
                  eval_shared_model=eval_shared_model
              )
          )
  
          # Verify our metrics are properly generated per slice.
          def check_result(got):
              self.assertEqual(3, len(got), "got: %s" % got)
              slices = {}
              for slice_key, metrics in got:
                  slices[slice_key] = metrics
              overall_slice = ()
              english_slice = (("language", "english"),)
              chinese_slice = (("language", "chinese"),)
  
              self.assertCountEqual(
                  list(slices),
                  [overall_slice, english_slice, chinese_slice],
              )
              # Overall slice has label/predictions sum = 24 and 12 elements.
              self.assertDictElementsAlmostEqual(
                  slices[overall_slice],
                  {
                      "tf_metric_mean": 2.0,
                      "py_func_total_label": 24.0,
                  },
              )
              # English slice has label/predictions sum = 5 and 4 elements.
              self.assertDictElementsAlmostEqual(
                  slices[english_slice],
                  {
                      "tf_metric_mean": 1.25,
                      "py_func_total_label": 5.0,
                  },
              )
              # Chinese slice has label/predictions sum = 6 and 4 elements.
              self.assertDictElementsAlmostEqual(
                  slices[chinese_slice],
                  {
                      "tf_metric_mean": 1.5,
                      "py_func_total_label": 6.0,
                  },
              )
  
          util.assert_that(metrics, check_result)
  
file_path: model-analysis/tensorflow_model_analysis/model_agnostic_eval/model_agnostic_evaluate_graph_test.py
imports: import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  from tensorflow_model_analysis.api import model_eval_lib
  from tensorflow_model_analysis.api import tfma_unit
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.eval_saved_model import testutil
  from tensorflow_model_analysis.evaluators import legacy_metrics_and_plots_evaluator
  from tensorflow_model_analysis.extractors import slice_key_extractor
  from tensorflow_model_analysis.model_agnostic_eval import model_agnostic_evaluate_graph
  from tensorflow_model_analysis.model_agnostic_eval import model_agnostic_extractor
  from tensorflow_model_analysis.model_agnostic_eval import model_agnostic_predict
  from tensorflow_model_analysis.slicer import slicer_lib
label_desc: Use image node_id 3 for calling the ModelAgnosticEvaluateGraphTest obj's underlying member method code with example usage: obj.testModelAgnosticConstructFn() without return types