code: def _ExtractPredictions(  # pylint: disable=invalid-name
      extracts: beam.PCollection[types.Extracts],
      cf_ptransforms: Dict[str, beam.PTransform],
      non_cf_ptransform: Optional[beam.PTransform],
  ) -> beam.PCollection[types.Extracts]:
      """Applies both CF and non-CF prediction ptransforms and merges results.
  
      Args:
        extracts: Incoming TFMA extracts.
        cf_ptransforms: A mapping from model name to
          _ExtractCounterfactualPredictions ptransforms
        non_cf_ptransform: Optionally, a ptransform responsible for computing the
          non-counterfactual predictions.
  
      Returns:
        A PCollection of extracts containing merged predictions from both
        counterfactual and non-counterfactual models.
      """
      if non_cf_ptransform:
          extracts = extracts | "PredictNonCF" >> non_cf_ptransform
      for model_name, cf_ptransform in cf_ptransforms.items():
          extracts = (
              extracts | f"PredictCF[{model_name}]" >> cf_ptransform
          )
      return extracts
  
file_path: model-analysis/tensorflow_model_analysis/extractors/counterfactual_predictions_extractor.py
imports: from typing import Dict, Iterable, Mapping, Optional
  import apache_beam
  import numpy
  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 predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
label_desc: Use image node_id 7 for calling a global function with example usage: _ExtractPredictions(extracts, cf_ptransforms, non_cf_ptransform) and returns: extracts