Upload orthrus_hf.py

Updated source code for Mamba-2 compatibility.

orthrus_hf.py

@@ -5,9 +5,16 @@ import torch
 import torch.nn as nn
 
 from functools import partial
-from mamba_ssm.modules.mamba_simple import Block, Mamba
 from transformers import PretrainedConfig, PreTrainedModel
 
+try:
+    from mamba_ssm.modules.mamba_simple import Mamba
+    from mamba_ssm.modules.block import Block
+except ImportError:
+    from mamba_ssm.modules.mamba_simple import Block, Mamba
+
+HAS_MLP = "mlp_cls" in Block.__init__.__code__.co_varnames
+
 
 class OrthrusConfig(PretrainedConfig):
     """HuggingFace config for pre-trained Orthrus model."""
@@ -108,11 +115,21 @@ class OrthrusPretrainedModel(PreTrainedModel):
         """
         mix_cls = partial(Mamba, layer_idx=layer_idx)
         norm_cls = nn.LayerNorm
-        block = Block(
-            d_model,
-            mix_cls,
-            norm_cls=norm_cls,
-        )
+
+        if HAS_MLP:
+            block = Block(
+                d_model,
+                mix_cls,
+                norm_cls=norm_cls,
+                mlp_cls=nn.Identity
+            )
+        else:
+            block = Block(
+                d_model,
+                mix_cls,
+                norm_cls=norm_cls
+            )
+
         block.layer_idx = layer_idx
         return block
 
@@ -166,7 +183,7 @@ class OrthrusPretrainedModel(PreTrainedModel):
         mean_tensor = self.mean_unpadded(out, lengths)
         return mean_tensor
 
-    def seq_to_oh(self, seq: str) -> torch.Tensor:
+    def seq_to_oh(self, seq: list[str]) -> torch.Tensor:
         """Convert nucleotide string into one-hot-encoding.
 
         The encoding uses ordering ["A", "C", "G", "T"].