code: def RunInference(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      inference_ptransform: beam.PTransform,
      output_batch_size: Optional[int] = None,
      output_keypath: Sequence[str] = (constants.PREDICTIONS_KEY,),
      prediction_log_keypath: Optional[Sequence[str]] = None,
  ) -> beam.pvalue.PCollection:
      """A PTransform that adds predictions and possibly other tensors to 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).
        inference_ptransform: Bulk inference ptransform used to generate
          predictions. This allows users to use different implementations depending
          on evironment or Beam runner (e.g. a cloud-friendly OSS implementation or
          an internal-specific implementation). These implementations should accept
          a pcollection consisting of tuples containing a key and a single example.
          The key may be anything and the example may be a tf.Example or serialized
          tf.Example.
        output_batch_size: Sets a static output batch size.
        output_keypath: A sequence of keys to be used as the path to traverse and
          insert the outputs in the extract.
        prediction_log_keypath: A sequence of keys to be used as the path to
          traverse and insert the prediction logs in the extract.
  
      Returns:
        PCollection of Extracts updated with the predictions and prediction logs.
      """
      extracts = (
          extracts
          # Extracts are fed in pre-batched, but BulkInference has specific
          # batch handling and batching requirements. To accomodate the API and
          # encapsulate the inference batching logic, we unbatch here. This function
          # returns new Extracts dicts and will not modify the input Extracts.
          | "SplitExtracts"
          >> beam.FlatMap(util.split_extracts, expand_zero_dims=False)
          # The BulkInference API allows for key forwarding. To avoid a join
          # after running inference, we forward the unbatched Extracts as a key.
          | "CreateInferenceInputTuple"
          >> beam.Map(_create_inference_input_tuple)
          | "RunInferencePerModel" >> inference_ptransform
          # Combine predictions back into the original Extracts.
          | "InsertPredictionsIntoExtracts"
          >> beam.Map(
              insert_predictions_into_extracts,
              output_keypath=output_keypath,
              prediction_log_keypath=prediction_log_keypath,
          )
      )
      # Beam batch will group single Extracts into a batch. Then
      # merge_extracts will flatten the batch into a single "batched"
      # extract.
      if output_batch_size is not None:
          batch_kwargs = {
              "min_batch_size": output_batch_size,
              "max_batch_size": output_batch_size,
          }
      else:
          # Default batch parameters.
          batch_kwargs = {}
      return (
          extracts
          | "BatchSingleExampleExtracts"
          >> beam.BatchElements(**batch_kwargs)
          | "MergeExtracts"
          >> beam.Map(util.merge_extracts, squeeze_two_dim_vector=False)
      )
  
file_path: model-analysis/tensorflow_model_analysis/extractors/inference_base.py
imports: from typing import Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import apache_beam
  import numpy
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tensorflow_model_analysis.utils import util
  from tensorflow.python.saved_model import loader_impl
  from tensorflow_serving.apis import prediction_log_pb2
label_desc: Use image node_id 5 for calling a global function with example usage: RunInference(extracts, inference_ptransform, output_batch_size, output_keypath, prediction_log_keypath) and returns: unknown