code: def create_squad_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),
      }
      if is_training:
          name_to_features["start_positions"] = tf.io.FixedLenFeature(
              [], tf.int64
          )
          name_to_features["end_positions"] = tf.io.FixedLenFeature(
              [], tf.int64
          )
      else:
          name_to_features["unique_ids"] = 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):
          """Dispatches record to features and labels."""
          x, y = {}, {}
          for name, tensor in record.items():
              if name in ("start_positions", "end_positions"):
                  y[name] = tensor
              elif name == "input_ids":
                  x["input_word_ids"] = tensor
              elif name == "segment_ids":
                  x["input_type_ids"] = tensor
              else:
                  x[name] = tensor
          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=True)
      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 5 for calling a global function with example usage: create_squad_dataset(file_path, seq_length, batch_size, is_training, input_pipeline_context) and returns: dataset, x,  y