def get_dataset(data_sources, num_epochs, batch_size, decode_record,
                static_padded_shapes, shuffle, drop_remainder):
  """Converts a dataset file to a dataset tensor."""
  # Get sharded files from sstable path.
  tf.logging.info('Reading from ' + str(data_sources))

  data_files = get_data_files(
      [source_name + '@*' for source_name in data_sources])
  tf.logging.info('First file: ' + data_files[0])

  num_of_files = len(data_files)
  tf.logging.info('# of files: ' + str(num_of_files))

  num_reader_threads = min(num_of_files, multiprocessing.cpu_count())

  ds_filenames = tf.data.Dataset.from_tensor_slices(data_files)

  if num_epochs is None:
    ds_filenames = ds_filenames.shuffle(buffer_size=num_of_files)
  dataset = ds_filenames.apply(
      tf.data.experimental.parallel_interleave(
          tf.data.TFRecordDataset,
          cycle_length=num_reader_threads,
          sloppy=num_epochs is None))

  capacity = 16 * batch_size
  dataset = dataset.map(
      decode_record, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  if static_padded_shapes is None:
    # Let all shapes be dynamic.
    padded_shapes = dict([(name, [None] * len(shape))
                          for name, shape in dataset.output_shapes.items()])
  else:
    padded_shapes = static_padded_shapes

  if shuffle:
    dataset = dataset.shuffle(capacity)

  # Loop inifinitely if training, just once otherwise.
  dataset = dataset.repeat(num_epochs)

  padded_values = dict([(name, get_padded_value(datatype))
                        for name, datatype in dataset.output_types.items()])

  dataset = dataset.padded_batch(
      batch_size,
      padded_shapes=padded_shapes,
      padding_values=padded_values,
      drop_remainder=drop_remainder)

  return dataset
Converts a dataset file to a dataset tensor.
language/xsp/model/input_utils.py