def build_metaqa_input_fn(
      self,
      batch_size,
      mode,
      epochs = 1,
      n_take = -1,
      shuffle = False):
    """Return tensors in required shape.

    Args:
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: epochs
      n_take: number of examples to pick in each epoch
      shuffle: if shuffle dataset

    Returns:
      tf.Dataset
    """
    if mode == 'train':
      data = self.train_data_metaqa
    elif mode == 'eval':
      data = self.test_data_metaqa
    else:
      raise ValueError('mode not recognized: %s' % mode)

    def metaqa_generator():
      """Generate metaqa question answer pair.

      Yields:
        a tuple
      """
      for one_data in data:
        question, question_entity_id, answer_fact_ids = one_data
        question = np.array(question + [-1] *
                            (self.max_question_len - len(question)))
        question_entity_sketch = \
            self.cm_context.get_sketch(xs=[question_entity_id])
        answers = np.zeros(self.num_facts, dtype=bool)
        answers[answer_fact_ids] = True

        yield (question, question_entity_id, question_entity_sketch, answers), 0

    def input_fn():
      """Create tf.dataset for estimator.

      Returns:
        tf dataset
      """
      tensor_shape = self.num_facts
      ds = tf.data.Dataset.from_generator(
          generator=metaqa_generator,
          output_types=((tf.int32, tf.int32, tf.float32, tf.float32), tf.int32),
          output_shapes=((tf.TensorShape([self.max_question_len]),
                          tf.TensorShape([]),
                          tf.TensorShape(
                              [self.cm_context.depth, self.cm_context.width]),
                          tf.TensorShape([tensor_shape])), tf.TensorShape([])))

      if n_take > 0:
        ds = ds.take(n_take)
      if shuffle:
        ds = ds.shuffle(SEED)
      ds = ds.repeat(epochs)
      ds = ds.batch(batch_size)
      return ds

    return input_fn
Return tensors in required shape.

    Args:
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: epochs
      n_take: number of examples to pick in each epoch
      shuffle: if shuffle dataset

    Returns:
      tf.Dataset
    
language/emql/data_loader.py