Batching behaviour of Pipelines with datasets.Dataset as input could be clarified/improved
Author: jack89roberts
Pull request: no

### System Info

- `transformers` version: 4.31.0
- Platform: macOS-13.4.1-arm64-arm-64bit
- Python version: 3.11.4
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@Narsil

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I wasn't sure whether to classify this as a bug or feature request so apologies if I've put it in the wrong place.

The wording of the documentation made me think batching should work if the input to a `Pipeline` is a `datasets.Dataset` instance, but this is not the case. It is the case for an `IterableDataset` instance.

Here's a script to show what I mean:

```python
from datasets import Dataset
from transformers import (
    Pipeline,
    ViTConfig,
    ViTForImageClassification,
    ViTImageProcessor,
)


# dummy custom pipeline that just passes through the input and prints what was passed to `_forward`
class MyPipeline(Pipeline):
    def _sanitize_parameters(self):
        return {}, {}, {}

    def preprocess(self, inputs):
        return inputs

    def _forward(self, model_inputs):
        print("_forward", model_inputs["idx"])
        return model_inputs

    def postprocess(self, model_outputs):
        return model_outputs


# make a small dummy model and processor to initialise the pipeline with
model = ViTForImageClassification(
    ViTConfig(
        hidden_size=1,
        num_hidden_layers=1,
        num_attention_heads=1,
        intermediate_size=1,
        image_size=1,
        patch_size=1,
        num_channels=1,
        encoder_stride=1,
    )
)
extractor = ViTImageProcessor()

pipeline = MyPipeline(model=model, image_processor=extractor)

# this will call _forward on the whole dataset in one go (despite setting a batch size)
print("Using datasets.Dataset as input:")
dataset = Dataset.from_dict({"idx": range(16)})
for i in pipeline(dataset, batch_size=4):
    continue
# output:
# _forward [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

# if dataset converted to IterableDataset the calls to _forward are batched as expected
print("Using datasets.IterableDataset as input:")
dataset = dataset.to_iterable_dataset()
for i in pipeline(dataset, batch_size=4):
    continue
# output:
# _forward [0, 1, 2, 3]
# _forward [4, 5, 6, 7]
# _forward [8, 9, 10, 11]
# _forward [12, 13, 14, 15]
```

### Expected behavior

I am happy to make a PR for any of these but unsure which would be most appropriate or if there are implementation details I'm missing.

#### Docs

I think the documentation could be clarified. I found this confusing:

https://github.com/huggingface/transformers/blob/2f8acfea1ca11fe3479fb379ccbded516d0cff57/docs/source/en/main_classes/pipelines.md?plain=1#L110-L113

Batching will only be used if the `Dataset` that is passed is one that has `isinstance(inputs, torch.utils.data.Dataset) == True`. This is not the case for `datasets.arrow_dataset.Dataset`, which generally seems to be the default meaning of `Dataset` within the `datasets` library.

Both `datasets.iterable_dataset.IterableDataset` and `transformers.pipelines.pt_utils.KeyDataset` do inherit from `torch.utils.data.Dataset`, so the pipeline will use batching if inputs is one of those.

A minimal change could be something like changing this:

> (so when passing lists or `Dataset` or `generator`)

to this:

> (so when passing lists or `torch.utils.data.Dataset` or `generator` or `datasets.IterableDataset` or `KeyDataset`).

Maybe the pipeline documentation should also explicitly say batching will not work with `datasets.Dataset` if this is the expected behaviour.

#### Pipeline base class

Adjustments that could be made to the base Pipeline class to clarify/change this:

##### 1)

Raise a warning/error if a `batch_size` has been set but the type of inputs means the inference will not actually be batched.

##### 2)

A one-line change to the base class should make them compatible with batching over `datasets.Dataset`. If this line:

https://github.com/huggingface/transformers/blob/e42587f596181396e1c4b63660abf0c736b10dae/src/transformers/pipelines/base.py#L1089

also included a check for `isinstance(inputs, datasets.Dataset)` (as well as a torch Dataset, which is what the above check is against) then the input dataset would later be converted to a `DataLoader` here: 

https://github.com/huggingface/transformers/blob/e42587f596181396e1c4b63660abf0c736b10dae/src/transformers/pipelines/base.py#L1054

This worked fine for my custom pipeline but I'm unsure if there are other consequences/reasons not to be inputting `datasets.Dataset` instances into pipelines.


Comments:
cc @Narsil

I'm pretty sure `datasets.Dataset` will not work out of the box, because `datasets.Dataset` uses different keys for different items (Same as `torch.Dataset` actually, but you're in control of those.)

Actually the easiest/most general way would be to use a generator for everything (you both have full control and works with any dataset) https://huggingface.co/docs/transformers/v4.32.0/en/main_classes/pipelines#pipeline-batching

