How to convert datasets.arrow_dataset.Dataset to torch.utils.data.Dataset
Author: y-rokutan
State: closed

Hi,

I'm trying to pretraine deep-speed model using HF arxiv dataset like:
```
train_ds = nlp.load_dataset('scientific_papers', 'arxiv')
train_ds.set_format(
        type="torch",
        columns=["input_ids", "attention_mask", "global_attention_mask", "labels"],
    )
engine, _, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=[p for p in model.parameters() if p.requires_grad],
    training_data=train_ds)
```
but deepspeed.initialize accepts torch.utils.data.Dataset only. How can I convert HF-style dataset to torch-style dataset?


Comments: Hi,

a HF dataset can be converted to a Torch Dataset with a simple wrapper as follows:
```python
from torch.utils.data import Dataset
 
class HFDataset(Dataset):
    def __init__(self, dset):
        self.dset = dset

    def __getitem__(self, idx):
        return self.dset[idx]

    def __len__(self):
        return len(self.dset)

train_ds = HFDataset(train_ds)
```
@lhoestq Since the Arrow Dataset already provides `__getitem__` and `__len__`, I think we could use the [virtual subclass](https://docs.python.org/3/library/abc.html#abc.ABCMeta.register) mechanism from the `abc` module to elegantly solve this issue. This mechanism would allow the Arrow Dataset to be used in place of the Torch Dataset because the `isinstance(instance of Arrow Dataset, TorchDataset)` check would return True (DeepSpeed has this check [here](https://github.com/microsoft/DeepSpeed/blob/ab5534fc4c0f8ca21ada321f9730d723aa31288b/deepspeed/runtime/engine.py#L823)).

And it requires a minimal change in the `arrow_dataset.py` file:
```python
if config.TORCH_AVAILABLE:
    from torch.utils.data import Dataset as TorchDataset
    TorchDataset.register(Dataset)
```

Interesting ! Thanks for sharing this @mariosasko . I like the idea
This looks like something we should add IMO

@mariosasko 
Thx for your code!
It perfectly works with a small modification for HF NLP dataset:
```
original_ds = nlp.load_dataset('scientific_papers', 'arxiv')
train_ds = HFDataset(train_ds['train']) # needs splitting
```

@lhoestq Sadly, from Python 3.7 onwards `torch.utils.data.Dataset` doesn't support the virtual subclass mechanism due to `typing.Generic` type no longer having `abc.ABCMeta` as its metaclass.

With that in mind, another option is to remove a direct type check (`isinstance(dataset, torch.utils.data.Dataset)`) in `deepspeed.initalize` and to rewrite the checks in a manner similar to `torch.utils.data.DataLoader` ([link](https://github.com/pytorch/pytorch/blob/b80c6f863f2327c712c478f67c248b94d66b65ac/torch/utils/data/dataloader.py#L197-L239)). This is exactly why the `DataLoader` works with arbitrary objects that provide `__getitem__` and `__len__` (and in our case, the `ArrowDataset`). By doing so, their code wouldn't be any stricter in comparison to the `DataLoader`.

So if you agree, I can open an issue in their repo and fix this if they like the idea.

That makes sense ! Feel free to open an issue on their repo and discuss this idea

@y-rokutan Hi, now if you install `deepspeed` from master (this feature will be available in the next official release), the code should work without subclassing. Let us know if you still have any issues.

Worth mentioning that any function that expects a `torch..Dataset` (like `torch..DataLoader`) will fail a mypy-esque typecheck if a `datasets.Dataset` is passed, even though it implements the interface correctly (I think). The virtual subclass idea was a good one- I wonder if there's another workaround given the Generic issue. What we're really talking about is something similar to the structural subtyping semantics that `typing.Protocol` defines. If `torch..DataLoader` accepted anything that supports `__getitem__` and `__len__` methods this would be much easier. Not sure if there's a way to do this without the wrapper from the perspective of `datasets`.