code: def test_load_dataset(self, label_type):
      input_path = os.path.join(self.get_temp_dir(), "train.tf_record")
      batch_size = 8
      list_size = 12
      seq_length = 128
      _create_fake_preprocessed_dataset(
          input_path, seq_length, label_type
      )
      label_spec = (
          "relevance",
          tf.io.FixedLenFeature(
              shape=[
                  1,
              ],
              dtype=label_type,
              default_value=-1,
          ),
      )
      data_config = tfrbert_task.TFRBertDataConfig(
          input_path=input_path,
          data_format=tfr_data.ELWC,
          list_size=list_size,
          shuffle_examples=False,
          seq_length=seq_length,
          global_batch_size=batch_size,
          mask_feature_name="example_list_mask",
          read_query_id=True,
          read_document_id=True,
      )
      dataset = tfrbert_task.TFRBertDataLoader(
          data_config, label_spec
      ).load()
      features, labels = next(iter(dataset))
      self.assertCountEqual(
          [
              "input_word_ids",
              "input_mask",
              "input_type_ids",
              "example_list_mask",
              "query_id",
              "document_id",
          ],
          features.keys(),
      )
      self.assertEqual(
          features["input_word_ids"].shape,
          (batch_size, list_size, seq_length),
      )
      self.assertEqual(
          features["input_mask"].shape,
          (batch_size, list_size, seq_length),
      )
      self.assertEqual(
          features["input_type_ids"].shape,
          (batch_size, list_size, seq_length),
      )
      self.assertEqual(
          features["example_list_mask"].shape, (batch_size, list_size)
      )
      self.assertEqual(
          features["query_id"].shape, (batch_size, list_size, 1)
      )
      self.assertEqual(
          features["document_id"].shape, (batch_size, list_size, 1)
      )
  
      self.assertEqual(labels.shape, (batch_size, list_size))
  
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 1 for calling the TFRBERTDataTest obj's underlying member method code with example usage: obj.test_load_dataset(label_type) without return types