huggingface / huggingface/datasets
Cant pass streaming dataset to dataloader after take()
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
## Describe the bug
I am trying to pass a streaming version of c4 to a dataloader, but it can't be passed after I call `dataset.take(n)`. Some functions such as `shuffle()` can be applied without breaking the dataloader but not take.
## Steps to reproduce the bug
```python
import datasets
import torch
dset = datasets.load_dataset(path='c4', name='en', split="train", streaming=True)
dset = dset.take(50_000)
dset = dset.with_format("torch")
num_workers = 8
batch_size = 512
loader = torch.utils.data.DataLoader(dataset=dset,
batch_size=batch_size,
num_workers=num_workers)
for batch in loader:
...
```
## Expected results
No error thrown when iterating over the dataloader
## Actual results
Original Traceback (most recent call last):
File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
data = fetcher.fetch(index)
File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/fetch.py", line 32, in fetch
data.append(next(self.dataset_iter))
File "/root/.local/lib/python3.9/site-packages/datasets/formatting/dataset_wrappers/torch_iterable_dataset.py", line 48, in __iter__
for key, example in self._iter_shard(shard_idx):
File "/root/.local/lib/python3.9/site-packages/datasets/iterable_dataset.py", line 586, in _iter_shard
yield from ex_iterable.shard_data_sources(shard_idx)
File "/root/.local/lib/python3.9/site-packages/datasets/iterable_dataset.py", line 60, in shard_data_sources
raise NotImplementedError(f"{type(self)} doesn't implement shard_data_sources yet")
NotImplementedError: doesn't implement shard_data_sources yet
## Environment info
- `datasets` version: 2.3.2
- Platform: Linux-5.4.0-120-generic-x86_64-with-glibc2.31
- Python version: 3.9.13
- PyArrow version: 8.0.0
- Pandas version: 1.4.3
Contributor guide
Assessment
This issue has not been assessed yet.