Upload MOJO

mojo.py

@@ -1,7 +1,7 @@
 import logging
 import math
-from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
 
 import numpy as np
 import torch
@@ -510,32 +510,40 @@ class LMHead(nn.Module):
         return out
 
 
-@dataclass
 class MOJOConfig(PretrainedConfig):  # noqa: N801
     model_type = "MOJO"
-    alphabet_size: dict[str, int] = field(
-        default_factory=lambda: {"rnaseq": 66, "methylation": 66}
-    )
-    token_embed_dim: int = 256
-    init_gene_embed_dim: int = 200
-    use_gene_embedding: bool = True
-    project_gene_embedding: bool = True
-    sequence_length: int = 17_116  # n_genes
-    fixed_sequence_length: int | None = None
-    num_downsamples: int = 8
-    conv_init_embed_dim: int = 512
-    stem_kernel_shape: int = 15
-    embed_dim: int = 512
-    filter_list: list[int] = field(default_factory=list)
-    num_attention_heads: int = 16
-    key_size: Optional[int] = None
-    ffn_embed_dim: int = 1_024
-    num_layers: int = 8
-    num_hidden_layers_head: int = 1
-
-    # return
-    embeddings_layers_to_save: tuple[int, ...] = field(default_factory=tuple)
-    attention_maps_to_save: list[tuple[int, int]] = field(default_factory=list)
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.alphabet_size = kwargs.get(
+            "alphabet_size", {"rnaseq": 66, "methylation": 66}
+        )
+        self.token_embed_dim = kwargs.get("token_embed_dim", 256)
+        self.init_gene_embed_dim = kwargs.get("init_gene_embed_dim", 200)
+        self.use_gene_embedding = kwargs.get("use_gene_embedding", True)
+        self.project_gene_embedding = kwargs.get("project_gene_embedding", True)
+        self.sequence_length = kwargs.get("sequence_length", 17_116)  # n_genes
+        self.fixed_sequence_length = kwargs.get("fixed_sequence_length", None)
+        self.num_downsamples = kwargs.get("num_downsamples", 8)
+        self.conv_init_embed_dim = kwargs.get("conv_init_embed_dim", 512)
+        self.stem_kernel_shape = kwargs.get("stem_kernel_shape", 15)
+        self.embed_dim = kwargs.get("embed_dim", 512)
+        self.filter_list = kwargs.get("filter_list", [])
+        self.num_attention_heads = kwargs.get("num_attention_heads", 16)
+        self.key_size = kwargs.get("key_size", None)
+        self.ffn_embed_dim = kwargs.get("ffn_embed_dim", 1_024)
+        self.num_layers = kwargs.get("num_layers", 8)
+        self.num_hidden_layers_head = kwargs.get("num_hidden_layers_head", 1)
+
+        # return
+        self.embeddings_layers_to_save: tuple[int, ...] = kwargs.get(
+            "embeddings_layers_to_save", ()
+        )
+        self.attention_maps_to_save: list[tuple[int, int]] = kwargs.get(
+            "attention_maps_to_save", []
+        )
+
+        self.__post_init__()
 
     def __post_init__(self):
         # Validate attention key size