code: def _ExtractMaterializedPredictions(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      eval_config: config_pb2.EvalConfig,
      output_keypath: Sequence[str],
  ) -> beam.pvalue.PCollection:
      """A PTransform that populates the predictions key in the extracts.
  
      Args:
        extracts: PCollection of extracts containing model inputs keyed by
          tfma.FEATURES_KEY (if model inputs are named) or tfma.INPUTS_KEY (if model
          takes raw tf.Examples as input).
        eval_config: Eval config.
        output_keypath: A list of keys that indicates the location to which the
          predictions are inserted.
  
      Returns:
        PCollection of Extracts updated with the predictions.
      """
  
      def rekey_predictions(  # pylint: disable=invalid-name
          batched_extracts: types.Extracts,
      ) -> types.Extracts:
          """Extract predictions from extracts containing features."""
          predictions = (
              model_util.get_feature_values_for_model_spec_field(
                  list(eval_config.model_specs),
                  "prediction_key",
                  "prediction_keys",
                  batched_extracts,
              )
          )
          if predictions is not None:
              return util.copy_and_set_by_keys(
                  batched_extracts, list(output_keypath), predictions
              )
          else:
              return batched_extracts
  
      return extracts | "RekeyPredictions" >> beam.Map(
          rekey_predictions
      )
  
file_path: model-analysis/tensorflow_model_analysis/extractors/materialized_predictions_extractor.py
imports: from typing import Sequence
  import apache_beam
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tensorflow_model_analysis.utils import util
label_desc: Use image node_id 2 for calling a global function with example usage: _ExtractMaterializedPredictions(extracts, eval_config, output_keypath) and returns: unknown, util, batched_extracts