The *only* reason for `KeyDataset` to exist, is because it can do the adaptation layer over `datasets.Dataset` and be sizeable (which generators aren't) and therefore it works better with tools like `tqdm`.

@ArthurZucker Should we just remove `KeyDataset` from the example and purely promote the generator one ?

Yes, I think it makes sense 👍🏻 

> I'm pretty sure `datasets.Dataset` will not work out of the box, because `datasets.Dataset` uses different keys for different items (Same as `torch.Dataset` actually, but you're in control of those.)

Thanks for having a look at this! I'm not familiar with the internals of the various dataset classes, but for reference here are a couple of examples where creating a data loader from a `datasets.Dataset` seems to work (this is what the pipeline class does for other input types currently, but not for `datasets.Dataset`):

```python
from datasets import Dataset, load_dataset
from torch.utils.data import DataLoader
from transformers.pipelines.base import pad_collate_fn

test_collate_fn = pad_collate_fn(None, "TEST")

dataset = Dataset.from_dict({"idx": range(16)})
dataloader = DataLoader(dataset, batch_size=4, collate_fn=test_collate_fn)
data_iter = iter(dataloader)
print(next(data_iter))

dataset = load_dataset("pcuenq/oxford-pets", split="train")
dataloader = DataLoader(dataset, batch_size=4, collate_fn=test_collate_fn)
data_iter = iter(dataloader)
print(next(data_iter))
```

@jack89roberts 

Pipeline expects a raw `string` not a dict of ` {"input": "myinput}` which all dataset contain.
And since there's no standard  key naming for dataset, we cannot infer those.

Many users seem to have this problem, so I think we should do something about it.

We could consider making the `Dataset` class a subclass of `torch.utils.data.Dataset` (as we do with the iterable version)  to pass the type check, but this change is not cheap as it requires `import torch`.

Alternatively, the pipeline code could check for the presence of `__getitem__` and `__len__`  instead of `isinstance(inputs, torch.utils.data.Dataset)`. This is what we suggested to fix a DeepSpeed limitation [here](https://github.com/huggingface/datasets/issues/2165) (the rationale was that `DataLoader` doesn't require map-style datasets to be an instance of Torch Dataset)

So, even though one can bypass this limitation by providing an HF dataset as a generator (`iter(dataset)`), allowing `pipeline(dataset)` would still be nice (for new users).

(This assumes the `dataset`'s columns match input names that the `pipeline` expects, e.g., `text` for `TextClassificationPipeline`)



I'd be happy to see a proof of concept working.

In my experience, the name matching just doesn't work on most datasets (and is even more confusing when it doesn't).

Yes, it doesn't always work, so [here](https://colab.research.google.com/drive/1r65SbdJWIZAeAadtRKTVhtj0_f9Zicnk?usp=sharing) is a notebook that suggests a solution.

WDYT?

PS: One more reason for these changes is that the `Trainer` is well-integrated with HF datasets, so it makes sense to do the same for the `Pipeline`.

I read a little bit, the working code in the notebook is pretty clear already I think `iter(dataset)`.

My main issue with "magic" is just that when it doesn't work (which should happen unfortunately quite a lot of the time since the names don't have a lot of reasons to match in general) then users will be confused.
"It work in the example, but it doesn't work on my dataset". And I'm not sure how we could convey easy simple steps for users to update their code to make it work.

```
ValueError: Incorrect format used for image. Should be an url linking to an image, a base64 string, a local path, or a PIL image.
```
For instance is pretty confusing imo. (I did send a URL linking to an image).

Again if you have some working PR I'm happy to take a look, but this is what I'll be looking for (How easy it is to know why the dataset is failing, and how easy is it to update from non working to working code)

The names can be aligned with `dataset.{rename_columns, remove_columns, select_columns}` (`KeyError` is raised if they don't match, so it should be clear what's the issue) 

> For instance is pretty confusing imo. (I did send a URL linking to an image).

I've updated the notebook with a call that works but does not scale.

To make my proposed solution clear, this is what the usage would look like:
* For the "single-input" pipelines (e.g., classification pipelines): `pipe(dataset["input_col"])` (no need for `KeyDataset`)
* For the "multi-input" pipelines (e.g., question answering pipelines): `pipe(dataset)` with the `dataset`  columns required to match the input arguments' names (e.g., "image" and "question" for `VisualQuestionAnswertingPipeline`)

I think it shouldn't be too hard to understand with some basic examples added to the `Pipeline` docs (e.g., replacing `KeyDataset` with `dataset["input_col"]` where possible)

I'll open a PR (should be pretty small) here as soon as `dataset["column"]` in `datasets`  is optimized so that it doesn't load everything in memory :).

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I'm getting bitten by this same issue.  I have a datasets.Dataset and discovering that transformers.pipelines.base does not do batching, leading to OOM.