Feature request

Add a feature to pass the key of the DatasetDict when performing map

Motivation

I often preprocess using map on DatasetDict.
Sometimes, I need to preprocess train and valid data differently depending on the task.
So, I thought it would be nice to pass the key (like train, valid) when performing map on DatasetDict.

What do you think?

Your contribution

I can submit a pull request to add the feature to pass the key of the DatasetDict when performing map.