mypy

Author: xadupre

onnx_diagnostic/export/api.py

@@ -65,9 +65,9 @@ def to_onnx(
         ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
         assert hasattr(mod, "config"), f"configuration is missing in model class {type(mod)}"
         assert not args, f"only kwargs can be defined with exporter={exporter!r}"
-        assert list(kwargs) == ["input_ids", "attention_mask", "past_key_values"], (
+        assert list(kwargs) == ["input_ids", "attention_mask", "past_key_values"], (  # type: ignore[arg-type]
             f"Only a specified set of inputs is supported for exporter={exporter!r}, "
-            f"but it is {list(kwargs)}"
+            f"but it is {list(kwargs)}"  # type: ignore[arg-type]
         )
         flat_inputs = flatten_object(kwargs, drop_keys=True)
         first = flat_inputs[0]

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -73,14 +73,19 @@ def _is_torchdynamo_exporting() -> bool:
     # Introduced in 4.52
     from transformers.masking_utils import (
         _ignore_causal_mask_sdpa,
-        _ignore_bidirectional_mask_sdpa,
         and_masks,
         bidirectional_mask_function,
         causal_mask_function,
         padding_mask_function,
         prepare_padding_mask,
     )
 
+    try:
+        # transformers>=5.0
+        from transformers.masking_utils import _ignore_bidirectional_mask_sdpa
+    except ImportError:
+        _ignore_bidirectional_mask_sdpa = None
+
     def patched__vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callable:
         """manual patch for function ``transformers.masking_utils._vmap_for_bhqkv``."""
         from ...helpers import string_type
@@ -187,7 +192,11 @@ def patched_sdpa_mask_recent_torch(
             padding_mask, q_length, kv_length, kv_offset, local_size
         ):
             return None
-        if allow_is_bidirectional_skip and _ignore_bidirectional_mask_sdpa(padding_mask):
+        if (
+            allow_is_bidirectional_skip
+            and _ignore_bidirectional_mask_sdpa
+            and _ignore_bidirectional_mask_sdpa(padding_mask)
+        ):
             return None
 
         if mask_function is bidirectional_mask_function: