def build_mixture_input_fn(
      self,
      batch_size,
      mode,
      epochs = 1,
      n_take = -1):
    """Build mixture input_fn of membership, intersection, union, set_follow.

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

    Returns:
      input_fn
    """
    if mode == 'train':
      set_members = self.train_data_membership
      set_pairs = self.train_data_set_pair
      set_follows = self.train_data_follow
    elif mode == 'eval':
      set_members = self.test_data_membership
      set_pairs = self.test_data_set_pair
      set_follows = self.test_data_follow
    else:
      raise ValueError('mode not recognized: %s' % mode)

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

      Yields:
        a tuple with one data for each task
      """
      data_per_epoch = DATA_PER_EPOCH if n_take == -1 else n_take
      for _ in range(data_per_epoch * epochs):
        # Get model results for the task of membership.
        m_objs = set_members[np.random.randint(len(set_members))]
        m_padded_ent_ids, m_labels = self.gen_one_membership_data(m_objs)

        # Get model results for the task of intersection and union.
        p_set1, p_set2 = set_pairs[np.random.randint(len(set_pairs))]
        p_candidate_set1, p_candidate_set2, p_union_labels = \
            self.gen_one_pair_data('union', p_set1, p_set2)
        p_intersection_labels = np.minimum(p_candidate_set1, p_candidate_set2)

        # Get model results for the task of set_follows.
        f_subj_factids, f_rel_factids = set_follows[np.random.randint(
            len(set_follows))]
        f_candidate_set1, f_candidate_set2, f_labels = self.gen_one_follow_data(
            f_subj_factids, f_rel_factids)

        yield (m_padded_ent_ids, m_labels, p_candidate_set1, p_candidate_set2,
               p_union_labels, p_intersection_labels, f_candidate_set1,
               f_candidate_set2, f_labels), 0.0

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

      Returns:
        tf dataset
      """
      ds = tf.data.Dataset.from_generator(
          generator=mixture_generator,
          output_types=((tf.int32, tf.float32, tf.float32, tf.float32,
                         tf.float32, tf.float32, tf.float32, tf.float32,
                         tf.float32), tf.float32),
          output_shapes=((tf.TensorShape([self.max_set]),
                          tf.TensorShape([self.num_entities]),
                          tf.TensorShape([self.num_entities]),
                          tf.TensorShape([self.num_entities]),
                          tf.TensorShape([self.num_entities]),
                          tf.TensorShape([self.num_entities]),
                          tf.TensorShape([self.num_facts]),
                          tf.TensorShape([self.num_facts]),
                          tf.TensorShape([self.num_facts])),
                         tf.TensorShape([])))

      ds = ds.batch(batch_size)
      return ds

    return input_fn
Build mixture input_fn of membership, intersection, union, set_follow.

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

    Returns:
      input_fn
    
language/emql/data_loader.py