code: def test_task(self):
      # Prepare check point and test data
      ckpt_path = self._create_bert_ckpt()
  
      input_path = os.path.join(self.get_temp_dir(), "train.tf_record")
      seq_length = 128
      _create_fake_preprocessed_dataset(
          input_path, seq_length, tf.float32
      )
  
      # Set up data config
      train_data_config = tfrbert_task.TFRBertDataConfig(
          input_path=input_path,
          is_training=True,
          global_batch_size=5,
          list_size=3,
          dataset_fn="tfrecord",
          seq_length=128,
      )
      validation_data_config = tfrbert_task.TFRBertDataConfig(
          input_path=input_path,
          is_training=False,
          global_batch_size=5,
          list_size=3,
          dataset_fn="tfrecord",
          seq_length=128,
          read_query_id=True,
          read_document_id=True,
      )
  
      # Set up task config
      task_config = tfrbert_task.TFRBertConfig(
          output_preds=True,
          init_checkpoint=ckpt_path,
          aggregated_metrics=True,
          train_data=train_data_config,
          validation_data=validation_data_config,
          model=tfrbert_task.TFRBertModelConfig(
              encoder=encoders.EncoderConfig(
                  bert=encoders.BertEncoderConfig(num_layers=1)
              )
          ),
      )
  
      # Set up TFRBertTask
      label_spec = (
          "label",
          tf.io.FixedLenFeature(
              shape=(1,), dtype=tf.int64, default_value=-1
          ),
      )
      task = tfrbert_task.TFRBertTask(
          task_config,
          label_spec=label_spec,
          dataset_fn=tf.data.TFRecordDataset,
          logging_dir=self._logging_dir,
      )
  
      # Test
      model = task.build_model()
      metrics = task.build_metrics()
      train_dataset = task.build_inputs(task_config.train_data)
      vali_dataset = task.build_inputs(task_config.validation_data)
      task.initialize(model)
      train_iterator = iter(train_dataset)
      vali_iterator = iter(vali_dataset)
      optimizer = tf.keras.optimizers.SGD(lr=0.1)
      task.train_step(
          next(train_iterator), model, optimizer, metrics=metrics
      )
      logs = task.validation_step(
          next(vali_iterator), model, metrics=metrics
      )
      logs = {x: (logs[x],) for x in logs}
      logs = task.aggregate_logs(step_outputs=logs)
      self.assertEqual(tf.constant(logs["query_id"]).shape, (1, 5, 3))
      self.assertEqual(
          tf.constant(logs["document_id"]).shape, (1, 5, 3)
      )
      self.assertEqual(
          tf.constant(logs[tfrbert_task._PREDICTION]).shape, (1, 5, 3)
      )
      self.assertEqual(
          tf.constant(logs[tfrbert_task._LABEL]).shape, (1, 5, 3)
      )
      metrics = task.reduce_aggregated_logs(logs)
  
imports: import os
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from official.nlp.configs import encoders
  from tensorflow_ranking.extension.premade import tfrbert_task
  from tensorflow_ranking.python import data
  from tensorflow_ranking.python.keras import model
  from tensorflow_serving.apis import input_pb2
label_desc: Use image node_id 3 for calling the TFRBertTaskTest obj's underlying member method code with example usage: obj.test_task() without return types