def build_input_fn(self,
                     name,
                     batch_size,
                     mode,
                     epochs = 1,
                     n_take = -1,
                     shuffle = False):
    """Construct tf.Dataset object as input.

    Args:
      name: name of training task
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: number of epochs to run
      n_take: number of random samples selected
      shuffle: if shuffle the dataset or not (train vs. test)

    Returns:
      input_fn
    """
    if name == 'membership':
      return self.build_membership_input_fn(batch_size, mode, epochs, n_take,
                                            shuffle)
    elif name == 'intersection' or name == 'union':
      return self.build_intersection_input_fn(name, batch_size, mode, epochs,
                                              n_take, shuffle)
    elif name == 'follow' or name == 'set_follow':
      return self.build_follow_input_fn(batch_size, mode, epochs, n_take,
                                        shuffle)
    elif name in ['metaqa2', 'metaqa3']:
      return self.build_metaqa_input_fn(batch_size, mode, epochs, n_take,
                                        shuffle)
    elif name == 'mixture':
      return self.build_mixture_input_fn(batch_size, mode, epochs, n_take)
    elif name == 'webqsp':
      return self.build_webqsp_input_fn(batch_size, mode, epochs, n_take,
                                        shuffle)
    elif name.startswith('query2box'):
      return self.build_query2box_input_fn(
          name, batch_size, mode)
    else:
      raise ValueError('name not recognized: %s' % name)
Construct tf.Dataset object as input.

    Args:
      name: name of training task
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: number of epochs to run
      n_take: number of random samples selected
      shuffle: if shuffle the dataset or not (train vs. test)

    Returns:
      input_fn
    
language/emql/data_loader.py