code: def _read_data(prefix, batch_size, augmentation=None):
      
      prefix = prefix.replace("\\", "/")
      folder = "/".join(prefix.split("/")[:-1])
      phase = prefix.split("/")[-1]
      config = json.load(open(prefix + '_config.json'))
      num_threads = config["num_threads"]
  
      filenames = [folder + "/" + f for f in listdir(folder) if isfile(join(folder, f)) and phase in f and not "config.json" in f]
  
      dataset = tf.data.TFRecordDataset(filenames=filenames, num_parallel_reads=num_threads)
      dataset = dataset.shuffle(buffer_size=10 * batch_size)
      dataset = dataset.repeat()
      dataset = dataset.map(map_func=_create_parser_fn(config, phase), num_parallel_calls=num_threads)
      if augmentation is not None:
          dataset = dataset.map(map_func=augmentation, num_parallel_calls=num_threads)
      dataset = dataset.batch(batch_size=batch_size)
      dataset = dataset.prefetch(buffer_size=1)
  
      return dataset
docs: Loads a dataset.
  
      :param prefix: The path prefix as defined in the write data method.
      :param batch_size: The batch size you want for the tensors.
      :param augmentation: An augmentation function.
      :return: A tensorflow.data.dataset object.