GitHub issue #2392: Update text classification template labels in DatasetInfo __post_init__
State: closed | Created: 2021-05-21T15:29:41Z

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