task text-classification to be supported in the ONNX export for bart

[bry00]

Feature request

(Re)implement text-classification task support in the ONNX export for bart models.

Motivation

Task "text-classification" was supported previously (optimum==1.3.0, transformers==4.21.3), but now (optimum==1.16.1, transformers==4.36.2) the same code raises Task "text-classification" was supported previously (optimum==1.3.0, transformers==4.21.3), but now (optimum==1.16.1, transformers==4.36.2) the same code raises.

Sample code:

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer, pipeline
from pprint import pp

sample_text = "This is great!"
sample_labels = ["positive", "neutral", "negative"]

model_name = "facebook/bart-large-mnli"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ORTModelForSequenceClassification.from_pretrained(model_name, from_transformers=True)
pl = pipeline("zero-shot-classification", model=model, tokenizer=tokenizer)
pred = pl(sample_text, sample_labels, multi_label=False)

pp(pred)

Results of the above code with optimum==1.3.0, transformers==4.21.3:

Downloading: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.15k/1.15k [00:00<00:00, 546kB/s]
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:230: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:237: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if attention_mask.size() != (bsz, 1, tgt_len, src_len):
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:269: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:911: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if input_shape[-1] > 1:
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:96: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
  mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min))
/Users/brybak/dev/valdi/valdi-services/samples/valdi-zero-shot-sample/.venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py:1516: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
  if len(torch.unique_consecutive(eos_mask.sum(1))) > 1:
{'sequence': 'This is great!',
 'labels': ['positive', 'neutral', 'negative'],
 'scores': [0.9919587969779968, 0.006466080900281668, 0.0015750662423670292]}

Current result with optimum==1.16.1, transformers==4.36.2:

The argument `from_transformers` is deprecated, and will be removed in optimum 2.0.  Use `export` instead
Framework not specified. Using pt to export to ONNX.
Traceback (most recent call last):
  File "/Users/brybak/dev/modeling/samples/zero/src/bart_classifier.py", line 10, in <module>
    model = ORTModelForSequenceClassification.from_pretrained(model_name, from_transformers=True)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/brybak/dev/modeling/samples/zero/.venv/lib/python3.11/site-packages/optimum/onnxruntime/modeling_ort.py", line 626, in from_pretrained
    return super().from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/brybak/dev/modeling/samples/zero/.venv/lib/python3.11/site-packages/optimum/modeling_base.py", line 372, in from_pretrained
    return from_pretrained_method(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/brybak/dev/modeling/samples/zero/.venv/lib/python3.11/site-packages/optimum/onnxruntime/modeling_ort.py", line 549, in _from_transformers
    main_export(
  File "/Users/brybak/dev/modeling/samples/zero/.venv/lib/python3.11/site-packages/optimum/exporters/onnx/__main__.py", line 356, in main_export
    raise ValueError(
ValueError: Asked to export a bart model for the task text-classification, but the Optimum ONNX exporter only supports the tasks feature-extraction, feature-extraction-with-past, text-generation, text-generation-with-past, text2text-generation, text2text-generation-with-past for bart. Please use a supported task. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the task text-classification to be supported in the ONNX export for bart.

Your contribution

Cannot contribute.

[bry00]

Python version: 3.11.7.

[fxmarty]

@bry00 thank you, fixed by https://github.com/huggingface/optimum/pull/1666

[tengomucho]

this seems to be fixed, closing.