Upload BulkRNABert

bulkrnabert.py

@@ -1,6 +1,5 @@
 import logging
-from dataclasses import dataclass, field
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 import torch
@@ -198,23 +197,31 @@ class SelfAttentionBlock(nn.Module):
         return output
 
 
-@dataclass
 class BulkRNABertConfig(PretrainedConfig):
     model_type = "BulkRNABert"
-    n_genes: int = 19_062
-    n_expressions_bins: int = 64
-    embed_dim: int = 256
-    init_gene_embed_dim: int = 200
-    use_gene_embedding: bool = True
-    project_gene_embedding: bool = True
-    num_attention_heads: int = 8
-    key_size: Optional[int] = None
-    ffn_embed_dim: int = 512
-    num_layers: int = 4
-
-    # 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.n_genes = kwargs.get("n_genes", 19_062)
+        self.n_expressions_bins = kwargs.get("n_expressions_bins", 64)
+        self.embed_dim = kwargs.get("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.num_attention_heads = kwargs.get("num_attention_heads", 8)
+        self.key_size = kwargs.get("key_size", None)
+        self.ffn_embed_dim = kwargs.get("ffn_embed_dim", 512)
+        self.num_layers = kwargs.get("num_layers", 4)
+
+        # 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