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
Create tf.dataset for estimator.

      Returns:
        tf dataset
      
language/emql/data_loader.py