Refactoring, add task fill-mask

.github/workflows/check-urls.yml

@@ -31,7 +31,7 @@ jobs:
         timeout: 2
         retry_count# : 2
         exclude_urls: https://github.com/pytorch/pytorch/pull/117009,https://github.com/huggingface/transformers/pull/29285,https://github.com/pytorch/pytorch/blob/a44f8894fa6d973693aab44a3dda079a168b05c1/torch/_decomp/decompositions.py#L1475
-        exclude_patterns: https://dumps.wikimedia.org/,https://github.com/pytorch/pytorch/pull/,https://github.com/pytorch/pytorch/blob/a44f8894fa6d973693aab44a3dda079a168b05c1/torch/_decomp/decompositions.py#L1475,https://huggingface.co/
+        exclude_patterns: https://dumps.wikimedia.org/,https://github.com/pytorch/pytorch/pull/,https://github.com/pytorch/pytorch/blob/a44f8894fa6d973693aab44a3dda079a168b05c1/torch/_decomp/decompositions.py#L1475,https://huggingface.co/,https://huggingface.co/
         # force_pass : true
 
     - name: urls-checker-docs
@@ -43,5 +43,5 @@ jobs:
         timeout: 2
         retry_count# : 2
         exclude_urls: https://hal.archives-,ouvertes.fr/hal-00990252/document,http://badge.fury.io/py/onnx-diagnostic,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://medium.com/@msouza.os/llm-from-scratch-with-pytorch-9f21808c6319,https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L5965,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311
-        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311,https://codecov.io/
+        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311,https://codecov.io/,https://huggingface.co/
         # force_pass : true

.github/workflows/documentation.yml

@@ -62,9 +62,11 @@ jobs:
 
       - name: tiny-llm example
         run: PYTHONPATH=. python _doc/examples/plot_export_tiny_llm.py
+        continue-on-error: true  # connectivity issues
 
       - name: tiny-llm bypass
         run: PYTHONPATH=. python _doc/examples/plot_export_tiny_llm_patched.py
+        continue-on-error: true  # connectivity issues
 
       - name: run tests bypassed
         run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms_bypassed.py

_doc/api/helpers/config_helper.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.config_helper
+=====================================
+
+.. automodule:: onnx_diagnostic.helpers.config_helper
+    :members:
+    :no-undoc-members:

_doc/api/helpers/index.rst

@@ -9,6 +9,7 @@ onnx_diagnostic.helpers
     args_helper
     bench_run
     cache_helper
+    config_helper
     helper
     memory_peak
     onnx_helper

_doc/api/index.rst

@@ -10,6 +10,7 @@ API of onnx_diagnostic
     export/index
     helpers/index
     reference/index
+    tasks/index
     torch_export_patches/index
     torch_models/index
     torch_onnx/index

_doc/api/tasks/automatic_speech_recognition.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.automatic_speech_recognition
+===================================================
+
+.. automodule:: onnx_diagnostic.tasks.automatic_speech_recognition
+    :members:
+    :no-undoc-members:

_doc/api/tasks/fill_mask.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.fill_mask
+================================
+
+.. automodule:: onnx_diagnostic.tasks.fill_mask
+    :members:
+    :no-undoc-members:

_doc/api/tasks/image_classification.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.image_classification
+===========================================
+
+.. automodule:: onnx_diagnostic.tasks.image_classification
+    :members:
+    :no-undoc-members:

_doc/api/tasks/image_text_to_text.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.image_text_to_text
+=========================================
+
+.. automodule:: onnx_diagnostic.tasks.image_text_to_text
+    :members:
+    :no-undoc-members:

_doc/api/tasks/index.rst

@@ -0,0 +1,18 @@
+onnx_diagnostic.tasks
+=====================
+
+.. toctree::
+    :maxdepth: 1
+    :caption: modules
+
+    automatic_speech_recognition
+    fill_mask
+    image_classification
+    image_text_to_text      
+    text_generation
+    text2text_generation
+    zero_shot_image_classification
+
+.. automodule:: onnx_diagnostic.tasks
+    :members:
+    :no-undoc-members:

