Update text classification template labels in DatasetInfo __post_init__
State: closed | Author: lewtun (MEMBER)

This PR implements the idea discussed in #2389 to update the `labels` of the `TextClassification` template in the `DatasetInfo.__post_init__`. The main reason for doing so is so avoid duplicating the label definitions in both `DatasetInfo.features` and `DatasetInfo.task_templates`.

To avoid storing state in `DatasetInfo.__post_init__`, the current implementation flushes `DatasetInfo.task_templates` before the features are cast in `Dataset.prepare_for_task` (thanks to @mariosasko for this idea!).

Here is an example of the current workflow:

```python
ds1 = load_dataset("./datasets/emotion/")
# cast features and flush templates
ds2 = ds1.prepare_for_task("text-classification")
assert ds2.info.task_templates is None
```

Note that if users want to pass a `TextClassification` template to `prepare_for_task`, we require them to set `TextClassification.labels` to match the dataset's features corresponding to `label_column`:

```python
ds1 = load_dataset("./datasets/emotion/")
# TextClassification.labels is None by default => invalid template
task = TextClassification(text_column="text", label_column="label")
# Raises ValueError
ds1.prepare_for_task(task)
# Specifying the labels => valid template
task = TextClassification(text_column="text", label_column="label", labels=['anger', 'fear', 'joy', 'love', 'sadness', 'surprise'])
ds1.prepare_for_task(task)
```

This PR also adds:

* New tests + fixed some old tests that weren't testing `assertRaises` properly
* A decorator to share docstrings across common functions. This allows us to document `DatasetDict.prepare_for_task` and `Dataset.prepare_for_task` in one place.
* Fixes to avoid side-effects from in-place replacements of `DatasetInfo.task_templates` in `DatasetInfo.__post_init__`. Thanks to @lhoestq for figuring this out!
* Removal of `FeaturesWithLazyClassLabel` since we now create a new instance of `TextClassification` in `DatasetInfo.__post_init__` and avoid the side-effects first pointed out by @mariosasko 

### PR Description from original WIP 

Hi @yjernite and @lhoestq, here's a first stab at the suggestion discussed in #2389 to update the `labels` of the `TextClassification` template in the `DatasetInfo.__post_init__`.

One problem I've spotted is that my current implementation introduces state into the `__post_init__`: 

* When we call `load_dataset`, `DatasetInfo.features` are the "raw" features without any casting so we can access the column names by the `label_column` specified in `TextClassification`
* When we call `Dataset.prepare_for_task` we run into a problem because the `DatasetInfo.features` are first cast into the new schema which triggers a `KeyError` when we update the infos [here](https://github.com/huggingface/datasets/blob/8b2a78520828e0cc13c14a31f413a5395ef25110/src/datasets/arrow_dataset.py#L1959).

Here's an explicit example of what I mean with the stack trace appended below:

```python
from datasets import load_dataset

# this works 
ds = load_dataset("emotion")
# we can verify the task template is correctly set
ds["train"].info.task_templates # returns [TextClassification(labels=('sadness', 'joy', 'love', 'anger', 'fear', 'surprise'), text_column='text', label_column='label')]
# but this fails because the _post_init__ is looking for the original column names
ds.prepare_for_task("text-classification")
```
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-4-54a43019b319> in <module>
----> 1 ds.prepare_for_task("text-classification")

~/git/datasets/src/datasets/dataset_dict.py in prepare_for_task(self, task)
    807         """
    808         self._check_values_type()
--> 809         return DatasetDict({k: dataset.prepare_for_task(task=task) for k, dataset in self.items()})

~/git/datasets/src/datasets/dataset_dict.py in <dictcomp>(.0)
    807         """
    808         self._check_values_type()
--> 809         return DatasetDict({k: dataset.prepare_for_task(task=task) for k, dataset in self.items()})

~/git/datasets/src/datasets/arrow_dataset.py in prepare_for_task(self, task)
   1421         dataset = self.remove_columns(columns_to_drop)
   1422         dataset = dataset.rename_columns(column_mapping)
-> 1423         dataset = dataset.cast(features=template.features)
   1424         return dataset
   1425 

~/git/datasets/src/datasets/arrow_dataset.py in cast(self, features, batch_size, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, num_proc)
    970         format = self.format
    971         dataset = self.with_format("arrow")
--> 972         dataset = dataset.map(
    973             lambda t: t.cast(schema),
    974             batched=True,

~/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)
   1583 
   1584         if num_proc is None or num_proc == 1:
-> 1585             return self._map_single(
   1586                 function=function,
   1587                 with_indices=with_indices,

~/git/datasets/src/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    173         }
    174         # apply actual function
--> 175         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    176         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    177         # re-apply format to the output

~/git/datasets/src/datasets/fingerprint.py in wrapper(*args, **kwargs)
    338             # Call actual function
    339 
--> 340             out = func(self, *args, **kwargs)
    341 
    342             # 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)
   1959         if update_data:
   1960             # Create new Dataset from buffer or file
-> 1961             info = self.info.copy()
   1962             info.features = writer._features
   1963             if buf_writer is None:

~/git/datasets/src/datasets/info.py in copy(self)
    274 
    275     def copy(self) -> "DatasetInfo":
--> 276         return self.__class__(**{k: copy.deepcopy(v) for k, v in self.__dict__.items()})
    277 
    278 

~/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)
    174                     # The reason is that Dataset.prepare_for_task calls Dataset.cast which converts the
    175                     # DatasetInfo.features to the new schema and thus template.label_column is no longer a valid key
