Data Map Trainer Callback

### Feature request It would be nice to have a callback for the trainer class which could create Data Maps. See the paper for more details https://arxiv.org/pdf/2009.10795. A Data Map measures how a model's prediction of specific training data change over the course of model training. The Callback should support: - Executing at each step or epoch - Should integrate directly with the Trainer class. - Should save the prediction of each training example as a matrix of the form [n_examples, n_labels] so that it can easily be stacked into [n_epochs, n_examples, n_labels]. Right now I'm saving things as a List[List[float]] but this might be sub optimal. It needs some way of getting the logged information later. Running [this colab notebook I made](https://colab.research.google.com/drive/1OyaQkFULDcCUVGEneyqFYK1_M1nunbUw?usp=sharing) will generate data map outputs for classification tasks using the Trainer in line with what I was thinking. Here is what I have so far that works will for multilabel and multiclass classification using transformers. ```python class DataMapCallback(TrainerCallback): """Trainer Callback to save DataMap data. Original Paper: https://arxiv.org/pdf/2009.10795.pdf. This callback saves the predictions of the model on each training example at the end of every epoch to callback_dir/{epoch}.json. """ def __init__( self, log_on: str = "epoch", callback_dir: str = ".", n_log_steps: Optional[int] = None, prediction_fn: Optional[Callable[[PreTrainedModel, DataLoader, TrainingArguments], List[List[float]]]] = None, ): self.callback_dir = callback_dir self.log_on = log_on self.log_count = 0 self.n_log_steps = n_log_steps self.prediction_fn = self._predict if prediction_fn is None else prediction_fn # Handle discrepencies in how we initialize the logging mode. if n_log_steps is not None and self.log_on != "step": raise Va…