split

Author: xadupre

_unittests/ut_tasks/test_tasks.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
@@ -36,22 +36,6 @@ def test_text_generation(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
-    @hide_stdout()
-    def test_image_classification(self):
-        mid = "hf-internal-testing/tiny-random-BeitForImageClassification"
-        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
-        self.assertEqual(data["task"], "image-classification")
-        self.assertIn((data["size"], data["n_weights"]), [(56880, 14220)])
-        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
-        model(**inputs)
-        model(**data["inputs2"])
-        if not has_transformers("4.51"):
-            raise unittest.SkipTest("_patch_make_causal_mask patch fails when an issue arises")
-        with bypass_export_some_errors(patch_transformers=True, verbose=10):
-            torch.export.export(
-                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
-            )
-
     @hide_stdout()
     def test_automatic_speech_recognition(self):
         mid = "openai/whisper-tiny"
@@ -124,22 +108,6 @@ def test_automatic_speech_recognition(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
-    @hide_stdout()
-    def test_image_text_to_text(self):
-        mid = "HuggingFaceM4/tiny-random-idefics"
-        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
-        self.assertEqual(data["task"], "image-text-to-text")
-        self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
-        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
-        model(**inputs)
-        model(**data["inputs2"])
-        if not has_torch("2.7.99"):
-            raise unittest.SkipTest("sym_max does not work with dynamic dimension")
-        with bypass_export_some_errors(patch_transformers=True, verbose=10):
-            torch.export.export(
-                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
-            )
-
     @hide_stdout()
     def test_fill_mask(self):
         mid = "google-bert/bert-base-multilingual-cased"
@@ -212,22 +180,6 @@ def test_falcon_mamba_dev(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
-    @hide_stdout()
-    def test_zero_shot_image_classification(self):
-        mid = "openai/clip-vit-base-patch16"
-        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
-        self.assertEqual(data["task"], "zero-shot-image-classification")
-        self.assertIn((data["size"], data["n_weights"]), [(188872708, 47218177)])
-        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
-        model(**inputs)
-        model(**data["inputs2"])
-        if not has_transformers("4.51"):
-            raise unittest.SkipTest("_patch_make_causal_mask patch fails when an issue arises")
-        with bypass_export_some_errors(patch_transformers=True, verbose=10):
-            torch.export.export(
-                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
-            )
-
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks_image_classification.py

@@ -0,0 +1,26 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_image_classification(self):
+        mid = "hf-internal-testing/tiny-random-BeitForImageClassification"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "image-classification")
+        self.assertIn((data["size"], data["n_weights"]), [(56880, 14220)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks_image_text_to_text.py

@@ -0,0 +1,30 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_image_text_to_text(self):
+        mid = "HuggingFaceM4/tiny-random-idefics"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "image-text-to-text")
+        self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        if not has_transformers("4.55"):
+            raise unittest.SkipTest("The model has control flow.")
+        if not has_torch("2.7.99"):
+            raise unittest.SkipTest("sym_max does not work with dynamic dimension")
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks_zero_shot_image_classification.py

@@ -0,0 +1,28 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_zero_shot_image_classification(self):
+        mid = "openai/clip-vit-base-patch16"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "zero-shot-image-classification")
+        self.assertIn((data["size"], data["n_weights"]), [(188872708, 47218177)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        if not has_transformers("4.51"):
+            raise unittest.SkipTest("_patch_make_causal_mask patch fails when an issue arises")
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)