Supports ModelBuilder (#111)

xadupre · web-flow · commit ac301d1810f3 · 2025-05-26T18:59:16.000+02:00
* Supports ModelBuilder

* pymp

* t

* mb

* mypy

* fix modelbuilder inputs

* fix issues

* test'

* page

* publish
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
@@ -21,9 +21,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
 
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,8 +4,9 @@ Change Logs
 0.6.0
 +++++
 
-* :pr:`108`: first version of an algorithm rendering small onnx graph in ascii,
-  patch for ``torch.vmap``
+* :pr:`111`: support ModelBuilder with command line validatz
+* :pr:`108`, :pr:`109`, :pr:`110`: first version of an algorithm rendering
+  small onnx graph in ascii, patch for ``torch.vmap``
 
 0.5.0
 +++++
diff --git a/_doc/api/helpers/index.rst b/_doc/api/helpers/index.rst
@@ -14,6 +14,7 @@ onnx_diagnostic.helpers
     helper
     memory_peak
     mini_onnx_builder
+    model_builder_helper
     onnx_helper
     ort_session
     rt_helper
diff --git a/_doc/api/helpers/model_builder_helper.rst b/_doc/api/helpers/model_builder_helper.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.model_builder_helper
+============================================
+
+.. automodule:: onnx_diagnostic.helpers.model_builder_helper
+    :members:
+    :no-undoc-members:
diff --git a/_doc/conf.py b/_doc/conf.py
@@ -116,6 +116,7 @@ def linkcode_resolve(domain, info):
     ("py:class", "Argument"),
     ("py:class", "default=sklearn.utils.metadata_routing.UNCHANGED"),
     ("py:class", "ModelProto"),
+    ("py:class", "Model"),
     ("py:class", "Module"),
     ("py:class", "np.ndarray"),
     ("py:class", "onnxscript.ir.Tuple"),
@@ -211,6 +212,7 @@ def linkcode_resolve(domain, info):
     "huggingface_hub": "https://github.com/huggingface/huggingface_hub",
     "Linux": "https://www.linux.org/",
     "ml_dtypes": "https://github.com/jax-ml/ml_dtypes",
+    "ModelBuilder": "https://onnxruntime.ai/docs/genai/howto/build-model.html",
     "monai": "https://monai.io/",
     "numpy": "https://numpy.org/",
     "onnx": "https://onnx.ai/onnx/",
diff --git a/_doc/status/exported_program_dynamic.rst b/_doc/status/exported_program_dynamic.rst
@@ -43,7 +43,11 @@ to the original model.
         print()
         print("::")
         print()
-        print(textwrap.indent(textwrap.dedent(inspect.getsource(cls_model.forward)), "    "))
+        src = inspect.getsource(cls_model.forward)
+        if src:
+            print(textwrap.indent(textwrap.dedent(src), "    "))
+        else:
+            print("    # code is missing")
         print()
         for exporter in (
             "export-strict",
@@ -74,7 +78,11 @@ to the original model.
                 print()
                 print("::")
                 print()
-                print(textwrap.indent(str(res["error"]), "    "))
+                err = str(res["error"])
+                if err:
+                    print(textwrap.indent(err, "    "))
+                else:
+                    print("    # no error found for the failure")
                 print()
                 obs.append(dict(case=case_ref, error="FAIL", exporter=expo))
 
diff --git a/_doc/status/patches_coverage.rst b/_doc/status/patches_coverage.rst
@@ -28,7 +28,7 @@ for transformers.
     import onnx_diagnostic.torch_export_patches.patches.patch_transformers as p
 
     for name, cls in p.__dict__.items():
-        if name.startswith("patched_"):
+        if name.startswith("patched_") and hasattr(cls, "_PATCHES_"):
             print(f"{cls._PATCHED_CLASS_.__name__}: {', '.join(cls._PATCHES_)}")
 
 Half Automated Rewrites for Control Flows
diff --git a/_unittests/ut_export/test_dynamic_shapes.py b/_unittests/ut_export/test_dynamic_shapes.py
@@ -6,6 +6,7 @@
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 from onnx_diagnostic.export import ModelInputs, CoupleInputsDynamicShapes
 from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
 
 
 class TestDynamicShapes(ExtTestCase):
@@ -512,7 +513,6 @@ def forward(self, cache, z):
         mi = ModelInputs(Model(), inputs)
         self.assertIn("DynamicCache", string_type(mi.inputs, with_shape=True))
         ds = mi.guess_dynamic_shapes(auto="dim")
-        print(ds)
         self.assertEqual(
             ds,
             (
@@ -845,6 +845,23 @@ def test_dynamic_cache_replace_by_string(self):
             as_string,
         )
 
+    @requires_transformers("4.51")
+    def test_unbatch_inputs(self):
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        cpl = CoupleInputsDynamicShapes(
+            None, data["inputs"], dynamic_shapes=data["dynamic_shapes"]
+        )
+        new_dims = cpl.change_dynamic_dimensions(
+            desired_values=dict(batch=1), only_desired=True
+        )
+        s = self.string_type(new_dims, with_shape=True)
+        self.assertEqual(
+            "dict(input_ids:T7s1x3,attention_mask:T7s1x33,position_ids:T7s1x3,"
+            "past_key_values:DynamicCache("
+            "key_cache=#1[T1s1x1x30x96], value_cache=#1[T1s1x1x30x96]))",
+            s,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_helpers/test_model_builder_helper.py b/_unittests/ut_helpers/test_model_builder_helper.py
@@ -0,0 +1,83 @@
+import os
+import unittest
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    requires_torch,
+    requires_transformers,
+    hide_stdout,
+)
+from onnx_diagnostic.helpers.model_builder_helper import (
+    download_model_builder_to_cache,
+    import_model_builder,
+    create_model_builder,
+    save_model_builder,
+)
+from onnx_diagnostic.torch_models.hghub import (
+    get_untrained_model_with_inputs,
+)
+from onnx_diagnostic.helpers.rt_helper import make_feeds
+
+
+class TestModelBuilderHelper(ExtTestCase):
+    # This is to limit impact on CI.
+    @requires_transformers("4.52")
+    @requires_torch("2.7.99")
+    def test_download_model_builder(self):
+        path = download_model_builder_to_cache()
+        self.assertExists(path)
+        builder = import_model_builder()
+        self.assertHasAttr(builder, "create_model")
+
+    # This is to limit impact on CI.
+    @requires_transformers("4.52")
+    @requires_torch("2.7.99")
+    @hide_stdout()
+    def test_model_builder_id(self):
+        # clear&&python ~/.cache/onnx-diagnostic/builder.py
+        # --model arnir0/Tiny-LLM -p fp16 -c dump_cache -e cpu -o dump_model
+        folder = self.get_dump_folder("test_model_builder_id")
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        onnx_model = create_model_builder(
+            data["configuration"],
+            data["model"],
+            precision="fp32",
+            execution_provider="cpu",
+            cache_dir=folder,
+            verbose=1,
+        )
+        self.assertGreater(len(onnx_model.nodes), 5)
+
+        proto = save_model_builder(onnx_model, verbose=1)
+        import onnxruntime
+
+        onnxruntime.InferenceSession(
+            proto.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+
+        # We need to start again.
+        onnx_model = create_model_builder(
+            data["configuration"],
+            data["model"],
+            precision="fp32",
+            execution_provider="cpu",
+            cache_dir=folder,
+            verbose=1,
+        )
+        save_model_builder(onnx_model, folder, verbose=1)
+        model_name = os.path.join(folder, "model.onnx")
+        self.assertExists(model_name)
+
+        feeds = make_feeds(proto, data["inputs"], use_numpy=True)
+        expected = data["model"](**data["inputs"])
+
+        sess = onnxruntime.InferenceSession(model_name, providers=["CPUExecutionProvider"])
+        try:
+            got = sess.run(None, feeds)
+        except onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument as e:
+            if "batch_size must be 1 when sequence_length > 1" in str(e):
+                raise unittest.SkipTest("batch_size must be 1 when sequence_length > 1")
+        self.assertEqualAny(expected, got)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/_unittests/ut_torch_models/test_test_helpers.py b/_unittests/ut_torch_models/test_test_helpers.py
@@ -9,6 +9,7 @@
     requires_torch,
     requires_experimental,
     requires_onnxscript,
+    requires_transformers,
 )
 from onnx_diagnostic.torch_models.test_helper import (
     get_inputs_for_task,
@@ -184,6 +185,25 @@ def test_filter_inputs(self):
         ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"], model=["a", "b"])
         self.assertEqual((ni, nd), (((None,), {"b": 4}), {"b": 30}))
 
+    @requires_torch("2.7")
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
+    @requires_transformers("4.51")
+    def test_validate_model_modelbuilder(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=10,
+            exporter="modelbuilder",
+            dump_folder="dump_test_validate_model_onnx_dynamo",
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        self.assertLess(summary["disc_onnx_ort_run_abs"], 1e-4)
+        onnx_filename = data["onnx_filename"]
+        self.assertExists(onnx_filename)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/onnx_diagnostic/export/dynamic_shapes.py b/onnx_diagnostic/export/dynamic_shapes.py
@@ -379,8 +379,9 @@ def _generic_walker_step(
         return torch.utils._pytree.tree_unflatten(res, spec)
 
     class ChangeDimensionProcessor:
-        def __init__(self, desired_values):
+        def __init__(self, desired_values, only_desired):
             self.mapping = desired_values or {}
+            self.only_desired = only_desired
 
         def _build_new_shape(
             self, shape: Tuple[int, ...], ds: Dict[int, Any]
@@ -397,14 +398,16 @@ def _build_new_shape(
                             torch.export.dynamic_shapes._Dim,
                         ),
                     ):
-                        d = str(ds[i])
+                        d = ds[i].__name__
                     elif not isinstance(ds[i], int):
                         raise NotImplementedError(f"Unable to handle type {ds[i]} in {ds}")
                     if d in self.mapping:
                         new_dim = self.mapping[d]
-                    else:
+                    elif not self.only_desired:
                         new_dim = shape[i] + 1
                         self.mapping[d] = new_dim
+                    else:
+                        new_dim = shape[i]
                     new_shape[i] = new_dim
             return tuple(new_shape)
 
@@ -447,7 +450,10 @@ def __call__(self, inputs, ds):
             return self._build_new_tensor(inputs, new_shape)
 
     def change_dynamic_dimensions(
-        self, desired_values: Optional[Dict[str, int]] = None, args_kwargs: bool = False
+        self,
+        desired_values: Optional[Dict[str, int]] = None,
+        args_kwargs: bool = False,
+        only_desired: bool = False,
     ):
         """
         A model exported with dynamic shapes is not necessarily dynamic
@@ -460,6 +466,8 @@ def change_dynamic_dimensions(
 
         :param desired_values: to fixed named dimension to have the desired value
         :param args_kwargs: return both args, kwargs even if empty
+        :param only_desired: if True, only change the dimension specified in
+            ``desired_values``
         :return: new inputs
 
         Example:
@@ -483,7 +491,8 @@ def change_dynamic_dimensions(
             print("-after:", string_type(new_kwargs, with_shape=True))
         """
         return self._generic_walker(
-            self.ChangeDimensionProcessor(desired_values), args_kwargs=args_kwargs
+            self.ChangeDimensionProcessor(desired_values, only_desired=only_desired),
+            args_kwargs=args_kwargs,
         )
 
 
diff --git a/onnx_diagnostic/ext_test_case.py b/onnx_diagnostic/ext_test_case.py
@@ -776,6 +776,13 @@ def get_dump_file(self, name: str, folder: Optional[str] = None) -> str:
             os.mkdir(folder)
         return os.path.join(folder, name)
 
+    def get_dump_folder(self, folder: str) -> str:
+        """Returns a folder."""
+        folder = os.path.join("dump_test", folder)
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+        return folder
+
     def dump_onnx(
         self,
         name: str,
@@ -813,6 +820,11 @@ def assertIn(self, tofind: str, text: str, msg: str = ""):
             msg or f"Unable to find the list of strings {tofind!r} in\n--\n{text}"
         )
 
+    def assertHasAttr(self, obj: Any, name: str):
+        assert hasattr(
+            obj, name
+        ), f"Unable to find attribute {name!r} in object type {type(obj)}"
+
     def assertSetContained(self, set1, set2):
         "Checks that ``set1`` is contained in ``set2``."
         set1 = set(set1)
diff --git a/onnx_diagnostic/helpers/model_builder_helper.py b/onnx_diagnostic/helpers/model_builder_helper.py
diff --git a/onnx_diagnostic/helpers/rt_helper.py b/onnx_diagnostic/helpers/rt_helper.py
diff --git a/onnx_diagnostic/torch_models/test_helper.py b/onnx_diagnostic/torch_models/test_helper.py
diff --git a/pyproject.toml b/pyproject.toml
