Upload SegmentNT

config.json

@@ -10,6 +10,22 @@
   },
   "emb_layer_norm_before": false,
   "esmfold_config": null,
+  "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"
+  ],
   "hidden_dropout_prob": 0.0,
   "hidden_size": 1024,
   "initializer_range": 0.02,
@@ -20,7 +36,6 @@
   "max_position_embeddings": 2050,
   "model_type": "esm",
   "num_attention_heads": 16,
-  "num_features": 14,
   "num_hidden_layers": 29,
   "num_layers_head": 2,
   "pad_token_id": 1,

modeling_segment_nt.py

@@ -1159,6 +1159,7 @@ class SegmentNT(EsmPreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.config = config
+        self.num_features = len(config.features)
 
         self.esm = EsmModel(config, add_pooling_layer=False)
 
@@ -1171,7 +1172,7 @@ class SegmentNT(EsmPreTrainedModel):
                 embed_dim * (2**i) for i in range(num_layers)
             ),
         )
-        self.fc = nn.Linear(in_features=embed_dim, out_features=6 * 2 * config.num_features)
+        self.fc = nn.Linear(in_features=embed_dim, out_features=6 * 2 * self.num_features)
         self.activation_fn = nn.SiLU()
 
         self.init_weights()
@@ -1232,7 +1233,7 @@ class SegmentNT(EsmPreTrainedModel):
         logits = self.fc(x)
 
         # Final reshape to have logits per nucleotides, per feature
-        logits = torch.reshape(logits, (x.shape[0], x.shape[1] * 6, self.config.num_features, 2))
+        logits = torch.reshape(logits, (x.shape[0], x.shape[1] * 6, self.num_features, 2))
 
         # Add logits to the ESM outputs
         outputs["logits"] = logits

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b3b9da828b4a2058d0c66c7b6e0249acd2e54b755c8fa300a8f7a5cb0ee19ebd
+oid sha256:deed8fda3b2b9a76ee11aade732666b8303f97f4ecd6ab8aa6e3e1a9de9e1aa5
 size 2237478985

segment_nt_config.py

@@ -15,7 +15,7 @@
 """ ESM model configuration"""
 
 from dataclasses import asdict, dataclass
-from typing import Optional
+from typing import List, Optional
 
 from transformers import PretrainedConfig, logging
 
@@ -99,6 +99,7 @@ class SegmentNTConfig(PretrainedConfig):
 
     def __init__(
         self,
+        features=None,
         vocab_size=None,
         mask_token_id=None,
         pad_token_id=None,
@@ -121,7 +122,6 @@ class SegmentNTConfig(PretrainedConfig):
         add_bias_fnn=True,
         rescaling_factor=None,
         num_layers_head=2,
-        num_features=14,
         **kwargs,
     ):
         super().__init__(
@@ -147,7 +147,7 @@ class SegmentNTConfig(PretrainedConfig):
         self.add_bias_fnn = add_bias_fnn
         # Arguments needed for Segment NT
         self.num_layers_head = num_layers_head
-        self.num_features = num_features
+        self.features = features
         self.rescaling_factor = rescaling_factor
         if is_folding_model:
             if esmfold_config is None: