Enable mambav2 compat

modeling_caduceus.py

@@ -2,21 +2,29 @@
 
 """
 
+import inspect
 import math
 from functools import partial
 from typing import Optional, Tuple, Union
 
 import torch
-from mamba_ssm.modules.mamba_simple import Mamba, Block
+from mamba_ssm.modules.mamba_simple import Mamba
+try:
+    from mamba_ssm.modules.mamba_simple import Block  # Legacy mambav1 file structure
+except ImportError:
+    from mamba_ssm.modules.block import Block  # mambav2 file structure
 from torch import nn
 from torch.nn import functional as F
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithNoAttention, MaskedLMOutput, SequenceClassifierOutput
 
 try:
-    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn  # Legacy mambav1 file structure
 except ImportError:
-    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+    try:
+        from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn  # mambav2 file structure
+    except ImportError:
+        RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
 
 from .configuration_caduceus import CaduceusConfig
 from .modeling_rcps import RCPSAddNormWrapper, RCPSEmbedding, RCPSLMHead, RCPSMambaBlock
@@ -54,13 +62,24 @@ def create_block(
         nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
     )
     block_cls = RCPSMambaBlock if rcps else Block
-    block = block_cls(
-        d_model,
-        mixer_cls,
-        norm_cls=norm_cls,
-        fused_add_norm=fused_add_norm,
-        residual_in_fp32=residual_in_fp32,
-    )
+    # mambav2 compatibility
+    if "mlp_cls" in inspect.signature(block_cls.__init__).parameters:
+        block = block_cls(
+            d_model,
+            mixer_cls,
+            mlp_cls=nn.Identity,
+            norm_cls=norm_cls,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+        )
+    else:
+        block = block_cls(
+            d_model,
+            mixer_cls,
+            norm_cls=norm_cls,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+        )
     block.layer_idx = layer_idx
     return block
 
@@ -264,15 +283,15 @@ def cross_entropy(logits, y, ignore_index=-100):
     return F.cross_entropy(logits, y, ignore_index=ignore_index)
 
 
-def weighted_cross_entropy(logits, y, loss_weight, ignore_index=-100):
+def weighted_cross_entropy(logits, y, loss_weights, ignore_index=-100):
     """Weighted cross entropy loss (discounts certain tokens, e.g., repeated base pairs in genome)."""
     logits = logits.view(-1, logits.shape[-1])
     y = y.view(-1)
     ce = F.cross_entropy(logits, y, ignore_index=ignore_index, reduction="none")
-    loss_weight = loss_weight.view(-1)
-    loss_weight[y == ignore_index] = 0.0
+    loss_weights = loss_weights.view(-1)
+    loss_weights[y == ignore_index] = 0.0
     # TODO: Follows GPN implementation, but should we remove weight normalization?
-    return (ce * (loss_weight / loss_weight.sum())).sum()
+    return (ce * (loss_weights / loss_weights.sum())).sum()
 
 
 class CaduceusPreTrainedModel(PreTrainedModel):
@@ -497,12 +516,6 @@ class CaduceusForSequenceClassification(CaduceusPreTrainedModel):
 
         # Initialize weights and apply final processing
         self.post_init()
-        self.init_scorer()
-
-    def init_scorer(self, initializer_range=0.02):
-        initializer_range = self.config.initializer_cfg.get("initializer_range", initializer_range) \
-            if self.config.initializer_cfg is not None else initializer_range
-        self.score.weight.data.normal_(std=initializer_range)
 
     def get_input_embeddings(self):
         return self.caduceus.backbone.embeddings.word_embeddings
@@ -530,7 +543,6 @@ class CaduceusForSequenceClassification(CaduceusPreTrainedModel):
         labels: Optional[torch.LongTensor] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        **kwargs,
     ) -> Union[Tuple, SequenceClassifierOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):