def build_intersection_input_fn(
      self,
      name,
      batch_size,
      mode,
      epochs = 1,
      n_take = -1,
      shuffle = False):
    """Build input function for set intersections.

    Args:
      name: 'intersection' or 'union'
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: number of epochs
      n_take: size of each epoch
      shuffle: if shuffle or not

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

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

      Yields:
        a tuple
      """
      for set1, set2 in set_pairs:
        yield self.gen_one_pair_data(name, set1, set2), 0

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

      Returns:
        tf dataset
      """
      tensor_shape = self.num_entities
      ds = tf.data.Dataset.from_generator(
          generator=set_intersection_generator,
          output_types=((tf.float32, tf.float32, tf.float32), tf.int32),
          output_shapes=((tf.TensorShape([tensor_shape]),
                          tf.TensorShape([tensor_shape]),
                          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
Build input function for set intersections.

    Args:
      name: 'intersection' or 'union'
      batch_size: batch_size
      mode: 'train' or 'test'
      epochs: number of epochs
      n_take: size of each epoch
      shuffle: if shuffle or not

    Returns:
      input_fn
    
language/emql/data_loader.py