code: def _ModelSignaturesInferenceWrapper(
      extracts: beam.pvalue.PCollection,
      model_specs: List[config_pb2.ModelSpec],
      eval_shared_model: types.MaybeMultipleEvalSharedModels,
      output_keypath: Sequence[str],
  ) -> 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).
        model_specs: Model specs each of which corresponds to each of the
          eval_shared_models.
        eval_shared_model: Shared model parameters keyed by model name.
        output_keypath: A sequence of keys to be used as the path to traverse and
          insert the outputs in the extract.
  
      Returns:
        PCollection of Extracts updated with the 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
      }
      signature_names = {}
      for model_spec in model_specs:
          model_name = "" if len(model_specs) == 1 else model_spec.name
          signature_names[model_name] = [model_spec.signature_name]
  
      return extracts | "Inference" >> beam.ParDo(
          model_util.ModelSignaturesDoFn(
              model_specs=model_specs,
              eval_shared_models=name_to_eval_shared_model,
              output_keypath=output_keypath,
              signature_names=signature_names,
              prefer_dict_outputs=False,
          )
      )
  
file_path: model-analysis/tensorflow_model_analysis/extractors/predictions_extractor.py
imports: from typing import List, Optional, Sequence
  from absl import logging
  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.extractors import materialized_predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
label_desc: Use image node_id 2 for calling a global function with example usage: _ModelSignaturesInferenceWrapper(extracts, model_specs, eval_shared_model, output_keypath) and returns: unknown