Streaming IterableDataset not working with translation pipeline
Opened by jlquinn (open)

### Describe the bug

I'm trying to use a streaming dataset for translation inference to avoid downloading the training data.

I'm using a pipeline and a dataset, and following the guidance in the tutorial.  

Instead I get an exception that IterableDataset has no len().

### Steps to reproduce the bug

CODE:
```
from transformers import pipeline
from transformers.pipelines.pt_utils import KeyDataset
from datasets import load_dataset
ds = load_dataset(path="wmt14", name="fr-en", split="test", streaming=True)
bs=1
mt = pipeline("translation_en_to_fr", model="t5-base", batch_size=bs)
#print(mt("hello"))    THIS WORKS                                                                                                                                                                       
ks = KeyDataset(ds, "translation")
print(f"{ks}")
xx= mt(ks)
for x in xx:
    print(x)
```

RUN:
```
(watnlp) [jlquinn@bertdev01 hf]$ python ende.t5.pipe.py 
2023-05-11 16:48:08.817572: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-05-11 16:48:08.821388: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-05-11 16:48:08.821407: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
<transformers.pipelines.pt_utils.KeyDataset object at 0x7f61ed5da9d0>
Traceback (most recent call last):
  File "/home/jlquinn/models/hf/ende.t5.pipe.py", line 11, in <module>
    for x in xx:
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 111, in __next__
    item = next(self.iterator)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 111, in __next__
    item = next(self.iterator)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 681, in __next__
    data = self._next_data()
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 720, in _next_data
    index = self._next_index()  # may raise StopIteration
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 671, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/torch/utils/data/sampler.py", line 247, in __iter__
    for idx in self.sampler:
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/torch/utils/data/sampler.py", line 76, in __iter__
    return iter(range(len(self.data_source)))
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 13, in __len__
    return len(self.dataset)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 289, in __len__
    return len(self.dataset)
TypeError: object of type 'IterableDataset' has no len()
```

### Expected behavior

I'm expecting french translations of the english test set to be printed.


### Environment info

Run on CPU with no GPU.

RHEL 8.7 x86_64
python 3.9.0
transformers 4.17.0
datasets 2.0.0
tokenizers 0.12.1

```
(watnlp) [jlquinn@bertdev01 hf]$ datasets-cli env

Copy-and-paste the text below in your GitHub issue.

- `datasets` version: 2.0.0
- Platform: Linux-4.18.0-372.19.1.el8_6.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.0
- PyArrow version: 8.0.0
- Pandas version: 1.4.4
```


Comments:
I wasn't sure to file this against transformers or datasets.

