bypass

Author: xadupre

_unittests/ut_torch_models/test_llm_phi2.py

@@ -1,14 +1,8 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import (
-    ExtTestCase,
-    ignore_warnings,
-    requires_transformers,
-    requires_python,
-)
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, requires_transformers
 from onnx_diagnostic.torch_models.llms import get_phi2
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 
 
 class TestLlmPhi(ExtTestCase):
@@ -29,23 +23,6 @@ def test_export_phi2_1(self):
         ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds)
         assert ep
 
-    @ignore_warnings(UserWarning)
-    @requires_python((3, 12))
-    def test_export_phi2_2_bypassed(self):
-        data = get_phi2(num_hidden_layers=2)
-        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
-        self.assertEqual(
-            {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
-        )
-        with bypass_export_some_errors(patch_transformers=True) as modificator:
-            inputs = modificator(inputs)
-            ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
-            assert ep
-        with bypass_export_some_errors(patch_transformers=True) as modificator:
-            inputs = modificator(inputs)
-            ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
-            assert ep
-
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_models/test_tiny_llms_bypassed.py

@@ -4,6 +4,7 @@
 from transformers.cache_utils import DynamicCache
 from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
+from onnx_diagnostic.torch_models.llms import get_phi2
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
@@ -53,6 +54,22 @@ def debug():
             got = ep.module()(**inputs)
             self.assertEqualArrayAny(expected, got)
 
+    @ignore_warnings(UserWarning)
+    def test_export_phi2_2_bypassed(self):
+        data = get_phi2(num_hidden_layers=2)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(
+            {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
+        )
+        with bypass_export_some_errors(patch_transformers=True) as modificator:
+            inputs = modificator(inputs)
+            ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+            assert ep
+        with bypass_export_some_errors(patch_transformers=True) as modificator:
+            inputs = modificator(inputs)
+            ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+            assert ep
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)