_doc/api/tasks/text2text_generation.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.text2text_generation
+===========================================
+
+.. automodule:: onnx_diagnostic.tasks.text2text_generation
+    :members:
+    :no-undoc-members:

_doc/api/tasks/text_generation.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.text_generation
+======================================
+
+.. automodule:: onnx_diagnostic.tasks.text_generation
+    :members:
+    :no-undoc-members:

_doc/api/tasks/zero_shot_image_classification.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.zero_shot_image_classification
+=====================================================
+
+.. automodule:: onnx_diagnostic.tasks.zero_shot_image_classification
+    :members:
+    :no-undoc-members:

_doc/conf.py

@@ -170,7 +170,9 @@
 }
 
 if int(os.environ.get("UNITTEST_GOING", "0")):
-    sphinx_gallery_conf["ignore_pattern"] = ".*((tiny_llm)|(dort)|(draft_mode)).*"
+    sphinx_gallery_conf["ignore_pattern"] = (
+        ".*((tiny_llm)|(dort)|(draft_mode)|(hub_codellama.py)).*"
+    )
 elif pv.Version(torch.__version__) < pv.Version("2.8"):
     sphinx_gallery_conf["ignore_pattern"] = ".*((_oe_)|(dort)|(draft_mode)).*"
 

_unittests/ut_helpers/test_config_helper.py

@@ -0,0 +1,20 @@
+import unittest
+import transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    requires_torch,
+    requires_transformers,
+)
+from onnx_diagnostic.helpers.config_helper import config_class_from_architecture
+
+
+class TestConfigHelper(ExtTestCase):
+    @requires_transformers("4.50")  # we limit to some versions of the CI
+    @requires_torch("2.7")
+    def test_config_class_from_architecture(self):
+        config = config_class_from_architecture("LlamaForCausalLM")
+        self.assertEqual(config, transformers.LlamaConfig)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks.py

