Unify repo: default 24l-512d at root, add size variants via subfolder

24l-512d/config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "genemamba",
+  "architectures": [
+    "GeneMambaModel"
+  ],
+  "vocab_size": 25426,
+  "max_position_embeddings": 2048,
+  "hidden_size": 512,
+  "num_hidden_layers": 24,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "initializer_range": 0.02,
+  "mamba_mode": "gate",
+  "embedding_pooling": "mean",
+  "num_labels": 2,
+  "pad_token_id": 1,
+  "eos_token_id": 2,
+  "bos_token_id": 0,
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "auto_map": {
+    "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+    "AutoModel": "modeling_genemamba.GeneMambaModel",
+    "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification"
+  }
+}

24l-512d/configuration_genemamba.py

@@ -0,0 +1,97 @@
+"""
+Configuration for GeneMamba model.
+Defines all hyperparameters and settings for the GeneMamba architecture.
+"""
+
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+class GeneMambaConfig(PretrainedConfig):
+    """
+    Configuration class for GeneMamba model.
+    
+    This class stores the configuration of a GeneMamba model, inheriting from PretrainedConfig.
+    It can be used to instantiate models from pretrained checkpoints or customize model initialization.
+    
+    Args:
+        vocab_size (int, optional, defaults to 25426):
+            Vocabulary size of the model. Number of gene tokens (Ensembl Gene IDs).
+        
+        hidden_size (int, optional, defaults to 512):
+            Dimensionality of the hidden/embedding layers (d_model in Mamba).
+        
+        num_hidden_layers (int, optional, defaults to 24):
+            Number of Mamba layers (mamba_layer).
+        
+        intermediate_size (int, optional, defaults to 2048):
+            Dimensionality of intermediate representations in MLP.
+        
+        max_position_embeddings (int, optional, defaults to 2048):
+            Maximum sequence length (seq_len).
+        
+        hidden_dropout_prob (float, optional, defaults to 0.1):
+            Dropout probability for hidden states.
+        
+        initializer_range (float, optional, defaults to 0.02):
+            Standard deviation of truncated normal initializer.
+        
+        mamba_mode (str, optional, defaults to "gate"):
+            Aggregation mode for bidirectional Mamba layers.
+            Options: "mean", "sum", "concat", "gate".
+        
+        embedding_pooling (str, optional, defaults to "mean"):
+            Method for pooling to get cell embedding.
+            Options: "CLS", "mean", "weighted".
+        
+        num_labels (int, optional, defaults to 2):
+            Number of labels for sequence classification tasks.
+        
+        pad_token_id (int, optional, defaults to 1):
+            Token ID for padding.
+        
+        bos_token_id (int, optional, defaults to None):
+            Token ID for beginning of sequence.
+        
+        eos_token_id (int, optional, defaults to None):
+            Token ID for end of sequence.
+    """
+    
+    model_type = "genemamba"
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+    
+    def __init__(
+        self,
+        vocab_size: int = 25426,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 2048,
+        max_position_embeddings: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        initializer_range: float = 0.02,
+        mamba_mode: str = "gate",
+        embedding_pooling: str = "mean",
+        num_labels: int = 2,
+        pad_token_id: int = 1,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.mamba_mode = mamba_mode
+        self.embedding_pooling = embedding_pooling
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id

24l-512d/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccb1fcb0ee4b3ea2013099b9b187455e160d3b66b76c606715231b70b13c2784
+size 262998656

24l-512d/modeling_genemamba.py

@@ -0,0 +1,395 @@
+"""
+PyTorch implementation of GeneMamba model for Hugging Face Transformers.
+Includes backbone model and task-specific heads for various downstream tasks.
+"""
+
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import normal_, constant_
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, ModelOutput
+from transformers.models.auto import register_model_for_auto_class
+
+from mamba_ssm import Mamba
+from mamba_ssm.ops.triton.layer_norm import RMSNorm
+
+from .configuration_genemamba import GeneMambaConfig
+from .modeling_outputs import GeneMambaModelOutput, GeneMambaSequenceClassifierOutput, GeneMambaMaskedLMOutput
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================
+# Core Architecture Components
+# ===========================
+
+class EncoderLayer(nn.Module):
+    """
+    Single Mamba encoder layer with residual connection.
+    Applies a Mamba2 or Mamba layer followed by addition with input.
+    
+    Args:
+        hidden_size (int): Dimension of hidden states.
+    """
+    
+    def __init__(self, hidden_size: int):
+        super(EncoderLayer, self).__init__()
+        self.mamba = Mamba(d_model=hidden_size, d_state=64, d_conv=4, expand=2)
+    
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+        
+        Returns:
+            torch.Tensor: Output after Mamba layer and residual connection.
+        """
+        output = self.mamba(X) + X
+        return output
+
+
+class MambaMixer(nn.Module):
+    """
+    Stack of Mamba encoder layers with bidirectional processing and aggregation.
+    Processes sequences in both forward and reverse directions, then aggregates.
+    
+    Args:
+        mode (str): Aggregation mode. Options: "mean", "sum", "concat", "gate".
+        hidden_size (int): Dimension of hidden states.
+        num_hidden_layers (int): Number of Mamba layers.
+    """
+    
+    def __init__(
+        self,
+        mode: str = "gate",
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24
+    ):
+        super(MambaMixer, self).__init__()
+        self.mode = mode
+        self.hidden_size = hidden_size
+        
+        # Create Mamba layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(hidden_size) for _ in range(num_hidden_layers)]
+        )
+        
+        # Aggregation modules for certain modes
+        if mode in ["concat", "gate"]:
+            self.aggr = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def flip_sequence(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Reverse a sequence based on actual length (ignoring padding).
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            mask (torch.Tensor, optional): Padding mask of shape (batch_size, seq_len).
+        
+        Returns:
+            torch.Tensor: Reversed tensor.
+        """
+        batch_size, seq_length, embedding_dim = X.size()
+        
+        if mask is None:
+            # Simple flip
+            return X.flip([1])
+        
+        # Flip based on actual sequence length (marked by mask)
+        lengths = (~mask).sum(dim=1)
+        pos_tensor = torch.arange(seq_length, device=X.device).unsqueeze(0).expand(batch_size, -1)
+        flip_mask = pos_tensor < lengths.unsqueeze(1)
+        reversed_positions = torch.where(
+            flip_mask,
+            lengths.unsqueeze(1) - 1 - pos_tensor,
+            pos_tensor
+        )
+        
+        X_reverse = torch.gather(X, 1, reversed_positions.unsqueeze(-1).expand(-1, -1, embedding_dim))
+        return X_reverse
+    
+    def forward(
+        self,
+        X: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Process sequence through bidirectional Mamba layers.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            padding_mask (torch.Tensor, optional): Padding mask.
+        
+        Returns:
+            torch.Tensor: Output after processing all layers and aggregation.
+        """
+        
+        for layer in self.layers:
+            # Flip sequence for reverse processing
+            X_flip = self.flip_sequence(X, padding_mask)
+            
+            # Forward and reverse passes
+            X_f = layer(X)
+            X_b = layer(X_flip)
+            
+            # Flip back the reverse output
+            X_b = self.flip_sequence(X_b, padding_mask)
+            
+            # Aggregate forward and reverse
+            if self.mode == "mean":
+                X = (X_f + X_b) / 2
+            elif self.mode == "sum":
+                X = X_f + X_b
+            elif self.mode == "concat":
+                X = torch.cat([X_f, X_b], dim=-1)
+                X = self.aggr(X)
+            elif self.mode == "gate":
+                z = torch.sigmoid(self.aggr(torch.cat([X_f, X_b], dim=-1)))
+                X = z * X_f + (1 - z) * X_b
+            else:
+                raise ValueError(f"Invalid aggregation mode: {self.mode}")
+        
+        return X
+
+
+# ===========================
+# Base Model Classes
+# ===========================
+
+class GeneMambaPreTrainedModel(PreTrainedModel):
+    """
+    Base class for all GeneMamba models.
+    Handles weight initialization and provides standard model interfaces.
+    """
+    
+    config_class = GeneMambaConfig
+    base_model_prefix = "genemamba"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize module weights."""
+        if isinstance(module, nn.Linear):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.bias is not None:
+                constant_(module.bias, 0.0)
+        elif isinstance(module, nn.Embedding):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            constant_(module.bias, 0.0)
+            constant_(module.weight, 1.0)
+
+
+class GeneMambaModel(GeneMambaPreTrainedModel):
+    """
+    GeneMamba backbone model - outputs cell embeddings and hidden states.
+    This is the core model used by task-specific heads.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Embedding layer
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Mamba layers with bidirectional aggregation
+        self.mamba_mixer = MambaMixer(
+            mode=config.mamba_mode,
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers
+        )
+        
+        # Final layer normalization
+        self.norm = RMSNorm(config.hidden_size)
+        
+        self.apply(self._init_weights)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return embedding layer."""
+        return self.embeddings
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        """Set embedding layer."""
+        self.embeddings = value
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaModelOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask of shape (batch_size, seq_len).
+            output_hidden_states (bool): Whether to output hidden states from all layers.
+        
+        Returns:
+            GeneMambaModelOutput: Contains last_hidden_state, pooled_embedding, etc.
+        """
+        # Get embeddings
+        hidden_states = self.embeddings(input_ids)
+        
+        # Pass through Mamba layers
+        hidden_states = self.mamba_mixer(hidden_states, attention_mask)
+        
+        # Apply final normalization
+        hidden_states = self.norm(hidden_states)
+        
+        # Compute pooled embedding (cell representation)
+        if self.config.embedding_pooling == "CLS":
+            # Use first token (CLS)
+            pooled_embedding = hidden_states[:, 0, :]
+        elif self.config.embedding_pooling == "mean":
+            # Mean pooling over sequence
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).expand(hidden_states.shape).float()
+                pooled_embedding = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
+            else:
+                pooled_embedding = hidden_states.mean(dim=1)
+        else:
+            raise ValueError(f"Unsupported embedding_pooling: {self.config.embedding_pooling}")
+        
+        return GeneMambaModelOutput(
+            last_hidden_state=hidden_states,
+            pooled_embedding=pooled_embedding,
+            hidden_states=hidden_states if output_hidden_states else None,
+            embedding_pooling=self.config.embedding_pooling,
+        )
+
+
+# ===========================
+# Task-Specific Models
+# ===========================
+
+@register_model_for_auto_class("AutoModel")
+class GeneMambaForMaskedLM(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for masked language modeling (MLM).
+    Suitable for pretraining and domain adaptation.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.genemamba = GeneMambaModel(config)
+        
+        # Language modeling head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaMaskedLMOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Target token ids for MLM loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaMaskedLMOutput: Contains logits and optional loss.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        logits = self.lm_head(outputs.last_hidden_state)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        return GeneMambaMaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+        )
+
+
+@register_model_for_auto_class("AutoModelForSequenceClassification")
+class GeneMambaForSequenceClassification(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for sequence classification tasks.
+    Ideal for cell type annotation, tissue classification, etc.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        self.genemamba = GeneMambaModel(config)
+        
+        # Classification head
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaSequenceClassifierOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Class labels for classification loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaSequenceClassifierOutput: Contains logits, optional loss, and embedding.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        pooled_embedding = outputs.pooled_embedding
+        logits = self.classifier(self.dropout(pooled_embedding))
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        
+        return GeneMambaSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            pooled_embedding=pooled_embedding,
+        )
+
+
+# Register tokenizer class
+register_model_for_auto_class("AutoModelForMaskedLM")(GeneMambaForMaskedLM)

24l-512d/modeling_outputs.py

@@ -0,0 +1,81 @@
+"""
+Custom ModelOutput classes for GeneMamba.
+Defines the output structure for different GeneMamba tasks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class GeneMambaModelOutput(ModelOutput):
+    """
+    Base output class for GeneMamba models.
+    
+    Attributes:
+        last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)):
+            Sequence of hidden-states at the output of the last layer of the model.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size)):
+            Cell/sequence-level embedding (pooled representation) used for downstream tasks.
+            This is the recommended embedding to use for classification, clustering, etc.
+        
+        embedding_pooling (str):
+            The pooling method used to generate pooled_embedding.
+    """
+    
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: torch.FloatTensor = None
+    embedding_pooling: str = "mean"
+
+
+@dataclass
+class GeneMambaSequenceClassifierOutput(ModelOutput):
+    """
+    Output class for GeneMamba sequence classification models.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            Classification loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, num_labels)):
+            Classification scores (before softmax).
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size), optional):
+            Cell embedding before classification head.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class GeneMambaMaskedLMOutput(ModelOutput):
+    """
+    Output class for GeneMamba masked language modeling.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            MLM loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size)):
+            Prediction scores of the language modeling head.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None

24l-512d/special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]"
+}

24l-512d/tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

24l-768d/config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "genemamba",
+  "architectures": [
+    "GeneMambaModel"
+  ],
+  "vocab_size": 25426,
+  "max_position_embeddings": 2048,
+  "hidden_size": 768,
+  "num_hidden_layers": 24,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "initializer_range": 0.02,
+  "mamba_mode": "gate",
+  "embedding_pooling": "mean",
+  "num_labels": 2,
+  "pad_token_id": 1,
+  "eos_token_id": 2,
+  "bos_token_id": 0,
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "auto_map": {
+    "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+    "AutoModel": "modeling_genemamba.GeneMambaModel",
+    "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification"
+  }
+}

24l-768d/configuration_genemamba.py

