Upload model

config.json

@@ -0,0 +1,37 @@
+{
+  "_name_or_path": "Rostlab/prot_albert",
+  "architectures": [
+    "ProtAlBertModel"
+  ],
+  "attention_probs_dropout_prob": 0,
+  "auto_map": {
+    "AutoModel": "prot_albert_model.ProtAlBertModel"
+  },
+  "bos_token_id": 2,
+  "classifier_dropout_prob": 0.1,
+  "down_scale_factor": 1,
+  "embedding_size": 128,
+  "eos_token_id": 3,
+  "gap_size": 0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0,
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "inner_group_num": 1,
+  "intermediate_size": 16384,
+  "layer_norm_eps": 1e-12,
+  "layers_to_keep": [],
+  "max_position_embeddings": 40000,
+  "model_type": "albert",
+  "net_structure_type": 0,
+  "num_attention_heads": 64,
+  "num_hidden_groups": 1,
+  "num_hidden_layers": 12,
+  "num_memory_blocks": 0,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.1",
+  "type_vocab_size": 2,
+  "vocab_size": 34
+}

model.safetensors

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

prot_albert_model.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+import torch.nn as nn
+from transformers import AlbertModel, AutoConfig
+from transformers import PreTrainedModel, AutoModel, AutoConfig
+
+
+class ProtAlBertModel(PreTrainedModel):
+    def __init__(self, config: Any, num_labels: int, hidden_size: int, model_name: str, *args, **kwargs):
+        """
+        Initialise the model.
+        :param: config class for the PreTrainedModel class
+        :param hidden_size: size of the hidden layer after the CLS token.
+        :param num_labels: the number of labels.
+        :param model_name: the name of the model.
+        """
+        super().__init__(config)
+        albert_config = AutoConfig.from_pretrained(model_name)
+        self.config = AutoConfig.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+        )
+        self.protbert = AlbertModel.from_pretrained(
+            model_name, config=self.config, trust_remote_code=True
+        )
+        self.last_layer = nn.Sequential(
+            nn.Dropout(0.1),
+            nn.LayerNorm(self.config.hidden_size),
+            nn.Linear(self.config.hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, num_labels),
+        )
+
+    def forward(self, x):
+        # Take the last embedding of the [CLS] token
+        z = self.protbert(**x).last_hidden_state[:, 0, :]
+        output = self.last_layer(z)
+        return {"logits": output}