Reactivate BART Onnx Export (#1666)

* Reactivate BART Onnx Export * Update tests/onnxruntime/test_modeling.py * Update tests/onnxruntime/test_modeling.py * Bart Reactivation - fxmarty review --------- Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com>
huggingface · Jan 30, 2024 · da6f9e2 · da6f9e2
1 parent 45c1c09
commit da6f9e2
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 6 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -629,6 +629,7 @@ def flatten_past_key_values(self, flattened_output, name, idx, t):
 
 class BartOnnxConfig(M2M100OnnxConfig):
     DEFAULT_ONNX_OPSET = 14  # Bart now uses F.scaled_dot_product_attention by default for torch>=2.1.1.
+    MIN_TORCH_VERSION = version.parse("2.1.2")
     pass
 
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -317,10 +317,8 @@ class TasksManager:
             "text-generation-with-past",
             "text2text-generation",
             "text2text-generation-with-past",
-            # text-classification and question-answering can be supported, but the ONNX export is currently broken due to a regression in PyTorch 2.1.
-            # Reference: https://github.com/pytorch/pytorch/issues/110597.
-            # "text-classification",
-            # "question-answering",
+            "text-classification",
+            "question-answering",
             onnx="BartOnnxConfig",
         ),
         # BEiT cannot be used with the masked image modeling autoclass, so this task is excluded here

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -1227,7 +1227,7 @@ def test_trust_remote_code(self):
 class ORTModelForQuestionAnsweringIntegrationTest(ORTModelTestMixin):
     SUPPORTED_ARCHITECTURES = [
         "albert",
-        # "bart",  # see tasks.py
+        "bart", 
         "bert",
         # "big_bird",
         # "bigbird_pegasus",
@@ -1592,7 +1592,7 @@ def test_compare_to_io_binding(self, model_arch):
 class ORTModelForSequenceClassificationIntegrationTest(ORTModelTestMixin):
     SUPPORTED_ARCHITECTURES = [
         "albert",
-        # "bart",  # see tasks.py
+        "bart", 
         "bert",
         # "big_bird",
         # "bigbird_pegasus",

diff --git a/tests/onnxruntime/test_optimization.py b/tests/onnxruntime/test_optimization.py
@@ -89,6 +89,7 @@ def tearDownClass(cls):
 class ORTOptimizerTest(unittest.TestCase):
     # Contribution note: Please add test models in alphabetical order. Find test models here: https://huggingface.co/hf-internal-testing.
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = (
+        (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-bart"),
         (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-bert"),
         # (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-big_bird"),
         (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-distilbert"),

diff --git a/tests/onnxruntime/test_quantization.py b/tests/onnxruntime/test_quantization.py
@@ -76,6 +76,7 @@ class ORTDynamicQuantizationTest(unittest.TestCase):
         (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-bert", 30),
         (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-roberta", 30),
         (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-distilbert", 30),
+        (ORTModelForSequenceClassification, "hf-internal-testing/tiny-random-bart", 32),
     )
 
     SUPPORTED_DECODER_ARCHITECTURES_WITH_EXPECTED_QUANTIZED_MATMULS = (