def build_webqsp_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_webqsp
    elif mode == 'eval':
      data = self.test_data_webqsp
    else:
      raise ValueError('mode not recognized: %s' % mode)

    def webqsp_generator():
      """Generate set intersections.

      Yields:
        a tuple of features and label
      """
      for (question_token_ids, segment_ids, question_mask, question_entity_id,
           constraint_entity_id, answer_ids) in data:
        # question
        question_entity_sketch = np.zeros((self.cm_depth, self.cm_width),
                                          dtype=np.float32)
        for i, j in enumerate(self.all_entity_sketches[question_entity_id]):
          question_entity_sketch[i, j] = 1.0
        # constraint
        constraint_entity_sketch = np.zeros((self.cm_depth, self.cm_width),
                                            dtype=np.float32)
        if constraint_entity_id != -1:
          for i, j in enumerate(self.all_entity_sketches[constraint_entity_id]):
            constraint_entity_sketch[i, j] = 1.0
        yield (question_token_ids, segment_ids, question_mask,
               question_entity_id, question_entity_sketch, constraint_entity_id,
               constraint_entity_sketch, answer_ids), 0

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

      Returns:
        tf dataset
      """
      ds = tf.data.Dataset.from_generator(
          generator=webqsp_generator,
          output_types=((tf.int32, tf.int32, tf.int32, tf.int32, tf.float32,
                         tf.int32, tf.float32, tf.int32), tf.int32),
          output_shapes=((tf.TensorShape([self.bert_tokenizer.max_seq_length]),
                          tf.TensorShape([self.bert_tokenizer.max_seq_length]),
                          tf.TensorShape([self.bert_tokenizer.max_seq_length]),
                          tf.TensorShape([]),
                          tf.TensorShape([self.cm_depth,
                                          self.cm_width]), tf.TensorShape([]),
                          tf.TensorShape([self.cm_depth, self.cm_width]),
                          tf.TensorShape([self.max_webqsp_ans])),
                         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)
      ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
      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