Zero-shot classification pipeline does not support bfloat16
Author: umarbutler
Pull request: no
Labels: name: Core: Pipeline

name: bug

name: bfloat16

### System Info

- `transformers` version: 4.44.2
- Platform: Windows-11-10.0.22631-SP0
- Python version: 3.12.4
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.1+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce RTX 4090

### 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

1. Run:
```python
import torch

from transformers import pipeline, AutoTokenizer, RobertaForSequenceClassification, AutoModelForSequenceClassification

model_name = 'MoritzLaurer/roberta-large-zeroshot-v2.0-c'

model = RobertaForSequenceClassification.from_pretrained(model_name, torch_dtype = torch.bfloat16).cuda().to_bettertransformer()
tokenizer = AutoTokenizer.from_pretrained(model_name)
zeroshot_classifier = pipeline("zero-shot-classification", model, tokenizer = tokenizer, device = 0)

output = zeroshot_classifier('Angela Merkel is a politician in Germany and leader of the CDU', ["politics", "economy", "entertainment", "environment"], hypothesis_template = 'This text is about {}', multi_label = False)
```
2. Observe the following error:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[5], [line 9](vscode-notebook-cell:?execution_count=5&line=9)
      [6](vscode-notebook-cell:?execution_count=5&line=6) tokenizer = AutoTokenizer.from_pretrained(model_name)
      [7](vscode-notebook-cell:?execution_count=5&line=7) zeroshot_classifier = pipeline("zero-shot-classification", model, tokenizer = tokenizer, device = 0)
----> [9](vscode-notebook-cell:?execution_count=5&line=9) output = zeroshot_classifier('Angela Merkel is a politician in Germany and leader of the CDU', ["politics", "economy", "entertainment", "environment"], hypothesis_template = 'This text is about {}', multi_label = False)

File d:\redcated\redcated\.venv\Lib\site-packages\transformers\pipelines\zero_shot_classification.py:206, in ZeroShotClassificationPipeline.__call__(self, sequences, *args, **kwargs)
    [203](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:203) else:
    [204](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:204)     raise ValueError(f"Unable to understand extra arguments {args}")
--> [206](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:206) return super().__call__(sequences, **kwargs)

File d:\redcated\redcated\.venv\Lib\site-packages\transformers\pipelines\base.py:1249, in Pipeline.__call__(self, inputs, num_workers, batch_size, *args, **kwargs)
   [1247](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1247)     return self.iterate(inputs, preprocess_params, forward_params, postprocess_params)
   [1248](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1248) elif self.framework == "pt" and isinstance(self, ChunkPipeline):
-> [1249](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1249)     return next(
   [1250](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1250)         iter(
   [1251](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1251)             self.get_iterator(
   [1252](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1252)                 [inputs], num_workers, batch_size, preprocess_params, forward_params, postprocess_params
   [1253](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1253)             )
   [1254](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1254)         )
   [1255](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1255)     )
   [1256](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1256) else:
   [1257](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/base.py:1257)     return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)

File d:\redcated\redcated\.venv\Lib\site-packages\transformers\pipelines\pt_utils.py:125, in PipelineIterator.__next__(self)
    [123](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:123) # We're out of items within a batch
    [124](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:124) item = next(self.iterator)
--> [125](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:125) processed = self.infer(item, **self.params)
    [126](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:126) # We now have a batch of "inferred things".
    [127](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:127) if self.loader_batch_size is not None:
    [128](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/pt_utils.py:128)     # Try to infer the size of the batch

File d:\redcated\redcated\.venv\Lib\site-packages\transformers\pipelines\zero_shot_classification.py:242, in ZeroShotClassificationPipeline.postprocess(self, model_outputs, multi_label)
    [240](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:240) candidate_labels = [outputs["candidate_label"] for outputs in model_outputs]
    [241](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:241) sequences = [outputs["sequence"] for outputs in model_outputs]
--> [242](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:242) logits = np.concatenate([output["logits"].numpy() for output in model_outputs])
    [243](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:243) N = logits.shape[0]
    [244](file:///D:/redcated/redcated/.venv/Lib/site-packages/transformers/pipelines/zero_shot_classification.py:244) n = len(candidate_labels)

TypeError: Got unsupported ScalarType BFloat16
```

### Expected behavior

No exception raised.

Comments:
Indeed pipelines aren't bfloat16-compatible right now; and they're not tested to be so. 

cc @Rocketknight1 for your info, in case it makes sense to start a community effort to support this

@LysandreJik The fix is actually super simple. Add `import torch` to `transformers/pipelines/zero_shot_classification.py`, then `logits = np.concatenate([output["logits"].numpy() for output in model_outputs])` -> `np.concatenate([output["logits"].astype(torch.float32).numpy() for output in model_outputs])`.

`torch.float32` is used as intermediary to turning `torch.bfloat16` into a numpy float.

We don't mandate support for float16/bfloat16, but I'm in favour of adding it where appropriate. @umarbutler would you be willing to open a PR to add those changes?

Sure :)