@@ -0,0 +1,97 @@
+"""
+Configuration for GeneMamba model.
+Defines all hyperparameters and settings for the GeneMamba architecture.
+"""
+
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+class GeneMambaConfig(PretrainedConfig):
+    """
+    Configuration class for GeneMamba model.
+    
+    This class stores the configuration of a GeneMamba model, inheriting from PretrainedConfig.
+    It can be used to instantiate models from pretrained checkpoints or customize model initialization.
+    
+    Args:
+        vocab_size (int, optional, defaults to 25426):
+            Vocabulary size of the model. Number of gene tokens (Ensembl Gene IDs).
+        
+        hidden_size (int, optional, defaults to 512):
+            Dimensionality of the hidden/embedding layers (d_model in Mamba).
+        
+        num_hidden_layers (int, optional, defaults to 24):
+            Number of Mamba layers (mamba_layer).
+        
+        intermediate_size (int, optional, defaults to 2048):
+            Dimensionality of intermediate representations in MLP.
+        
+        max_position_embeddings (int, optional, defaults to 2048):
+            Maximum sequence length (seq_len).
+        
+        hidden_dropout_prob (float, optional, defaults to 0.1):
+            Dropout probability for hidden states.
+        
+        initializer_range (float, optional, defaults to 0.02):
+            Standard deviation of truncated normal initializer.
+        
+        mamba_mode (str, optional, defaults to "gate"):
+            Aggregation mode for bidirectional Mamba layers.
+            Options: "mean", "sum", "concat", "gate".
+        
+        embedding_pooling (str, optional, defaults to "mean"):
+            Method for pooling to get cell embedding.
+            Options: "CLS", "mean", "weighted".
+        
+        num_labels (int, optional, defaults to 2):
+            Number of labels for sequence classification tasks.
+        
+        pad_token_id (int, optional, defaults to 1):
+            Token ID for padding.
+        
+        bos_token_id (int, optional, defaults to None):
+            Token ID for beginning of sequence.
+        
+        eos_token_id (int, optional, defaults to None):
+            Token ID for end of sequence.
+    """
+    
+    model_type = "genemamba"
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+    
+    def __init__(
+        self,
+        vocab_size: int = 25426,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 2048,
+        max_position_embeddings: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        initializer_range: float = 0.02,
+        mamba_mode: str = "gate",
+        embedding_pooling: str = "mean",
+        num_labels: int = 2,
+        pad_token_id: int = 1,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.mamba_mode = mamba_mode
+        self.embedding_pooling = embedding_pooling
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id

24l-768d/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b423a3555eecacc88ff587c1d3f689a2caa05ede0a01d09dbaae175f23a2e7e1
+size 508241792

24l-768d/modeling_genemamba.py

@@ -0,0 +1,395 @@
+"""
+PyTorch implementation of GeneMamba model for Hugging Face Transformers.
+Includes backbone model and task-specific heads for various downstream tasks.
+"""
+
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import normal_, constant_
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, ModelOutput
+from transformers.models.auto import register_model_for_auto_class
+
+from mamba_ssm import Mamba
+from mamba_ssm.ops.triton.layer_norm import RMSNorm
+
+from .configuration_genemamba import GeneMambaConfig
+from .modeling_outputs import GeneMambaModelOutput, GeneMambaSequenceClassifierOutput, GeneMambaMaskedLMOutput
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================
+# Core Architecture Components
+# ===========================
+
+class EncoderLayer(nn.Module):
+    """
+    Single Mamba encoder layer with residual connection.
+    Applies a Mamba2 or Mamba layer followed by addition with input.
+    
+    Args:
+        hidden_size (int): Dimension of hidden states.
+    """
+    
+    def __init__(self, hidden_size: int):
+        super(EncoderLayer, self).__init__()
+        self.mamba = Mamba(d_model=hidden_size, d_state=64, d_conv=4, expand=2)
+    
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+        
+        Returns:
+            torch.Tensor: Output after Mamba layer and residual connection.
+        """
+        output = self.mamba(X) + X
+        return output
+
+
+class MambaMixer(nn.Module):
+    """
+    Stack of Mamba encoder layers with bidirectional processing and aggregation.
+    Processes sequences in both forward and reverse directions, then aggregates.
+    
+    Args:
+        mode (str): Aggregation mode. Options: "mean", "sum", "concat", "gate".
+        hidden_size (int): Dimension of hidden states.
+        num_hidden_layers (int): Number of Mamba layers.
+    """
+    
+    def __init__(
+        self,
+        mode: str = "gate",
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24
+    ):
+        super(MambaMixer, self).__init__()
+        self.mode = mode
+        self.hidden_size = hidden_size
+        
+        # Create Mamba layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(hidden_size) for _ in range(num_hidden_layers)]
+        )
+        
+        # Aggregation modules for certain modes
+        if mode in ["concat", "gate"]:
+            self.aggr = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def flip_sequence(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Reverse a sequence based on actual length (ignoring padding).
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            mask (torch.Tensor, optional): Padding mask of shape (batch_size, seq_len).
+        
+        Returns:
+            torch.Tensor: Reversed tensor.
+        """
+        batch_size, seq_length, embedding_dim = X.size()
+        
+        if mask is None:
+            # Simple flip
+            return X.flip([1])
+        
+        # Flip based on actual sequence length (marked by mask)
+        lengths = (~mask).sum(dim=1)
+        pos_tensor = torch.arange(seq_length, device=X.device).unsqueeze(0).expand(batch_size, -1)
+        flip_mask = pos_tensor < lengths.unsqueeze(1)
+        reversed_positions = torch.where(
+            flip_mask,
+            lengths.unsqueeze(1) - 1 - pos_tensor,
+            pos_tensor
+        )
+        
+        X_reverse = torch.gather(X, 1, reversed_positions.unsqueeze(-1).expand(-1, -1, embedding_dim))
+        return X_reverse
+    
+    def forward(
+        self,
+        X: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Process sequence through bidirectional Mamba layers.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            padding_mask (torch.Tensor, optional): Padding mask.
+        
+        Returns:
+            torch.Tensor: Output after processing all layers and aggregation.
+        """
+        
+        for layer in self.layers:
+            # Flip sequence for reverse processing
+            X_flip = self.flip_sequence(X, padding_mask)
+            
+            # Forward and reverse passes
+            X_f = layer(X)
+            X_b = layer(X_flip)
+            
+            # Flip back the reverse output
+            X_b = self.flip_sequence(X_b, padding_mask)
+            
+            # Aggregate forward and reverse
+            if self.mode == "mean":
+                X = (X_f + X_b) / 2
+            elif self.mode == "sum":
+                X = X_f + X_b
+            elif self.mode == "concat":
+                X = torch.cat([X_f, X_b], dim=-1)
+                X = self.aggr(X)
+            elif self.mode == "gate":
+                z = torch.sigmoid(self.aggr(torch.cat([X_f, X_b], dim=-1)))
+                X = z * X_f + (1 - z) * X_b
+            else:
+                raise ValueError(f"Invalid aggregation mode: {self.mode}")
+        
+        return X
+
+
+# ===========================
+# Base Model Classes
+# ===========================
+
+class GeneMambaPreTrainedModel(PreTrainedModel):
+    """
+    Base class for all GeneMamba models.
+    Handles weight initialization and provides standard model interfaces.
+    """
+    
+    config_class = GeneMambaConfig
+    base_model_prefix = "genemamba"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize module weights."""
+        if isinstance(module, nn.Linear):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.bias is not None:
+                constant_(module.bias, 0.0)
+        elif isinstance(module, nn.Embedding):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            constant_(module.bias, 0.0)
+            constant_(module.weight, 1.0)
+
+
+class GeneMambaModel(GeneMambaPreTrainedModel):
+    """
+    GeneMamba backbone model - outputs cell embeddings and hidden states.
+    This is the core model used by task-specific heads.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Embedding layer
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Mamba layers with bidirectional aggregation
+        self.mamba_mixer = MambaMixer(
+            mode=config.mamba_mode,
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers
+        )
+        
+        # Final layer normalization
+        self.norm = RMSNorm(config.hidden_size)
+        
+        self.apply(self._init_weights)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return embedding layer."""
+        return self.embeddings
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        """Set embedding layer."""
+        self.embeddings = value
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaModelOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask of shape (batch_size, seq_len).
+            output_hidden_states (bool): Whether to output hidden states from all layers.
+        
+        Returns:
+            GeneMambaModelOutput: Contains last_hidden_state, pooled_embedding, etc.
+        """
+        # Get embeddings
+        hidden_states = self.embeddings(input_ids)
+        
+        # Pass through Mamba layers
+        hidden_states = self.mamba_mixer(hidden_states, attention_mask)
+        
+        # Apply final normalization
+        hidden_states = self.norm(hidden_states)
+        
+        # Compute pooled embedding (cell representation)
+        if self.config.embedding_pooling == "CLS":
+            # Use first token (CLS)
+            pooled_embedding = hidden_states[:, 0, :]
+        elif self.config.embedding_pooling == "mean":
+            # Mean pooling over sequence
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).expand(hidden_states.shape).float()
+                pooled_embedding = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
+            else:
+                pooled_embedding = hidden_states.mean(dim=1)
+        else:
+            raise ValueError(f"Unsupported embedding_pooling: {self.config.embedding_pooling}")
+        
+        return GeneMambaModelOutput(
+            last_hidden_state=hidden_states,
+            pooled_embedding=pooled_embedding,
+            hidden_states=hidden_states if output_hidden_states else None,
+            embedding_pooling=self.config.embedding_pooling,
+        )
+
+
+# ===========================
+# Task-Specific Models
+# ===========================
+
+@register_model_for_auto_class("AutoModel")
+class GeneMambaForMaskedLM(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for masked language modeling (MLM).
+    Suitable for pretraining and domain adaptation.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.genemamba = GeneMambaModel(config)
+        
+        # Language modeling head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaMaskedLMOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Target token ids for MLM loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaMaskedLMOutput: Contains logits and optional loss.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        logits = self.lm_head(outputs.last_hidden_state)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        return GeneMambaMaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+        )
+
+
+@register_model_for_auto_class("AutoModelForSequenceClassification")
+class GeneMambaForSequenceClassification(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for sequence classification tasks.
+    Ideal for cell type annotation, tissue classification, etc.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        self.genemamba = GeneMambaModel(config)
+        
+        # Classification head
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaSequenceClassifierOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Class labels for classification loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaSequenceClassifierOutput: Contains logits, optional loss, and embedding.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        pooled_embedding = outputs.pooled_embedding
+        logits = self.classifier(self.dropout(pooled_embedding))
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        
+        return GeneMambaSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            pooled_embedding=pooled_embedding,
+        )
+
+
+# Register tokenizer class
+register_model_for_auto_class("AutoModelForMaskedLM")(GeneMambaForMaskedLM)

24l-768d/modeling_outputs.py

@@ -0,0 +1,81 @@
+"""
+Custom ModelOutput classes for GeneMamba.
+Defines the output structure for different GeneMamba tasks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class GeneMambaModelOutput(ModelOutput):
+    """
+    Base output class for GeneMamba models.
+    
+    Attributes:
+        last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)):
+            Sequence of hidden-states at the output of the last layer of the model.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size)):
+            Cell/sequence-level embedding (pooled representation) used for downstream tasks.
+            This is the recommended embedding to use for classification, clustering, etc.
+        
+        embedding_pooling (str):
+            The pooling method used to generate pooled_embedding.
+    """
+    
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: torch.FloatTensor = None
+    embedding_pooling: str = "mean"
+
+
+@dataclass
+class GeneMambaSequenceClassifierOutput(ModelOutput):
+    """
+    Output class for GeneMamba sequence classification models.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            Classification loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, num_labels)):
+            Classification scores (before softmax).
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size), optional):
+            Cell embedding before classification head.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class GeneMambaMaskedLMOutput(ModelOutput):
+    """
+    Output class for GeneMamba masked language modeling.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            MLM loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size)):
+            Prediction scores of the language modeling head.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None

24l-768d/special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]"
+}

24l-768d/tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

48l-512d/config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "genemamba",
+  "architectures": [
+    "GeneMambaModel"
+  ],
+  "vocab_size": 25426,
+  "max_position_embeddings": 2048,
+  "hidden_size": 512,
+  "num_hidden_layers": 48,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "initializer_range": 0.02,
+  "mamba_mode": "gate",
+  "embedding_pooling": "mean",
+  "num_labels": 2,
+  "pad_token_id": 1,
+  "eos_token_id": 2,
+  "bos_token_id": 0,
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "auto_map": {
+    "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+    "AutoModel": "modeling_genemamba.GeneMambaModel",
+    "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification"
+  }
+}

48l-512d/configuration_genemamba.py

@@ -0,0 +1,97 @@
+"""
+Configuration for GeneMamba model.
+Defines all hyperparameters and settings for the GeneMamba architecture.
+"""
+
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+class GeneMambaConfig(PretrainedConfig):
+    """
+    Configuration class for GeneMamba model.
+    
+    This class stores the configuration of a GeneMamba model, inheriting from PretrainedConfig.
+    It can be used to instantiate models from pretrained checkpoints or customize model initialization.
+    
+    Args:
+        vocab_size (int, optional, defaults to 25426):
+            Vocabulary size of the model. Number of gene tokens (Ensembl Gene IDs).
+        
+        hidden_size (int, optional, defaults to 512):
+            Dimensionality of the hidden/embedding layers (d_model in Mamba).
+        
+        num_hidden_layers (int, optional, defaults to 24):
+            Number of Mamba layers (mamba_layer).
+        
+        intermediate_size (int, optional, defaults to 2048):
+            Dimensionality of intermediate representations in MLP.
+        
+        max_position_embeddings (int, optional, defaults to 2048):
+            Maximum sequence length (seq_len).
+        
+        hidden_dropout_prob (float, optional, defaults to 0.1):
+            Dropout probability for hidden states.
+        
+        initializer_range (float, optional, defaults to 0.02):
+            Standard deviation of truncated normal initializer.
+        
+        mamba_mode (str, optional, defaults to "gate"):
+            Aggregation mode for bidirectional Mamba layers.
+            Options: "mean", "sum", "concat", "gate".
+        
+        embedding_pooling (str, optional, defaults to "mean"):
+            Method for pooling to get cell embedding.
+            Options: "CLS", "mean", "weighted".
+        
+        num_labels (int, optional, defaults to 2):
+            Number of labels for sequence classification tasks.
+        
+        pad_token_id (int, optional, defaults to 1):
+            Token ID for padding.
+        
+        bos_token_id (int, optional, defaults to None):
+            Token ID for beginning of sequence.
+        
+        eos_token_id (int, optional, defaults to None):
+            Token ID for end of sequence.
+    """
+    
+    model_type = "genemamba"
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+    
+    def __init__(
+        self,
+        vocab_size: int = 25426,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 2048,
+        max_position_embeddings: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        initializer_range: float = 0.02,
+        mamba_mode: str = "gate",
+        embedding_pooling: str = "mean",
+        num_labels: int = 2,
+        pad_token_id: int = 1,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.mamba_mode = mamba_mode
+        self.embedding_pooling = embedding_pooling
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id

48l-512d/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a715342c6cc00b20161a05941d9d181cca73c7ecc9cae17fd3a04bf92590a7d
+size 421748360

48l-512d/modeling_genemamba.py

@@ -0,0 +1,395 @@
+"""
+PyTorch implementation of GeneMamba model for Hugging Face Transformers.
+Includes backbone model and task-specific heads for various downstream tasks.
+"""
+
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import normal_, constant_
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, ModelOutput
+from transformers.models.auto import register_model_for_auto_class
+
+from mamba_ssm import Mamba
+from mamba_ssm.ops.triton.layer_norm import RMSNorm
+
+from .configuration_genemamba import GeneMambaConfig
+from .modeling_outputs import GeneMambaModelOutput, GeneMambaSequenceClassifierOutput, GeneMambaMaskedLMOutput
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================
+# Core Architecture Components
+# ===========================
+
+class EncoderLayer(nn.Module):
+    """
+    Single Mamba encoder layer with residual connection.
+    Applies a Mamba2 or Mamba layer followed by addition with input.
+    
+    Args:
+        hidden_size (int): Dimension of hidden states.
+    """
+    
+    def __init__(self, hidden_size: int):
+        super(EncoderLayer, self).__init__()
+        self.mamba = Mamba(d_model=hidden_size, d_state=64, d_conv=4, expand=2)
+    
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+        
+        Returns:
+            torch.Tensor: Output after Mamba layer and residual connection.
+        """
+        output = self.mamba(X) + X
+        return output
+
+
+class MambaMixer(nn.Module):
+    """
+    Stack of Mamba encoder layers with bidirectional processing and aggregation.
+    Processes sequences in both forward and reverse directions, then aggregates.
+    
+    Args:
+        mode (str): Aggregation mode. Options: "mean", "sum", "concat", "gate".
+        hidden_size (int): Dimension of hidden states.
+        num_hidden_layers (int): Number of Mamba layers.
+    """
+    
+    def __init__(
+        self,
+        mode: str = "gate",
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24
+    ):
+        super(MambaMixer, self).__init__()
+        self.mode = mode
+        self.hidden_size = hidden_size
+        
+        # Create Mamba layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(hidden_size) for _ in range(num_hidden_layers)]
+        )
+        
+        # Aggregation modules for certain modes
+        if mode in ["concat", "gate"]:
+            self.aggr = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def flip_sequence(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Reverse a sequence based on actual length (ignoring padding).
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            mask (torch.Tensor, optional): Padding mask of shape (batch_size, seq_len).
+        
+        Returns:
+            torch.Tensor: Reversed tensor.
+        """
+        batch_size, seq_length, embedding_dim = X.size()
+        
+        if mask is None:
+            # Simple flip
+            return X.flip([1])
+        
+        # Flip based on actual sequence length (marked by mask)
+        lengths = (~mask).sum(dim=1)
+        pos_tensor = torch.arange(seq_length, device=X.device).unsqueeze(0).expand(batch_size, -1)
+        flip_mask = pos_tensor < lengths.unsqueeze(1)
+        reversed_positions = torch.where(
+            flip_mask,
+            lengths.unsqueeze(1) - 1 - pos_tensor,
+            pos_tensor
+        )
+        
+        X_reverse = torch.gather(X, 1, reversed_positions.unsqueeze(-1).expand(-1, -1, embedding_dim))
+        return X_reverse
+    
+    def forward(
+        self,
+        X: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Process sequence through bidirectional Mamba layers.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            padding_mask (torch.Tensor, optional): Padding mask.
+        
+        Returns:
+            torch.Tensor: Output after processing all layers and aggregation.
+        """
+        
+        for layer in self.layers:
+            # Flip sequence for reverse processing
+            X_flip = self.flip_sequence(X, padding_mask)
+            
+            # Forward and reverse passes
+            X_f = layer(X)
+            X_b = layer(X_flip)
+            
+            # Flip back the reverse output
+            X_b = self.flip_sequence(X_b, padding_mask)
+            
+            # Aggregate forward and reverse
+            if self.mode == "mean":
+                X = (X_f + X_b) / 2
+            elif self.mode == "sum":
+                X = X_f + X_b
+            elif self.mode == "concat":
+                X = torch.cat([X_f, X_b], dim=-1)
+                X = self.aggr(X)
+            elif self.mode == "gate":
+                z = torch.sigmoid(self.aggr(torch.cat([X_f, X_b], dim=-1)))
+                X = z * X_f + (1 - z) * X_b
+            else:
+                raise ValueError(f"Invalid aggregation mode: {self.mode}")
+        
+        return X
+
+
+# ===========================
+# Base Model Classes
+# ===========================
+
+class GeneMambaPreTrainedModel(PreTrainedModel):
+    """
+    Base class for all GeneMamba models.
+    Handles weight initialization and provides standard model interfaces.
+    """
+    
+    config_class = GeneMambaConfig
+    base_model_prefix = "genemamba"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize module weights."""
+        if isinstance(module, nn.Linear):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.bias is not None:
+                constant_(module.bias, 0.0)
+        elif isinstance(module, nn.Embedding):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            constant_(module.bias, 0.0)
+            constant_(module.weight, 1.0)
+
+
+class GeneMambaModel(GeneMambaPreTrainedModel):
+    """
+    GeneMamba backbone model - outputs cell embeddings and hidden states.
+    This is the core model used by task-specific heads.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Embedding layer
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Mamba layers with bidirectional aggregation
+        self.mamba_mixer = MambaMixer(
+            mode=config.mamba_mode,
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers
+        )
+        
+        # Final layer normalization
+        self.norm = RMSNorm(config.hidden_size)
+        
+        self.apply(self._init_weights)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return embedding layer."""
+        return self.embeddings
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        """Set embedding layer."""
+        self.embeddings = value
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaModelOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask of shape (batch_size, seq_len).
+            output_hidden_states (bool): Whether to output hidden states from all layers.
+        
+        Returns:
+            GeneMambaModelOutput: Contains last_hidden_state, pooled_embedding, etc.
+        """
+        # Get embeddings
+        hidden_states = self.embeddings(input_ids)
+        
+        # Pass through Mamba layers
+        hidden_states = self.mamba_mixer(hidden_states, attention_mask)
+        
+        # Apply final normalization
+        hidden_states = self.norm(hidden_states)
+        
+        # Compute pooled embedding (cell representation)
+        if self.config.embedding_pooling == "CLS":
+            # Use first token (CLS)
+            pooled_embedding = hidden_states[:, 0, :]
+        elif self.config.embedding_pooling == "mean":
+            # Mean pooling over sequence
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).expand(hidden_states.shape).float()
+                pooled_embedding = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
+            else:
+                pooled_embedding = hidden_states.mean(dim=1)
+        else:
+            raise ValueError(f"Unsupported embedding_pooling: {self.config.embedding_pooling}")
+        
+        return GeneMambaModelOutput(
+            last_hidden_state=hidden_states,
+            pooled_embedding=pooled_embedding,
+            hidden_states=hidden_states if output_hidden_states else None,
+            embedding_pooling=self.config.embedding_pooling,
+        )
+
+
+# ===========================
+# Task-Specific Models
+# ===========================
+
+@register_model_for_auto_class("AutoModel")
+class GeneMambaForMaskedLM(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for masked language modeling (MLM).
+    Suitable for pretraining and domain adaptation.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.genemamba = GeneMambaModel(config)
+        
+        # Language modeling head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaMaskedLMOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Target token ids for MLM loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaMaskedLMOutput: Contains logits and optional loss.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        logits = self.lm_head(outputs.last_hidden_state)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        return GeneMambaMaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+        )
+
+
+@register_model_for_auto_class("AutoModelForSequenceClassification")
+class GeneMambaForSequenceClassification(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for sequence classification tasks.
+    Ideal for cell type annotation, tissue classification, etc.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        self.genemamba = GeneMambaModel(config)
+        
+        # Classification head
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaSequenceClassifierOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Class labels for classification loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaSequenceClassifierOutput: Contains logits, optional loss, and embedding.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        pooled_embedding = outputs.pooled_embedding
+        logits = self.classifier(self.dropout(pooled_embedding))
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        
+        return GeneMambaSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            pooled_embedding=pooled_embedding,
+        )
+
+
+# Register tokenizer class
+register_model_for_auto_class("AutoModelForMaskedLM")(GeneMambaForMaskedLM)

48l-512d/modeling_outputs.py

@@ -0,0 +1,81 @@
+"""
+Custom ModelOutput classes for GeneMamba.
+Defines the output structure for different GeneMamba tasks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class GeneMambaModelOutput(ModelOutput):
+    """
+    Base output class for GeneMamba models.
+    
+    Attributes:
+        last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)):
+            Sequence of hidden-states at the output of the last layer of the model.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size)):
+            Cell/sequence-level embedding (pooled representation) used for downstream tasks.
+            This is the recommended embedding to use for classification, clustering, etc.
+        
+        embedding_pooling (str):
+            The pooling method used to generate pooled_embedding.
+    """
+    
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: torch.FloatTensor = None
+    embedding_pooling: str = "mean"
+
+
+@dataclass
+class GeneMambaSequenceClassifierOutput(ModelOutput):
+    """
+    Output class for GeneMamba sequence classification models.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            Classification loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, num_labels)):
+            Classification scores (before softmax).
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size), optional):
+            Cell embedding before classification head.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class GeneMambaMaskedLMOutput(ModelOutput):
+    """
+    Output class for GeneMamba masked language modeling.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            MLM loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size)):
+            Prediction scores of the language modeling head.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None

