huggingface / huggingface/datasets

Cant pass streaming dataset to dataloader after take()

Open
#4,678 1 comment 1 reaction 0 assignees View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.