Removing label column in a text classification dataset yields to errors
State: closed | Author: sgugger (CONTRIBUTOR)
Labels: name: bug

## Describe the bug

This looks like #3059 but it's not linked to the cache this time. Removing the `label` column from a text classification dataset and then performing any processing will result in an error.

To reproduce:
```py
from datasets import load_dataset
from transformers import AutoTokenizer

raw_datasets = load_dataset("imdb")
raw_datasets = raw_datasets.remove_columns("label")

model_checkpoint = "distilbert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
context_length = 128

def tokenize_pad_and_truncate(texts):
    return tokenizer(texts["text"], truncation=True, padding="max_length", max_length=context_length)

tokenized_datasets = raw_datasets.map(tokenize_pad_and_truncate, batched=True)
```


Traceback:

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-1-ba61bb32f786> in <module>
     12     return tokenizer(texts["text"], truncation=True, padding="max_length", max_length=context_length)
     13 
---> 14 tokenized_datasets = raw_datasets.map(tokenize_pad_and_truncate, batched=True)

~/git/datasets/src/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)
    500                     desc=desc,
    501                 )
--> 502                 for k, dataset in self.items()
    503             }
    504         )

~/git/datasets/src/datasets/dataset_dict.py in <dictcomp>(.0)
    500                     desc=desc,
    501                 )
--> 502                 for k, dataset in self.items()
    503             }
    504         )

~/git/datasets/src/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   2051                 new_fingerprint=new_fingerprint,
   2052                 disable_tqdm=disable_tqdm,
-> 2053                 desc=desc,
   2054             )
   2055         else:

~/git/datasets/src/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    501             self: "Dataset" = kwargs.pop("self")
    502         # apply actual function
--> 503         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    504         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    505         for dataset in datasets:

~/git/datasets/src/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    468         }
    469         # apply actual function
--> 470         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    471         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    472         # re-apply format to the output

~/git/datasets/src/datasets/fingerprint.py in wrapper(*args, **kwargs)
    404             # Call actual function
    405 
--> 406             out = func(self, *args, **kwargs)
    407 
    408             # Update fingerprint of in-place transforms + update in-place history of transforms

~/git/datasets/src/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, disable_tqdm, desc, cache_only)
   2243             if os.path.exists(cache_file_name) and load_from_cache_file:
   2244                 logger.warning("Loading cached processed dataset at %s", cache_file_name)
-> 2245                 info = self.info.copy()
   2246                 info.features = features
   2247                 info.task_templates = None

~/git/datasets/src/datasets/info.py in copy(self)
    278 
    279     def copy(self) -> "DatasetInfo":
--> 280         return self.__class__(**{k: copy.deepcopy(v) for k, v in self.__dict__.items()})
    281 
    282 

~/git/datasets/src/datasets/info.py in __init__(self, description, citation, homepage, license, features, post_processed, supervised_keys, task_templates, builder_name, config_name, version, splits, download_checksums, download_size, post_processing_size, dataset_size, size_in_bytes)

~/git/datasets/src/datasets/info.py in __post_init__(self)
    177                 for idx, template in enumerate(self.task_templates):
    178                     if isinstance(template, TextClassification):
--> 179                         labels = self.features[template.label_column].names
    180                         self.task_templates[idx] = TextClassification(
    181                             text_column=template.text_column, label_column=template.label_column, labels=labels

KeyError: 'label'
```