48l-512d/special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]"
+}

48l-512d/tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

48l-768d/config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "genemamba",
+  "architectures": [
+    "GeneMambaModel"
+  ],
+  "vocab_size": 25426,
+  "max_position_embeddings": 2048,
+  "hidden_size": 768,
+  "num_hidden_layers": 48,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "initializer_range": 0.02,
+  "mamba_mode": "gate",
+  "embedding_pooling": "mean",
+  "num_labels": 2,
+  "pad_token_id": 1,
+  "eos_token_id": 2,
+  "bos_token_id": 0,
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "auto_map": {
+    "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+    "AutoModel": "modeling_genemamba.GeneMambaModel",
+    "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification"
+  }
+}

48l-768d/configuration_genemamba.py

@@ -0,0 +1,97 @@
+"""
+Configuration for GeneMamba model.
+Defines all hyperparameters and settings for the GeneMamba architecture.
+"""
+
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+class GeneMambaConfig(PretrainedConfig):
+    """
+    Configuration class for GeneMamba model.
+    
+    This class stores the configuration of a GeneMamba model, inheriting from PretrainedConfig.
+    It can be used to instantiate models from pretrained checkpoints or customize model initialization.
+    
+    Args:
+        vocab_size (int, optional, defaults to 25426):
+            Vocabulary size of the model. Number of gene tokens (Ensembl Gene IDs).
+        
+        hidden_size (int, optional, defaults to 512):
+            Dimensionality of the hidden/embedding layers (d_model in Mamba).
+        
+        num_hidden_layers (int, optional, defaults to 24):
+            Number of Mamba layers (mamba_layer).
+        
+        intermediate_size (int, optional, defaults to 2048):
+            Dimensionality of intermediate representations in MLP.
+        
+        max_position_embeddings (int, optional, defaults to 2048):
+            Maximum sequence length (seq_len).
+        
+        hidden_dropout_prob (float, optional, defaults to 0.1):
+            Dropout probability for hidden states.
+        
+        initializer_range (float, optional, defaults to 0.02):
+            Standard deviation of truncated normal initializer.
+        
+        mamba_mode (str, optional, defaults to "gate"):
+            Aggregation mode for bidirectional Mamba layers.
+            Options: "mean", "sum", "concat", "gate".
+        
+        embedding_pooling (str, optional, defaults to "mean"):
+            Method for pooling to get cell embedding.
+            Options: "CLS", "mean", "weighted".
+        
+        num_labels (int, optional, defaults to 2):
+            Number of labels for sequence classification tasks.
+        
+        pad_token_id (int, optional, defaults to 1):
+            Token ID for padding.
+        
+        bos_token_id (int, optional, defaults to None):
+            Token ID for beginning of sequence.
+        
+        eos_token_id (int, optional, defaults to None):
+            Token ID for end of sequence.
+    """
+    
+    model_type = "genemamba"
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+    
+    def __init__(
+        self,
+        vocab_size: int = 25426,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 2048,
+        max_position_embeddings: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        initializer_range: float = 0.02,
+        mamba_mode: str = "gate",
+        embedding_pooling: str = "mean",
+        num_labels: int = 2,
+        pad_token_id: int = 1,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.mamba_mode = mamba_mode
+        self.embedding_pooling = embedding_pooling
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id

48l-768d/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:728514a211350e69937d73398dffa4c6bbb7f59366fb6c8b39f27437a6a5af77
+size 860161160

48l-768d/modeling_genemamba.py

@@ -0,0 +1,395 @@
+"""
+PyTorch implementation of GeneMamba model for Hugging Face Transformers.
+Includes backbone model and task-specific heads for various downstream tasks.
+"""
+
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import normal_, constant_
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, ModelOutput
+from transformers.models.auto import register_model_for_auto_class
+
+from mamba_ssm import Mamba
+from mamba_ssm.ops.triton.layer_norm import RMSNorm
+
+from .configuration_genemamba import GeneMambaConfig
+from .modeling_outputs import GeneMambaModelOutput, GeneMambaSequenceClassifierOutput, GeneMambaMaskedLMOutput
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================
+# Core Architecture Components
+# ===========================
+
+class EncoderLayer(nn.Module):
+    """
+    Single Mamba encoder layer with residual connection.
+    Applies a Mamba2 or Mamba layer followed by addition with input.
+    
+    Args:
+        hidden_size (int): Dimension of hidden states.
+    """
+    
+    def __init__(self, hidden_size: int):
+        super(EncoderLayer, self).__init__()
+        self.mamba = Mamba(d_model=hidden_size, d_state=64, d_conv=4, expand=2)
+    
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+        
+        Returns:
+            torch.Tensor: Output after Mamba layer and residual connection.
+        """
+        output = self.mamba(X) + X
+        return output
+
+
+class MambaMixer(nn.Module):
+    """
+    Stack of Mamba encoder layers with bidirectional processing and aggregation.
+    Processes sequences in both forward and reverse directions, then aggregates.
+    
+    Args:
+        mode (str): Aggregation mode. Options: "mean", "sum", "concat", "gate".
+        hidden_size (int): Dimension of hidden states.
+        num_hidden_layers (int): Number of Mamba layers.
+    """
+    
+    def __init__(
+        self,
+        mode: str = "gate",
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24
+    ):
+        super(MambaMixer, self).__init__()
+        self.mode = mode
+        self.hidden_size = hidden_size
+        
+        # Create Mamba layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(hidden_size) for _ in range(num_hidden_layers)]
+        )
+        
+        # Aggregation modules for certain modes
+        if mode in ["concat", "gate"]:
+            self.aggr = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def flip_sequence(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Reverse a sequence based on actual length (ignoring padding).
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            mask (torch.Tensor, optional): Padding mask of shape (batch_size, seq_len).
+        
+        Returns:
+            torch.Tensor: Reversed tensor.
+        """
+        batch_size, seq_length, embedding_dim = X.size()
+        
+        if mask is None:
+            # Simple flip
+            return X.flip([1])
+        
+        # Flip based on actual sequence length (marked by mask)
+        lengths = (~mask).sum(dim=1)
+        pos_tensor = torch.arange(seq_length, device=X.device).unsqueeze(0).expand(batch_size, -1)
+        flip_mask = pos_tensor < lengths.unsqueeze(1)
+        reversed_positions = torch.where(
+            flip_mask,
+            lengths.unsqueeze(1) - 1 - pos_tensor,
+            pos_tensor
+        )
+        
+        X_reverse = torch.gather(X, 1, reversed_positions.unsqueeze(-1).expand(-1, -1, embedding_dim))
+        return X_reverse
+    
+    def forward(
+        self,
+        X: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Process sequence through bidirectional Mamba layers.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            padding_mask (torch.Tensor, optional): Padding mask.
+        
+        Returns:
+            torch.Tensor: Output after processing all layers and aggregation.
+        """
+        
+        for layer in self.layers:
+            # Flip sequence for reverse processing
+            X_flip = self.flip_sequence(X, padding_mask)
+            
+            # Forward and reverse passes
+            X_f = layer(X)
+            X_b = layer(X_flip)
+            
+            # Flip back the reverse output
+            X_b = self.flip_sequence(X_b, padding_mask)
+            
+            # Aggregate forward and reverse
+            if self.mode == "mean":
+                X = (X_f + X_b) / 2
+            elif self.mode == "sum":
+                X = X_f + X_b
+            elif self.mode == "concat":
+                X = torch.cat([X_f, X_b], dim=-1)
+                X = self.aggr(X)
+            elif self.mode == "gate":
+                z = torch.sigmoid(self.aggr(torch.cat([X_f, X_b], dim=-1)))
+                X = z * X_f + (1 - z) * X_b
+            else:
+                raise ValueError(f"Invalid aggregation mode: {self.mode}")
+        
+        return X
+
+
+# ===========================
+# Base Model Classes
+# ===========================
+
+class GeneMambaPreTrainedModel(PreTrainedModel):
+    """
+    Base class for all GeneMamba models.
+    Handles weight initialization and provides standard model interfaces.
+    """
+    
+    config_class = GeneMambaConfig
+    base_model_prefix = "genemamba"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize module weights."""
+        if isinstance(module, nn.Linear):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.bias is not None:
+                constant_(module.bias, 0.0)
+        elif isinstance(module, nn.Embedding):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            constant_(module.bias, 0.0)
+            constant_(module.weight, 1.0)
+
+
+class GeneMambaModel(GeneMambaPreTrainedModel):
+    """
+    GeneMamba backbone model - outputs cell embeddings and hidden states.
+    This is the core model used by task-specific heads.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Embedding layer
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Mamba layers with bidirectional aggregation
+        self.mamba_mixer = MambaMixer(
+            mode=config.mamba_mode,
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers
+        )
+        
+        # Final layer normalization
+        self.norm = RMSNorm(config.hidden_size)
+        
+        self.apply(self._init_weights)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return embedding layer."""
+        return self.embeddings
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        """Set embedding layer."""
+        self.embeddings = value
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaModelOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask of shape (batch_size, seq_len).
+            output_hidden_states (bool): Whether to output hidden states from all layers.
+        
+        Returns:
+            GeneMambaModelOutput: Contains last_hidden_state, pooled_embedding, etc.
+        """
+        # Get embeddings
+        hidden_states = self.embeddings(input_ids)
+        
+        # Pass through Mamba layers
+        hidden_states = self.mamba_mixer(hidden_states, attention_mask)
+        
+        # Apply final normalization
+        hidden_states = self.norm(hidden_states)
+        
+        # Compute pooled embedding (cell representation)
+        if self.config.embedding_pooling == "CLS":
+            # Use first token (CLS)
+            pooled_embedding = hidden_states[:, 0, :]
+        elif self.config.embedding_pooling == "mean":
+            # Mean pooling over sequence
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).expand(hidden_states.shape).float()
+                pooled_embedding = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
+            else:
+                pooled_embedding = hidden_states.mean(dim=1)
+        else:
+            raise ValueError(f"Unsupported embedding_pooling: {self.config.embedding_pooling}")
+        
+        return GeneMambaModelOutput(
+            last_hidden_state=hidden_states,
+            pooled_embedding=pooled_embedding,
+            hidden_states=hidden_states if output_hidden_states else None,
+            embedding_pooling=self.config.embedding_pooling,
+        )
+
+
+# ===========================
+# Task-Specific Models
+# ===========================
+
+@register_model_for_auto_class("AutoModel")
+class GeneMambaForMaskedLM(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for masked language modeling (MLM).
+    Suitable for pretraining and domain adaptation.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.genemamba = GeneMambaModel(config)
+        
+        # Language modeling head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaMaskedLMOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Target token ids for MLM loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaMaskedLMOutput: Contains logits and optional loss.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        logits = self.lm_head(outputs.last_hidden_state)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        return GeneMambaMaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+        )
+
+
+@register_model_for_auto_class("AutoModelForSequenceClassification")
+class GeneMambaForSequenceClassification(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for sequence classification tasks.
+    Ideal for cell type annotation, tissue classification, etc.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        self.genemamba = GeneMambaModel(config)
+        
+        # Classification head
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaSequenceClassifierOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Class labels for classification loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaSequenceClassifierOutput: Contains logits, optional loss, and embedding.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        pooled_embedding = outputs.pooled_embedding
+        logits = self.classifier(self.dropout(pooled_embedding))
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        
+        return GeneMambaSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            pooled_embedding=pooled_embedding,
+        )
+
+
+# Register tokenizer class
+register_model_for_auto_class("AutoModelForMaskedLM")(GeneMambaForMaskedLM)

48l-768d/modeling_outputs.py

@@ -0,0 +1,81 @@
+"""
+Custom ModelOutput classes for GeneMamba.
+Defines the output structure for different GeneMamba tasks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class GeneMambaModelOutput(ModelOutput):
+    """
+    Base output class for GeneMamba models.
+    
+    Attributes:
+        last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)):
+            Sequence of hidden-states at the output of the last layer of the model.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size)):
+            Cell/sequence-level embedding (pooled representation) used for downstream tasks.
+            This is the recommended embedding to use for classification, clustering, etc.
+        
+        embedding_pooling (str):
+            The pooling method used to generate pooled_embedding.
+    """
+    
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: torch.FloatTensor = None
+    embedding_pooling: str = "mean"
+
+
+@dataclass
+class GeneMambaSequenceClassifierOutput(ModelOutput):
+    """
+    Output class for GeneMamba sequence classification models.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            Classification loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, num_labels)):
+            Classification scores (before softmax).
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size), optional):
+            Cell embedding before classification head.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class GeneMambaMaskedLMOutput(ModelOutput):
+    """
+    Output class for GeneMamba masked language modeling.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            MLM loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size)):
+            Prediction scores of the language modeling head.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None

