Using batching with pipeline and transformers
Author: arunasank
Pull request: no
Labels: name: Core: Pipeline

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-5.15.0-107-generic-x86_64-with-glibc2.35
- Python version: 3.11.8
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.2
- Accelerate version: 0.29.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- 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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker @Narsil @stevhliu 

### Information

- [X] 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

```python
import pandas as pd
import utils
import datasets
import bitsandbytes
from torch.utils.data import DataLoader
from accelerate import infer_auto_device_map
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from transformers.pipelines.pt_utils import KeyDataset
from tqdm.auto import tqdm
import os
import torch
assert torch.cuda.is_available()

BATCH_SIZE = 2
os.environ['TRANSFORMERS_CACHE'] = '<redacted>'
os.environ['HF_TOKEN']='<redacted>'

dataset_filename: str = "<redacted>"

# Only for Llama-70B-chat
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load dataset
dataset = utils.load_from_jsonl(dataset_filename)
dataframe = pd.DataFrame(dataset[:4])
dataset = datasets.Dataset.from_pandas(dataframe)

pipe = pipeline("text-generation", model="meta-llama/Llama-2-70b-chat-hf", device_map='auto', max_new_tokens=1024, model_kwargs={"cache_dir": os.environ['TRANSFORMERS_CACHE'], "quantization_config": nf4_config})
pipe.tokenizer.pad_token_id = pipe.model.config.eos_token_id

responses = []
for d in tqdm(pipe(KeyDataset(dataset, 'prompt'), batch_size=BATCH_SIZE)):
    print('here!')
    responses.extend(d)

# Save responses
with open('llama-70-responses.jsonl', 'a') as f:
    for response in responses:
        f.write(response['generated_text'] + "\n")

print("Processing complete.")
```

1. Code hangs at `for d in tqdm(pipe(KeyDataset(dataset, 'prompt'), batch_size=BATCH_SIZE)):` and doesn't progress. No errors are thrown.
2. Code works when `batch_size` is not provided in the above line, but tqdm only shows progress on completion.
3. Alternatively, I tried to provide the `batch_size` with this line `pipe = pipeline("text-generation", model="meta-llama/Llama-2-70b-chat-hf", device_map='auto', max_new_tokens=1024, batch_size=BATCH_SIZE, model_kwargs={"cache_dir": os.environ['TRANSFORMERS_CACHE'], "quantization_config": nf4_config})` but it doesn't work.
4. Example from https://huggingface.co/docs/transformers/en/main_classes/pipelines#pipeline-batching does not work with Llama-2-79b-chat-hf

### Expected behavior

I would expect the code to work with batching.
Clear documentation on batching when using pipelines would be appreciated.

Comments:
cc @Rocketknight1 

~~UPDATE: I got the code to work with batching. It's just very hard to gauge progress because tqdm does not report progress until the whole pipeline has finished the task. Even so, times are not reported, making it impossible to to sub-sample and determine time either.~~

~~TL;DR - 1, 3, and 4 are resolved and 2 still remains as an issue.~~

UPDATE: It doesn't work! It works for small datasets and smaller batch sizes, but fails silently with no errors on a large dataset with a reasonable batch size. For example, the script uses BATCH_SIZE =2 and a dataset with 4 inputs. This took 4 minutes to run with Llama-2-chat on an A100. When I considered the full dataset, which is ~7200 inputs with a batch size of 4, it was hanging on the `for d in tqdm(pipe(KeyDataset(dataset, 'prompt'), batch_size=BATCH_SIZE)):` line silently, with no errors. 

@arunasank have you tried using the **map** method ?

```python
def fn(batch) : 
     batch["new_col"] = pipe(batch["prompt"]) 
     return batch

dataset = dataset.map(fn, batched = True , batch_size = BATCH_SIZE )
responses = dataset["new_col"]
```


Yes, I did that with/without setting `batch-size` on the pipeline method and it didn't work!

Hi @arunasank, have you tried other models and other methods of dataset creation to see if the issue recurs? It'd be helpful if you could figure out some minimal specific reproducer code for this issue, so that we could run it here and figure out what's going on!

Hi @arunasank , I am also troubled by the problem of pipeline progress bar. When processing a large dataset, the program is not hanging actually. The model is still inferring. After the inference of whole dataset is completed, the progress bar will be updated to the end. Because the dataset is so large, it looks like it's hanging without any errors. This seems to be a bug within the pipeline. I think the idea provided by @not-lain is worth to try. Thanks for @not-lain. Use dataset.map to implement a progress bar. Or something similar, you can implement a similar logic yourself like this:
```python
for i in tqdm(range(0, len(dataset), batch_size)):
    batch_data = dataset[i:i + args.batch_size]
    # Code to process your data.
```

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.

This still needs to be addressed, I think: the tqdm bar doesn't update until inference is complete for all items in the dataset, then the tqdm bar updates all at once. 

For example, the following code snippet exhibits this behavior:
```
data_iter = KeyDataset(ds, key_name)

generator = pipe(
    data_iter,
    batch_size=batch_size,
    max_new_tokens=max_new_tokens,
    do_sample=True,
    temperature=temperature if temperature > 0 else None,
)

completions: List[str] = []
for item in tqdm(generator, total=len(prompts), desc=progress_desc, disable=not show_progress):
    completions.append(item[0]["generated_text"])
```

As suggested above, a hacky approach is to call pipe() on each batch individually, but this produces the following warning:

> You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset

I'm not sure what the actual effect is on performance, but it seems reasonable that this would be slower than using a KeyDataset natively rather than manually breaking it into batches.

cc @Rocketknight1 

Hi @rmovva, I can't reproduce the issue! I turned your snippet into a complete script, but when I ran it, the TQDM bar filled normally for me. If you're still seeing the bug after updating to the latest versions of `transformers` and `datasets`, can you please send a **complete** reproducer script that I can run without any changes to see the issue?

```python                                                                                                           
from datasets import load_dataset
from transformers import pipeline
from tqdm import tqdm
from transformers.pipelines.pt_utils import KeyDataset
import torch

dataset = load_dataset("nyu-mll/glue", "cola")['train']

data_iter = KeyDataset(dataset, "sentence")

pipe = pipeline("text-generation", "google/gemma-3-1b-pt", torch_dtype=torch.bfloat16, device_map="cuda")

generator = pipe(
    data_iter,
    batch_size=4,
    max_new_tokens=12,
    do_sample=True,
)

completions = []
for item in tqdm(generator):
    completions.append(item[0]["generated_text"])
```