Upload model

config.json

@@ -1,4 +1,5 @@
 {
+  "angles": "BACKBONE",
   "architectures": [
     "RNATorsionBERTModel"
   ],
@@ -9,7 +10,6 @@
   "hidden_size": 1024,
   "k": 3,
   "model_type": "rna_torsionbert",
-  "num_classes": 18,
   "torch_dtype": "float32",
-  "transformers_version": "4.29.0"
+  "transformers_version": "4.40.1"
 }

enums.py

@@ -0,0 +1,22 @@
+BACKBONE = [
+    "alpha",
+    "beta",
+    "gamma",
+    "delta",
+    "epsilon",
+    "zeta",
+    "chi",
+    "eta",
+    "theta",
+    "eta'",
+    "theta'",
+    "v0",
+    "v1",
+    "v2",
+    "v3",
+    "v4"
+]
+
+ANGLES_TO_LIST = {
+    "BACKBONE": BACKBONE
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a48154892c8697ebb962ba1e31ecb4bedad20c0de4e0836661955b29f54309
+size 347688872

rna_torsionbert_config.py

@@ -3,14 +3,15 @@ from transformers import PretrainedConfig
 
 class RNATorsionBertConfig(PretrainedConfig):
     model_type = "rna_torsionbert"
-    def __init__(self, k: int = 3, num_classes: int = 18, hidden_size: int = 1024, **kwargs):
+    def __init__(self, k: int = 3, angles: str = "BACKBONE", hidden_size: int = 1024, **kwargs):
         """
         Initialise the model.
         :param k: the k-mer size.
-        :param num_classes: the number of labels.
+        :param angles: types of angles to use. "BACKBONE" for eight torsional angles +
+                                                            two pseudo-torsional angles
         :param hidden_size: size of the hidden layer after BERT hidden states.
         """
         self.k = k
-        self.num_classes = num_classes
+        self.angles = angles
         self.hidden_size = hidden_size
         super().__init__(**kwargs)

rna_torsionbert_model.py

@@ -2,6 +2,7 @@ from torch import nn
 from transformers import PreTrainedModel, AutoModel, AutoConfig
 
 from .rna_torsionbert_config import RNATorsionBertConfig
+from .enums import ANGLES_TO_LIST
 
 
 class RNATorsionBERTModel(PreTrainedModel):
@@ -10,6 +11,7 @@ class RNATorsionBERTModel(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.init_model(config.k)
+        self.num_labels = 2 * len(ANGLES_TO_LIST.get(config.angles, []))
         self.dnabert = AutoModel.from_pretrained(
             self.model_name, config=self.dnabert_config, trust_remote_code=True
         )
@@ -17,7 +19,7 @@ class RNATorsionBERTModel(PreTrainedModel):
             nn.LayerNorm(self.dnabert_config.hidden_size),
             nn.Linear(self.dnabert_config.hidden_size, config.hidden_size),
             nn.GELU(),
-            nn.Linear(config.hidden_size, config.num_classes),
+            nn.Linear(config.hidden_size, self.num_labels),
         )
         self.activation = nn.Tanh()