code: def CounterfactualPredictionsExtractor(  # pylint: disable=invalid-name
      eval_shared_models: types.MaybeMultipleEvalSharedModels,
      eval_config: config_pb2.EvalConfig,
      cf_configs: Mapping[str, CounterfactualConfig],
  ) -> extractor.Extractor:
      """Creates a CF predictions extractor by wrapping the PredictionsExtractor.
  
      Example usage:
  
        eval_config = tfma.EvalConfig(model_specs=[
           tfma.ModelSpec(name='orig', is_baseline=True),
           tfma.ModelSpec(name='cf')])
        eval_shared_models = {
          'orig': eval_shared_model,
          'cf' eval_shared_model}
        cf_configs = {'cf': {'x_cf': 'x'}}
        extractors = tfma.default_extractors(eval_shared_models, eval_config,
            custom_predictions_extractor=CounterfactualPredictionsExtractor(
                eval_shared_models,cf_configs))
        tfma.run_model_analysis(eval_shared_models, eval_config,
            extractors=etractors)
  
      Args:
        eval_shared_models: The set of eval_shared_models for which to generate
          predictions. If a model is to be computed with original inputs and CF
          inputs, it should be provided twice, with distinct names. The name of the
          model to be computed with CF inputs should match the name provided in
          cf_configs as well as the ModelSpec.name in the provided EvalConfig.
        eval_config: The EvalConfig for this evaluation. If a model is to be
          computed with original inputs and CF inputs, it should correspond to two
          ModelSpecs with distinct names. The CF model name should match the name
          provided in cf_configs as well as the EvalSharedModel.model_name.
        cf_configs: A mapping from a model name to the CF config which should be
          used to preprocess its inputs. Any models in eval_shared_models not
          specified will have their predictions computed on the original input
  
      Returns:
        A tfma.Extractor which performs counterfactual inference along with non-
        counterfactual inference.
  
      Raises:
        ValueError if eval_shared_models is empty.
      """
      (
          eval_shared_models,
          cf_configs,
      ) = _validate_and_update_models_and_configs(
          eval_shared_models, cf_configs
      )
      cf_ptransforms = {}
      non_cf_models = []
      for model in eval_shared_models:
          cf_config = cf_configs.get(model.model_name, None)
          if cf_config:
              # filter EvalConfig so that it matches single EvalSavedModel
              cf_eval_config = _filter_model_specs(eval_config, [model])
              # TODO(b/258850519): Refactor default_extractors logic to expose new api
              # for constructing the default predictions extractor and call it here.
              predictions_ptransform = (
                  predictions_extractor.PredictionsExtractor(
                      eval_shared_model=model,
                      eval_config=cf_eval_config,
                      output_keypath=(
                          constants.PREDICTIONS_KEY,
                          model.model_name,
                      ),
                  ).ptransform
              )
              cf_ptransforms[
                  model.model_name
              ] = _ExtractCounterfactualPredictions(  # pylint: disable=no-value-for-parameter
                  config=cf_config,
                  predictions_ptransform=predictions_ptransform,
              )
          else:
              non_cf_models.append(model)
      non_cf_eval_config = _filter_model_specs(
          eval_config, non_cf_models
      )
      if non_cf_models:
          output_keypath = (constants.PREDICTIONS_KEY,)
          if len(non_cf_models) == 1:
              output_keypath = output_keypath + (
                  non_cf_models[0].model_name,
              )
          non_cf_ptransform = (
              predictions_extractor.PredictionsExtractor(
                  eval_shared_model=non_cf_models,
                  eval_config=non_cf_eval_config,
                  output_keypath=output_keypath,
              ).ptransform
          )
      else:
          non_cf_ptransform = None
      return extractor.Extractor(
          stage_name=_COUNTERFACTUAL_PREDICTIONS_EXTRACTOR_NAME,
          ptransform=_ExtractPredictions(  # pylint: disable=no-value-for-parameter
              cf_ptransforms=cf_ptransforms,
              non_cf_ptransform=non_cf_ptransform,
          ),
      )
  
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 1 for calling a global function with example usage: CounterfactualPredictionsExtractor(eval_shared_models, eval_config, cf_configs) and returns: extractor