code: def _TFMAPredict(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      eval_shared_models: Dict[str, types.EvalSharedModel],
      desired_batch_size: Optional[int] = None,
      materialize: Optional[bool] = True,
      eval_config: Optional[config_pb2.EvalConfig] = None,
  ) -> beam.pvalue.PCollection:
      """A PTransform that adds predictions to Extracts.
  
      Args:
        extracts: PCollection of Extracts containing a serialized example to be fed
          to the model.
        eval_shared_models: Shared model parameters keyed by model name.
        desired_batch_size: Optional. Desired batch size for prediction.
        materialize: True to call the FeatureExtractor to add MaterializedColumn
          entries for the features, predictions, and labels.
        eval_config: Eval config.
  
      Returns:
        PCollection of Extracts, where the extracts contains the features,
        predictions, labels retrieved.
      """
      if not eval_config:
          batch_args = {}
  
          # TODO(b/143484017): Consider removing this option if autotuning is better
          # able to handle batch size selection.
          if desired_batch_size:
              batch_args = dict(
                  min_batch_size=desired_batch_size,
                  max_batch_size=desired_batch_size,
              )
  
          extracts = extracts | "Batch" >> beam.BatchElements(
              **batch_args
          )
      else:
          extracts = extracts | "UnwrapBatchedExtract" >> beam.Map(
              _unwrap_batched_extract
          )
  
      # We don't actually need to add the add_metrics_callbacks to do Predict,
      # but because if we want to share the model between Predict and subsequent
      # stages (i.e. we use same shared handle for this and subsequent stages),
      # then if we don't add the metrics callbacks here, they won't be present
      # in the model in the later stages if we reuse the model from this stage.
      extracts = extracts | "Predict" >> beam.ParDo(
          _TFMAPredictionDoFn(
              eval_shared_models=eval_shared_models,
              eval_config=eval_config,
          )
      )
  
      if materialize and not eval_config:
          additional_fetches = []
          for m in eval_shared_models.values():
              if m.additional_fetches:
                  additional_fetches.extend(m.additional_fetches)
          return (
              extracts
              | "ExtractFeatures"
              >> legacy_feature_extractor._ExtractFeatures(  # pylint: disable=protected-access
                  additional_extracts=additional_fetches or None
              )
          )
  
      return extracts
  
imports: import collections
  import copy
  from typing import Any, Dict, List, Optional, Sequence
  import apache_beam
  import numpy
  import pyarrow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.eval_saved_model import constants
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.extractors import legacy_feature_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
label_desc: Use image node_id 5 for calling a global function with example usage: _TFMAPredict(extracts, eval_shared_models, desired_batch_size, materialize, eval_config) and returns: extracts, unknown