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' % (FLAGS.rootdir, filename)

  def dset_fn():
    """Construct a tf dataset.

    Returns:
      a tf dataset
    """

    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)

    def feature_dict_mapper(q, s):
      feature_dict = {
          question_name: tf.strings.regex_replace(q, r'\[([^]]+)\]', ''),
          'seeds': s,
      }
      return feature_dict

    question_name = get_text_module_input_name()
    return dset.map(lambda _, q, s, ans: (feature_dict_mapper(q, 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/metaqa/metaqa.py