code: def test_pipeline_with_multi_task(self):
      data_dir = self.create_tempdir()
      data_file = os.path.join(data_dir, "elwc.tfrecord")
      if tf.io.gfile.exists(data_file):
          tf.io.gfile.remove(data_file)
  
      with tf.io.TFRecordWriter(data_file) as writer:
          for _ in range(256):
              writer.write(ELWC.SerializeToString())
  
      model_dir = os.path.join(data_dir, "model")
  
      dataset_hparams = pipeline.DatasetHparams(
          train_input_pattern=data_file,
          valid_input_pattern=data_file,
          train_batch_size=128,
          valid_batch_size=128,
          list_size=2,
          dataset_reader=tf.data.TFRecordDataset,
          convert_labels_to_binary=False,
      )
      pipeline_hparams = pipeline.PipelineHparams(
          model_dir=model_dir,
          num_epochs=2,
          steps_per_epoch=5,
          validation_steps=2,
          learning_rate=0.01,
          loss={
              "task1": "softmax_loss",
              "task2": "pairwise_logistic_loss",
          },
          loss_weights={"task1": 1.0, "task2": 2.0},
          export_best_model=True,
          strategy="MirroredStrategy",
      )
  
      context_feature_spec = {
          "cf_1": tf.io.FixedLenFeature(
              shape=(1,), dtype=tf.float32, default_value=0.0
          ),
      }
      example_feature_spec = {
          "custom_features_{}".format(i + 1): tf.io.FixedLenFeature(
              shape=(1,), dtype=tf.float32, default_value=0.0
          )
          for i in range(3)
      }
      label_spec = (
          _LABEL_FEATURE,
          tf.io.FixedLenFeature(
              shape=(1,), dtype=tf.float32, default_value=_PADDING_LABEL
          ),
      )
      label_spec = {"task1": label_spec, "task2": label_spec}
      weight_spec = (
          "weight",
          tf.io.FixedLenFeature(
              shape=(1,), dtype=tf.float32, default_value=1.0
          ),
      )
  
      model_builder = model_lib.ModelBuilder(
          input_creator=model_lib.FeatureSpecInputCreator(
              context_feature_spec, example_feature_spec
          ),
          preprocessor=model_lib.PreprocessorWithSpec({}),
          scorer=DummyMultiTaskScorer(),
          mask_feature_name=_MASK,
          name="multi_task_model",
      )
  
      ranking_pipeline = pipeline.MultiTaskPipeline(
          model_builder,
          dataset_builder=pipeline.MultiLabelDatasetBuilder(
              context_feature_spec,
              example_feature_spec,
              _MASK,
              label_spec,
              dataset_hparams,
              sample_weight_spec=weight_spec,
          ),
          hparams=pipeline_hparams,
      )
  
      ranking_pipeline.train_and_validate(verbose=1)
  
      latest_model_path = os.path.join(model_dir, "export/latest_model")
      self.assertTrue(
          tf.saved_model.contains_saved_model(latest_model_path)
      )
  
      latest_model = tf.saved_model.load(export_dir=latest_model_path)
      listwise_predictor = latest_model.signatures[
          tf.saved_model.PREDICT_METHOD_NAME
      ]
      listwise_logits = listwise_predictor(
          tf.convert_to_tensor([ELWC.SerializeToString()] * 2)
      )["task1"]
      self.assertAllEqual([2, 2], listwise_logits.get_shape().as_list())
  
      pointwise_predictor = latest_model.signatures[
          tf.saved_model.REGRESS_METHOD_NAME
      ]
      pointwise_logits = pointwise_predictor(
          tf.convert_to_tensor(
              [
                  EXAMPLE_PROTO_1.SerializeToString(),
                  EXAMPLE_PROTO_2.SerializeToString(),
              ]
          )
      )["task1"]
      self.assertAllEqual([2], pointwise_logits.get_shape().as_list())
  
      self.assertAllClose(pointwise_logits, listwise_logits[0])
  
file_path: ranking/tensorflow_ranking/python/keras/pipeline_test.py
imports: import os
  from typing import Dict
  from absl.testing import parameterized
  import tensorflow
  from tensorflow_ranking.python.keras import model
  from tensorflow_ranking.python.keras import pipeline
  from google.protobuf import text_format
  from tensorflow_serving.apis import input_pb2
label_desc: Use image node_id 2 for calling the PipelineTest obj's underlying member method code with example usage: obj.test_pipeline_with_multi_task() without return types