def get_dataset(examples_path, mask_rate, bert_hub_module_path, query_seq_len,
                block_seq_len, num_block_records, num_input_threads):
  """An input function satisfying the tf.estimator API."""
  # The input file is not sharded. We can still get the randomization and
  # efficiency benefits of sharded inputs by doing multiple reads concurrently
  # but starting at different points.
  skips = perturbed_chunks(num_block_records, num_input_threads)
  tf.logging.info("Concurrent reads of %d records: %s",
                  num_block_records, skips)
  dataset = tf.data.Dataset.from_tensor_slices(tf.constant(skips, tf.int64))

  def _skipped_dataset(skip):
    """Get skipped dataset."""
    dataset = tf.data.TFRecordDataset(
        examples_path, buffer_size=16 * 1024 * 1024)
    dataset = dataset.repeat()
    dataset = dataset.skip(skip)
    dataset = dataset.map(
        functools.partial(
            get_retrieval_examples,
            mask_rate=mask_rate,
            bert_hub_module_path=bert_hub_module_path,
            query_seq_len=query_seq_len,
            block_seq_len=block_seq_len))
    dataset = dataset.filter(lambda d: d.pop("keep_example"))
    return dataset

  dataset = dataset.apply(tf.data.experimental.parallel_interleave(
      _skipped_dataset,
      sloppy=True,
      cycle_length=num_input_threads))
  return dataset
An input function satisfying the tf.estimator API.
language/orqa/datasets/ict_dataset.py