| """ |
| 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 |
|
|
