| from typing import List |
|
|
| from einops import rearrange |
| from enformer_pytorch import Enformer |
| from transformers import PretrainedConfig, PreTrainedModel |
|
|
| from genomics_research.segmentnt.layers.torch.segmentation_head import TorchUNetHead |
|
|
| FEATURES = [ |
| "protein_coding_gene", |
| "lncRNA", |
| "exon", |
| "intron", |
| "splice_donor", |
| "splice_acceptor", |
| "5UTR", |
| "3UTR", |
| "CTCF-bound", |
| "polyA_signal", |
| "enhancer_Tissue_specific", |
| "enhancer_Tissue_invariant", |
| "promoter_Tissue_specific", |
| "promoter_Tissue_invariant", |
| ] |
|
|
|
|
| class SegmentEnformerConfig(PretrainedConfig): |
| model_type = "segment_enformer" |
|
|
| def __init__( |
| self, |
| features: List[str] = FEATURES, |
| embed_dim: int = 1536, |
| dim_divisible_by: int = 128, |
| **kwargs |
| ): |
| self.features = features |
| self.embed_dim = embed_dim |
| self.dim_divisible_by = dim_divisible_by |
|
|
| super().__init__(**kwargs) |
|
|
|
|
| class SegmentEnformer(PreTrainedModel): |
| config_class = SegmentEnformerConfig |
|
|
| def __init__(self, config: SegmentEnformerConfig): |
| super().__init__(config=config) |
|
|
| enformer = Enformer.from_pretrained("EleutherAI/enformer-official-rough") |
|
|
| self.stem = enformer.stem |
| self.conv_tower = enformer.conv_tower |
| self.transformer = enformer.transformer |
|
|
| self.unet_head = TorchUNetHead( |
| features=config.features, |
| embed_dimension=config.embed_dim, |
| nucl_per_token=config.dim_divisible_by, |
| remove_cls_token=False, |
| ) |
|
|
| def __call__(self, x): |
| x = rearrange(x, "b n d -> b d n") |
| x = self.stem(x) |
|
|
| x = self.conv_tower(x) |
|
|
| x = rearrange(x, "b d n -> b n d") |
| x = self.transformer(x) |
|
|
| x = rearrange(x, "b n d -> b d n") |
| x = self.unet_head(x) |
|
|
| return x |
|
|
