def _get_file_names(folder):
  """Returns the file names expected to exist in the input_dir."""
  assert folder in ['cifar-10', 'cifar-100']

  file_names = {}
  if folder == 'cifar-10':
    file_names['train'] = ['data_batch_%d' % i for i in range(1, 5)]
    file_names['validation'] = ['data_batch_5']
    file_names['train_all'] = ['data_batch_%d' % i for i in range(1, 6)]
    file_names['test'] = ['test_batch']
  else:
    file_names['train_all'] = ['train']
    file_names['test'] = ['test']
    # Split in `convert_to_tfrecord` function
    file_names['train'] = ['train']
    file_names['validation'] = ['train']
  return file_names