@@ -0,0 +1,105 @@
+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
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_text2text_generation(self):
+        mid = "sshleifer/tiny-marian-en-de"
+        # mid = "Salesforce/codet5-small"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertIn((data["size"], data["n_weights"]), [(473928, 118482)])
+        model, inputs = data["model"], data["inputs"]
+        raise unittest.SkipTest(f"not working for {mid!r}")
+        model(**inputs)
+
+    @hide_stdout()
+    def test_automatic_speech_recognition(self):
+        mid = "openai/whisper-tiny"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertIn((data["size"], data["n_weights"]), [(132115968, 33028992)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        Dim = torch.export.Dim
+        self.maxDiff = None
+        self.assertIn("{0:Dim(batch),1:Dim(seq_length)}", self.string_type(ds))
+        self.assertEqualAny(
+            {
+                "decoder_input_ids": {
+                    0: Dim("batch", min=1, max=1024),
+                    1: Dim("seq_length", min=1, max=4096),
+                },
+                "cache_position": {0: Dim("seq_length", min=1, max=4096)},
+                "encoder_outputs": [{0: Dim("batch", min=1, max=1024)}],
+                "past_key_values": [
+                    [
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                    ],
+                    [
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                    ],
+                ],
+            },
+            ds,
+        )
+        model(**inputs)
+        self.assertEqual(
+            "#1[T1r3]",
+            self.string_type(torch.utils._pytree.tree_flatten(inputs["encoder_outputs"])[0]),
+        )
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            flat = torch.utils._pytree.tree_flatten(inputs["past_key_values"])[0]
+            self.assertIsInstance(flat, list)
+            self.assertIsInstance(flat[0], torch.Tensor)
+            self.assertEqual(
+                "#8[T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4]",
+                self.string_type(flat),
+            )
+            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            flat = torch.utils._pytree.tree_flatten(inputs["past_key_values"])[0]
+            self.assertIsInstance(flat, list)
+            self.assertIsInstance(flat[0], torch.Tensor)
+            self.assertEqual(
+                "#8[T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4]",
+                self.string_type(flat),
+            )
+            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+
+    @hide_stdout()
+    def test_imagetext2text_generation(self):
+        mid = "HuggingFaceM4/tiny-random-idefics"
+        # mid = "Salesforce/codet5-small"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
+        model, inputs = data["model"], data["inputs"]
+        model(**inputs)
+
+    @hide_stdout()
+    def test_fill_mask(self):
+        mid = "google-bert/bert-base-multilingual-cased"
+        # mid = "Salesforce/codet5-small"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertIn((data["size"], data["n_weights"]), [(428383212, 107095803)])
+        model, inputs = data["model"], data["inputs"]
+        model(**inputs)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_torch_models/try_tasks.py → _unittests/ut_tasks/try_tasks.py

@@ -7,7 +7,7 @@
 class TestHuggingFaceHubModel(ExtTestCase):
     @never_test()
     def test_image_classification(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k image_c
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k image_c
 
         from transformers import ViTImageProcessor, ViTModel
         from PIL import Image
@@ -27,7 +27,7 @@ def test_image_classification(self):
 
     @never_test()
     def test_image_classification_resnet(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k resnet
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k resnet
 
         from transformers import ViTImageProcessor, ViTModel
         from PIL import Image
@@ -47,7 +47,7 @@ def test_image_classification_resnet(self):
 
     @never_test()
     def test_zero_shot_image_classification(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k zero
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k zero
         from PIL import Image
         import requests
         from transformers import CLIPProcessor, CLIPModel
@@ -74,7 +74,7 @@ def test_zero_shot_image_classification(self):
 
     @never_test()
     def test_text2text_generation(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k text2t
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k text2t
 
         import torch
         from transformers import RobertaTokenizer, T5ForConditionalGeneration
@@ -100,7 +100,7 @@ def test_text2text_generation(self):
 
     @never_test()
     def test_imagetext2text_generation(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k etext2t
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k etext2t
         # https://huggingface.co/docs/transformers/main/en/tasks/idefics
 
         import torch
@@ -131,7 +131,7 @@ def test_imagetext2text_generation(self):
 
     @never_test()
     def test_automatic_speech_recognition(self):
-        # clear&&NEVERTEST=1 python _unittests/ut_torch_models/try_tasks.py -k automatic_speech
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k automatic_speech
         # https://huggingface.co/openai/whisper-tiny
 
         from transformers import WhisperProcessor, WhisperForConditionalGeneration
@@ -195,6 +195,22 @@ def test_automatic_speech_recognition(self):
         transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
         print("--", transcription)
 
+    @never_test()
+    def test_fill_mask(self):
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k fill
+        # https://huggingface.co/google-bert/bert-base-multilingual-cased
+
+        from transformers import BertTokenizer, BertModel
+
+        tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")
+        model = BertModel.from_pretrained("bert-base-multilingual-cased")
+        text = "Replace me by any text you'd like."
+        encoded_input = tokenizer(text, return_tensors="pt")
+        print()
+        print("-- inputs", string_type(encoded_input, with_shape=True, with_min_max=True))
+        output = model(**encoded_input)
+        print("-- outputs", string_type(output, with_shape=True, with_min_max=True))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_models/test_hghub_api.py

@@ -100,7 +100,7 @@ def test_hf_all_models(self):
 
     def test_load_architecture_task(self):
         data = load_architecture_task()
-        print(set(data.values()))
+        self.assertNotEmpty(set(data.values()))
 
     def test_task_from_tags(self):
         _tags = [
@@ -132,10 +132,10 @@ def test_model_testings_and_architectures(self):
 
     def test__ccached_config_64(self):
         from onnx_diagnostic.torch_models.hghub.hub_data_cached_configs import (
-            _cached_hf_internal_testing_tiny_random_beitforimageclassification,
+            _ccached_hf_internal_testing_tiny_random_beitforimageclassification,
         )
 
-        conf = _cached_hf_internal_testing_tiny_random_beitforimageclassification()
+        conf = _ccached_hf_internal_testing_tiny_random_beitforimageclassification()
         self.assertEqual(conf.auxiliary_channels, 256)