Update modeling_mos_mamba.py

modeling_mos_mamba.py

@@ -26,22 +26,35 @@ from torch.nn import CrossEntropyLoss
 from transformers.activations import ACT2FN
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import ModelOutput
-from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available
+# from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available
 from .configuration_mos_mamba import MoSMambaConfig
 
 import torch.nn.functional as F
 
 
-if is_mamba_ssm_available():
+# if is_mamba_ssm_available():
+#     from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
+#     from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+# else:
+#     selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None
+
+# if is_causal_conv1d_available():
+#     from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+# else:
+#     causal_conv1d_update, causal_conv1d_fn = None, None
+
+
+try:
     from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
     from mamba_ssm.ops.triton.selective_state_update import selective_state_update
-else:
+except:
     selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None
 
-if is_causal_conv1d_available():
+try:
     from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
+except:
     causal_conv1d_update, causal_conv1d_fn = None, None
+    
 
 is_fast_path_available = all(
     (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)