def build_query2box_input_fn(
      self,
      name,
      batch_size,
      mode):
    """Build input function.

    Build input function and return the inputs in their required
    formats:
        1c: ent, rel
        2c: ent, rel1, rel2
        3c: ent, rel1, rel2, rel3
        2i or 2u: ent1, rel1, ent2, rel2
        3i: ent1, rel1, ent2, rel2, ent3, rel3
        ic or uc: ent1, rel1, ent2, rel2, rel3
        ci: ent1, rel1, rel2, ent2, rel3

    Args:
      name: name of the experiment.
      batch_size: batch size
      mode: 'train' or 'eval'

    Returns:
      tf.data.Dataset

    """

    task = name.split('_')[-1]

    if mode == 'eval':
      data = self.test_data_query2box
    else:
      raise ValueError('Only test mode allowed for query2box: mode = %s' % mode)

    def convert_entid(q2b_entid):
      # KB might be loaded in different ways such that entity ids
      # won't match. So this function converts q2b ids to our
      # entity ids.
      return self.entity2id[self.q2b_id2entity[q2b_entid]]

    def convert_relid(q2b_relid):
      return self.relation2id[self.q2b_id2relation[q2b_relid]]

    def query2box_generator():
      """Convert query2box features to EmQL features.

      We simply hard-code their input format for parsing, since there's not
      a simple way to automatically serialized to our input format.

      Yields:
        a tuple of features and label
      """
      for query in data:
        if task == '1c':
          assert len(query) == 1
          ent, rels = query[0]
          yield (convert_entid(ent), convert_relid(rels[0])), 0
        elif task == '2c':
          assert len(query) == 1
          ent, rels = query[0]
          yield (convert_entid(ent),
                 convert_relid(rels[0]),
                 convert_relid(rels[1])), 0
        elif task == '3c':
          assert len(query) == 1
          ent, rels = query[0]
          yield (convert_entid(ent),
                 convert_relid(rels[0]),
                 convert_relid(rels[1]),
                 convert_relid(rels[2])), 0
        elif task == '2i' or task == '2u':
          assert len(query) == 2
          ent1, rels1 = query[0]
          ent2, rels2 = query[1]
          yield (convert_entid(ent1), convert_relid(rels1[0]),
                 convert_entid(ent2), convert_relid(rels2[0])), 0
        elif task == '3i':
          assert len(query) == 3
          ent1, rels1 = query[0]
          ent2, rels2 = query[1]
          ent3, rels3 = query[2]
          yield (convert_entid(ent1), convert_relid(rels1[0]),
                 convert_entid(ent2), convert_relid(rels2[0]),
                 convert_entid(ent3), convert_relid(rels3[0])), 0
        elif task == 'ic' or task == 'uc':
          assert len(query) == 3
          ent1, rels1 = query[0]
          ent2, rels2 = query[1]
          rel3 = query[2]
          yield (convert_entid(ent1), convert_relid(rels1[0]),
                 convert_entid(ent2), convert_relid(rels2[0]),
                 convert_relid(rel3)), 0
        elif task == 'ci':
          assert len(query) == 2
          ent1, rels1 = query[0]
          ent2, rels2 = query[1]
          yield (convert_entid(ent1), convert_relid(rels1[0]),
                 convert_relid(rels1[1]),
                 convert_entid(ent2), convert_relid(rels2[0])), 0
        else:
          raise ValueError

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

      Returns:
        tf dataset
      """
      # output_types
      if task == '1c':
        output_types = tuple([tf.int32] * 2)
        output_shapes = tuple([tf.TensorShape([])] * 2)
      elif task == '2c':
        output_types = tuple([tf.int32] * 3)
        output_shapes = tuple([tf.TensorShape([])] * 3)
      elif task == '3c':
        output_types = tuple([tf.int32] * 4)
        output_shapes = tuple([tf.TensorShape([])] * 4)
      elif task == '2i' or task == '2u':
        output_types = tuple([tf.int32] * 4)
        output_shapes = tuple([tf.TensorShape([])] * 4)
      elif task == '3i':
        output_types = tuple([tf.int32] * 6)
        output_shapes = tuple([tf.TensorShape([])] * 6)
      elif task == 'ic' or task == 'uc':
        output_types = tuple([tf.int32] * 5)
        output_shapes = tuple([tf.TensorShape([])] * 5)
      elif task == 'ci':
        output_types = tuple([tf.int32] * 5)
        output_shapes = tuple([tf.TensorShape([])] * 5)
      else:
        raise ValueError

      ds = tf.data.Dataset.from_generator(
          generator=query2box_generator,
          output_types=(output_types, tf.int32),
          output_shapes=(output_shapes, tf.TensorShape([])))

      ds = ds.batch(batch_size)
      ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
      return ds

    return input_fn
Build input function.

    Build input function and return the inputs in their required
    formats:
        1c: ent, rel
        2c: ent, rel1, rel2
        3c: ent, rel1, rel2, rel3
        2i or 2u: ent1, rel1, ent2, rel2
        3i: ent1, rel1, ent2, rel2, ent3, rel3
        ic or uc: ent1, rel1, ent2, rel2, rel3
        ci: ent1, rel1, rel2, ent2, rel3

    Args:
      name: name of the experiment.
      batch_size: batch size
      mode: 'train' or 'eval'

    Returns:
      tf.data.Dataset

    
language/emql/data_loader.py