[`KeyDataset`](https://github.com/huggingface/transformers/blob/7f8b909189547944617741d8d3c6c84504701693/src/transformers/pipelines/pt_utils.py#L296) doesn't support iterable datasets, so you either need to implement a version that does (and also indexing nested (translation) fields):

```python
from torch.utils.data import Dataset, IterableDataset

def build_key_fetcher(key: str):
    def _key_fetcher(item):
        for sub_key in key.split("."):
            item = item[sub_key]
        return item
    return _key_fetcher

class KeyDataset(Dataset):
    def __new__(cls, dataset: Dataset, key: str):
        cls = _KeyIterableDataset if isinstance(dataset, IterableDataset) else _KeyMapDataset
        self = object.__new__(cls)
        self.dataset = dataset
        self.key = key
        self._key_fetcher = build_key_fetcher(key)
        return self

class _KeyMapDataset(KeyDataset):
    def __getitem__(self, i):
        return self._key_fetcher(self.dataset[i])
    
    def __len__(self):
        return len(self.dataset)


class _KeyIterableDataset(KeyDataset):
    def __iter__(self):
        for ex in self.dataset:
            yield self._key_fetcher(ex)

ks = KeyDataset(ds, "translation.en")
```

or use `IterableDataset`'s `map`:
```python
def fetch_en_translation(ex):
    return {"en": ex["translation"]["en"]}
ks = ds.map(fetch_en_translation, remove_columns=ds.column_names) 
```

cc @sgugger: Perhaps the `KeyDataset` + PyTorch `IterableDataset` case should be supported by Transformers

@mariosasko The map snippet didn't quite work, but gave me enough of a clue to get it working.  The following snippet does work:
```
def en_translation(x):
    return {"en":x['translation']['en']}
ks = ds.map(en_translation, remove_columns=['translation'])
test=[]
for x in iter(ks):
    test.append(x['en'])
xx= mt(test)
for x in xx:
    print(x)
```

I tried just returning `x['translation']['en`]` in the helper function instead of the dict, but that didn't give me an iterator over strings that pipeline would work with either.


The snippet as is gives the following error:
```
Traceback (most recent call last):
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/pdb.py", line 1704, in main
    pdb._runscript(mainpyfile)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/pdb.py", line 1573, in _runscript
    self.run(statement)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/bdb.py", line 580, in run
    exec(cmd, globals, locals)
  File "<string>", line 1, in <module>
  File "/home/jlquinn/models/hf/ende.t5.pipe.py", line 1, in <module>
    from transformers import pipeline
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/text2text_generation.py", line 335, in __call__
    return super().__call__(*args, **kwargs)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/text2text_generation.py", line 138, in __call__
    result = super().__call__(*args, **kwargs)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/base.py", line 1027, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/base.py", line 1033, in run_single
    model_inputs = self.preprocess(inputs, **preprocess_params)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/text2text_generation.py", line 287, in preprocess
    return super()._parse_and_tokenize(*args, truncation=truncation)
  File "/home/jlquinn/miniconda3/envs/watnlp/lib/python3.9/site-packages/transformers/pipelines/text2text_generation.py", line 100, in _parse_and_tokenize
    raise ValueError(
ValueError:  `args[0]`: <datasets.iterable_dataset.IterableDataset object at 0x7f5fd38ef1c0> have the wrong format. The should be either of type `str` or type `list`
Uncaught exception. Entering post mortem debugging
Running 'cont' or 'step' will restart the program
```


So perhaps there's no bug exactly, but I would love to see two things:  1) improve the documentation to better understand what's really getting returned.  2) update the example provided of using transformer pipeline with a dataset to include the oddball case that translation appears to be.

cc @Narsil 

Hi,

for the original snippet, the issue is that `streaming` datasets are not countable (they have no len) and therefore `KeyDataset` cannot work with them ( KeyDataset is a dataset and therefore requires a length).

I modified slightly the original snippet to make it work:

```python
from transformers import pipeline
from transformers.pipelines.pt_utils import KeyDataset
from datasets import load_dataset

ds = load_dataset(path="wmt14", name="fr-en", split="test", streaming=True)
bs = 1
mt = pipeline(
    "translation_en_to_fr", model="hf-internal-testing/tiny-random-T5ForConditionalGeneration", batch_size=bs
)


def ks(ds):
    for item in ds:
        yield item["translation"]["en"]


# print(f"{ks}")
xx = mt(ks(ds))
for x in xx:
    print(x)
```

This is what the first example in the docs suggests to use (as it's the most flexible): https://huggingface.co/docs/transformers/v4.29.1/en/pipeline_tutorial#using-pipelines-on-a-dataset

`KeyDataset` really exists only to get a `sized` dataset to work nicer with `tqdm` for instance.

@sgugger should we update the docs to remove `KeyDataset` entirely ? (We can add a note to pass manually the length of the data to tqdm so that the progress bar option can still be easy to use ?)


Maybe moving `KeyDataset` later on in the guide and specify it's mostly for streaming then? Or is it also necessary for batch_size>1 (which is what the current doc implies)?

Hmm

Iterator (`yield`) :
- Not countable
- Super flexible
- Cannot use `num_workers>1` (threading requires indexing at the correct location, iterators require to iterate in order,so each thread would iterate over the full thing being genuinely a bad idea)
- Can batch
- tqdm doesn't show a nice progress bar (it has no total)

KeyDataset (Or any PyTorch like Dataset returning the correct object for the pipeline):
- Countable
- Less flexible (not applicable to datasets with streaming), can only work on single keys. But should be easy to read and write your own (like @mariosasko did)
- Works with  `num_workers > 1` (Every worker can fetch exactly what's needed)
- Can batch 
- tqdm shows a nice progress bar

In the docs, if we update all the examples to use iterators, and include an example with

```
for item in tqdm.tqdm(pipe(iterator(), total=len(dataset))))
```

We can save the biggest feature that doesn't work out of the box with iterators which is the tqdm progress bar.

`num_workers>1` we can mention it, but it tends to be an issues only on CPU intensive loads, like image (and maybe audio)