def build_membership_input_fn(
      self,
      batch_size,
      mode,
      epochs = 1,
      n_take = -1,
      shuffle = False):
    """Build set membership input_fn.

    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:
      input_fn
    """
    if mode == 'train':
      data = self.train_data_membership
    elif mode == 'eval':
      data = self.test_data_membership
    else:
      raise ValueError('mode not recognized: %s' % mode)

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

      Yields:
        a tuple
      """
      if shuffle:
        random_ids = np.random.permutation(len(data))
      else:
        random_ids = np.arange(len(data))

      for data_id in random_ids:
        objs = data[data_id]
        yield self.gen_one_membership_data(objs), 0

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

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

      if n_take > 0:
        ds = ds.take(n_take)
      ds = ds.repeat(epochs)
      ds = ds.batch(batch_size)
      return ds

    return input_fn
Build set membership input_fn.

    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:
      input_fn
    
language/emql/data_loader.py