def load_dataset(shard_paths, randomize):
  """Loads a dataset from a set of sharded TFRecord files."""
  tf.logging.info('Reading data from these shards:')
  for shard_path in shard_paths:
    tf.logging.info(shard_path)

  dataset = tf.data.Dataset.from_tensor_slices(tf.constant(shard_paths))

  def tfrecord_dataset(path):
    buffer_size = 16 * 1024 * 1024  # Max number of bytes to store.
    return tf.data.TFRecordDataset(
        path, compression_type='GZIP', buffer_size=buffer_size)

  # `sloppy` mode means that the interleaving is not exact. This adds more
  # randomness to the training pipeline.
  dataset = dataset.apply(
      tf.data.experimental.parallel_interleave(
          map_func=tfrecord_dataset,
          sloppy=randomize,
          cycle_length=len(shard_paths)))

  if randomize:
    dataset = dataset.shuffle(buffer_size=10000)

  dataset = dataset.prefetch(16 * 1024)
  return dataset
Loads a dataset from a set of sharded TFRecord files.
language/realm/example_generator.py