code: def _ExtractPredictions(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      eval_shared_models: Dict[str, tfma.EvalSharedModel],
  ) -> beam.pvalue.PCollection:
      """A PTransform that adds predictions and possibly other tensors to extracts.
  
      Args:
        extracts: PCollection of extracts with inputs keyed by tfma.INPUTS_KEY.
        eval_shared_models: Shared model parameters keyed by model name.
  
      Returns:
        PCollection of Extracts updated with the predictions.
      """
      return extracts | "Predict" >> beam.ParDo(
          _TFMAPredictionDoFn(eval_shared_models)
      )
  
file_path: tfx/tfx/examples/penguin/experimental/sklearn_predict_extractor.py
imports: import copy
  import os
  import pickle
  from typing import Dict, Iterable, List
  import apache_beam
  import tensorflow
  import tensorflow_model_analysis
  from tfx_bsl.tfxio import tensor_adapter
label_desc: Use image node_id 2 for calling a global function with example usage: _ExtractPredictions(extracts, eval_shared_models) and returns: unknown