fix mambacache import

Author: xadupre

_unittests/ut_torch_export_patches/test_onnx_export_errors.py

@@ -22,7 +22,11 @@ class TestOnnxExportErrors(ExtTestCase):
     def test_pytree_flatten_mamba_cache(self):
         import torch
         import torch.utils._pytree as py_pytree
-        from transformers.cache_utils import MambaCache
+
+        try:
+            from transformers.models.mamba.cache_mamba import MambaCache
+        except ImportError:
+            from transformers.cache_utils import MambaCache
 
         class _config:
             def __init__(self):

onnx_diagnostic/helpers/cache_helper.py

@@ -4,6 +4,11 @@
 import transformers
 import transformers.cache_utils
 
+try:
+    from transformers.models.mamba.cache_mamba import MambaCache
+except ImportError:
+    from transformers.cache_utils import MambaCache
+
 
 def flatten_unflatten_for_dynamic_shapes(
     obj: Any,
@@ -242,10 +247,8 @@ def make_encoder_decoder_cache(
     )
 
 
-def make_mamba_cache(
-    key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
-) -> transformers.cache_utils.MambaCache:
-    "Creates a :class:`transformers.cache_utils.MambaCache`."
+def make_mamba_cache(key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]]) -> MambaCache:
+    "Creates a ``MambaCache``."
     dtype = key_value_pairs[0][0].dtype
 
     class _config:
@@ -256,7 +259,7 @@ def __init__(self):
             self.num_hidden_layers = len(key_value_pairs)
             self.dtype = dtype
 
-    cache = transformers.cache_utils.MambaCache(
+    cache = MambaCache(
         _config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
         device=key_value_pairs[0][0].device,
@@ -286,7 +289,7 @@ def __init__(self):
 
 def make_sliding_window_cache(
     key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
-) -> transformers.cache_utils.MambaCache:
+) -> transformers.cache_utils.SlidingWindowCache:
     "Creates a :class:`transformers.cache_utils.SlidingWindowCache`."
 
     class _config:

onnx_diagnostic/tasks/text_generation.py

@@ -1,6 +1,5 @@
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 import torch
-import transformers
 from ..helpers.cache_helper import (
     make_dynamic_cache,
     make_mamba_cache,
@@ -95,9 +94,14 @@ def get_inputs(
     cache_length = "cache_length"  # torch.export.Dim("cache_length", min=1, max=4096)
 
     if config is not None and config.__class__.__name__ == "FalconMambaConfig":
+        try:
+            from transformers.models.mamba.cache_mamba import MambaCache
+        except ImportError:
+            from transformers.cache_utils import MambaCache
+
         assert cls_cache in (
             "MambaCache",
-            transformers.cache_utils.MambaCache,
+            MambaCache,
         ), f"Unexpected value for cls_cache={cls_cache} and config={config}"
         seq_length_multiple = 8
         sequence_length = (

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -6,12 +6,16 @@
 import transformers
 from transformers.cache_utils import (
     DynamicCache,
-    MambaCache,
     EncoderDecoderCache,
     SlidingWindowCache,
     StaticCache,
 )
 
+try:
+    from transformers.models.mamba.cache_mamba import MambaCache
+except ImportError:
+    from transformers.cache_utils import MambaCache
+
 from ..helpers import string_type
 from .serialization import _lower_name_with_
 

onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py

@@ -3,11 +3,15 @@
 import transformers
 from transformers.cache_utils import (
     DynamicCache,
-    MambaCache,
     EncoderDecoderCache,
     SlidingWindowCache,
     StaticCache,
 )
+
+try:
+    from transformers.models.mamba.cache_mamba import MambaCache
+except ImportError:
+    from transformers.cache_utils import MambaCache
 from transformers.modeling_outputs import BaseModelOutput
 from ...helpers.cache_helper import make_static_cache
 from . import make_serialization_function_for_dataclass