Rename convert_pytorch to export

mfuntowicz · mfuntowicz · commit 4eebfda73684 · 2021-07-06T20:28:32.000+02:00
diff --git a/src/transformers/onnx/__init__.py b/src/transformers/onnx/__init__.py
@@ -14,4 +14,5 @@
 # limitations under the License.
 
 from .config import EXTERNAL_DATA_FORMAT_SIZE_LIMIT, OnnxConfig, OnnxConfigWithPast
+from .convert import export, validate_model_outputs
 from .utils import ParameterFormat, compute_serialized_parameters_size
diff --git a/src/transformers/onnx/__main__.py b/src/transformers/onnx/__main__.py
@@ -29,7 +29,7 @@
 
 from .. import is_torch_available
 from ..utils import logging
-from .convert import convert_pytorch, validate_model_outputs
+from .convert import export, validate_model_outputs
 
 
 if is_torch_available():
@@ -142,7 +142,7 @@ def main():
             f"At least  {onnx_config.default_onnx_opset} is required."
         )
 
-    onnx_inputs, onnx_outputs = convert_pytorch(tokenizer, model, onnx_config, args.opset, args.output)
+    onnx_inputs, onnx_outputs = export(tokenizer, model, onnx_config, args.opset, args.output)
 
     validate_model_outputs(onnx_config, tokenizer, model, args.output, onnx_outputs, args.atol)
     logger.info(f"All good, model saved at: {args.output.as_posix()}")
diff --git a/src/transformers/onnx/convert.py b/src/transformers/onnx/convert.py
@@ -62,7 +62,7 @@ def check_onnxruntime_requirements(minimum_version: Version):
         )
 
 
-def convert_pytorch(
+def export(
     tokenizer: PreTrainedTokenizer, model: PreTrainedModel, config: OnnxConfig, opset: int, output: Path
 ) -> Tuple[List[str], List[str]]:
     """
diff --git a/tests/test_onnx_v2.py b/tests/test_onnx_v2.py
@@ -24,9 +24,8 @@
 from transformers.models.roberta import RobertaOnnxConfig
 from transformers.models.t5 import T5OnnxConfig
 from transformers.models.xlm_roberta import XLMRobertaOnnxConfig
-from transformers.onnx import EXTERNAL_DATA_FORMAT_SIZE_LIMIT, OnnxConfig, ParameterFormat
+from transformers.onnx import EXTERNAL_DATA_FORMAT_SIZE_LIMIT, OnnxConfig, ParameterFormat, validate_model_outputs
 from transformers.onnx.config import DEFAULT_ONNX_OPSET, OnnxConfigWithPast
-from transformers.onnx.convert import validate_model_outputs
 from transformers.onnx.utils import (
     compute_effective_axis_dimension,
     compute_serialized_parameters_size,
@@ -204,7 +203,7 @@ class OnnxExportTestCaseV2(TestCase):
     @slow
     @require_torch
     def test_pytorch_export_default(self):
-        from transformers.onnx.convert import convert_pytorch
+        from transformers.onnx import export
 
         for name, model, model_class, config_class, onnx_config_class in PYTORCH_EXPORT_DEFAULT_MODELS:
             with self.subTest(name):
@@ -215,7 +214,7 @@ def test_pytorch_export_default(self):
                 onnx_config = onnx_config_class.default(model.config)
 
                 with NamedTemporaryFile("w") as output:
-                    onnx_inputs, onnx_outputs = convert_pytorch(
+                    onnx_inputs, onnx_outputs = export(
                         tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, Path(output.name)
                     )
 
@@ -227,7 +226,7 @@ def test_pytorch_export_default(self):
     @slow
     @require_torch
     def test_pytorch_export_with_past(self):
-        from transformers.onnx.convert import convert_pytorch
+        from transformers.onnx import export
 
         for name, model, model_class, config_class, onnx_config_class in PYTORCH_EXPORT_WITH_PAST_MODELS:
             with self.subTest(name):
@@ -244,7 +243,7 @@ def test_pytorch_export_with_past(self):
 
                 with NamedTemporaryFile("w") as output:
                     output = Path(output.name)
-                    onnx_inputs, onnx_outputs = convert_pytorch(
+                    onnx_inputs, onnx_outputs = export(
                         tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, output
                     )
 

Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,7 @@ def check_onnxruntime_requirements(minimum_version: Version):`
`62`	`62`	`)`
`63`	`63`
`64`	`64`
`65`		`-def convert_pytorch(`
	`65`	`+def export(`
`66`	`66`	`tokenizer: PreTrainedTokenizer, model: PreTrainedModel, config: OnnxConfig, opset: int, output: Path`
`67`	`67`	`) -> Tuple[List[str], List[str]]:`
`68`	`68`	`"""`