Fixes text2text generation, add summarization (#104)

* improve task text2text

* fix issues

* add summarization

* more robust

* insert clone

* fix a few things

* assert

* fix issues

* change

* fix

* fix

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.5.0
 +++++
 
+* :pr:`104`: add summarization task, add rewrite to command line validate
+* :pr:`101`: first draft to rewrite loops
 * :pr:`100`: implements a context to automatically rewrite methods or function with control flows
 * :pr:`96`: implements ``is_stealing``, ``steal_append`` to complement ``steal_forward``
 * :pr:`95`: fixzq Scan implementation for ``OnnxruntimeEvaluator``

_doc/api/tasks/index.rst

@@ -43,6 +43,7 @@ Or:
     mixture_of_expert
     object_detection
     sentence_similarity
+    summarization
     text_classification
     text_generation
     text2text_generation

_doc/api/tasks/summarization.rst

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

_unittests/ut_tasks/test_tasks.py

@@ -1,6 +1,11 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    has_transformers,
+    requires_transformers,
+)
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
@@ -123,7 +128,7 @@ def test_fill_mask(self):
             )
 
     @hide_stdout()
-    def test_feature_extraction(self):
+    def test_feature_extraction_bart_base(self):
         mid = "facebook/bart-base"
         data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
         self.assertEqual(data["task"], "feature-extraction")
@@ -136,6 +141,35 @@ def test_feature_extraction(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
+    @hide_stdout()
+    def test_feature_extraction_tiny_bart(self):
+        mid = "hf-tiny-model-private/tiny-random-PLBartForConditionalGeneration"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text2text-generation")
+        self.assertIn((data["size"], data["n_weights"]), [(3243392, 810848)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        with torch_export_patches(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+    @requires_transformers("4.51.999")
+    @hide_stdout()
+    def test_summarization(self):
+        mid = "facebook/bart-large-cnn"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "summarization")
+        self.assertIn((data["size"], data["n_weights"]), [(1625161728, 406290432)])
+        model, inputs, _ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        # with torch_export_patches(patch_transformers=True, verbose=10):
+        #    torch.export.export(
+        #        model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+        #    )
+
     @hide_stdout()
     def test_text_classification(self):
         mid = "Intel/bert-base-uncased-mrpc"

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -153,6 +153,30 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
+    def test_check_syntax_assign_noelse(self):
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+
+                def branch_cond_then_1(x):
+                    x = torch.abs(x) + 1
+                    return x
+
+                def branch_cond_else_1(x):
+                    return x.clone()
+
+                x = torch.cond(x.sum() > 0, branch_cond_then_1, branch_cond_else_1, [x])
+                return x + y
+
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        expected, expected_ = Model()(x, y), Model()(-x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        ep = torch.export.export(Model(), (x, y), dynamic_shapes=ds)
+        self.assertIn("cond", [str(getattr(n, "target", "?")) for n in ep.graph.nodes])
+        self.assertEqualAny(expected, ep.module()(x, y))
+        self.assertEqualAny(expected_, ep.module()(-x, y))
+
     def test_rewrite_test_in_forward_assign_noelse(self):
 
         class Model(torch.nn.Module):

_unittests/ut_torch_models/test_hghub_mode_rewrite.py

@@ -0,0 +1,13 @@
+import unittest
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.torch_models.hghub.hub_data import code_needing_rewriting
+
+
+class TestHuggingFaceHubModelRewrite(ExtTestCase):
+
+    def test_code_needing_rewriting(self):
+        self.assertEqual(1, len(code_needing_rewriting("BartForConditionalGeneration")))
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_torch_models/test_test_helpers.py

@@ -15,6 +15,7 @@
     validate_model,
     filter_inputs,
     run_ort_fusion,
+    empty,
 )
 from onnx_diagnostic.tasks import supported_tasks
 
@@ -32,6 +33,9 @@ def test_get_inputs_for_task(self):
             self.assertIn("dynamic_shapes", data)
             copy.deepcopy(data["inputs"])
 
+    def test_empty(self):
+        self.assertFalse(empty("float16"))
+
     @hide_stdout()
     def test_validate_model(self):
         mid = "arnir0/Tiny-LLM"

onnx_diagnostic/_command_lines_parser.py

@@ -309,12 +309,17 @@ def get_parser_validate() -> ArgumentParser:
         help="catches exception, report them in the summary",
     )
     parser.add_argument(
-        "-p",
         "--patch",
         default=True,
         action=BooleanOptionalAction,
         help="applies patches before exporting",
     )
+    parser.add_argument(
+        "--rewrite",
+        default=True,
+        action=BooleanOptionalAction,
+        help="applies rewrite before exporting",
+    )
     parser.add_argument(
         "--stop-if-static",
         default=0,
@@ -411,6 +416,7 @@ def _cmd_validate(argv: List[Any]):
             dtype=args.dtype,
             device=args.device,
             patch=args.patch,
+            rewrite=args.rewrite,
             stop_if_static=args.stop_if_static,
             optimization=args.opt,
             exporter=args.export,

onnx_diagnostic/helpers/torch_helper.py

@@ -689,9 +689,22 @@ def forward(self, input_ids):
     raise NotImplementedError(f"cls_name={cls_name}")
 
 
-def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
+def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
     """Applies torch.to if applicable. Goes recursively."""
-    if isinstance(value, (torch.nn.Module, torch.Tensor)):
+    if isinstance(value, (torch.nn.Module, torch.Tensor)) and value.__class__.__name__ not in {
+        "DynamicCache",
+        "EncoderDecoderCache",
+    }:
+        if (
+            (
+                isinstance(to_value, torch.dtype)
+                or to_value in {"float16", "bfloat16", "float32", "float64"}
+            )
+            and hasattr(value, "dtype")
+            and value.dtype in {torch.int32, torch.int64, torch.int8, torch.int16}
+        ):
+            # int vector should not be changed.
+            return value
         return value.to(to_value)
     if isinstance(value, list):
         return [to_any(t, to_value) for t in value]
@@ -701,8 +714,6 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
         return {to_any(t, to_value) for t in value}
     if isinstance(value, dict):
         return {k: to_any(t, to_value) for k, t in value.items()}
-    if hasattr(value, "to"):
-        return value.to(to_value)
     if value.__class__.__name__ == "DynamicCache":
         return make_dynamic_cache(
             list(
@@ -712,11 +723,23 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
                 )
             )
         )
+    if value.__class__.__name__ == "EncoderDecoderCache":
+        return make_encoder_decoder_cache(
+            to_any(value.self_attention_cache, to_value),
+            to_any(value.cross_attention_cache, to_value),
+        )
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         args, spec = torch.utils._pytree.tree_flatten(value)
         new_args = to_any(args, to_value)
         return torch.utils._pytree.tree_unflatten(new_args, spec)
 
+    if hasattr(value, "to"):
+        return value.to(to_value)
+
+    assert "Cache" not in value.__class__.__name__, (
+        f"Class {value.__class__.__name__!r} should be registered "
+        f"to be able to change the type in every tensor it contains."
+    )
     assert not isinstance(value, Iterable), f"Unsupported type {type(value)}"
     return value
 

onnx_diagnostic/tasks/__init__.py

@@ -8,6 +8,7 @@
     mixture_of_expert,
     object_detection,
     sentence_similarity,
+    summarization,
     text_classification,
     text_generation,
     text2text_generation,
@@ -23,6 +24,7 @@
     mixture_of_expert,
     object_detection,
     sentence_similarity,
+    summarization,
     text_classification,
     text_generation,
     text2text_generation,