code: def create_classifier_dataset(
      file_path,
      seq_length,
      batch_size,
      is_training=True,
      input_pipeline_context=None,
  ):
      """Creates input dataset from (tf)records files for train/eval."""
      name_to_features = {
          "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
          "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
          "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
          "label_ids": tf.io.FixedLenFeature([], tf.int64),
          "is_real_example": tf.io.FixedLenFeature([], tf.int64),
      }
      dataset = single_file_dataset(file_path, name_to_features)
  
      # The dataset is always sharded by number of hosts.
      # num_input_pipelines is the number of hosts rather than number of cores.
      if (
          input_pipeline_context
          and input_pipeline_context.num_input_pipelines > 1
      ):
          dataset = dataset.shard(
              input_pipeline_context.num_input_pipelines,
              input_pipeline_context.input_pipeline_id,
          )
  
      def _select_data_from_record(record):
          x = {
              "input_word_ids": record["input_ids"],
              "input_mask": record["input_mask"],
              "input_type_ids": record["segment_ids"],
          }
          y = record["label_ids"]
          return (x, y)
  
      dataset = dataset.map(_select_data_from_record)
  
      if is_training:
          dataset = dataset.shuffle(100)
          dataset = dataset.repeat()
  
      dataset = dataset.batch(batch_size, drop_remainder=is_training)
      dataset = dataset.prefetch(1024)
      return dataset
  
file_path: autodist/examples/benchmark/utils/input_pipeline.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: create_classifier_dataset(file_path, seq_length, batch_size, is_training, input_pipeline_context) and returns: dataset, x,  y