class_docstr: Test class with extra methods for unit-testing TFMA models / metrics.
class_docstr_tok: ["Test","class","with","extra","methods","for","unit-testing","TFMA","models","\/","metrics","."]
docstr: Checks metrics computed using Beam.
  
  A more general version of assertMetricsComputedWithBeamAre. Note that the
  caller is responsible for setting up and running the Beam pipeline.
  
  Example usage:
    def add_metrics(features, predictions, labels):
     metric_ops = {
       'mse': tf.metrics.mean_squared_error(labels, predictions['logits']),
       'mae': tf.metrics.mean_absolute_error(labels, predictions['logits']),
    }
    return metric_ops
  
    with beam.Pipeline() as pipeline:
      expected_slice_metrics = {
          (): {
            'mae': 0.1,
            'mse': 0.2,
            tfma.post_export_metrics.metric_keys.AUC:
              tfma.test.BoundedValue(lower_bound=0.5)
          },
          (('age', 10),): {
            'mae': 0.2,
            'mse': 0.3,
            tfma.post_export_metrics.metric_keys.AUC:
              tfma.test.BoundedValue(lower_bound=0.5)
          },
      }
      examples = pipeline | 'ReadExamples' >> beam.io.ReadFromTFRecord(path)
      self.assertGeneralMetricsComputedWithBeamAre(
        eval_saved_model_path=path,
        examples_pcollection=examples,
        slice_spec=[tfma.slicer.SingleSliceSpec(),
                    tfma.slicer.SingleSliceSpec(columns=['age'])],
        add_metrics_callbacks=[
          add_metrics, tfma.post_export_metrics.auc()],
        expected_slice_metrics=expected_slice_metrics)
  
  Args:
    eval_saved_model_path: Path to the directory containing the
      EvalSavedModel.
    examples_pcollection: A PCollection of serialized example bytes.
    slice_spec: List of slice specifications.
    add_metrics_callbacks: Callbacks for adding additional metrics.
    expected_slice_metrics: Dictionary of dictionaries describing the expected
      metrics for each slice. The outer dictionary map slice keys to the
      expected metrics for that slice.
