class_docstr: A helper class for testing interop with TFX pipelines.
class_docstr_tok: ["A","helper","class","for","testing","interop","with","TFX","pipelines","."]
docstr: Runs a sample TFMA job and stores output.
  
  This uses a trivial inputs (two examples, with prediction/label = 0/1 and
  1/1). and writes metrics and plots to the specified path.
  
  Args:
    tfma_path: The path to save the TFMA output to.
    output_file_format: The format to save TFMA output to.
    See [TFMA API Docs](https://www.tensorflow.org/tfx/model_analysis/api_docs/python/tfma/writers/MetricsPlotsAndValidationsWriter) # pylint: disable=line-too-long
        for the most up-to-date reference. If the empty string, 'tfrecord'
        will be used.
    add_metrics_callbacks: TFMA metric callbacks to compute.
    See [TFMA API Docs](https://www.tensorflow.org/tfx/model_analysis/api_docs/python/tfma/post_export_metrics) # pylint: disable=line-too-long
        for examples.
    store: The MLMD store to save the TFMA output artifact.
docstr_tok: ["Runs","a","sample","TFMA","job","and","stores","output",".","This","uses","a","trivial","inputs","(","two","examples",",","with","prediction\/label","=","0\/1","and","1\/1",")",".","and","writes","metrics","and","plots","to","the","specified","path",".","Args",":","tfma_path",":","The","path","to","save","the","TFMA","output","to",".","output_file_format",":","The","format","to","save","TFMA","output","to",".","See","[","TFMA","API","Docs","]","(","https",":","\/\/www.tensorflow.org\/tfx\/model_analysis\/api_docs\/python\/tfma\/writers\/MetricsPlotsAndValidationsWriter",")","#","pylint",":","disable=line-too-long","for","the","most","up-to-date","reference",".","If","the","empty","string",",","'tfrecord'","will","be","used",".","add_metrics_callbacks",":","TFMA","metric","callbacks","to","compute",".","See","[","TFMA","API","Docs","]","(","https",":","\/\/www.tensorflow.org\/tfx\/model_analysis\/api_docs\/python\/tfma\/post_export_metrics",")","#","pylint",":","disable=line-too-long","for","examples",".","store",":","The","MLMD","store","to","save","the","TFMA","output","artifact","."]
code: def _write_tfma(
      self,
      tfma_path: str,
      output_file_format: str,
      add_metrics_callbacks: List[Callable[..., Any]],
      store: Optional[mlmd.MetadataStore] = None,
  ) -> None:
      """Runs a sample TFMA job and stores output.
  
      This uses a trivial inputs (two examples, with prediction/label = 0/1 and
      1/1). and writes metrics and plots to the specified path.
  
      Args:
        tfma_path: The path to save the TFMA output to.
        output_file_format: The format to save TFMA output to.
        See [TFMA API Docs](https://www.tensorflow.org/tfx/model_analysis/api_docs/python/tfma/writers/MetricsPlotsAndValidationsWriter) # pylint: disable=line-too-long
            for the most up-to-date reference. If the empty string, 'tfrecord'
            will be used.
        add_metrics_callbacks: TFMA metric callbacks to compute.
        See [TFMA API Docs](https://www.tensorflow.org/tfx/model_analysis/api_docs/python/tfma/post_export_metrics) # pylint: disable=line-too-long
            for examples.
        store: The MLMD store to save the TFMA output artifact.
      """
      if not output_file_format:
          output_file_format = "tfrecord"
      (
          _,
          eval_saved_model_path,
      ) = fixed_prediction_estimator.simple_fixed_prediction_estimator(
          export_path=None,
          eval_export_path=os.path.join(self.tmpdir, "eval_export_dir"),
      )
      eval_config = tfma.EvalConfig(model_specs=[tfma.ModelSpec()])
      eval_shared_model = self.createTestEvalSharedModel(
          eval_saved_model_path=eval_saved_model_path,
          add_metrics_callbacks=add_metrics_callbacks,
      )
      extractors = [
          tfma.extractors.legacy_predict_extractor.PredictExtractor(
              eval_shared_model, eval_config=eval_config
          ),
          tfma.extractors.unbatch_extractor.UnbatchExtractor(),
          tfma.extractors.slice_key_extractor.SliceKeyExtractor(),
      ]
      evaluators = [
          tfma.evaluators.legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator(
              eval_shared_model
          )
      ]
      writers = [
          tfma.writers.MetricsPlotsAndValidationsWriter(
              output_paths={
                  "metrics": os.path.join(tfma_path, "metrics"),
                  "plots": os.path.join(tfma_path, "plots"),
              },
              output_file_format=output_file_format,
              eval_config=eval_config,
              add_metrics_callbacks=eval_shared_model.add_metrics_callbacks,
          ),
      ]
  
      tfx_io = raw_tf_record.RawBeamRecordTFXIO(
          physical_format="inmemory",
          raw_record_column_name="__raw_record__",
          telemetry_descriptors=["TFMATest"],
      )
      with beam.Pipeline() as pipeline:
          example1 = self._makeExample(prediction=0.0, label=1.0)
          example2 = self._makeExample(prediction=1.0, label=1.0)
          _ = (
              pipeline
              | "Create"
              >> beam.Create(
                  [
                      example1.SerializeToString(),
                      example2.SerializeToString(),
                  ]
              )
              | "BatchExamples" >> tfx_io.BeamSource()
              | "ExtractEvaluateAndWriteResults"
              >> tfma.ExtractEvaluateAndWriteResults(
                  eval_config=eval_config,
                  eval_shared_model=eval_shared_model,
                  extractors=extractors,
                  evaluators=evaluators,
                  writers=writers,
              )
          )
  
      if store:
          self._put_artifact(store, _TFX_METRICS_TYPE, tfma_path)
  
file_path: model-card-toolkit/model_card_toolkit/utils/testdata/tfxtest.py
imports: import os
  from typing import Any, Callable, List, Optional
  import apache_beam
  import ml_metadata
  import tensorflow_model_analysis
  from absl import flags
  from ml_metadata.proto import metadata_store_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
  from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator
  from tfx_bsl.tfxio import raw_tf_record
  from model_card_toolkit.utils import io_utils
  from model_card_toolkit.utils.tf_utils import _TFX_METRICS_TYPE, _TFX_STATS_TYPE
label_desc: Use image node_id 4 for calling the TfxTest obj's underlying member method code with example usage: obj._write_tfma(tfma_path, output_file_format, add_metrics_callbacks, store) without return types