Add transformers tool test to pipeline

cmake/onnxruntime_python.cmake

@@ -67,7 +67,7 @@ if (onnxruntime_ENABLE_TRAINING)
   target_link_libraries(onnxruntime_pybind11_state PRIVATE onnxruntime_training)
 endif()
 
-target_link_libraries(onnxruntime_pybind11_state PRIVATE 
+target_link_libraries(onnxruntime_pybind11_state PRIVATE
     onnxruntime_session
     ${onnxruntime_libs}
     ${PROVIDERS_MIGRAPHX}
@@ -219,6 +219,12 @@ if (onnxruntime_BUILD_UNIT_TESTS)
   file(GLOB onnxruntime_python_dhp_parallel_test_srcs CONFIGURE_DEPENDS
       "${ORTTRAINING_SOURCE_DIR}/test/python/dhp_parallel/*.py"
   )
+  file(GLOB onnxruntime_python_transformers_test_srcs CONFIGURE_DEPENDS
+      "${ONNXRUNTIME_ROOT}/test/python/transformers/*.py"
+  )
+  file(GLOB onnxruntime_python_transformers_testdata_srcs CONFIGURE_DEPENDS
+      "${ONNXRUNTIME_ROOT}/test/python/transformers/test_data/models/*.onnx"
+  )
 endif()
 
 file(GLOB onnxruntime_python_tools_srcs CONFIGURE_DEPENDS
@@ -278,6 +284,8 @@ add_custom_command(
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/checkpoint
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/dhp_parallel
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/quantization
+  COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers
+  COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models
   COMMAND ${CMAKE_COMMAND} -E copy
       ${ONNXRUNTIME_ROOT}/__init__.py
       $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/
@@ -343,7 +351,7 @@ add_custom_command(
       $<TARGET_FILE_DIR:${build_output_target}>
 )
 
-if (NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_EXTENDED_MINIMAL_BUILD 
+if (NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_EXTENDED_MINIMAL_BUILD
                                   AND NOT onnxruntime_ENABLE_TRAINING
                                   AND NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin|iOS"
                                   AND NOT (CMAKE_SYSTEM_NAME STREQUAL "Android")
@@ -371,6 +379,12 @@ if (onnxruntime_BUILD_UNIT_TESTS)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_dhp_parallel_test_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/dhp_parallel/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_transformers_test_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/transformers/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_transformers_testdata_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models/
   )
 endif()
 

onnxruntime/python/tools/transformers/benchmark_gpt2.py

@@ -18,7 +18,7 @@
 import onnx
 from packaging import version
 from transformers import AutoConfig
-from gpt2_helper import DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS
+from gpt2_helper import Gpt2Helper, DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS
 from gpt2_beamsearch_helper import Gpt2HelperFactory, MODEL_CLASSES
 from quantize_helper import QuantizeHelper
 from benchmark_helper import create_onnxruntime_session, setup_logger, prepare_environment, Precision

onnxruntime/python/tools/transformers/test/test_attention_fusion.py → onnxruntime/test/python/transformers/test_attention_fusion.py

@@ -12,8 +12,7 @@
 
 # set path so that we could import from parent directory
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
-from optimizer import optimize_model
-
+from onnxruntime.transformers.optimizer import optimize_model
 
 class TestFusion(unittest.TestCase):
     def test_attention_fusion_pruned_model(self):
@@ -24,7 +23,7 @@ def test_attention_fusion_pruned_model(self):
         optimized_model = optimize_model(model_path)
         os.remove(model_path)
 
-        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'fusion',
+        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'models',
                                            'pruned_attention_opt.onnx')
         expected = onnx.load(expected_model_path)
         self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
@@ -38,7 +37,7 @@ def test_attention_fusion_reverse_add_order(self):
         os.remove(model_path)
 
         # reverse add input order will get same optimized model
-        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'fusion',
+        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'models',
                                            'pruned_attention_opt.onnx')
         expected = onnx.load(expected_model_path)
         self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
@@ -51,11 +50,10 @@ def test_3d_attention_fusion_tf2onnx_model(self):
         optimized_model = optimize_model(model_path, model_type='bert_tf', num_heads=4, hidden_size=16)
         os.remove(model_path)
 
-        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'fusion',
+        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'models',
                                            'bert_3d_attention_opt.onnx')
         expected = onnx.load(expected_model_path)
         self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
 
-
 if __name__ == '__main__':
     unittest.main()

onnxruntime/python/tools/transformers/test/test_gelu_fusions.py → onnxruntime/test/python/transformers/test_gelu_fusions.py

@@ -27,12 +27,8 @@ def forward(self, x):
         return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
 
 
-test_cases = [
-    ('huggingface', 'Gelu', HuggingfaceGelu),
-    ('huggingface', 'FastGelu', HuggingfaceFastGelu),
-    ('megatron', 'Gelu', MegatronGelu),
-    ('megatron', 'FastGelu', MegatronFastGelu)
-]
+test_cases = [('huggingface', 'Gelu', HuggingfaceGelu), ('huggingface', 'FastGelu', HuggingfaceFastGelu),
+              ('megatron', 'Gelu', MegatronGelu), ('megatron', 'FastGelu', MegatronFastGelu)]
 
 
 class TestGeluFusions(unittest.TestCase):
@@ -46,7 +42,7 @@ def verify_node_count(self, bert_model, expected_node_count, test_name):
 
     def test_fusions(self):
         sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
-        from optimizer import optimize_model
+        from onnxruntime.transformers.optimizer import optimize_model
 
         for test_case in test_cases:
             source, operator, model_class = test_case

onnxruntime/python/tools/transformers/test/test_gpt2.py → onnxruntime/test/python/transformers/test_gpt2.py

@@ -16,7 +16,7 @@
 
 class TestGpt2(unittest.TestCase):
     def run_benchmark_gpt2(self, arguments: str):
-        from benchmark_gpt2 import parse_arguments, main
+        from onnxruntime.transformers.benchmark_gpt2 import parse_arguments, main
         args = parse_arguments(arguments.split())
         csv_filename = main(args)
         self.assertTrue(os.path.exists(csv_filename))