def input_fn(params):
  """An input function satisfying the tf.estimator API."""
  compression_type = "GZIP" if FLAGS.examples_path.endswith(".gz") else ""
  examples_paths = sorted(tf.io.gfile.glob(FLAGS.examples_path))
  dataset = tf.data.TFRecordDataset(
      examples_paths,
      compression_type=compression_type,
      buffer_size=16 * 1024 * 1024)
  dataset = dataset.map(parse_examples, num_parallel_calls=FLAGS.num_threads)

  # Repeat just one extra batch to make sure we don't lose any remainder
  # blocks that might be dropped.
  dataset = dataset.concatenate(dataset.take(params["batch_size"]))

  dataset = dataset.batch(params["batch_size"], drop_remainder=True)
  dataset = dataset.prefetch(10)
  return dataset
An input function satisfying the tf.estimator API.
language/multivec/predict/encode_blocks.py