Make `with_format` correctly indicate that a `Dataset` is compatible with PyTorch's `Dataset` class

[benglewis]

Feature request

Currently datasets does not correctly indicate to the Python type-checker (e.g. pyright / Pylance) that the output of with_format is compatible with PyTorch's Dataloader since it does not indicate that the HuggingFace Dataset is compatible with the PyTorch Dataset class. It would be great if we could get the typing to work nicely.

Motivation

To avoid casting types in our Python code.

Your contribution

I would be happy to contribute a PR if this is something that may be accepted and could work with the current approach.
This doesn't have to be for just PyTorch, but I imagine that the same thing would be useful for tensorflow and such, but we only have a need for PyTorch at this stage.

[lhoestq]

Does the torch DataLoader really require the dataset to be a subclass of torch.utils.data.Dataset ? Or is there a simpler type we could use ?

PS: also note that a dataset without with_format() can also be used in a torch DataLoader . Calling with_format("torch") simply makes the output of the dataset torch Tensors in an efficient way.