def simple_tf_dataset(context,
                      tuple_input,
                      x_type,
                      y_type,
                      normalize_outputs=False,
                      batch_size=1,
                      shuffle_buffer_size=1000,
                      feature_key=None,
                      field_separator="\t"):
  """A dataset with just two columns, x and y.

  Args:
    context: a NeuralQueryContext
    tuple_input: passed to util.tuple_dataset
    x_type: type of entities x
    y_type: type of entities y1,...,yk
    normalize_outputs: make the encoding of {y1,...,yk} sum to 1
    batch_size: size of minibatches
    shuffle_buffer_size: if zero, do not shuffle the dataset. Otherwise, this is
      passed in as argument to shuffle
    feature_key: if not None, wrap the x part of the minibatch in a dictionary
      with the given key
    field_separator: passed in to dataset.tuple_dataset

  Returns:
    a tf.data.Dataset formed by wrapping the generator
  """
  dset = dataset.tuple_dataset(
      context,
      tuple_input, [x_type, y_type],
      normalize_outputs=normalize_outputs,
      field_separator=field_separator)
  if shuffle_buffer_size > 0:
    dset = dset.shuffle(shuffle_buffer_size)
  dset = dset.batch(batch_size)
  if feature_key is None:
    return dset
  else:
    wrap_x_in_dict = lambda x, y: ({feature_key: x}, y)
    return dset.map(wrap_x_in_dict)
A dataset with just two columns, x and y.

  Args:
    context: a NeuralQueryContext
    tuple_input: passed to util.tuple_dataset
    x_type: type of entities x
    y_type: type of entities y1,...,yk
    normalize_outputs: make the encoding of {y1,...,yk} sum to 1
    batch_size: size of minibatches
    shuffle_buffer_size: if zero, do not shuffle the dataset. Otherwise, this is
      passed in as argument to shuffle
    feature_key: if not None, wrap the x part of the minibatch in a dictionary
      with the given key
    field_separator: passed in to dataset.tuple_dataset

  Returns:
    a tf.data.Dataset formed by wrapping the generator
  
language/nql/nql/util_test.py