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))
Construct a tf dataset.

    Returns:
      a tf dataset
    
language/nql/demos/metaqa/metaqa.py