import tensorflow as tf

    if single_file:
      dataset_path = os.path.join(dataset_path, 'train_annotated.json')
  else:
    dataset_path = os.path.join(dataset_path, 'dev_annotated.json')

  def load_dataset():
    dataset = []
    if single_file:
      # Opening with GFile allows to use remotely stored files, e.g.
      # in a gs bucket.
      dataset_handle = tf.io.gfile.GFile(dataset_path, 'r')
      for line in dataset_handle:
        dataset.append(json.loads(line))
    else:
      all_files = tf.io.gfile.listdir(dataset_path)
      for filename in all_files:
        if 'json' in filename:
          print('Loading data from file {}'.format(filename))
          with tf.io.gfile.GFile(os.path.join(dataset_path, filename)) as f:
            for line in f:
              dataset.append(json.loads(line))
    print('The total size of the dataset {}'.format(len(dataset)))
    return dataset[:int(len(dataset) * percentile)]

  def drop_annotated_yield_examples(generator=None):
    del generator
    while True:
      passages = set()
