File size: 2,815 Bytes
3bbb8d3 187fa63 3bbb8d3 187fa63 a393106 187fa63 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from dataclasses import dataclass
from typing import Tuple
from transformers import PretrainedConfig
@dataclass
class sCTConfig(PretrainedConfig): # noqa: N801
model_type = "sCT"
def __init__(self, **kwargs): # type: ignore
self.alphabet_size = kwargs.get("alphabet_size", 7)
self.pad_token_id = kwargs.get("pad_token_id", 5)
self.mask_token_id = kwargs.get("mask_token_id", 6)
self.cell_len = kwargs.get("cell_len", 19968)
self.num_downsamples = kwargs.get("num_downsamples", 8)
self.attention_heads = kwargs.get("attention_heads", 16)
self.key_size = kwargs.get("key_size", None)
self.token_embed_dim = kwargs.get("token_embed_dim", 16)
self.embed_dim = kwargs.get("embed_dim", 1024)
self.ffn_embed_dim = kwargs.get("ffn_embed_dim", 2048)
self.num_layers = kwargs.get("num_layers", 4)
self.layer_norm_eps = kwargs.get("layer_norm_eps", 1e-5)
self.interpolation_method = kwargs.get("interpolation_method", "nearest")
# bad hack to satisfy cellnt_celltype_annotation.py:312
self.max_positions: int = kwargs.get("max_positions", 20480)
self.num_cells: int = kwargs.get("num_cells", 50)
self.num_hidden_layers_head: int = kwargs.get("num_hidden_layers_head", 1)
self.use_skip_connection: bool = kwargs.get("use_skip_connection", True)
# logging
self.use_gradient_checkpointing: bool = False
# 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", []
)
# Spatial info configuration
self.use_spatial_information: bool = kwargs.get(
"use_spatial_information", False
)
self.num_scales: int = kwargs.get("num_scales", 10)
self.sigma_min: float = kwargs.get("sigma_min", 1.0)
self.sigma_max: float = kwargs.get("sigma_max", 10.0)
super().__init__(**kwargs)
def __post_init__(self) -> None: # type: ignore # noqa: N807
"""
Checks that the given values are compatible.
"""
if self.key_size is None:
if not self.embed_dim % self.attention_heads == 0:
raise ValueError(
f"When no key size is provided, the embedding dimension"
f"should be divisible by the number of heads, however "
f"provided embedding dimension is {self.embed_dim} and "
f"the number of heads is {self.attention_heads}."
)
self.key_size = self.embed_dim // self.attention_heads
|