onnx-diagnostic 0.6.3__py3-none-any.whl → 0.7.0__py3-none-any.whl

onnx_diagnostic/tasks/text_generation.py

@@ -5,6 +5,7 @@ from ..helpers.cache_helper import (
     make_dynamic_cache,
     make_mamba_cache,
     make_sliding_window_cache,
+    make_static_cache,
 )
 from ..helpers.config_helper import update_config, check_hasattr, _pick
 
@@ -151,52 +152,98 @@ def get_inputs(
             assert config, "head_dim is None, the value cannot be set without a configuration"
             head_dim = config.hidden_size // config.num_attention_heads
 
-        shapes = {
-            "input_ids": {0: batch, 1: seq_length},
-            "attention_mask": {
-                0: batch,
-                1: "cache+seq",  # cache_length + seq_length
-            },
-            "position_ids": {
-                0: batch,
-                1: "cache+seq",  # cache_length + seq_length
-            },
-            "past_key_values": [
-                [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
-                [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
-            ],
+        cache_name = (
+            cls_cache
+            if cls_cache is None or isinstance(cls_cache, str)
+            else cls_cache.__name__
+        )
+        make_caches = {
+            "DynamicCache": make_dynamic_cache,
+            "SlidingWindowCache": make_sliding_window_cache,
+            "StaticCache": make_static_cache,
         }
-
-        make_cache = (
-            make_sliding_window_cache
-            if cls_cache in ("SlidingWindowCache", transformers.cache_utils.SlidingWindowCache)
-            else make_dynamic_cache
+        assert cache_name is None or cache_name in make_caches, (
+            f"Unable to handle cls_cache={cache_name!r}, it should be in "
+            f"{sorted(make_caches)}"
         )
+        make_cache = make_dynamic_cache if cache_name is None else make_caches[cache_name]
+        is_static = cache_name == "StaticCache"
 
-        inputs = dict(
-            input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length2)).to(
-                torch.int64
-            ),
-            attention_mask=torch.ones((batch_size, sequence_length + sequence_length2)).to(
-                torch.int64
-            ),
-            position_ids=torch.arange(sequence_length, sequence_length + sequence_length2)
-            .to(torch.int64)
-            .expand((batch_size, -1)),
-            past_key_values=make_cache(
-                [
-                    (
-                        torch.randn(
-                            batch_size, num_key_value_heads, sequence_length, head_dim
-                        ),
-                        torch.randn(
-                            batch_size, num_key_value_heads, sequence_length, head_dim
-                        ),
-                    )
-                    for i in range(num_hidden_layers)
-                ]
-            ),
-        )
+        if is_static:
+            # static
+            shapes = {
+                "input_ids": {0: batch, 1: seq_length},
+                "attention_mask": {0: batch, 2: "seq"},
+                "cache_position": {0: "seq"},
+                "past_key_values": [
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                ],
+            }
+            inputs = dict(
+                input_ids=torch.randint(
+                    0, dummy_max_token_id, (batch_size, sequence_length2)
+                ).to(torch.int64),
+                attention_mask=torch.ones(
+                    (batch_size, num_key_value_heads, sequence_length2, head_dim)
+                ).to(torch.bool),
+                cache_position=torch.arange(sequence_length2).to(torch.int64),
+                past_key_values=make_cache(
+                    [
+                        (
+                            torch.randn(
+                                batch_size, num_key_value_heads, sequence_length, head_dim
+                            ),
+                            torch.randn(
+                                batch_size, num_key_value_heads, sequence_length, head_dim
+                            ),
+                        )
+                        for i in range(num_hidden_layers)
+                    ]
+                ),
+            )
+        else:
+            # dynamic
+            shapes = {
+                "input_ids": {0: batch, 1: seq_length},
+                "attention_mask": {
+                    0: batch,
+                    1: "cache+seq",  # cache_length + seq_length
+                },
+                "position_ids": {
+                    0: batch,
+                    1: "cache+seq",  # cache_length + seq_length
+                },
+                "past_key_values": [
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                ],
+            }
+
+            inputs = dict(
+                input_ids=torch.randint(
+                    0, dummy_max_token_id, (batch_size, sequence_length2)
+                ).to(torch.int64),
+                attention_mask=torch.ones((batch_size, sequence_length + sequence_length2)).to(
+                    torch.int64
+                ),
+                position_ids=torch.arange(sequence_length, sequence_length + sequence_length2)
+                .to(torch.int64)
+                .expand((batch_size, -1)),
+                past_key_values=make_cache(
+                    [
+                        (
+                            torch.randn(
+                                batch_size, num_key_value_heads, sequence_length, head_dim
+                            ),
+                            torch.randn(
+                                batch_size, num_key_value_heads, sequence_length, head_dim
+                            ),
+                        )
+                        for i in range(num_hidden_layers)
+                    ]
+                ),
+            )
         res = dict(inputs=inputs, dynamic_shapes=shapes)
     if add_second_input:
         res["inputs2"] = get_inputs(

onnx_diagnostic/tasks/zero_shot_image_classification.py

@@ -55,6 +55,9 @@ def get_inputs(
     # attention_mask:T7s2x7
     # pixel_values:T1s2x3x224x224
     """
+    assert (
+        "cls_cache" not in kwargs
+    ), f"Not yet implemented for cls_cache={kwargs['cls_cache']!r}."
     assert isinstance(
         input_width, int
     ), f"Unexpected type for input_width {type(input_width)}{config}"

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -1,5 +1,8 @@
+import functools
+import importlib
 import contextlib
-from typing import Any, Callable, Dict, List, Optional
+import re
+from typing import Any, Callable, Dict, List, Optional, Tuple
 from .onnx_export_serialization import (
     register_cache_serialization,
     unregister_cache_serialization,
@@ -7,6 +10,41 @@ from .onnx_export_serialization import (
 from .patches import patch_transformers as patch_transformers_list
 
 
+def get_function(name: str) -> Tuple[type, Callable]:
+    """Returns the module and the function based on its name."""
+    spl = name.split(".")
+    module_name = ".".join(spl[:-1])
+    fname = spl[-1]
+    mod = importlib.import_module(module_name)
+    return mod, getattr(mod, fname)
+
+
+@functools.lru_cache
+def get_patches(mod, verbose: int = 0) -> Tuple[str, List[Any]]:
+    """Returns the list of patches to make for a specific module."""
+    to_patch = []
+    for k in dir(mod):
+        if k.startswith("patched_"):
+            v = getattr(mod, k)
+            if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
+                to_patch.append(v)
+            else:
+                # a function
+                doc = v.__doc__.lstrip()
+                if doc.startswith("manual patch"):
+                    continue
+                reg = re.compile("[[]patch:([a-z_A-Z.]+)[]]")
+                fall = reg.findall(doc)
+                assert (
+                    len(fall) == 1
+                ), f"Unable to find patching information for {v} in \n{doc}"
+                fmod, f = get_function(fall[0])
+                to_patch.append({"module": fmod, "function": f, "patch": v})
+
+    name = mod.__name__
+    return name, to_patch
+
+
 def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Callable]]:
     """
     Applies all patches defined in classes prefixed by ``patched_``
@@ -23,16 +61,21 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
         to_patch = mod
         name = "list"
     else:
-        to_patch = []
-        for k in dir(mod):
-            if k.startswith("patched_"):
-                v = getattr(mod, k)
-                if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
-                    to_patch.append(v)
-        name = mod.__name__
+        name, to_patch = get_patches(mod, verbose)
 
     res = {}
     for cls in to_patch:
+        if isinstance(cls, dict):
+            # a function
+            keep = {}
+            original = cls["module"]
+            f = cls["function"]
+            res[f] = f
+            if verbose:
+                print(f"[patch_module_or_classes] function: {original.__name__}.{f.__name__}")
+            setattr(original, f.__name__, cls["patch"])
+            continue
+
         original = cls._PATCHED_CLASS_
         methods = cls._PATCHES_
         if verbose:
@@ -57,26 +100,36 @@ def unpatch_module_or_classes(mod, info: Dict[type, Dict[type, Callable]], verbo
         to_patch = mod
         name = "list"
     else:
-        to_patch = []
-        for k in dir(mod):
-            if k.startswith("patched_"):
-                v = getattr(mod, k)
-                if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
-                    to_patch.append(v)
-        name = mod.__name__
-    set_patch = set(to_patch)
+        name, to_patch = get_patches(mod, verbose)
+
+    set_patch_cls = {i for i in to_patch if not isinstance(i, dict)}
+    dict_patch_fct = {i["function"]: i for i in to_patch if isinstance(i, dict)}
 
     for cls, methods in info.items():
-        assert cls in set_patch, f"No patch registered for {cls} in {mod} (found {set_patch})"
+        if cls in set_patch_cls:
+            if verbose:
+                print(
+                    f"[unpatch_module_or_classes] {name}.{cls.__name__}: {', '.join(methods)}"
+                )
+            original = cls._PATCHED_CLASS_
+            for n, v in methods.items():
+                if v is None:
+                    # The method did not exist. We remove it.
+                    delattr(original, n)
+                else:
+                    setattr(original, n, v)
+            continue
+        assert cls in dict_patch_fct, (
+            f"No patch registered for {cls} in {mod} "
+            f"(found {set_patch_cls} and {set(dict_patch_fct)})"
+        )
+        patch = dict_patch_fct[cls]
         if verbose:
-            print(f"[unpatch_module_or_classes] {name}.{cls.__name__}: {', '.join(methods)}")
-        original = cls._PATCHED_CLASS_
-        for n, v in methods.items():
-            if v is None:
-                # The method did not exist. We remove it.
-                delattr(original, n)
-            else:
-                setattr(original, n, v)
+            print(
+                f"[unpatch_module_or_classes] function "
+                f"{patch['module'].__name__}.{cls.__name__}"
+            )
+        setattr(patch["module"], cls.__name__, patch["function"])
 
 
 @contextlib.contextmanager

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -9,9 +9,11 @@ from transformers.cache_utils import (
     MambaCache,
     EncoderDecoderCache,
     SlidingWindowCache,
+    StaticCache,
 )
 from transformers.modeling_outputs import BaseModelOutput
 from ..helpers import string_type
+from ..helpers.cache_helper import make_static_cache
 
 
 PATCH_OF_PATCHES: Set[Any] = set()
@@ -175,6 +177,13 @@ def serialization_functions(verbose: int = 0) -> Dict[str, Union[Callable, int]]
             flatten_with_keys_sliding_window_cache,
             verbose=verbose,
         ),
+        StaticCache=register_class_serialization(
+            StaticCache,
+            flatten_static_cache,
+            unflatten_static_cache,
+            flatten_with_keys_static_cache,
+            verbose=verbose,
+        ),
     )
 
 
@@ -309,6 +318,34 @@ def unflatten_dynamic_cache(
     return cache
 
 
+##############
+# DynamicCache
+##############
+
+
+def flatten_static_cache(
+    cache: StaticCache,
+) -> Tuple[List[Any], torch.utils._pytree.Context]:
+    """Serializes a :class:`transformers.cache_utils.StaticCache` with python objects."""
+    flat = [("key_cache", cache.key_cache), ("value_cache", cache.value_cache)]
+    return [f[1] for f in flat], [f[0] for f in flat]
+
+
+def flatten_with_keys_static_cache(
+    cache: StaticCache,
+) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
+    """Serializes a :class:`transformers.cache_utils.StaticCache` with python objects."""
+    values, context = flatten_static_cache(cache)
+    return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
+
+
+def unflatten_static_cache(
+    values: List[Any], context: torch.utils._pytree.Context, output_type=None
+) -> StaticCache:
+    """Restores a :class:`transformers.cache_utils.StaticCache` from python objects."""
+    return make_static_cache(list(zip(values[0], values[1])))
+
+
 ####################
 # SlidingWindowCache
 ####################