docstr_tok: ["Checks","metrics","computed","using","Beam",".","A","more","general","version","of","assertMetricsComputedWithBeamAre",".","Note","that","the","caller","is","responsible","for","setting","up","and","running","the","Beam","pipeline",".","Example","usage",":","def","add_metrics","(","features",",","predictions",",","labels",")",":","metric_ops","=","{","'mse","'",":","tf.metrics.mean_squared_error","(","labels",",","predictions","[","'logits","'","]",")",",","'mae","'",":","tf.metrics.mean_absolute_error","(","labels",",","predictions","[","'logits","'","]",")",",","}","return","metric_ops","with","beam.Pipeline","(",")","as","pipeline",":","expected_slice_metrics","=","{","(",")",":","{","'mae","'",":","0.1",",","'mse","'",":","0.2",",","tfma.post_export_metrics.metric_keys.AUC",":","tfma.test.BoundedValue","(","lower_bound=0.5",")","}",",","(","(","'age","'",",","10",")",",",")",":","{","'mae","'",":","0.2",",","'mse","'",":","0.3",",","tfma.post_export_metrics.metric_keys.AUC",":","tfma.test.BoundedValue","(","lower_bound=0.5",")","}",",","}","examples","=","pipeline","|","'ReadExamples","'",">",">","beam.io.ReadFromTFRecord","(","path",")","self.assertGeneralMetricsComputedWithBeamAre","(","eval_saved_model_path=path",",","examples_pcollection=examples",",","slice_spec=","[","tfma.slicer.SingleSliceSpec","(",")",",","tfma.slicer.SingleSliceSpec","(","columns=","[","'age","'","]",")","]",",","add_metrics_callbacks=","[","add_metrics",",","tfma.post_export_metrics.auc","(",")","]",",","expected_slice_metrics=expected_slice_metrics",")","Args",":","eval_saved_model_path",":","Path","to","the","directory","containing","the","EvalSavedModel",".","examples_pcollection",":","A","PCollection","of","serialized","example","bytes",".","slice_spec",":","List","of","slice","specifications",".","add_metrics_callbacks",":","Callbacks","for","adding","additional","metrics",".","expected_slice_metrics",":","Dictionary","of","dictionaries","describing","the","expected","metrics","for","each","slice",".","The","outer","dictionary","map","slice","keys","to","the","expected","metrics","for","that","slice","."]
code: def assertGeneralMetricsComputedWithBeamAre(
      self,
      eval_saved_model_path: str,
      examples_pcollection: beam.pvalue.PCollection,
      slice_spec: List[slicer.SingleSliceSpec],
      add_metrics_callbacks: List[types.AddMetricsCallbackType],
      expected_slice_metrics: Dict[Any, Dict[str, Any]],
  ):
      """Checks metrics computed using Beam.
  
      A more general version of assertMetricsComputedWithBeamAre. Note that the
      caller is responsible for setting up and running the Beam pipeline.
  
      Example usage:
        def add_metrics(features, predictions, labels):
         metric_ops = {
           'mse': tf.metrics.mean_squared_error(labels, predictions['logits']),
           'mae': tf.metrics.mean_absolute_error(labels, predictions['logits']),
        }
        return metric_ops
  
        with beam.Pipeline() as pipeline:
          expected_slice_metrics = {
              (): {
                'mae': 0.1,
                'mse': 0.2,
                tfma.post_export_metrics.metric_keys.AUC:
                  tfma.test.BoundedValue(lower_bound=0.5)
              },
              (('age', 10),): {
                'mae': 0.2,
                'mse': 0.3,
                tfma.post_export_metrics.metric_keys.AUC:
                  tfma.test.BoundedValue(lower_bound=0.5)
              },
          }
          examples = pipeline | 'ReadExamples' >> beam.io.ReadFromTFRecord(path)
          self.assertGeneralMetricsComputedWithBeamAre(
            eval_saved_model_path=path,
            examples_pcollection=examples,
            slice_spec=[tfma.slicer.SingleSliceSpec(),
                        tfma.slicer.SingleSliceSpec(columns=['age'])],
            add_metrics_callbacks=[
              add_metrics, tfma.post_export_metrics.auc()],
            expected_slice_metrics=expected_slice_metrics)
  
      Args:
        eval_saved_model_path: Path to the directory containing the
          EvalSavedModel.
        examples_pcollection: A PCollection of serialized example bytes.
        slice_spec: List of slice specifications.
        add_metrics_callbacks: Callbacks for adding additional metrics.
        expected_slice_metrics: Dictionary of dictionaries describing the expected
          metrics for each slice. The outer dictionary map slice keys to the
          expected metrics for that slice.
      """
  
      def check_metrics(got):
          """Check metrics callback."""
          try:
              slices = {}
              for slice_key, value in got:
                  slices[slice_key] = value
              self.assertCountEqual(
                  list(slices), list(expected_slice_metrics)
              )
              for (
                  slice_key,
                  expected_metrics,
              ) in expected_slice_metrics.items():
                  self.assertDictElementsWithinBounds(
                      got_values_dict=slices[slice_key],
                      expected_values_dict=expected_metrics,
                  )
          except AssertionError as err:
              raise beam_util.BeamAssertException(err)
  
      slicing_specs = None
      if slice_spec:
          slicing_specs = [s.to_proto() for s in slice_spec]
      eval_config = config_pb2.EvalConfig(slicing_specs=slicing_specs)
      eval_shared_model = self.createTestEvalSharedModel(
          eval_saved_model_path=eval_saved_model_path,
          add_metrics_callbacks=add_metrics_callbacks,
      )
      extractors = model_eval_lib.default_extractors(
          eval_config=eval_config, eval_shared_model=eval_shared_model
      )
  
      tfx_io = raw_tf_record.RawBeamRecordTFXIO(
          physical_format="inmemory",
          raw_record_column_name=constants.ARROW_INPUT_COLUMN,
          telemetry_descriptors=["TFMATest"],
      )
      # pylint: disable=no-value-for-parameter
      (metrics, _), _ = (
          examples_pcollection
          | "BatchExamples" >> tfx_io.BeamSource()
          | "InputsToExtracts"
          >> model_eval_lib.BatchedInputsToExtracts()
          | "Extract" >> Extract(extractors=extractors)
          | "ComputeMetricsAndPlots"
          >> legacy_metrics_and_plots_evaluator._ComputeMetricsAndPlots(  # pylint: disable=protected-access
              eval_shared_model=eval_shared_model
          )
      )
      # pylint: enable=no-value-for-parameter
  
      beam_util.assert_that(metrics, check_metrics)
  
imports: from typing import Any, Dict, List, Optional
  import apache_beam
  from apache_beam.testing import util
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import model_eval_lib
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.eval_saved_model import load
  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 extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.slicer import slicer_lib
  from tfx_bsl.tfxio import raw_tf_record
label_desc: Use image node_id 8 for calling the TestCase obj's underlying member method code with example usage: obj.assertGeneralMetricsComputedWithBeamAre(eval_saved_model_path, examples_pcollection, slice_spec, add_metrics_callbacks, expected_slice_metrics) without return types