Upload model

config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "RNATorsionBERTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "rna_torsionbert_config.RNATorsionBertConfig",
+    "AutoModel": "rna_torsionbert_model.RNATorsionBERTModel"
+  },
+  "hidden_size": 1024,
+  "k": 3,
+  "model_type": "rna_torsionbert",
+  "num_classes": 18,
+  "torch_dtype": "float32",
+  "transformers_version": "4.29.0"
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e770ee16a98d493a917a3b445c84ae2b4c22e5a94fa129e0bc9d222107596b9
+size 347678917

rna_torsionbert_config.py

@@ -0,0 +1,16 @@
+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):
+        """
+        Initialise the model.
+        :param k: the k-mer size.
+        :param num_classes: the number of labels.
+        :param hidden_size: size of the hidden layer after BERT hidden states.
+        """
+        self.k = k
+        self.num_classes = num_classes
+        self.hidden_size = hidden_size
+        super().__init__(**kwargs)

rna_torsionbert_model.py

@@ -0,0 +1,38 @@
+from torch import nn
+from transformers import PreTrainedModel, AutoModel, AutoConfig
+
+from .rna_torsionbert_config import RNATorsionBertConfig
+
+
+class RNATorsionBERTModel(PreTrainedModel):
+    config_class = RNATorsionBertConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.init_model(config.k)
+        self.dnabert = AutoModel.from_pretrained(
+            self.model_name, config=self.dnabert_config, trust_remote_code=True
+        )
+        self.regressor = nn.Sequential(
+            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.Softmax(dim=-1)
+        )
+
+    def init_model(self, k: int):
+        model_name = f"zhihan1996/DNA_bert_{k}"
+        revisions = {3: "ed28178", 4: "c8499f0", 5: "c296157", 6: "a79a8fd"}
+        dnabert_config = AutoConfig.from_pretrained(
+            model_name,
+            revision=revisions[k],
+            trust_remote_code=True,
+        )
+        self.dnabert_config = dnabert_config
+        self.model_name = model_name
+
+    def forward(self, tensor):
+        z = self.dnabert(**tensor).last_hidden_state
+        output = self.regressor(z)
+        return {"logits": output}