code: def testPredictionsExtractorWithMultiClassModel(self, signature_name):
      temp_export_dir = self._getExportDir()
      num_classes = 3
      export_dir, _ = dnn_classifier.simple_dnn_classifier(
          temp_export_dir, None, n_classes=num_classes
      )
  
      eval_config = config_pb2.EvalConfig(
          model_specs=[
              config_pb2.ModelSpec(signature_name=signature_name)
          ]
      )
      eval_shared_model = self.createTestEvalSharedModel(
          eval_saved_model_path=export_dir,
          tags=[tf.saved_model.SERVING],
      )
      (
          tfx_io,
          feature_extractor,
      ) = self._create_tfxio_and_feature_extractor(
          eval_config,
          text_format.Parse(
              """
          feature {
            name: "age"
            type: FLOAT
          }
          feature {
            name: "langauge"
            type: BYTES
          }
          feature {
            name: "label"
            type: INT
          }
          """,
              schema_pb2.Schema(),
          ),
      )
  
      examples = [
          self._makeExample(age=1.0, language="english", label=0),
          self._makeExample(age=2.0, language="chinese", label=1),
          self._makeExample(age=3.0, language="english", label=2),
          self._makeExample(age=4.0, language="chinese", label=1),
      ]
      num_examples = len(examples)
  
      prediction_extractor = predictions_extractor.PredictionsExtractor(
          eval_config=eval_config, eval_shared_model=eval_shared_model
      )
  
      with beam.Pipeline() as pipeline:
          # pylint: disable=no-value-for-parameter
          result = (
              pipeline
              | "Create"
              >> beam.Create(
                  [e.SerializeToString() for e in examples],
                  reshuffle=False,
              )
              | "BatchExamples"
              >> tfx_io.BeamSource(batch_size=num_examples)
              | "InputsToExtracts"
              >> model_eval_lib.BatchedInputsToExtracts()
              | feature_extractor.stage_name
              >> feature_extractor.ptransform
              | prediction_extractor.stage_name
              >> prediction_extractor.ptransform
          )
          # pylint: enable=no-value-for-parameter
  
          def check_result(got):
              try:
                  self.assertLen(got, 1)
                  # We can't verify the actual predictions, but we can verify the keys.
                  self.assertIn(constants.PREDICTIONS_KEY, got[0])
                  # Prediction API cases. We default '' to 'predict'.
                  if signature_name in ("", "predict"):
                      for pred_key in ("probabilities", "all_classes"):
                          self.assertIn(
                              pred_key,
                              got[0][constants.PREDICTIONS_KEY],
                          )
                      self.assertEqual(
                          (num_examples, num_classes),
                          got[0][constants.PREDICTIONS_KEY][
                              "probabilities"
                          ].shape,
                      )
                  # Classification API cases. The classification signature is also the
                  # 'serving_default' signature for this model.
                  if signature_name in (
                      "serving_default",
                      "classification",
                  ):
                      for pred_key in ("classes", "scores"):
                          self.assertIn(
                              pred_key,
                              got[0][constants.PREDICTIONS_KEY],
                          )
                      self.assertEqual(
                          (num_examples, num_classes),
                          got[0][constants.PREDICTIONS_KEY][
                              "classes"
                          ].shape,
                      )
                      self.assertEqual(
                          (num_examples, num_classes),
                          got[0][constants.PREDICTIONS_KEY][
                              "scores"
                          ].shape,
                      )
  
              except AssertionError as err:
                  raise util.BeamAssertException(err)
  
          util.assert_that(result, check_result)
  
imports: import os
  from absl.testing import parameterized
  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.eval_saved_model.example_trainers import batch_size_limited_classifier
  from tensorflow_model_analysis.eval_saved_model.example_trainers import dnn_classifier
  from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator_extra_fields
  from tensorflow_model_analysis.eval_saved_model.example_trainers import multi_head
  from tensorflow_model_analysis.extractors import features_extractor
  from tensorflow_model_analysis.extractors import predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tfx_bsl.tfxio import tensor_adapter
  from tfx_bsl.tfxio import test_util
  from google.protobuf import text_format
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 6 for calling the PredictionsExtractorTest obj's underlying member method code with example usage: obj.testPredictionsExtractorWithMultiClassModel(signature_name) without return types