code: def TfxBslPredictionsExtractor(
      eval_config: config_pb2.EvalConfig,
      eval_shared_model: types.MaybeMultipleEvalSharedModels,
      output_batch_size: Optional[int] = None,
      output_keypath: Iterable[str] = (constants.PREDICTIONS_KEY,),
  ) -> extractor.Extractor:
      """Creates an extractor for performing predictions over a batch.
  
      The extractor's PTransform loads and runs the serving saved_model(s) against
      every Extracts yielding a copy of the incoming Extracts with an additional
      Extracts added for the predictions keyed by tfma.PREDICTIONS_KEY. The model
      inputs are searched for under tfma.FEATURES_KEY (keras only) or tfma.INPUT_KEY
      (if tfma.FEATURES_KEY is not set or the model is non-keras). If multiple
      models are used the predictions will be stored in a dict keyed by model name.
  
      Note that the prediction_key in the ModelSpecs also serves as a key into the
      dict of the prediction's output.
  
      Args:
        eval_config: Eval config.
        eval_shared_model: Shared model (single-model evaluation) or list of shared
          models (multi-model evaluation).
        output_batch_size: Sets a static output batch size for bulk inference. Note:
          this only affects the rebatched output batch size to set inference batch
          size set ModelSpec.inference_batch_size.
        output_keypath: A sequence of keys to be used as the path to traverse and
          insert the outputs in the extract.
  
      Returns:
        Extractor for extracting predictions.
      """
      eval_shared_models = (
          model_util.verify_and_update_eval_shared_models(
              eval_shared_model
          )
      )
      # This should never happen, but verify_and_update_eval_shared_models can
      # theoretically return None or empty iterables.
      if not eval_shared_models:
          raise ValueError(
              "No valid model(s) were provided. Please ensure that "
              "EvalConfig.ModelSpec is correctly configured to enable "
              "using the PredictionsExtractor."
          )
  
      name_to_eval_shared_model = {
          m.model_name: m for m in eval_shared_models
      }
      model_specs = []
      for model_spec in eval_config.model_specs:
          if not model_spec.signature_name:
              eval_shared_model = model_util.get_eval_shared_model(
                  model_spec.name, name_to_eval_shared_model
              )
              model_spec = copy.copy(model_spec)
              # Select a default signature. Note that this may differ from the
              # 'serving_default' signature.
              model_spec.signature_name = (
                  model_util.get_default_signature_name_from_model_path(
                      eval_shared_model.model_path
                  )
              )
          model_specs.append(model_spec)
  
      tfx_bsl_inference_ptransform = inference_base.RunInference(
          inference_ptransform=TfxBslInferenceWrapper(
              model_specs, name_to_eval_shared_model
          ),
          output_batch_size=output_batch_size,
          output_keypath=output_keypath,
      )
      # pylint: disable=no-value-for-parameter
      return extractor.Extractor(
          stage_name=predictions_extractor.PREDICTIONS_EXTRACTOR_STAGE_NAME,
          ptransform=tfx_bsl_inference_ptransform,
      )
  
file_path: model-analysis/tensorflow_model_analysis/extractors/tfx_bsl_predictions_extractor.py
imports: import copy
  from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
  import apache_beam
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.extractors import inference_base
  from tensorflow_model_analysis.extractors import predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tfx_bsl.public.beam import run_inference
  from tfx_bsl.public.proto import model_spec_pb2
  from tensorflow_serving.apis import prediction_log_pb2
label_desc: Use image node_id 1 for calling a global function with example usage: TfxBslPredictionsExtractor(eval_config, eval_shared_model, output_batch_size, output_keypath) and returns: extractor