def get_dataset(params, mode, file_pattern, vocab):
  """Gets a tf.data.Dataset representing the captions data."""
  data_files = tf.io.gfile.glob(file_pattern)
  if not data_files:
    tf.logging.fatal("Found no input files matching %s", file_pattern)
  else:
    tf.logging.info("Prefetching values from %d files matching %s",
                    len(data_files), file_pattern)

  # Create dataset from files.
  dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))

  # Shuffle.
  if mode == tf_estimator.ModeKeys.TRAIN:
    dataset = dataset.shuffle(buffer_size=len(data_files))

  # Repeat.
  if mode != tf_estimator.ModeKeys.PREDICT:
    dataset = dataset.repeat()

  # Load TFRecords from files.
  dataset = dataset.interleave(
      functools.partial(tf.data.TFRecordDataset, buffer_size=16 * 1024 * 1024),
      cycle_length=len(data_files),
      num_parallel_calls=len(data_files))

  # Shuffle.
  if mode == tf_estimator.ModeKeys.TRAIN:
    dataset = dataset.shuffle(buffer_size=1000)

  # Decode TFRecords.
  dataset = dataset.map(
      parse_example, num_parallel_calls=params["num_input_threads"])

  # Preprocess.
  dataset = dataset.map(
      functools.partial(
          preprocess_mapper,
          params=params,
          lookup_table=vocab.get_string_lookup_table(),
          vocab=vocab,
          mode=mode),
      num_parallel_calls=params["num_input_threads"])

  return dataset
Gets a tf.data.Dataset representing the captions data.
language/capwap/datasets/captions_dataset.py