Update WellcomeBertMesh with transformers based trained model

model.py

@@ -1,4 +1,4 @@
-from transformers import AutoModel, AutoConfig, PreTrainedModel
+from transformers import AutoModel, PreTrainedModel, BertConfig
 import torch
 
 
@@ -16,34 +16,33 @@ class MultiLabelAttention(torch.nn.Module):
 
 
 class BertMesh(PreTrainedModel):
+    config_class = BertConfig
+
     def __init__(
         self,
         config,
-        pretrained_model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract",
-        num_labels=28761,
-        hidden_size=1024,
-        dropout=0,
-        multilabel_attention=True,
     ):
         super().__init__(config=config)
-        self.config.auto_map = {"AutoModel": "transformers_model.BertMesh"}
-        self.pretrained_model = pretrained_model
-        self.num_labels = num_labels
-        self.hidden_size = hidden_size
-        self.dropout = dropout
-        self.multilabel_attention = multilabel_attention
-
-        self.bert = AutoModel.from_pretrained(pretrained_model)  # 768
+        self.config.auto_map = {"AutoModel": "model.BertMesh"}
+        self.pretrained_model = self.config.pretrained_model
+        self.num_labels = self.config.num_labels
+        self.hidden_size = getattr(self.config, "hidden_size", 512)
+        self.dropout = getattr(self.config, "dropout", 0.1)
+        self.multilabel_attention = getattr(self.config, "multilabel_attention", False)
+
+        self.bert = AutoModel.from_pretrained(self.pretrained_model)  # 768
         self.multilabel_attention_layer = MultiLabelAttention(
-            768, num_labels
+            768, self.num_labels
         )  # num_labels, 768
-        self.linear_1 = torch.nn.Linear(768, hidden_size)  # num_labels, 512
-        self.linear_2 = torch.nn.Linear(hidden_size, 1)  # num_labels, 1
-        self.linear_out = torch.nn.Linear(hidden_size, num_labels)
+        self.linear_1 = torch.nn.Linear(768, self.hidden_size)  # num_labels, 512
+        self.linear_2 = torch.nn.Linear(self.hidden_size, 1)  # num_labels, 1
+        self.linear_out = torch.nn.Linear(self.hidden_size, self.num_labels)
         self.dropout_layer = torch.nn.Dropout(self.dropout)
 
-    def forward(self, input_ids, token_type_ids=None, attention_mask=None):
-        input_ids = torch.tensor(input_ids)
+    def forward(self, input_ids, **kwargs):
+        if type(input_ids) is list:
+            # coming from tokenizer
+            input_ids = torch.tensor(input_ids)
         if self.multilabel_attention:
             hidden_states = self.bert(input_ids=input_ids)[0]
             attention_outs = self.multilabel_attention_layer(hidden_states)
@@ -57,6 +56,3 @@ class BertMesh(PreTrainedModel):
             outs = self.dropout_layer(outs)
             outs = torch.sigmoid(self.linear_out(outs))
         return outs
-
-    def _init_weights(self, module):
-        pass

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9588cca776ff05298d33d4785a1053cafddc231e3ebc069bc8137ed28c647fa0
-size 647426403
+oid sha256:1c80db3a392fe08b3faa111d46e48fef56eb2c0efe862f0a80cc7fe4da55baea
+size 647442531