Enable ONNX export test for supported model.

mfuntowicz · mfuntowicz · commit 4311e8a6f619 · 2021-06-28T20:08:09.000+02:00
diff --git a/tests/test_onnx_v2.py b/tests/test_onnx_v2.py
@@ -1,17 +1,27 @@
 from pathlib import Path
+from tempfile import NamedTemporaryFile
 from unittest import TestCase
 from unittest.mock import patch
 
-from transformers.models.bert.configuration_bert import BertOnnxConfig
+from transformers import AutoTokenizer, is_torch_available, AlbertConfig, DistilBertConfig, LongformerConfig, \
+    RobertaConfig, XLMRobertaConfig
+from transformers.models.albert import AlbertOnnxConfig
+
+from transformers.models.bert.configuration_bert import BertOnnxConfig, BertConfig
+from transformers.models.distilbert import DistilBertOnnxConfig
+from transformers.models.longformer import LongformerOnnxConfig
+from transformers.models.roberta import RobertaOnnxConfig
+from transformers.models.xlm_roberta import XLMRobertaOnnxConfig
 from transformers.onnx import EXTERNAL_DATA_FORMAT_SIZE_LIMIT, OnnxConfig, ParameterFormat
 
 # from transformers.onnx.convert import convert_pytorch
+from transformers.onnx.config import DEFAULT_ONNX_OPSET
 from transformers.onnx.utils import (
     compute_effective_axis_dimension,
     compute_serialized_parameters_size,
     flatten_output_collection_property,
 )
-from transformers.testing_utils import require_onnx
+from transformers.testing_utils import require_onnx, slow, require_torch
 
 
 @require_onnx
@@ -69,13 +79,35 @@ def test_use_external_data_format(self):
         self.assertTrue(OnnxConfig.use_external_data_format((TWO_GB_LIMIT + 1) // ParameterFormat.Float.size))
 
 
-class OnnxExportTestCaseV2(TestCase):
-    EXPORT_DEFAULT_MODELS = {
-        ("BERT", "bert-base-cased", BertOnnxConfig),
+if is_torch_available():
+    from transformers import AlbertModel, BertModel, DistilBertModel, LongformerModel, RobertaModel, XLMRobertaModel
+    PYTORCH_EXPORT_DEFAULT_MODELS = {
+        ("ALBERT", "albert-base-v2", AlbertModel, AlbertConfig, AlbertOnnxConfig),
+        ("BERT", "bert-base-cased", BertModel, BertConfig, BertOnnxConfig),
+        ("DistilBERT", "distilbert-base-cased", DistilBertModel, DistilBertConfig, DistilBertOnnxConfig),
+        # ("LongFormer", "longformer-base-4096", LongformerModel, LongformerConfig, LongformerOnnxConfig),
+        ("Roberta", "roberta-base", RobertaModel, RobertaConfig, RobertaOnnxConfig),
+        ("XLM-Roberta", "roberta-base", XLMRobertaModel, XLMRobertaConfig, XLMRobertaOnnxConfig),
     }
 
-    def export_default(self):
-        pass
 
-    def export_with_past(self):
+class OnnxExportTestCaseV2(TestCase):
+    @slow
+    @require_torch
+    def test_pytorch_export_default(self):
+        from transformers.onnx.convert import convert_pytorch
+
+        for name, model, model_class, config_class, onnx_config_class in PYTORCH_EXPORT_DEFAULT_MODELS:
+
+            with self.subTest(name):
+                tokenizer = AutoTokenizer.from_pretrained(model)
+                model = model_class(config_class())
+                onnx_config = onnx_config_class.default(model.config)
+
+                with NamedTemporaryFile("w") as output:
+                    convert_pytorch(tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, Path(output.name))
+
+    @slow
+    @require_torch
+    def test_pytorch_export_with_past(self):
         pass
