fix import issues

Author: xadupre

_unittests/ut_torch_export_patches/test_onnx_export_errors.py

@@ -24,7 +24,7 @@ def test_pytree_flatten_mamba_cache(self):
         import torch.utils._pytree as py_pytree
 
         try:
-            from transformers.models.mamba.cache_mamba import MambaCache
+            from transformers.models.mamba.modeling_mamba import MambaCache
         except ImportError:
             from transformers.cache_utils import MambaCache
 

onnx_diagnostic/helpers/cache_helper.py

@@ -5,7 +5,7 @@
 import transformers.cache_utils
 
 try:
-    from transformers.models.mamba.cache_mamba import MambaCache
+    from transformers.models.mamba.modeling_mamba import MambaCache
 except ImportError:
     from transformers.cache_utils import MambaCache
 

onnx_diagnostic/tasks/text_generation.py

@@ -95,7 +95,7 @@ def get_inputs(
 
     if config is not None and config.__class__.__name__ == "FalconMambaConfig":
         try:
-            from transformers.models.mamba.cache_mamba import MambaCache
+            from transformers.models.mamba.modeling_mamba import MambaCache
         except ImportError:
             from transformers.cache_utils import MambaCache
 

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -12,7 +12,7 @@
 )
 
 try:
-    from transformers.models.mamba.cache_mamba import MambaCache
+    from transformers.models.mamba.modeling_mamba import MambaCache
 except ImportError:
     from transformers.cache_utils import MambaCache
 

onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py

@@ -9,7 +9,7 @@
 )
 
 try:
-    from transformers.models.mamba.cache_mamba import MambaCache
+    from transformers.models.mamba.modeling_mamba import MambaCache
 except ImportError:
     from transformers.cache_utils import MambaCache
 from transformers.modeling_outputs import BaseModelOutput