def make_dset_fn(context, filename, shuffle=False, epochs=1, n_take=0):
  """Return a data_fn usable for an estimator.

  Args:
    context: a NeuralQueryContext
    filename: name of a file in FLAGS.rootdir
    shuffle: if True, shuffle the data
    epochs: number of times to repeat the data
    n_take: size of subset of examples for training/testing

  Returns:
    a function f so that f() returns a tf.data.Dataset
  """
  data_specs = [str, str, 'entity_t', 'entity_t']
  full_filename = '%s/%s/%s' % (FLAGS.rootdir, FLAGS.task, filename)
  logging.info('dataset from %s minibatch %d epochs %d shuffle %r take %d',
               filename, FLAGS.minibatch_size, epochs, shuffle, n_take)

  def dset_fn():
    dset = dataset.tuple_dataset(context, full_filename, data_specs)
    if shuffle:
      dset = dset.shuffle(1000, seed=SEED)
    if n_take > 0:
      dset = dset.take(n_take)
    dset = dset.batch(FLAGS.minibatch_size)
    dset = dset.repeat(epochs)
    return dset.map(lambda _, q, s, ans: ({'seeds': s}, ans))

  return dset_fn
Return a data_fn usable for an estimator.

  Args:
    context: a NeuralQueryContext
    filename: name of a file in FLAGS.rootdir
    shuffle: if True, shuffle the data
    epochs: number of times to repeat the data
    n_take: size of subset of examples for training/testing

  Returns:
    a function f so that f() returns a tf.data.Dataset
  
language/nql/demos/nell995/nell995.py