code: def testInputExtractorMultiModel(self):
      model_spec1 = config_pb2.ModelSpec(
          name="model1",
          label_key="label",
          example_weight_key="example_weight",
          prediction_key="fixed_float",
      )
      model_spec2 = config_pb2.ModelSpec(
          name="model2",
          label_keys={"output1": "label1", "output2": "label2"},
          example_weight_keys={
              "output1": "example_weight1",
              "output2": "example_weight2",
          },
          prediction_keys={
              "output1": "fixed_float",
              "output2": "fixed_float",
          },
      )
      extractor = input_extractor.InputExtractor(
          eval_config=config_pb2.EvalConfig(
              model_specs=[model_spec1, model_spec2]
          )
      )
  
      examples = [
          self._makeExample(
              label=1.0,
              label1=1.0,
              label2=0.0,
              example_weight=0.5,
              example_weight1=0.5,
              example_weight2=0.5,
              fixed_int=1,
              fixed_float=1.0,
              fixed_string="fixed_string1",
          ),
          self._makeExample(
              label=1.0,
              label1=1.0,
              label2=1.0,
              example_weight=0.0,
              example_weight1=0.0,
              example_weight2=1.0,
              fixed_int=1,
              fixed_float=2.0,
              fixed_string="fixed_string2",
          ),
      ]
  
      with beam.Pipeline() as pipeline:
          # pylint: disable=no-value-for-parameter
          result = (
              pipeline
              | "Create"
              >> beam.Create(
                  [e.SerializeToString() for e in examples],
                  reshuffle=False,
              )
              | "InputsToExtracts" >> model_eval_lib.InputsToExtracts()
              | extractor.stage_name >> extractor.ptransform
          )
  
          # pylint: enable=no-value-for-parameter
  
          def check_result(got):
              try:
                  self.assertLen(got, 2)
                  self.assertDictElementsAlmostEqual(
                      got[0][constants.FEATURES_KEY],
                      {
                          "fixed_int": np.array([1]),
                          "label": np.array([1.0]),
                          "label1": np.array([1.0]),
                          "label2": np.array([0.0]),
                          "example_weight": np.array([0.5]),
                          "example_weight1": np.array([0.5]),
                          "example_weight2": np.array([0.5]),
                      },
                  )
                  self.assertEqual(
                      got[0][constants.FEATURES_KEY]["fixed_string"],
                      np.array([b"fixed_string1"]),
                  )
                  for model_name in ("model1", "model2"):
                      self.assertIn(
                          model_name, got[0][constants.LABELS_KEY]
                      )
                      self.assertIn(
                          model_name,
                          got[0][constants.EXAMPLE_WEIGHTS_KEY],
                      )
                      self.assertIn(
                          model_name, got[0][constants.PREDICTIONS_KEY]
                      )
                  self.assertAlmostEqual(
                      got[0][constants.LABELS_KEY]["model1"],
                      np.array([1.0]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[0][constants.LABELS_KEY]["model2"],
                      {
                          "output1": np.array([1.0]),
                          "output2": np.array([0.0]),
                      },
                  )
                  self.assertAlmostEqual(
                      got[0][constants.EXAMPLE_WEIGHTS_KEY]["model1"],
                      np.array([0.5]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[0][constants.EXAMPLE_WEIGHTS_KEY]["model2"],
                      {
                          "output1": np.array([0.5]),
                          "output2": np.array([0.5]),
                      },
                  )
                  self.assertAlmostEqual(
                      got[0][constants.PREDICTIONS_KEY]["model1"],
                      np.array([1.0]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[0][constants.PREDICTIONS_KEY]["model2"],
                      {
                          "output1": np.array([1.0]),
                          "output2": np.array([1.0]),
                      },
                  )
  
                  self.assertDictElementsAlmostEqual(
                      got[1][constants.FEATURES_KEY],
                      {
                          "fixed_int": np.array([1]),
                          "label": np.array([1.0]),
                          "label1": np.array([1.0]),
                          "label2": np.array([1.0]),
                          "example_weight": np.array([0.0]),
                          "example_weight1": np.array([0.0]),
                          "example_weight2": np.array([1.0]),
                      },
                  )
                  self.assertEqual(
                      got[1][constants.FEATURES_KEY]["fixed_string"],
                      np.array([b"fixed_string2"]),
                  )
                  for model_name in ("model1", "model2"):
                      self.assertIn(
                          model_name, got[1][constants.LABELS_KEY]
                      )
                      self.assertIn(
                          model_name,
                          got[1][constants.EXAMPLE_WEIGHTS_KEY],
                      )
                      self.assertIn(
                          model_name, got[1][constants.PREDICTIONS_KEY]
                      )
                  self.assertAlmostEqual(
                      got[1][constants.LABELS_KEY]["model1"],
                      np.array([1.0]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[1][constants.LABELS_KEY]["model2"],
                      {
                          "output1": np.array([1.0]),
                          "output2": np.array([1.0]),
                      },
                  )
                  self.assertAlmostEqual(
                      got[1][constants.EXAMPLE_WEIGHTS_KEY]["model1"],
                      np.array([0.0]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[1][constants.EXAMPLE_WEIGHTS_KEY]["model2"],
                      {
                          "output1": np.array([0.0]),
                          "output2": np.array([1.0]),
                      },
                  )
                  self.assertAlmostEqual(
                      got[1][constants.PREDICTIONS_KEY]["model1"],
                      np.array([2.0]),
                  )
                  self.assertDictElementsAlmostEqual(
                      got[1][constants.PREDICTIONS_KEY]["model2"],
                      {
                          "output1": np.array([2.0]),
                          "output2": np.array([2.0]),
                      },
                  )
  
              except AssertionError as err:
                  raise util.BeamAssertException(err)
  
          util.assert_that(result, check_result, label="result")
  
file_path: model-analysis/tensorflow_model_analysis/extractors/legacy_input_extractor_test.py
imports: import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import model_eval_lib
  from tensorflow_model_analysis.eval_saved_model import testutil
  from tensorflow_model_analysis.extractors import legacy_input_extractor
  from tensorflow_model_analysis.proto import config_pb2
label_desc: Use image node_id 3 for calling the InputExtractorTest obj's underlying member method code with example usage: obj.testInputExtractorMultiModel() without return types