Upload SegmentEnformer

segment_enformer.py

@@ -1,10 +1,13 @@
-from typing import List
+from typing import Any, Dict, List
 
+import torch
 from einops import rearrange
 from enformer_pytorch import Enformer
 from transformers import PretrainedConfig, PreTrainedModel
 
-from genomics_research.segmentnt.layers.torch.segmentation_head import TorchUNetHead
+from genomics_research.segmentnt.porting_to_pytorch.layers.segmentation_head import (
+    TorchUNetHead,
+)
 
 FEATURES = [
     "protein_coding_gene",
@@ -32,8 +35,8 @@ class SegmentEnformerConfig(PretrainedConfig):
         features: List[str] = FEATURES,
         embed_dim: int = 1536,
         dim_divisible_by: int = 128,
-        **kwargs
-    ):
+        **kwargs: Dict[str, Any]
+    ) -> None:
         self.features = features
         self.embed_dim = embed_dim
         self.dim_divisible_by = dim_divisible_by
@@ -44,7 +47,7 @@ class SegmentEnformerConfig(PretrainedConfig):
 class SegmentEnformer(PreTrainedModel):
     config_class = SegmentEnformerConfig
 
-    def __init__(self, config: SegmentEnformerConfig):
+    def __init__(self, config: SegmentEnformerConfig) -> None:
         super().__init__(config=config)
 
         enformer = Enformer.from_pretrained("EleutherAI/enformer-official-rough")
@@ -60,7 +63,7 @@ class SegmentEnformer(PreTrainedModel):
             remove_cls_token=False,
         )
 
-    def __call__(self, x):
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
         x = rearrange(x, "b n d -> b d n")
         x = self.stem(x)