--> 176                     object.__setattr__(template, "labels", tuple(self.features[template.label_column].names))
    177                     template.label_schema["labels"] = ClassLabel(names=template.labels)
    178                     self.task_templates[idx] = template

KeyError: 'label'
```

What do you think? I did this a bit quickly, so maybe I'm overlooking something obvious :) One thing would be to only update the labels of the task template on load, but this seems a bit hacky IMO

Comments: If I'm not mistaken, one way to fix this would be to drop the task templates when copying the info by inserting `dataset.info.task_templates = None` before the `Dataset.cast` call in `Dataset.prepare_for_task`. Moreover, we should do this change independently of the KeyError being raised because currently the following is possible:
```python
dset = load_dataset("some_dataset") # let's say 'some_dataset' supports text classification and question answering
dset_tc = dset.prepare_for_task("text-classification")
dset_tc.preprare_for_task("question-answering") # this should raise an error because the schema is no longer valid for this task; currently this fails on 'rename_columns'
```
I see 2 options:
1. to drop the task templates after the first `Dataset.prepare_for_task` call
2. to save only the tasks compatible with the new schema after Dataset.prepare_for_task` (but then we have to update the column names of the compatible tasks to make sure the column mapping is still valid) 

> If I'm not mistaken, one way to fix this would be to drop the task templates when copying the info by inserting `dataset.info.task_templates = None` before the `Dataset.cast` call in `Dataset.prepare_for_task`. Moreover, we should do this change independently of the KeyError being raised because currently the following is possible:
> 
> ```python
> dset = load_dataset("some_dataset") # let's say 'some_dataset' supports text classification and question answering
> dset_tc = dset.prepare_for_task("text-classification")
> dset_tc.preprare_for_task("question-answering") # this should raise an error because the schema is no longer valid for this task; currently this fails on 'rename_columns'
> ```
> 
> I see 2 options:
> 
> 1. to drop the task templates after the first `Dataset.prepare_for_task` call
> 2. to save only the tasks compatible with the new schema after Dataset.prepare_for_task` (but then we have to update the column names of the compatible tasks to make sure the column mapping is still valid)

thanks for the great idea @mariosasko and for spotting the problem with sequential task preparation! i am in favour of your option (1) since it is simple and saves us from having to keep track of the column mappings across multiple steps. 

i've implemented the change and refactored the tests to account for the new approach (including a new test that the templates are flushed after we call `prepare_for_task`). perhaps the slightly inelegant aspect here is that if we want to allow the user to set `labels` in the `TextClassification` template, then we have two places (`DatasetInfo.__post_init__` and `TextClassification.__post_init__`) where we need to update `label_schema`. 

on the other hand, dropping `labels` from the `TextClassification` signature would have the nice effect that users only have to think about column names when defining their tasks.

in any case, i think it would be a good idea to merge #2376 soon as the current PR is touching a lot of the same places in the codebase 😄 


cc @SBrandeis who might also be interested in this feature :)

Tests are failing only because the `emotion` dataset card doesn't pass our dataset card validator (tags are missing), you can ignore this since it's unrelated to this PR.

@lhoestq @SBrandeis i've fixed the tests and think this is now in a good state for another review :)

Maybe @SBrandeis you can also take a look to make sure you're fine with it ?