Zero-shot classification pipeline does not support bfloat16

[umarbutler]

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
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

1. Run:

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.

[LysandreJik]

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

[umarbutler]

@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.

[Rocketknight1]

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?

[umarbutler]

Sure :)