48l-768d/special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]"
+}

48l-768d/tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

README.md

@@ -0,0 +1,133 @@
+---
+library_name: transformers
+tags:
+  - genomics
+  - single-cell
+  - mamba
+  - biology
+pipeline_tag: feature-extraction
+---
+
+# GeneMamba
+
+This repository contains a **default GeneMamba model** plus full usage assets:
+- default model weights at repository root (**24l-512d**)
+- custom modeling/config files for `trust_remote_code=True`
+- preprocessing example from `h5ad` to `input_ids`
+- tokenizer assets and id mapping files
+
+Additional model sizes are provided as subfolders:
+- `24l-512d` (same architecture class as default)
+- `24l-768d`
+- `48l-512d`
+- `48l-768d`
+
+## 1) Input format (very important)
+
+GeneMamba input is **ranked gene token IDs** per cell:
+1. Start from one cell expression vector
+2. Keep genes with expression > 0
+3. Sort genes by expression descending
+4. Convert each gene ID (Ensembl, e.g. `ENSG00000000003`) to token ID
+5. Use resulting list as `input_ids`
+
+Each sample is one list of integers:
+
+```python
+{"input_ids": [145, 2088, 531, 91, ...]}
+```
+
+For batch input, shape is typically `(batch_size, seq_len)` after padding/truncation.
+
+## 2) Where tokenizer and id mapping come from
+
+- Main tokenizer used for model inference: `tokenizer.json`
+- Original full tokenizer table: `tokenizer_assets/gene_tokenizer.json`
+- Gene symbol -> token id mapping: `tokenizer_assets/symbol2id.pkl`
+- Token id -> gene symbol mapping: `tokenizer_assets/id2symbol.pkl`
+
+Special tokens:
+- `[UNK]` = 0
+- `[PAD]` = 1
+
+## 3) Preprocess your data
+
+See script:
+- `examples/00_preprocess_to_input_ids.py`
+
+Example:
+
+```bash
+python examples/00_preprocess_to_input_ids.py \
+  --h5ad /path/to/your_data.h5ad \
+  --tokenizer_json tokenizer.json \
+  --output_arrow ./my_data/sorted_gene_token_ids.arrow
+```
+
+This output Arrow file has one column: `input_ids`.
+
+## 4) Load model and extract embedding
+
+### Default load (24l-512d)
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained(
+  "mineself2016/GeneMamba",
+    trust_remote_code=True
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+  "mineself2016/GeneMamba",
+    trust_remote_code=True
+)
+```
+
+### Load other sizes (via `subfolder`)
+
+```python
+from transformers import AutoModel
+
+model_24l_768d = AutoModel.from_pretrained(
+  "mineself2016/GeneMamba",
+  subfolder="24l-768d",
+  trust_remote_code=True,
+)
+
+model_48l_512d = AutoModel.from_pretrained(
+  "mineself2016/GeneMamba",
+  subfolder="48l-512d",
+  trust_remote_code=True,
+)
+
+model_48l_768d = AutoModel.from_pretrained(
+  "mineself2016/GeneMamba",
+  subfolder="48l-768d",
+  trust_remote_code=True,
+)
+```
+
+More complete example:
+- `examples/01_extract_embeddings.py`
+
+## 6) Downstream task examples (added)
+
+See:
+- `examples/downstream/README.md`
+
+Included downstream tasks:
+- cell type annotation fine-tuning
+- zero-shot embedding + logistic regression
+- batch integration proxy evaluation
+- original legacy downstream scripts from `gene_mamba/analysis/cell_type_annotation`
+
+## 7) Source of preprocessing logic
+
+The preprocessing/tokenization pipeline is aligned with assets from:
+- `/project/zhiwei/cq5/PythonWorkSpace/gene_mamba`
+
+Key references used:
+- tokenizer: `gene_tokenizer.json`
+- mappings: `symbol2id.pkl`, `id2symbol.pkl`
+- dataset build logic (Arrow + `input_ids`): `utils.py` (`build_dataset`)

config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "genemamba",
+  "architectures": [
+    "GeneMambaModel"
+  ],
+  "vocab_size": 25426,
+  "max_position_embeddings": 2048,
+  "hidden_size": 512,
+  "num_hidden_layers": 24,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "initializer_range": 0.02,
+  "mamba_mode": "gate",
+  "embedding_pooling": "mean",
+  "num_labels": 2,
+  "pad_token_id": 1,
+  "eos_token_id": 2,
+  "bos_token_id": 0,
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "auto_map": {
+    "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+    "AutoModel": "modeling_genemamba.GeneMambaModel",
+    "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification"
+  }
+}

configuration_genemamba.py

@@ -0,0 +1,97 @@
+"""
+Configuration for GeneMamba model.
+Defines all hyperparameters and settings for the GeneMamba architecture.
+"""
+
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+class GeneMambaConfig(PretrainedConfig):
+    """
+    Configuration class for GeneMamba model.
+    
+    This class stores the configuration of a GeneMamba model, inheriting from PretrainedConfig.
+    It can be used to instantiate models from pretrained checkpoints or customize model initialization.
+    
+    Args:
+        vocab_size (int, optional, defaults to 25426):
+            Vocabulary size of the model. Number of gene tokens (Ensembl Gene IDs).
+        
+        hidden_size (int, optional, defaults to 512):
+            Dimensionality of the hidden/embedding layers (d_model in Mamba).
+        
+        num_hidden_layers (int, optional, defaults to 24):
+            Number of Mamba layers (mamba_layer).
+        
+        intermediate_size (int, optional, defaults to 2048):
+            Dimensionality of intermediate representations in MLP.
+        
+        max_position_embeddings (int, optional, defaults to 2048):
+            Maximum sequence length (seq_len).
+        
+        hidden_dropout_prob (float, optional, defaults to 0.1):
+            Dropout probability for hidden states.
+        
+        initializer_range (float, optional, defaults to 0.02):
+            Standard deviation of truncated normal initializer.
+        
+        mamba_mode (str, optional, defaults to "gate"):
+            Aggregation mode for bidirectional Mamba layers.
+            Options: "mean", "sum", "concat", "gate".
+        
+        embedding_pooling (str, optional, defaults to "mean"):
+            Method for pooling to get cell embedding.
+            Options: "CLS", "mean", "weighted".
+        
+        num_labels (int, optional, defaults to 2):
+            Number of labels for sequence classification tasks.
+        
+        pad_token_id (int, optional, defaults to 1):
+            Token ID for padding.
+        
+        bos_token_id (int, optional, defaults to None):
+            Token ID for beginning of sequence.
+        
+        eos_token_id (int, optional, defaults to None):
+            Token ID for end of sequence.
+    """
+    
+    model_type = "genemamba"
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+    
+    def __init__(
+        self,
+        vocab_size: int = 25426,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 2048,
+        max_position_embeddings: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        initializer_range: float = 0.02,
+        mamba_mode: str = "gate",
+        embedding_pooling: str = "mean",
+        num_labels: int = 2,
+        pad_token_id: int = 1,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.mamba_mode = mamba_mode
+        self.embedding_pooling = embedding_pooling
+        self.num_labels = num_labels
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id

examples/00_preprocess_to_input_ids.py

@@ -0,0 +1,75 @@
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import scanpy as sc
+import pyarrow as pa
+
+
+def load_vocab(tokenizer_json_path: str):
+    with open(tokenizer_json_path, "r") as f:
+        tokenizer = json.load(f)
+    vocab = tokenizer["model"]["vocab"]
+    pad_id = vocab.get("[PAD]", 1)
+    unk_id = vocab.get("[UNK]", 0)
+    return vocab, pad_id, unk_id
+
+
+def ranked_gene_ids_for_cell(expr_values, gene_names, vocab):
+    nonzero_idx = np.where(expr_values > 0)[0]
+    if len(nonzero_idx) == 0:
+        return []
+
+    genes = np.array(gene_names)[nonzero_idx]
+    values = expr_values[nonzero_idx]
+
+    order = np.argsort(-values)
+    ranked_genes = genes[order]
+
+    token_ids = [vocab[g] for g in ranked_genes if g in vocab]
+    return token_ids
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert h5ad to GeneMamba input_ids (Arrow)")
+    parser.add_argument("--h5ad", required=True, help="Input h5ad file")
+    parser.add_argument("--tokenizer_json", required=True, help="Path to tokenizer.json or gene_tokenizer.json")
+    parser.add_argument("--output_arrow", required=True, help="Output arrow file path")
+    parser.add_argument("--max_cells", type=int, default=None, help="Optional: process first N cells only")
+    args = parser.parse_args()
+
+    adata = sc.read_h5ad(args.h5ad)
+    vocab, _, _ = load_vocab(args.tokenizer_json)
+
+    gene_names = list(adata.var_names)
+    n_cells = adata.n_obs if args.max_cells is None else min(args.max_cells, adata.n_obs)
+
+    rows = []
+    X = adata.X
+
+    for i in range(n_cells):
+        row = X[i]
+        if hasattr(row, "toarray"):
+            expr = row.toarray().ravel()
+        else:
+            expr = np.asarray(row).ravel()
+
+        token_ids = ranked_gene_ids_for_cell(expr, gene_names, vocab)
+        rows.append(token_ids)
+
+    df = pd.DataFrame({"input_ids": rows})
+    table = pa.Table.from_pandas(df)
+
+    output_path = Path(args.output_arrow)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with pa.OSFile(str(output_path), "wb") as sink:
+        with pa.ipc.new_stream(sink, table.schema) as writer:
+            writer.write_table(table)
+
+    print(f"Saved {len(rows)} cells to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/01_extract_embeddings.py

@@ -0,0 +1,150 @@
+"""
+Phase 1: Extract Cell Embeddings
+Demonstrates how to load GeneMamba and extract cell embeddings for downstream analysis.
+
+Usage:
+    python examples/1_extract_embeddings.py
+"""
+
+import torch
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+
+
+def main():
+    print("=" * 80)
+    print("GeneMamba Phase 1: Extract Cell Embeddings")
+    print("=" * 80)
+    
+    # ============================================================
+    # Step 1: Load pretrained model and tokenizer
+    # ============================================================
+    print("\n[Step 1] Loading model and tokenizer...")
+    
+    # For this example, we use a local model path
+    # In practice, you would use: "username/GeneMamba-24l-512d"
+    model_name = "GeneMamba-24l-512d"  # Change to HF Hub path when available
+    
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            local_files_only=True  # Try local first
+        )
+        model = AutoModel.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            local_files_only=True
+        )
+    except Exception as e:
+        print(f"Note: Could not load from '{model_name}': {e}")
+        print("Using mock data for demonstration...")
+        
+        # For demonstration without actual checkpoint
+        from configuration_genemamba import GeneMambaConfig
+        from modeling_genemamba import GeneMambaModel
+        
+        config = GeneMambaConfig(
+            vocab_size=25426,
+            hidden_size=512,
+            num_hidden_layers=24,
+            embedding_pooling="mean",
+        )
+        model = GeneMambaModel(config)
+        tokenizer = None
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    model.eval()
+    
+    print(f"✓ Model loaded on device: {device}")
+    print(f"✓ Model config: hidden_size={model.config.hidden_size}, "
+          f"num_layers={model.config.num_hidden_layers}")
+    
+    # ============================================================
+    # Step 2: Prepare simulated single-cell data
+    # ============================================================
+    print("\n[Step 2] Preparing sample data...")
+    
+    batch_size = 8
+    seq_len = 2048
+    vocab_size = 25426
+    
+    # Simulate ranked gene sequences
+    # In practice, this would come from your scRNA-seq data
+    # Genes should be ranked by expression (highest first)
+    input_ids = torch.randint(2, vocab_size, (batch_size, seq_len)).to(device)
+    
+    print(f"✓ Created sample input:")
+    print(f"  - Batch size: {batch_size}")
+    print(f"  - Sequence length: {seq_len}")
+    print(f"  - Input shape: {input_ids.shape}")
+    
+    # ============================================================
+    # Step 3: Inference - Extract embeddings
+    # ============================================================
+    print("\n[Step 3] Extracting cell embeddings...")
+    
+    with torch.no_grad():
+        outputs = model(input_ids, output_hidden_states=False)
+    
+    # Get the pooled embedding (cell representation)
+    cell_embeddings = outputs.pooled_embedding
+    
+    print(f"✓ Extraction complete!")
+    print(f"  - Cell embeddings shape: {cell_embeddings.shape}")
+    print(f"  - Pooling method used: {outputs.embedding_pooling}")
+    print(f"  - Embedding type: {cell_embeddings.dtype}")
+    
+    # ============================================================
+    # Step 4: Example downstream analyses
+    # ============================================================
+    print("\n[Step 4] Example downstream uses...")
+    
+    # Example 1: Clustering (KMeans)
+    from sklearn.cluster import KMeans
+    n_clusters = 3
+    kmeans = KMeans(n_clusters=n_clusters, n_init=10)
+    clusters = kmeans.fit_predict(cell_embeddings.cpu().numpy())
+    print(f"✓ Clustering: Assigned {len(np.unique(clusters))} clusters")
+    
+    # Example 2: Dimensionality reduction (PCA)
+    from sklearn.decomposition import PCA
+    pca = PCA(n_components=2)
+    embedding_2d = pca.fit_transform(cell_embeddings.cpu().numpy())
+    print(f"✓ PCA reduction: {cell_embeddings.shape} → {embedding_2d.shape}")
+    
+    # Example 3: Similarity search
+    # Find the most similar cell to the first cell
+    similarities = torch.nn.functional.cosine_similarity(
+        cell_embeddings[0:1],
+        cell_embeddings
+    )
+    most_similar_idx = torch.argmax(similarities).item()
+    print(f"✓ Similarity search: Most similar cell to cell 0 is cell {most_similar_idx} "
+          f"(similarity: {similarities[most_similar_idx]:.4f})")
+    
+    # Example 4: Statistics
+    print("\n[Step 5] Embedding statistics:")
+    print(f"  - Mean: {cell_embeddings.mean(dim=0).norm():.4f}")
+    print(f"  - Std: {cell_embeddings.std(dim=0).mean():.4f}")
+    print(f"  - Min: {cell_embeddings.min():.4f}")
+    print(f"  - Max: {cell_embeddings.max():.4f}")
+    
+    # ============================================================
+    # Step 6: Save embeddings (optional)
+    # ============================================================
+    print("\n[Step 6] Saving embeddings...")
+    
+    np.save("cell_embeddings.npy", cell_embeddings.cpu().numpy())
+    print("✓ Embeddings saved to 'cell_embeddings.npy'")
+    
+    print("\n" + "=" * 80)
+    print("Phase 1 Complete!")
+    print("=" * 80)
+    
+    return model, cell_embeddings
+
+
+if __name__ == "__main__":
+    model, embeddings = main()

examples/downstream/10_finetune_classification.py

@@ -0,0 +1,248 @@
+"""
+Phase 2: Downstream Task - Fine-tune for Classification
+Demonstrates cell type annotation and other sequence classification tasks.
+
+Usage:
+    python examples/2_finetune_classification.py
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments
+
+
+class GeneExpressionDataset(Dataset):
+    """
+    Simple dataset for gene expression classification.
+    In practice, this would load from h5ad or other single-cell formats.
+    """
+    
+    def __init__(self, input_ids, labels, max_length=2048):
+        self.input_ids = input_ids
+        self.labels = labels
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.input_ids)
+    
+    def __getitem__(self, idx):
+        input_id = self.input_ids[idx]
+        label = self.labels[idx]
+        
+        return {
+            "input_ids": torch.tensor(input_id, dtype=torch.long),
+            "labels": torch.tensor(label, dtype=torch.long),
+        }
+
+
+def create_mock_data(n_samples=1000, n_features=2048, n_classes=5):
+    """Create mock single-cell data for demonstration."""
+    
+    print("Creating mock dataset...")
+    
+    # Create random ranked gene sequences
+    input_ids = np.random.randint(2, 25426, (n_samples, n_features))
+    
+    # Create random labels (e.g., cell types)
+    labels = np.random.randint(0, n_classes, n_samples)
+    
+    # Split into train/val/test
+    train_size = int(0.7 * n_samples)
+    val_size = int(0.15 * n_samples)
+    
+    train_ids = input_ids[:train_size]
+    train_labels = labels[:train_size]
+    
+    val_ids = input_ids[train_size:train_size + val_size]
+    val_labels = labels[train_size:train_size + val_size]
+    
+    test_ids = input_ids[train_size + val_size:]
+    test_labels = labels[train_size + val_size:]
+    
+    print(f"✓ Dataset created:")
+    print(f"  - Train: {len(train_ids)} samples")
+    print(f"  - Val: {len(val_ids)} samples")
+    print(f"  - Test: {len(test_ids)} samples")
+    print(f"  - Classes: {n_classes}")
+    
+    return (
+        GeneExpressionDataset(train_ids, train_labels),
+        GeneExpressionDataset(val_ids, val_labels),
+        GeneExpressionDataset(test_ids, test_labels),
+    )
+
+
+def main():
+    print("=" * 80)
+    print("GeneMamba Phase 2: Downstream Classification")
+    print("=" * 80)
+    
+    # ============================================================
+    # Step 1: Load pretrained model with classification head
+    # ============================================================
+    print("\n[Step 1] Loading pretrained model with classification head...")
+    
+    num_classes = 5
+    
+    try:
+        model = AutoModelForSequenceClassification.from_pretrained(
+            "GeneMamba-24l-512d",
+            num_labels=num_classes,
+            trust_remote_code=True,
+            local_files_only=True,
+        )
+    except Exception as e:
+        print(f"Note: Could not load from hub ({e})")
+        print("Using local initialization...")
+        
+        # Initialize locally
+        from configuration_genemamba import GeneMambaConfig
+        from modeling_genemamba import GeneMambaForSequenceClassification
+        
+        config = GeneMambaConfig(
+            vocab_size=25426,
+            hidden_size=512,
+            num_hidden_layers=24,
+            num_labels=num_classes,
+        )
+        model = GeneMambaForSequenceClassification(config)
+    
+    print(f"✓ Model loaded")
+    print(f"  - Classification head: input={model.config.hidden_size} → output={num_classes}")
+    
+    # ============================================================
+    # Step 2: Prepare data
+    # ============================================================
+    print("\n[Step 2] Preparing dataset...")
+    
+    train_dataset, val_dataset, test_dataset = create_mock_data(
+        n_samples=1000,
+        n_features=2048,
+        n_classes=num_classes,
+    )
+    
+    # ============================================================
+    # Step 3: Set up training arguments
+    # ============================================================
+    print("\n[Step 3] Setting up training...")
+    
+    output_dir = "./classification_results"
+    
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=3,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=100,
+        logging_steps=50,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="accuracy",
+        report_to="none",  # Disable W&B logging
+        seed=42,
+    )
+    
+    print(f"✓ Training config:")
+    print(f"  - Output dir: {output_dir}")
+    print(f"  - Epochs: {training_args.num_train_epochs}")
+    print(f"  - Batch size: {training_args.per_device_train_batch_size}")
+    print(f"  - Learning rate: {training_args.learning_rate}")
+    
+    # ============================================================
+    # Step 4: Train using Trainer
+    # ============================================================
+    print("\n[Step 4] Training model...")
+    
+    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+    
+    def compute_metrics(eval_pred):
+        """Compute evaluation metrics."""
+        predictions, labels = eval_pred
+        predictions = np.argmax(predictions, axis=1)
+        
+        return {
+            "accuracy": accuracy_score(labels, predictions),
+            "f1": f1_score(labels, predictions, average="weighted", zero_division=0),
+            "precision": precision_score(labels, predictions, average="weighted", zero_division=0),
+            "recall": recall_score(labels, predictions, average="weighted", zero_division=0),
+        }
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        compute_metrics=compute_metrics,
+    )
+    
+    train_result = trainer.train()
+    
+    print(f"✓ Training complete!")
+    print(f"  - Final training loss: {train_result.training_loss:.4f}")
+    
+    # ============================================================
+    # Step 5: Evaluate on test set
+    # ============================================================
+    print("\n[Step 5] Evaluating on test set...")
+    
+    test_results = trainer.evaluate(test_dataset)
+    
+    print(f"✓ Test Results:")
+    for metric, value in test_results.items():
+        if isinstance(value, float):
+            print(f"  - {metric}: {value:.4f}")
+    
+    # ============================================================
+    # Step 6: Make predictions
+    # ============================================================
+    print("\n[Step 6] Making predictions...")
+    
+    predictions = trainer.predict(test_dataset)
+    predicted_classes = np.argmax(predictions.predictions, axis=1)
+    
+    print(f"✓ Predictions made:")
+    print(f"  - Predicted classes: {len(predicted_classes)} samples")
+    print(f"  - Class distribution: {np.bincount(predicted_classes)}")
+    
+    # ============================================================
+    # Step 7: Save model
+    # ============================================================
+    print("\n[Step 7] Saving model...")
+    
+    save_dir = "./my_genemamba_classifier"
+    model.save_pretrained(save_dir)
+    print(f"✓ Model saved to '{save_dir}'")
+    
+    # ============================================================
+    # Step 8: Load and test saved model
+    # ============================================================
+    print("\n[Step 8] Testing model reloading...")
+    
+    loaded_model = AutoModelForSequenceClassification.from_pretrained(
+        save_dir,
+        trust_remote_code=True,
+    )
+    loaded_model.eval()
+    
+    # Test on a single batch
+    with torch.no_grad():
+        sample_input = torch.randint(2, 25426, (1, 2048))
+        output = loaded_model(sample_input)
+        logits = output.logits
+        prediction = torch.argmax(logits, dim=1)
+    
+    print(f"✓ Loaded model test prediction: class {prediction.item()}")
+    
+    print("\n" + "=" * 80)
+    print("Phase 2 Complete! Model ready for deployment.")
+    print("=" * 80)
+    
+    return model, trainer
+
+
+if __name__ == "__main__":
+    model, trainer = main()

examples/downstream/11_zero_shot_logreg.py

@@ -0,0 +1,98 @@
+"""
+Zero-shot downstream baseline:
+1) Extract frozen GeneMamba embeddings
+2) Train LogisticRegression on train split
+3) Evaluate on test split
+
+Expected h5ad columns:
+- obs['celltype']
+- obs['partition'] with values in {'train', 'test'}
+"""
+
+import argparse
+import numpy as np
+import scanpy as sc
+import torch
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+from sklearn.preprocessing import LabelEncoder
+from transformers import AutoModel
+
+
+def build_ranked_input_ids(adata, symbol2id, seq_len=2048, pad_id=1):
+    gene_names = np.array(adata.var_names)
+    X = adata.X
+    out = np.full((adata.n_obs, seq_len), pad_id, dtype=np.int64)
+
+    for i in range(adata.n_obs):
+        row = X[i]
+        if hasattr(row, "toarray"):
+            expr = row.toarray().ravel()
+        else:
+            expr = np.asarray(row).ravel()
+
+        nz = np.where(expr > 0)[0]
+        if len(nz) == 0:
+            continue
+
+        genes = gene_names[nz]
+        vals = expr[nz]
+        order = np.argsort(-vals)
+        ranked_genes = genes[order]
+        ids = [symbol2id[g] for g in ranked_genes if g in symbol2id][:seq_len]
+        out[i, : len(ids)] = ids
+
+    return out
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", required=True)
+    parser.add_argument("--h5ad", required=True)
+    parser.add_argument("--symbol2id_npy", default=None, help="Optional .npy dumped dict path")
+    parser.add_argument("--seq_len", type=int, default=2048)
+    parser.add_argument("--batch_size", type=int, default=64)
+    args = parser.parse_args()
+
+    adata = sc.read_h5ad(args.h5ad)
+    assert "celltype" in adata.obs, "h5ad must include obs['celltype']"
+    assert "partition" in adata.obs, "h5ad must include obs['partition']"
+
+    if args.symbol2id_npy is None:
+        raise ValueError("Please provide --symbol2id_npy (dict saved by np.save(..., allow_pickle=True))")
+
+    symbol2id = np.load(args.symbol2id_npy, allow_pickle=True).item()
+
+    input_ids = build_ranked_input_ids(adata, symbol2id, seq_len=args.seq_len)
+    labels = LabelEncoder().fit_transform(adata.obs["celltype"].values)
+
+    model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True)
+    model.eval().cuda()
+
+    embeds = []
+    with torch.no_grad():
+        for s in range(0, input_ids.shape[0], args.batch_size):
+            batch = torch.tensor(input_ids[s : s + args.batch_size], dtype=torch.long, device="cuda")
+            out = model(batch)
+            embeds.append(out.pooled_embedding.detach().cpu().numpy())
+    embeds = np.concatenate(embeds, axis=0)
+
+    train_mask = adata.obs["partition"].values == "train"
+    test_mask = adata.obs["partition"].values == "test"
+
+    X_train, y_train = embeds[train_mask], labels[train_mask]
+    X_test, y_test = embeds[test_mask], labels[test_mask]
+
+    clf = LogisticRegression(max_iter=2000)
+    clf.fit(X_train, y_train)
+    pred = clf.predict(X_test)
+
+    print("accuracy:", accuracy_score(y_test, pred))
+    print("micro_f1:", f1_score(y_test, pred, average="micro"))
+    print("macro_f1:", f1_score(y_test, pred, average="macro"))
+    print("precision_weighted:", precision_score(y_test, pred, average="weighted", zero_division=0))
+    print("recall_weighted:", recall_score(y_test, pred, average="weighted", zero_division=0))
+
+
+if __name__ == "__main__":
+    main()

examples/downstream/12_batch_integration_eval.py

@@ -0,0 +1,79 @@
+"""
+Batch integration downstream example:
+- Extract embeddings with frozen GeneMamba
+- Evaluate simple batch mixing score proxy (silhouette by batch)
+
+Expected h5ad columns:
+- obs['batch']
+"""
+
+import argparse
+import numpy as np
+import scanpy as sc
+import torch
+from sklearn.metrics import silhouette_score
+from sklearn.preprocessing import LabelEncoder
+from transformers import AutoModel
+
+
+def build_ranked_input_ids(adata, symbol2id, seq_len=2048, pad_id=1):
+    gene_names = np.array(adata.var_names)
+    X = adata.X
+    out = np.full((adata.n_obs, seq_len), pad_id, dtype=np.int64)
+
+    for i in range(adata.n_obs):
+        row = X[i]
+        if hasattr(row, "toarray"):
+            expr = row.toarray().ravel()
+        else:
+            expr = np.asarray(row).ravel()
+
+        nz = np.where(expr > 0)[0]
+        if len(nz) == 0:
+            continue
+
+        genes = gene_names[nz]
+        vals = expr[nz]
+        order = np.argsort(-vals)
+        ranked_genes = genes[order]
+        ids = [symbol2id[g] for g in ranked_genes if g in symbol2id][:seq_len]
+        out[i, : len(ids)] = ids
+
+    return out
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", required=True)
+    parser.add_argument("--h5ad", required=True)
+    parser.add_argument("--symbol2id_npy", required=True)
+    parser.add_argument("--seq_len", type=int, default=2048)
+    parser.add_argument("--batch_size", type=int, default=64)
+    args = parser.parse_args()
+
+    adata = sc.read_h5ad(args.h5ad)
+    assert "batch" in adata.obs, "h5ad must include obs['batch']"
+
+    symbol2id = np.load(args.symbol2id_npy, allow_pickle=True).item()
+    input_ids = build_ranked_input_ids(adata, symbol2id, seq_len=args.seq_len)
+
+    model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True)
+    model.eval().cuda()
+
+    embeds = []
+    with torch.no_grad():
+        for s in range(0, input_ids.shape[0], args.batch_size):
+            batch = torch.tensor(input_ids[s : s + args.batch_size], dtype=torch.long, device="cuda")
+            out = model(batch)
+            embeds.append(out.pooled_embedding.detach().cpu().numpy())
+    embeds = np.concatenate(embeds, axis=0)
+
+    batch_labels = LabelEncoder().fit_transform(adata.obs["batch"].values)
+    score = silhouette_score(embeds, batch_labels, metric="euclidean")
+
+    print("silhouette_by_batch:", score)
+    print("(Closer to 0 typically indicates better batch mixing than very high positive values.)")
+
+
+if __name__ == "__main__":
+    main()

examples/downstream/20_continue_pretraining_reference.py

@@ -0,0 +1,265 @@
+"""
+Phase 3: Continue Pretraining
+Demonstrates how to continue pretraining GeneMamba on your own data using masked LM objective.
+
+Usage:
+    python examples/3_continue_pretraining.py
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+)
+
+
+class PretrainingDataset(Dataset):
+    """
+    Dataset for pretraining/continued pretraining.
+    Loads sequences and their lengths.
+    """
+    
+    def __init__(self, input_ids_list, max_length=2048):
+        self.input_ids_list = input_ids_list
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.input_ids_list)
+    
+    def __getitem__(self, idx):
+        input_ids = self.input_ids_list[idx]
+        
+        # Pad or truncate to max_length
+        if len(input_ids) >= self.max_length:
+            input_ids = input_ids[:self.max_length]
+        else:
+            input_ids = np.pad(
+                input_ids,
+                (0, self.max_length - len(input_ids)),
+                constant_values=1  # Pad token ID
+            )
+        
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+        }
+
+
+def create_mock_pretraining_data(n_sequences=5000, seq_len=2048):
+    """Create mock single-cell sequences for pretraining."""
+    
+    print("Creating mock pretraining dataset...")
+    
+    # Create ranked gene sequences
+    # In practice, these would come from your scRNA-seq data
+    sequences = []
+    for _ in range(n_sequences):
+        # Random ranked sequence
+        seq = np.random.randint(2, 25426, seq_len)
+        sequences.append(seq)
+    
+    print(f"✓ Created {n_sequences} sequences of length {seq_len}")
+    
+    return sequences
+
+
+def main():
+    print("=" * 80)
+    print("GeneMamba Phase 3: Continue Pretraining")
+    print("=" * 80)
+    
+    # ============================================================
+    # Step 1: Load pretrained model for masked LM
+    # ============================================================
+    print("\n[Step 1] Loading model for masked LM...")
+    
+    try:
+        model = AutoModelForMaskedLM.from_pretrained(
+            "GeneMamba-24l-512d",
+            trust_remote_code=True,
+            local_files_only=True,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            "GeneMamba-24l-512d",
+            trust_remote_code=True,
+            local_files_only=True,
+        )
+    except Exception as e:
+        print(f"Note: Could not load from hub ({e})")
+        print("Using local initialization...")
+        
+        # Initialize locally
+        from configuration_genemamba import GeneMambaConfig
+        from modeling_genemamba import GeneMambaForMaskedLM
+        
+        config = GeneMambaConfig(
+            vocab_size=25426,
+            hidden_size=512,
+            num_hidden_layers=24,
+        )
+        model = GeneMambaForMaskedLM(config)
+        tokenizer = None
+    
+    print(f"✓ Model loaded")
+    print(f"  - Architecture: {model.config.num_hidden_layers} layers, "
+          f"hidden_size={model.config.hidden_size}")
+    
+    # ============================================================
+    # Step 2: Prepare pretraining data
+    # ============================================================
+    print("\n[Step 2] Preparing pretraining dataset...")
+    
+    sequences = create_mock_pretraining_data(n_sequences=5000, seq_len=2048)
+    
+    # Split train/eval
+    train_size = int(0.9 * len(sequences))
+    train_sequences = sequences[:train_size]
+    eval_sequences = sequences[train_size:]
+    
+    train_dataset = PretrainingDataset(train_sequences)
+    eval_dataset = PretrainingDataset(eval_sequences)
+    
+    print(f"✓ Datasets created:")
+    print(f"  - Training: {len(train_dataset)} samples")
+    print(f"  - Evaluation: {len(eval_dataset)} samples")
+    
+    # ============================================================
+    # Step 3: Set up data collator for MLM
+    # ============================================================
+    print("\n[Step 3] Setting up data collator...")
+    
+    if tokenizer is not None:
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=True,
+            mlm_probability=0.15,  # Mask 15% of tokens
+        )
+    else:
+        # Custom collator if no tokenizer available
+        class CustomDataCollator:
+            def __call__(self, batch):
+                input_ids = torch.stack([item["input_ids"] for item in batch])
+                
+                # Create masked labels (for MLM loss)
+                labels = input_ids.clone()
+                mask = torch.rand(input_ids.shape) < 0.15
+                
+                # Set input to [MASK] token (id=0)
+                input_ids[mask] = 0
+                
+                # Set labels to -100 where not masked (loss ignores these)
+                labels[~mask] = -100
+                
+                return {"input_ids": input_ids, "labels": labels}
+        
+        data_collator = CustomDataCollator()
+    
+    print(f"✓ Data collator ready (MLM probability: 0.15)")
+    
+    # ============================================================
+    # Step 4: Set up training arguments
+    # ============================================================
+    print("\n[Step 4] Setting up training...")
+    
+    output_dir = "./pretrain_results"
+    
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=2,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=500,
+        logging_steps=100,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        report_to="none",  # Disable W&B
+        seed=42,
+    )
+    
+    print(f"✓ Training config:")
+    print(f"  - Output dir: {output_dir}")
+    print(f"  - Epochs: {training_args.num_train_epochs}")
+    print(f"  - Batch size: {training_args.per_device_train_batch_size}")
+    print(f"  - Learning rate: {training_args.learning_rate}")
+    print(f"  - MLM masking: 15%")
+    
+    # ============================================================
+    # Step 5: Train
+    # ============================================================
+    print("\n[Step 5] Starting continued pretraining...")
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=data_collator,
+    )
+    
+    train_result = trainer.train()
+    
+    print(f"✓ Training complete!")
+    print(f"  - Final training loss: {train_result.training_loss:.4f}")
+    
+    # ============================================================
+    # Step 6: Evaluate
+    # ============================================================
+    print("\n[Step 6] Evaluating on held-out set...")
+    
+    eval_results = trainer.evaluate()
+    
+    print(f"✓ Evaluation Results:")
+    for metric, value in eval_results.items():
+        if isinstance(value, (int, float)):
+            print(f"  - {metric}: {value:.4f}")
+    
+    # ============================================================
+    # Step 7: Save model
+    # ============================================================
+    print("\n[Step 7] Saving continued pretrained model...")
+    
+    save_dir = "./genemamba_continued_pretrain"
+    model.save_pretrained(save_dir)
+    if tokenizer is not None:
+        tokenizer.save_pretrained(save_dir)
+    
+    print(f"✓ Model saved to '{save_dir}'")
+    
+    # ============================================================
+    # Step 8: Test model inference
+    # ============================================================
+    print("\n[Step 8] Testing inference on masked input...")
+    
+    model.eval()
+    
+    # Create sample input with masked tokens
+    sample_input = torch.randint(2, 25426, (1, 2048))
+    sample_input[0, :10] = 0  # Mask first 10 tokens
+    
+    with torch.no_grad():
+        outputs = model(sample_input)
+        logits = outputs.logits
+        predictions = torch.argmax(logits, dim=-1)
+    
+    print(f"✓ Sample predictions generated")
+    print(f"  - Input shape: {sample_input.shape}")
+    print(f"  - Output logits shape: {logits.shape}")
+    print(f"  - Top predicted genes (tokens): {predictions[0, :10].tolist()}")
+    
+    print("\n" + "=" * 80)
+    print("Phase 3 Complete! Model ready for downstream tasks or further training.")
+    print("=" * 80)
+    
+    return model, trainer
+
+
+if __name__ == "__main__":
+    model, trainer = main()

examples/downstream/21_pretrain_from_scratch_reference.py

@@ -0,0 +1,280 @@
+"""
+Phase 4: Train from Scratch
+Demonstrates how to initialize and train a GeneMamba model from scratch.
+
+Usage:
+    python examples/4_pretrain_from_scratch.py
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import (
+    AutoConfig,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+)
+
+
+class PretrainingDataset(Dataset):
+    """Dataset for pretraining."""
+    
+    def __init__(self, input_ids_list, max_length=2048):
+        self.input_ids_list = input_ids_list
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.input_ids_list)
+    
+    def __getitem__(self, idx):
+        input_ids = self.input_ids_list[idx]
+        
+        # Pad or truncate
+        if len(input_ids) >= self.max_length:
+            input_ids = input_ids[:self.max_length]
+        else:
+            input_ids = np.pad(
+                input_ids,
+                (0, self.max_length - len(input_ids)),
+                constant_values=1
+            )
+        
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+        }
+
+
+def create_mock_pretraining_data(n_sequences=5000, seq_len=2048):
+    """Create mock pretraining data."""
+    
+    print("Creating mock pretraining dataset for from-scratch training...")
+    
+    sequences = []
+    for _ in range(n_sequences):
+        seq = np.random.randint(2, 25426, seq_len)
+        sequences.append(seq)
+    
+    print(f"✓ Created {n_sequences} sequences")
+    
+    return sequences
+
+
+def main():
+    print("=" * 80)
+    print("GeneMamba Phase 4: Train from Scratch")
+    print("=" * 80)
+    
+    # ============================================================
+    # Step 1: Create config from scratch
+    # ============================================================
+    print("\n[Step 1] Creating model configuration...")
+    
+    from configuration_genemamba import GeneMambaConfig
+    from modeling_genemamba import GeneMambaForMaskedLM
+    
+    config = GeneMambaConfig(
+        vocab_size=25426,
+        hidden_size=256,  # Smaller for faster demo
+        num_hidden_layers=12,  # Reduced for demo
+        intermediate_size=1024,
+        max_position_embeddings=2048,
+        mamba_mode="gate",
+        embedding_pooling="mean",
+        num_labels=2,
+        hidden_dropout_prob=0.1,
+        initializer_range=0.02,
+    )
+    
+    print(f"✓ Config created:")
+    print(f"  - Model type: {config.model_type}")
+    print(f"  - Hidden size: {config.hidden_size}")
+    print(f"  - Num layers: {config.num_hidden_layers}")
+    print(f"  - Vocab size: {config.vocab_size}")
+    print(f"  - Mode: {config.mamba_mode}")
+    
+    # ============================================================
+    # Step 2: Initialize model from config
+    # ============================================================
+    print("\n[Step 2] Initializing model from config...")
+    
+    model = GeneMambaForMaskedLM(config)
+    
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+    print(f"✓ Model initialized:")
+    print(f"  - Total parameters: {total_params / 1e6:.2f}M")
+    print(f"  - Trainable parameters: {trainable_params / 1e6:.2f}M")
+    
+    # ============================================================
+    # Step 3: Prepare data
+    # ============================================================
+    print("\n[Step 3] Preparing training data...")
+    
+    sequences = create_mock_pretraining_data(n_sequences=5000, seq_len=2048)
+    
+    # Split
+    train_size = int(0.8 * len(sequences))
+    train_sequences = sequences[:train_size]
+    eval_sequences = sequences[train_size:]
+    
+    train_dataset = PretrainingDataset(train_sequences)
+    eval_dataset = PretrainingDataset(eval_sequences)
+    
+    print(f"✓ Datasets created:")
+    print(f"  - Train: {len(train_dataset)}")
+    print(f"  - Eval: {len(eval_dataset)}")
+    
+    # ============================================================
+    # Step 4: Data collator for MLM
+    # ============================================================
+    print("\n[Step 4] Setting up data collator...")
+    
+    class CustomDataCollator:
+        """Custom collator for MLM without tokenizer."""
+        
+        def __call__(self, batch):
+            input_ids = torch.stack([item["input_ids"] for item in batch])
+            
+            # Create labels for MLM
+            labels = input_ids.clone()
+            
+            # Mask 15% of tokens
+            mask = torch.rand(input_ids.shape) < 0.15
+            input_ids[mask] = 0  # [MASK] token
+            
+            # Don't compute loss on non-masked tokens
+            labels[~mask] = -100
+            
+            return {"input_ids": input_ids, "labels": labels}
+    
+    data_collator = CustomDataCollator()
+    print(f"✓ Data collator ready")
+    
+    # ============================================================
+    # Step 5: Training arguments
+    # ============================================================
+    print("\n[Step 5] Setting up training...")
+    
+    output_dir = "./from_scratch_pretrain"
+    
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=5,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        learning_rate=5e-4,
+        weight_decay=0.01,
+        warmup_steps=500,
+        logging_steps=50,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        report_to="none",
+        seed=42,
+        optim="adamw_torch",
+        gradient_accumulation_steps=1,
+        max_grad_norm=1.0,
+    )
+    
+    print(f"✓ Training config:")
+    print(f"  - Output: {output_dir}")
+    print(f"  - Epochs: {training_args.num_train_epochs}")
+    print(f"  - Batch size: {training_args.per_device_train_batch_size}")
+    print(f"  - Learning rate: {training_args.learning_rate}")
+    
+    # ============================================================
+    # Step 6: Train
+    # ============================================================
+    print("\n[Step 6] Starting training from scratch...")
+    print("(This may take a while. In practice, use more GPUs/data for real pretraining)")
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=data_collator,
+    )
+    
+    train_result = trainer.train()
+    
+    print(f"✓ Training complete!")
+    print(f"  - Final training loss: {train_result.training_loss:.4f}")
+    
+    # ============================================================
+    # Step 7: Evaluate
+    # ============================================================
+    print("\n[Step 7] Evaluating...")
+    
+    eval_results = trainer.evaluate()
+    
+    print(f"✓ Evaluation Results:")
+    for metric, value in eval_results.items():
+        if isinstance(value, (int, float)):
+            print(f"  - {metric}: {value:.4f}")
+    
+    # ============================================================
+    # Step 8: Save model and config
+    # ============================================================
+    print("\n[Step 8] Saving model...")
+    
+    save_dir = "./my_genemamba_from_scratch"
+    model.save_pretrained(save_dir)
+    config.save_pretrained(save_dir)
+    
+    print(f"✓ Model and config saved to '{save_dir}'")
+    print(f"  Files created:")
+    print(f"    - config.json")
+    print(f"    - model.safetensors (or pytorch_model.bin)")
+    
+    # ============================================================
+    # Step 9: Reload and verify
+    # ============================================================
+    print("\n[Step 9] Reloading model from checkpoint...")
+    
+    from transformers import AutoModelForMaskedLM
+    
+    loaded_model = AutoModelForMaskedLM.from_pretrained(
+        save_dir,
+        trust_remote_code=True,
+    )
+    
+    loaded_model.eval()
+    
+    # Test inference
+    with torch.no_grad():
+        sample_input = torch.randint(2, 25426, (2, 2048))
+        sample_input[:, :10] = 0  # Mask first 10 tokens
+        
+        outputs = loaded_model(sample_input)
+        logits = outputs.logits
+    
+    print(f"✓ Model reloaded and tested!")
+    print(f"  - Input shape: {sample_input.shape}")
+    print(f"  - Logits shape: {logits.shape}")
+    
+    # ============================================================
+    # Step 10: Optional - Convert to different format
+    # ============================================================
+    print("\n[Step 10] Model ready for conversion/deployment!")
+    print(f"✓ You can now:")
+    print(f"  1. Push to Hugging Face Hub:")
+    print(f"     model.push_to_hub('your-username/GeneMamba-custom')")
+    print(f"  2. Use with downstream tasks:")
+    print(f"     AutoModelForSequenceClassification.from_pretrained('{save_dir}', num_labels=N)")
+    print(f"  3. Extract embeddings:")
+    print(f"     AutoModel.from_pretrained('{save_dir}')")
+    
+    print("\n" + "=" * 80)
+    print("Phase 4 Complete! Model trained from scratch and ready to use.")
+    print("=" * 80)
+    
+    return model, trainer, config
+
+
+if __name__ == "__main__":
+    model, trainer, config = main()

examples/downstream/README.md

@@ -0,0 +1,35 @@
+# Downstream Examples
+
+This folder now contains both **ready-to-run** examples and **legacy scripts** from the original GeneMamba project.
+
+## Ready-to-run scripts
+
+- `10_finetune_classification.py`  
+  Fine-tune `AutoModelForSequenceClassification` for cell-type annotation.
+
+- `11_zero_shot_logreg.py`  
+  Freeze GeneMamba, extract `pooled_embedding`, train LogisticRegression on train split, evaluate on test split.
+
+- `12_batch_integration_eval.py`  
+  Batch integration proxy evaluation using silhouette score by `obs['batch']`.
+
+## Reference training scripts
+
+- `20_continue_pretraining_reference.py`
+- `21_pretrain_from_scratch_reference.py`
+
+## Legacy scripts from original repo
+
+- `legacy_from_gene_mamba/mamba2_classification_finetune_with_label.py`
+- `legacy_from_gene_mamba/mamba2_classification_finetune_without_label.py`
+- `legacy_from_gene_mamba/mamba2_classification_finetune_without_label_zero_shot.py`
+
+## Required h5ad conventions
+
+For downstream compatibility, standardize columns in `adata.obs`:
+
+- `celltype` for label
+- `batch` for batch id
+- `partition` in `{train, test}` for train/test split
+
+This matches conventions described in the original `dataset/downstream/README.md`.

examples/downstream/legacy_from_gene_mamba/mamba2_classification_finetune_with_label.py

@@ -0,0 +1,378 @@
+# %%
+import torch
+from transformers import Trainer
+import os
+
+import pyarrow as pa 
+import pandas as pd
+import numpy as np
+
+from matplotlib import pyplot as plt
+
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, TrainingArguments
+
+import argparse
+
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from transformers import AutoTokenizer, TrainingArguments, MambaForCausalLM
+
+from dotmap import DotMap
+
+import sys
+import os
+import torch
+
+# from trange import trange
+
+sys.path.append("/project/zhiwei/cq5/PythonWorkSpace/gene_mamba")
+from models import Classifier, GeneMamba, GeneMambaForCellAnnotation, GeneMambaForGeneClassification, GeneMamba2, GeneMamba2ForCellClassification
+from utils import permute_genes_by_expression
+from utils2 import standardize_columns
+
+import importlib
+importlib.reload(sys.modules['models'])
+importlib.reload(sys.modules['utils'])
+importlib.reload(sys.modules['utils2'])
+
+
+# %%
+DATA_PATH = "/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/dataset/downstream/"
+# CHECKPOINT_PATH = "/project/zhiwei/cq5/LLM_checkpoints/GeneMamba/GeneMamba2_48l_512d/1/3m/checkpoint-31250"
+TOKENIZER_PATH = "/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/gene_tokenizer.json"
+SAVE_PATH = "/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/dataset/embeddings/cell"
+
+# %%
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset_name", type=str)
+parser.add_argument("--ckpt_path", type = str)
+parser.add_argument("--seq_len", type=int, default=2048)
+parser.add_argument("--batch_size", type=int, default=24)
+parser.add_argument("--num_epochs", type=int, default=5)
+parser.add_argument("--test_size", type = float, default=0.1)
+parser.add_argument("--split", type=lambda x: x.lower() in ["true", "1", "yes"], default=False,)
+
+args = parser.parse_args()
+
+
+# args = DotMap({
+#     "dataset_name": "ms",
+#     "seq_len": 512,
+#     "batch_size": 24,
+#     "num_epochs": 5,
+#     "test_size": 0.1
+# })
+
+
+#%%
+CHECKPOINT_PATH = args.ckpt_path
+model_name = CHECKPOINT_PATH.split("/")[-4]
+mamba_layer = int(model_name.split("_")[1][:-1])
+d_model = int(model_name.split("_")[2][:-1])
+
+
+# make the sub directories to save the results
+SAVE_PATH = os.path.join(SAVE_PATH, model_name)
+sub_directories = ["predictions", "metrics", "figures", "repr"]
+for sub_dir in sub_directories:
+    os.makedirs(os.path.join(SAVE_PATH, sub_dir), exist_ok=True)
+
+
+# %%
+import scanpy as sc
+
+# Load the .h5ad file
+dataset_name = args.dataset_name
+# adata = sc.read_h5ad(os.path.join(DATA_PATH ,f'{dataset_name}.h5ad'))
+
+adata = None
+
+if args.split:
+    adata = sc.read_h5ad(os.path.join(DATA_PATH ,f'split/{dataset_name}_split.h5ad'))
+    print(f"Read data from {dataset_name}_split.h5ad")
+    dataset_name = dataset_name + "_split"
+else:
+    adata = sc.read_h5ad(os.path.join(DATA_PATH ,f'processed/{dataset_name}_processed.h5ad'))
+    print(f"Read data from {dataset_name}_processed.h5ad")
+
+# Display basic information about the data
+print(adata)
+
+# %%
+# adata = standardize_columns(adata, dataset_name)
+# assert "batch" in adata.obs.columns and "celltype" in adata.obs.columns
+
+# %%
+from sklearn.preprocessing import LabelEncoder
+
+y_names = np.array(adata.obs['celltype'].values.tolist())
+
+label_encoder = LabelEncoder()
+y = label_encoder.fit_transform(y_names)
+
+num_class = len(label_encoder.classes_)
+
+# %%
+from transformers import PretrainedConfig
+
+config = PretrainedConfig.from_dict({
+    "d_model": d_model,
+    "mamba_layer": mamba_layer,
+})
+
+
+# %%
+model_cell_cls = GeneMamba2ForCellClassification(config, model_path=CHECKPOINT_PATH, tokenizer_path = TOKENIZER_PATH, args=None, output_dim_cls = num_class, hidden_dim= 512, num_layers_cls = 4)
+
+# %%
+permuted_gene_ids = permute_genes_by_expression(adata, dataset_name, model_cell_cls.tokenizer, model_cell_cls.symbol2id)
+permuted_gene_ids
+
+# %%
+seq_len = args.seq_len
+
+input_data = permuted_gene_ids[:, :seq_len]
+
+# %%
+model_cell_cls.tokenizer.cls_token_id
+
+# %%
+torch.tensor([model_cell_cls.tokenizer.cls_token_id for _ in range(input_data.shape[0])])
+
+# %%
+model_cell_cls.tokenizer.cls_token_id
+
+# %%
+input_data.shape[0]
+
+# %%
+input_data
+
+# %%
+# add the cls token to the input data
+input_data = np.hstack([np.array([model_cell_cls.tokenizer.cls_token_id for _ in range(input_data.shape[0])]).reshape(-1, 1), input_data])
+input_data
+
+# %%
+input_data.shape
+
+#%%
+from sklearn.model_selection import train_test_split
+import numpy as np
+
+def manual_stratified_split(X, y, test_size=0.1, random_state=None):
+    # separate the samples for each class
+    unique_classes = np.unique(y)
+    X_train, X_test, y_train, y_test = [], [], [], []
+    
+    for cls in unique_classes:
+        cls_indices = np.where(y == cls)[0]
+        
+        if len(cls_indices) > 1:
+
+            cls_train, cls_test = train_test_split(cls_indices, test_size=test_size, random_state=random_state)
+        else:
+            # if a class has only one sample, put it in the training set
+            cls_train, cls_test = cls_indices, []
+        
+        X_train.extend(X[cls_train])
+        y_train.extend(y[cls_train])
+        X_test.extend(X[cls_test])
+        y_test.extend(y[cls_test])
+    
+    return np.array(X_train), np.array(X_test), np.array(y_train), np.array(y_test)
+
+# %%
+# from sklearn.model_selection import train_test_split
+
+# X_train, X_test, y_train, y_test = manual_stratified_split(input_data, y, test_size=args.test_size, random_state=42)
+
+
+#%%
+# train and test split is done and stored in the adata.obs["partition"] column, so we can extract the train and test data from there
+
+X_train = input_data[adata.obs["partition"] == "train"]
+X_test = input_data[adata.obs["partition"] == "test"]
+y_train = y[adata.obs["partition"] == "train"]
+y_test = y[adata.obs["partition"] == "test"]
+
+X_train.shape, X_test.shape, y_train.shape, y_test.shape
+
+# %%
+from torch.utils.data import DataLoader, Dataset
+
+class GeneDataset(Dataset):
+    def __init__(self, data, y):
+        self.data = data
+        self.labels = y
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx], self.labels[idx]
+
+train_dataset = GeneDataset(X_train, y_train)
+test_dataset = GeneDataset(X_test, y_test)
+all_dataset = GeneDataset(input_data, y)
+
+train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
+all_loader = DataLoader(all_dataset, batch_size=args.batch_size, shuffle=False)
+
+# %%
+from sklearn.metrics import classification_report, accuracy_score, f1_score, precision_score, recall_score, roc_auc_score
+
+# %%
+def compute_metrics(y_pred, y_prob, y_true):
+
+    metrics = {
+        "accuracy": accuracy_score(y_true, y_pred),
+        "Micro-F1 score": f1_score(y_true, y_pred, average='micro'),
+        "Macro-F1 score": f1_score(y_true, y_pred, average='macro'),
+        "precision": precision_score(y_true, y_pred, average='weighted'),
+        "recall": recall_score(y_true, y_pred, average='weighted'),
+        # "auc_roc": roc_auc_score(y_true, y_prob, multi_class = 'ovr'),
+    }
+    return metrics
+
+# %%
+epochs = args.num_epochs
+optimizer = torch.optim.Adam(model_cell_cls.parameters(), lr=1e-4)
+loss = torch.nn.CrossEntropyLoss()
+
+for epoch in range(epochs):
+    model_cell_cls.train()
+    for i, batch in enumerate(train_loader):
+        data = batch[0]
+        target = batch[1]
+        data = data.to(model_cell_cls.device)
+        target = target.to(model_cell_cls.device)
+        model_cell_cls = model_cell_cls.to(model_cell_cls.device)
+        
+        optimizer.zero_grad()
+        output = model_cell_cls(data, None)
+        loss_val = loss(output, target)
+        loss_val.backward()
+        optimizer.step()
+        if i % 10 == 0:
+            print(f"Epoch {epoch}, Iteration {i}, Loss: {loss_val}")
+
+    model_cell_cls.eval()
+    with torch.no_grad():
+        # add code to compute the metrics
+        pred_prob = []
+        pred_label = []
+        targets = []
+        cell_repr = []
+
+        for i, batch in enumerate(test_loader):
+            data = batch[0]
+            target = batch[1]
+            data = data.to(model_cell_cls.device)
+            target = target.to(model_cell_cls.device)
+            model_cell_cls = model_cell_cls.to(model_cell_cls.device)
+
+            output, output_test_repr = model_cell_cls(data, None, return_cls = True)
+            cell_repr.append(output_test_repr.cpu().numpy())
+            
+            # calculate the probability from the output
+            pred_prob.append(torch.nn.functional.softmax(output, dim=1).cpu().numpy())
+
+            _, predicted = torch.max(output, 1)
+            pred_label.append(predicted.cpu().numpy())
+            targets.append(target.cpu().numpy())
+        
+        pred_prob = np.concatenate(pred_prob)
+        pred_label = np.concatenate(pred_label)
+        targets = np.concatenate(targets)
+        cell_repr = np.concatenate(cell_repr)
+
+        # break
+        # save the predictions
+        np.save(os.path.join(SAVE_PATH, f"predictions/pred_prob_{dataset_name}_{epoch}.npy"), pred_prob)
+        np.save(os.path.join(SAVE_PATH, f"predictions/pred_label_{dataset_name}_{epoch}.npy"), pred_label)
+        np.save(os.path.join(SAVE_PATH, f"predictions/targets_{dataset_name}_{epoch}.npy"), targets)
+        
+        metrics = compute_metrics(pred_label, pred_prob, targets)
+        
+        with open(os.path.join(SAVE_PATH, f"metrics/metrics_{dataset_name}_{epoch}.txt"), "w") as f:
+            print(metrics, file=f)
+            print(metrics)
+        
+
+        # draw scatter plot for the first two components
+        from sklearn.decomposition import PCA
+
+        pca = PCA(n_components=2)
+        pca_result = pca.fit_transform(cell_repr)
+
+        plt.figure(figsize=(8, 8))
+
+        plt.scatter(pca_result[:, 0], pca_result[:, 1], c = targets)
+        plt.savefig(os.path.join(SAVE_PATH, f"figures/scatter_{dataset_name}_{epoch}.png"))
+        # plt.show()
+
+
+# %%
+model_cell_cls.eval()
+
+def cell_embeddings(data_loader, model_cell_cls):
+    cell_repr = []
+
+    for i, batch in enumerate(data_loader):
+        data = batch[0]
+        target = batch[1]
+        data = data.to(model_cell_cls.device)
+        target = target.to(model_cell_cls.device)
+        model_cell_cls = model_cell_cls.to(model_cell_cls.device)
+
+        output, output_test_repr = model_cell_cls(data, None, return_cls = True)
+        cell_repr.append(output_test_repr.detach().cpu().numpy())
+        if i % 10 == 0:
+            print(f"Processed {i} batches")
+
+    cell_repr = np.concatenate(cell_repr)
+    return cell_repr
+
+
+test_cell_repr = cell_embeddings(test_loader, model_cell_cls)
+save_path_test = os.path.join(SAVE_PATH, f"repr/{dataset_name}_test_cell_repr.npy")
+np.save(save_path_test, test_cell_repr)
+del test_cell_repr
+
+
+train_cell_repr = cell_embeddings(train_loader, model_cell_cls)
+save_path_train = os.path.join(SAVE_PATH, f"repr/{dataset_name}_train_cell_repr.npy")
+np.save(save_path_train, train_cell_repr)
+del train_cell_repr
+
+all_cell_repr = cell_embeddings(all_loader, model_cell_cls)
+save_path_all = os.path.join(SAVE_PATH, f"repr/{dataset_name}_cell_repr.npy")
+np.save(save_path_all, all_cell_repr)
+del all_cell_repr
+
+
+# %%
+# original_data = adata.X.toarray()
+# original_data.shape
+
+# %%
+# draw the scatter figure on the original data
+# from sklearn.decomposition import PCA
+
+# pca = PCA(n_components=2)
+# pca_result = pca.fit_transform(original_data)
+
+# plt.figure(figsize=(8, 8))
+
+# plt.scatter(pca_result[:, 0], pca_result[:, 1], c = y)
+# plt.show()
+
+
+# %%
+
+
+

examples/downstream/legacy_from_gene_mamba/mamba2_classification_finetune_without_label.py

@@ -0,0 +1,161 @@
+# %%
+import torch
+from transformers import Trainer
+import os
+
+import pyarrow as pa 
+import pandas as pd
+import numpy as np
+
+from matplotlib import pyplot as plt
+
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, TrainingArguments
+
+import argparse
+
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from transformers import AutoTokenizer, TrainingArguments, MambaForCausalLM
+
+from dotmap import DotMap
+
+import sys
+import os
+import torch
+
+sys.path.append("/project/zhiwei/cq5/PythonWorkSpace/gene_mamba")
+
+from models import Classifier, GeneMamba, GeneMambaForCellAnnotation, GeneMambaForGeneClassification, GeneMamba2, GeneMamba2ForCellClassification
+from utils import permute_genes_by_expression, build_downstream_dataset
+
+
+import importlib
+importlib.reload(sys.modules['models'])
+importlib.reload(sys.modules['utils'])
+
+# %%
+import scanpy as sc
+
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset_name", type=str)
+
+args2 = parser.parse_args()
+
+# Load the .h5ad file
+dataset_name = args2.dataset_name
+
+
+assert dataset_name in ["pbmc12k", "perirhinal_cortex", "covid19"]
+
+adata = sc.read_h5ad(f'/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/dataset/downstream/processed/{dataset_name}_processed.h5ad')
+
+assert "celltype" in adata.obs
+
+print(adata)
+
+# %%
+from sklearn.preprocessing import LabelEncoder
+
+y_names = np.array(adata.obs['celltype'].values.tolist())
+
+label_encoder = LabelEncoder()
+y = label_encoder.fit_transform(y_names)
+
+num_class = len(label_encoder.classes_)
+
+# %%
+from transformers import PretrainedConfig
+
+config = PretrainedConfig.from_dict({
+    "d_model": 512,
+    "mamba_layer": 24,
+})
+
+
+# %%
+model = GeneMamba2(config, model_path="/project/zhiwei/cq5/LLM_checkpoints/GeneMamba/GeneMamba2_24l_512d/1/16m/checkpoint-31250", tokenizer_path="/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/gene_tokenizer.json", args=None)
+
+# %%
+permuted_gene_ids = permute_genes_by_expression(adata, dataset_name, model.tokenizer, model.symbol2id)
+permuted_gene_ids
+
+# %%
+num_samples = permuted_gene_ids.shape[0]
+num_avaliable_gpu = torch.cuda.device_count()
+
+# %%
+from dotmap import DotMap
+
+args = DotMap(
+    {
+        # "model": "state-spaces/mamba-130m-hf",
+        # "tokenizer": "state-spaces/mamba-130m-hf",
+        "learning_rate": 5e-5,
+        "batch_size": 16,
+        "gradient_accumulation_steps": 1,
+        "optim": "adamw_torch",
+        # "data_path": "/home/cong/study/codeSpace/VSCodeSpace/PythonWorkSpace/TCRPrediction/mamba_transformer/smiles_data.txt",
+        # "num_epochs": args2.num_epochs,
+        "seq_len": 2048,
+        "num_samples": num_samples,
+        "num_gpus": num_avaliable_gpu,
+        "output_dir": "/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/analysis/cell_type_annotation/fine-tuned/debug",
+    }
+)
+
+# %%
+input_data = permuted_gene_ids[:, :args.seq_len]
+
+# %%
+input_data.shape
+
+#%%
+# check if cls_token in the tokenizer:
+if model.tokenizer.cls_token_id is None:
+    model.tokenizer.add_special_tokens({'cls_token': '[CLS]'})
+
+#%%
+input_data = np.hstack([np.array([model.tokenizer.cls_token_id for _ in range(input_data.shape[0])]).reshape(-1, 1), input_data])
+
+#%%
+input_data.shape
+
+# %%
+sample_dataset = build_downstream_dataset(input_data, model.tokenizer)
+sample_dataset
+
+# input_data = np.hstack([np.array([model.tokenizer.cls_token_id for _ in range(input_data.shape[0])]).reshape(-1, 1), input_data])
+# input_data
+
+# %%
+args=TrainingArguments(
+    learning_rate=args.learning_rate,
+    num_train_epochs = 4,
+    per_device_train_batch_size=args.batch_size,
+    gradient_accumulation_steps=args.gradient_accumulation_steps,
+    optim=args.optim,
+    output_dir=os.path.join(args.output_dir, dataset_name),
+    # output_dir=f"/scratch/zhiwei/cq5/logs/mamba/test/context_length",
+    # logging_dir=f"{args.output_dir}/{args.num_epochs}/{args.num_samples // 1000000 + args.bulk_id}m_logging",
+    logging_steps=args.num_samples // args.batch_size // 10,
+    save_steps=args.num_samples // args.batch_size // 10,
+)
+
+
+# %%
+model.finetune(sample_dataset, args)
+
+# %%
+# ckpt_pth = get_last_checkpoint(os.path.join(args.output_dir, dataset_name))
+# ckpt_pth
+
+# #%%
+# model = GeneMamba2(config, model_path=, tokenizer_path="/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/gene_tokenizer.json", args=None)
+
+#%%
+
+
+
+

examples/downstream/legacy_from_gene_mamba/mamba2_classification_finetune_without_label_zero_shot.py

@@ -0,0 +1,197 @@
+# %%
+import torch
+from transformers import Trainer
+import os
+
+import pyarrow as pa 
+import pandas as pd
+import numpy as np
+
+from matplotlib import pyplot as plt
+
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, TrainingArguments
+
+import argparse
+
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from transformers import AutoTokenizer, TrainingArguments, MambaForCausalLM
+
+from dotmap import DotMap
+
+import sys
+import os
+import torch
+
+sys.path.append("/project/zhiwei/cq5/PythonWorkSpace/gene_mamba")
+
+from models import Classifier, GeneMamba, GeneMambaForCellAnnotation, GeneMambaForGeneClassification, GeneMamba2, GeneMamba2ForCellClassification
+from utils import permute_genes_by_expression, build_downstream_dataset, get_last_checkpoint
+
+
+import importlib
+importlib.reload(sys.modules['models'])
+importlib.reload(sys.modules['utils'])
+
+# %%
+import scanpy as sc
+
+# import argparse
+
+# parser = argparse.ArgumentParser()
+# parser.add_argument("--dataset_name", type=str)
+
+# args2 = parser.parse_args()
+
+# dataset_name = args2.dataset_name
+
+dataset_name = "pbmc12k"
+
+assert dataset_name in ["pbmc12k", "perirhinal_cortex", "covid19"]
+
+adata = sc.read_h5ad(f'/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/dataset/downstream/processed/{dataset_name}_processed.h5ad')
+
+assert "celltype" in adata.obs
+
+print(adata)
+
+# %%
+from transformers import PretrainedConfig
+
+config = PretrainedConfig.from_dict({
+    "d_model": 512,
+    "mamba_layer": 24,
+})
+
+
+# %%
+model = GeneMamba2(config, model_path="/project/zhiwei/cq5/LLM_checkpoints/GeneMamba/GeneMamba2_24l_512d/1/16m/checkpoint-31250", tokenizer_path="/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/gene_tokenizer.json", args=None)
+
+# %%
+permuted_gene_ids = permute_genes_by_expression(adata, dataset_name, model.tokenizer, model.symbol2id)
+permuted_gene_ids
+
+# %%
+num_samples = permuted_gene_ids.shape[0]
+num_avaliable_gpu = torch.cuda.device_count()
+
+# %%
+from dotmap import DotMap
+
+args = DotMap(
+    {
+        # "model": "state-spaces/mamba-130m-hf",
+        # "tokenizer": "state-spaces/mamba-130m-hf",
+        "learning_rate": 5e-5,
+        "batch_size": 16,
+        "gradient_accumulation_steps": 1,
+        "optim": "adamw_torch",
+        # "data_path": "/home/cong/study/codeSpace/VSCodeSpace/PythonWorkSpace/TCRPrediction/mamba_transformer/smiles_data.txt",
+        # "num_epochs": args2.num_epochs,
+        "seq_len": 2048,
+        "num_samples": num_samples,
+        "num_gpus": num_avaliable_gpu,
+        "output_dir": "/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/analysis/cell_type_annotation/fine-tuned",
+    }
+)
+
+
+#%%
+model = GeneMamba2(config, model_path="/project/zhiwei/cq5/LLM_checkpoints/GeneMamba/GeneMamba2_24l_512d/1/16m/checkpoint-31250", tokenizer_path="/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/gene_tokenizer.json", args=None)
+
+model.resize_token_embeddings()
+
+#%%
+def get_last_checkpoint(output_dir):
+    checkpoints = os.listdir(output_dir)
+    checkpoints = [ckpt for ckpt in checkpoints if "checkpoint" in ckpt]
+    checkpoints = [int(ckpt.split("-")[1]) for ckpt in checkpoints]
+    checkpoints = sorted(checkpoints)
+    last_checkpoint = checkpoints[-1]
+    last_checkpoint = os.path.join(output_dir, f"checkpoint-{last_checkpoint}")
+    return last_checkpoint
+
+ckpt_pth = f"/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/analysis/cell_type_annotation/fine-tuned/{dataset_name}"
+
+last_checkpoint = get_last_checkpoint(ckpt_pth)
+state_dict_pth = os.path.join(last_checkpoint, "model.safetensors")
+
+print(state_dict_pth)
+
+#%%
+from safetensors.torch import load_file
+
+state_dict = load_file(state_dict_pth)
+
+model.model.load_state_dict(state_dict)
+
+
+# %%
+input_data = permuted_gene_ids[:, :args.seq_len]
+
+# %%
+input_data.shape
+
+#%%
+# check if cls_token in the tokenizer:
+if model.tokenizer.cls_token_id is None:
+    model.tokenizer.add_special_tokens({'cls_token': '[CLS]'})
+
+#%%
+input_data = np.hstack([np.array([model.tokenizer.cls_token_id for _ in range(input_data.shape[0])]).reshape(-1, 1), input_data])
+
+#%%
+input_data.shape
+
+
+#%%
+from torch.utils.data import DataLoader, Dataset
+
+class GeneDataset(Dataset):
+    def __init__(self, data):
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+
+#%%
+all_dataset = GeneDataset(input_data)
+all_loader = DataLoader(all_dataset, batch_size = args.batch_size, shuffle=False)
+
+# %%
+def cell_embeddings(data_loader, model):
+    
+    cell_repr = []
+
+    for i, batch in enumerate(data_loader):
+        batch = batch.to(model.device)
+        outputs = model(batch)
+
+
+        cls_representation = outputs.hidden_states[:, 0, :]
+        cell_repr.append(cls_representation.detach().cpu().numpy())
+
+        if i % 10 == 0:
+            print(f"Processed {i} batches")
+
+    cell_repr = np.concatenate(cell_repr)
+    return cell_repr
+
+# %%
+model = model.to("cuda")
+model.eval()
+
+# %%
+cell_repr = cell_embeddings(all_loader, model)
+cell_repr.shape
+
+
+# cell_repr = np.concatenate(cell_repr)
+# %%
+np.save(f"/project/zhiwei/cq5/PythonWorkSpace/gene_mamba/analysis/cell_type_annotation/embeddings/fine-tuned/{dataset_name}_cell_repr.npy", cell_repr)
+
+# %%

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccb1fcb0ee4b3ea2013099b9b187455e160d3b66b76c606715231b70b13c2784
+size 262998656

modeling_genemamba.py

@@ -0,0 +1,395 @@
+"""
+PyTorch implementation of GeneMamba model for Hugging Face Transformers.
+Includes backbone model and task-specific heads for various downstream tasks.
+"""
+
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import normal_, constant_
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, ModelOutput
+from transformers.models.auto import register_model_for_auto_class
+
+from mamba_ssm import Mamba
+from mamba_ssm.ops.triton.layer_norm import RMSNorm
+
+from .configuration_genemamba import GeneMambaConfig
+from .modeling_outputs import GeneMambaModelOutput, GeneMambaSequenceClassifierOutput, GeneMambaMaskedLMOutput
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================
+# Core Architecture Components
+# ===========================
+
+class EncoderLayer(nn.Module):
+    """
+    Single Mamba encoder layer with residual connection.
+    Applies a Mamba2 or Mamba layer followed by addition with input.
+    
+    Args:
+        hidden_size (int): Dimension of hidden states.
+    """
+    
+    def __init__(self, hidden_size: int):
+        super(EncoderLayer, self).__init__()
+        self.mamba = Mamba(d_model=hidden_size, d_state=64, d_conv=4, expand=2)
+    
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+        
+        Returns:
+            torch.Tensor: Output after Mamba layer and residual connection.
+        """
+        output = self.mamba(X) + X
+        return output
+
+
+class MambaMixer(nn.Module):
+    """
+    Stack of Mamba encoder layers with bidirectional processing and aggregation.
+    Processes sequences in both forward and reverse directions, then aggregates.
+    
+    Args:
+        mode (str): Aggregation mode. Options: "mean", "sum", "concat", "gate".
+        hidden_size (int): Dimension of hidden states.
+        num_hidden_layers (int): Number of Mamba layers.
+    """
+    
+    def __init__(
+        self,
+        mode: str = "gate",
+        hidden_size: int = 512,
+        num_hidden_layers: int = 24
+    ):
+        super(MambaMixer, self).__init__()
+        self.mode = mode
+        self.hidden_size = hidden_size
+        
+        # Create Mamba layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(hidden_size) for _ in range(num_hidden_layers)]
+        )
+        
+        # Aggregation modules for certain modes
+        if mode in ["concat", "gate"]:
+            self.aggr = nn.Linear(hidden_size * 2, hidden_size)
+    
+    def flip_sequence(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Reverse a sequence based on actual length (ignoring padding).
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            mask (torch.Tensor, optional): Padding mask of shape (batch_size, seq_len).
+        
+        Returns:
+            torch.Tensor: Reversed tensor.
+        """
+        batch_size, seq_length, embedding_dim = X.size()
+        
+        if mask is None:
+            # Simple flip
+            return X.flip([1])
+        
+        # Flip based on actual sequence length (marked by mask)
+        lengths = (~mask).sum(dim=1)
+        pos_tensor = torch.arange(seq_length, device=X.device).unsqueeze(0).expand(batch_size, -1)
+        flip_mask = pos_tensor < lengths.unsqueeze(1)
+        reversed_positions = torch.where(
+            flip_mask,
+            lengths.unsqueeze(1) - 1 - pos_tensor,
+            pos_tensor
+        )
+        
+        X_reverse = torch.gather(X, 1, reversed_positions.unsqueeze(-1).expand(-1, -1, embedding_dim))
+        return X_reverse
+    
+    def forward(
+        self,
+        X: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Process sequence through bidirectional Mamba layers.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, seq_len, hidden_size).
+            padding_mask (torch.Tensor, optional): Padding mask.
+        
+        Returns:
+            torch.Tensor: Output after processing all layers and aggregation.
+        """
+        
+        for layer in self.layers:
+            # Flip sequence for reverse processing
+            X_flip = self.flip_sequence(X, padding_mask)
+            
+            # Forward and reverse passes
+            X_f = layer(X)
+            X_b = layer(X_flip)
+            
+            # Flip back the reverse output
+            X_b = self.flip_sequence(X_b, padding_mask)
+            
+            # Aggregate forward and reverse
+            if self.mode == "mean":
+                X = (X_f + X_b) / 2
+            elif self.mode == "sum":
+                X = X_f + X_b
+            elif self.mode == "concat":
+                X = torch.cat([X_f, X_b], dim=-1)
+                X = self.aggr(X)
+            elif self.mode == "gate":
+                z = torch.sigmoid(self.aggr(torch.cat([X_f, X_b], dim=-1)))
+                X = z * X_f + (1 - z) * X_b
+            else:
+                raise ValueError(f"Invalid aggregation mode: {self.mode}")
+        
+        return X
+
+
+# ===========================
+# Base Model Classes
+# ===========================
+
+class GeneMambaPreTrainedModel(PreTrainedModel):
+    """
+    Base class for all GeneMamba models.
+    Handles weight initialization and provides standard model interfaces.
+    """
+    
+    config_class = GeneMambaConfig
+    base_model_prefix = "genemamba"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize module weights."""
+        if isinstance(module, nn.Linear):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.bias is not None:
+                constant_(module.bias, 0.0)
+        elif isinstance(module, nn.Embedding):
+            normal_(module.weight, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            constant_(module.bias, 0.0)
+            constant_(module.weight, 1.0)
+
+
+class GeneMambaModel(GeneMambaPreTrainedModel):
+    """
+    GeneMamba backbone model - outputs cell embeddings and hidden states.
+    This is the core model used by task-specific heads.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Embedding layer
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Mamba layers with bidirectional aggregation
+        self.mamba_mixer = MambaMixer(
+            mode=config.mamba_mode,
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers
+        )
+        
+        # Final layer normalization
+        self.norm = RMSNorm(config.hidden_size)
+        
+        self.apply(self._init_weights)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return embedding layer."""
+        return self.embeddings
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        """Set embedding layer."""
+        self.embeddings = value
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaModelOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask of shape (batch_size, seq_len).
+            output_hidden_states (bool): Whether to output hidden states from all layers.
+        
+        Returns:
+            GeneMambaModelOutput: Contains last_hidden_state, pooled_embedding, etc.
+        """
+        # Get embeddings
+        hidden_states = self.embeddings(input_ids)
+        
+        # Pass through Mamba layers
+        hidden_states = self.mamba_mixer(hidden_states, attention_mask)
+        
+        # Apply final normalization
+        hidden_states = self.norm(hidden_states)
+        
+        # Compute pooled embedding (cell representation)
+        if self.config.embedding_pooling == "CLS":
+            # Use first token (CLS)
+            pooled_embedding = hidden_states[:, 0, :]
+        elif self.config.embedding_pooling == "mean":
+            # Mean pooling over sequence
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).expand(hidden_states.shape).float()
+                pooled_embedding = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
+            else:
+                pooled_embedding = hidden_states.mean(dim=1)
+        else:
+            raise ValueError(f"Unsupported embedding_pooling: {self.config.embedding_pooling}")
+        
+        return GeneMambaModelOutput(
+            last_hidden_state=hidden_states,
+            pooled_embedding=pooled_embedding,
+            hidden_states=hidden_states if output_hidden_states else None,
+            embedding_pooling=self.config.embedding_pooling,
+        )
+
+
+# ===========================
+# Task-Specific Models
+# ===========================
+
+@register_model_for_auto_class("AutoModel")
+class GeneMambaForMaskedLM(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for masked language modeling (MLM).
+    Suitable for pretraining and domain adaptation.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.genemamba = GeneMambaModel(config)
+        
+        # Language modeling head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaMaskedLMOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Target token ids for MLM loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaMaskedLMOutput: Contains logits and optional loss.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        logits = self.lm_head(outputs.last_hidden_state)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        return GeneMambaMaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+        )
+
+
+@register_model_for_auto_class("AutoModelForSequenceClassification")
+class GeneMambaForSequenceClassification(GeneMambaPreTrainedModel):
+    """
+    GeneMamba model for sequence classification tasks.
+    Ideal for cell type annotation, tissue classification, etc.
+    
+    Args:
+        config (GeneMambaConfig): Model configuration class.
+    """
+    
+    def __init__(self, config: GeneMambaConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        
+        self.genemamba = GeneMambaModel(config)
+        
+        # Classification head
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        
+        self.apply(self._init_weights)
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> GeneMambaSequenceClassifierOutput:
+        """
+        Args:
+            input_ids (torch.Tensor): Token indices of shape (batch_size, seq_len).
+            attention_mask (torch.Tensor, optional): Attention mask.
+            labels (torch.Tensor, optional): Class labels for classification loss.
+            output_hidden_states (bool): Whether to output hidden states.
+        
+        Returns:
+            GeneMambaSequenceClassifierOutput: Contains logits, optional loss, and embedding.
+        """
+        outputs = self.genemamba(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        
+        pooled_embedding = outputs.pooled_embedding
+        logits = self.classifier(self.dropout(pooled_embedding))
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        
+        return GeneMambaSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            pooled_embedding=pooled_embedding,
+        )
+
+
+# Register tokenizer class
+register_model_for_auto_class("AutoModelForMaskedLM")(GeneMambaForMaskedLM)

modeling_outputs.py

@@ -0,0 +1,81 @@
+"""
+Custom ModelOutput classes for GeneMamba.
+Defines the output structure for different GeneMamba tasks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class GeneMambaModelOutput(ModelOutput):
+    """
+    Base output class for GeneMamba models.
+    
+    Attributes:
+        last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)):
+            Sequence of hidden-states at the output of the last layer of the model.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size)):
+            Cell/sequence-level embedding (pooled representation) used for downstream tasks.
+            This is the recommended embedding to use for classification, clustering, etc.
+        
+        embedding_pooling (str):
+            The pooling method used to generate pooled_embedding.
+    """
+    
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: torch.FloatTensor = None
+    embedding_pooling: str = "mean"
+
+
+@dataclass
+class GeneMambaSequenceClassifierOutput(ModelOutput):
+    """
+    Output class for GeneMamba sequence classification models.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            Classification loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, num_labels)):
+            Classification scores (before softmax).
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+        
+        pooled_embedding (torch.FloatTensor of shape (batch_size, hidden_size), optional):
+            Cell embedding before classification head.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    pooled_embedding: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class GeneMambaMaskedLMOutput(ModelOutput):
+    """
+    Output class for GeneMamba masked language modeling.
+    
+    Attributes:
+        loss (torch.FloatTensor of shape (), optional):
+            MLM loss (if labels were provided).
+        
+        logits (torch.FloatTensor of shape (batch_size, sequence_length, vocab_size)):
+            Prediction scores of the language modeling head.
+        
+        hidden_states (tuple(torch.FloatTensor), optional):
+            Hidden-states of the model at the output of each layer.
+    """
+    
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None

special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]"
+}