❓ On-the-fly tokenization with datasets, tokenizers, and torch Datasets and Dataloaders

Hi there,

I have a question regarding "on-the-fly" tokenization. This question was elicited by reading the "How to train a new language model from scratch using Transformers and Tokenizers" [here](https://huggingface.co/blog/how-to-train). Towards the end there is this sentence: "If your dataset is very large, you can opt to load and tokenize examples on the fly, rather than as a preprocessing step". I've tried coming up with a solution that would combine both `datasets` and `tokenizers`, but did not manage to find a good pattern.

I guess the solution would entail wrapping a dataset into a Pytorch dataset.

As a concrete example from the [docs](https://huggingface.co/transformers/custom_datasets.html)

```python
import torch

class SquadDataset(torch.utils.data.Dataset):
    def __init__(self, encodings):
        # instead of doing this beforehand, I'd like to do tokenization on the fly
        self.encodings = encodings 

    def __getitem__(self, idx):
        return {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}

    def __len__(self):
        return len(self.encodings.input_ids)

train_dataset = SquadDataset(train_encodings)
```

How would one implement this with "on-the-fly" tokenization exploiting the vectorized capabilities of tokenizers?


----

Edit: I have come up with this solution. It does what I want, but I feel it's not very elegant

```python
class CustomPytorchDataset(Dataset):
    def __init__(self):
        self.dataset = some_hf_dataset(...)
        self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

    def __getitem__(self, batch_idx):
        instance = self.dataset[text_col][batch_idx]
        tokenized_text = self.tokenizer(instance, truncation=True, padding=True)
        return tokenized_text

    def __len__(self):
        return len(self.dataset)

    @staticmethod
    def collate_fn(batch):
        # batch is a list, however it will always contain 1 item because we should not use the
        # batch_size argument as batch_size is controlled by the sampler
        return {k: torch.tensor(v) for k, v in batch[0].items()}

torch_ds = CustomPytorchDataset()

# NOTE: batch_sampler returns list of integers and since here we have SequentialSampler
# it returns: [1, 2, 3], [4, 5, 6], etc. - check calling `list(batch_sampler)`
batch_sampler = BatchSampler(SequentialSampler(torch_ds), batch_size=3, drop_last=True)

# NOTE: no `batch_size` as now the it is controlled by the sampler!
dl = DataLoader(dataset=torch_ds, sampler=batch_sampler, collate_fn=torch_ds.collate_fn)
```

Comments:
We're working on adding on-the-fly transforms in datasets.
Currently the only on-the-fly functions that can be applied are in `set_format` in which we transform the data in either numpy/torch/tf tensors or pandas.
For example
```python
dataset.set_format("torch")
```
applies `torch.Tensor` to the dataset entries on-the-fly.

We plan to extend this to user-defined formatting transforms.
For example
```python
dataset.set_format(transform=tokenize)
```

What do you think ?

You can now use `set_transform` to define custom formatting transforms.