Upload model

bert_classification_model.py

@@ -0,0 +1,24 @@
+import torch
+
+from typing import Tuple
+from transformers import PreTrainedModel, BertModel
+from torch import nn
+
+
+class BertClassificationModel(PreTrainedModel):
+
+    def __init__(self, config, num_main_segment=None, num_sub_segment=None):
+        super(BertClassificationModel, self).__init__(config=config)
+        self.num_main_segment = num_main_segment if num_main_segment else config.num_main_segment
+        self.num_sub_segment = num_sub_segment if num_sub_segment else config.num_sub_segment
+        self.bert = BertModel.from_pretrained("bert-base-multilingual-uncased")
+        self.dropout = nn.Dropout(0.1)
+        self.hidden_2 = nn.Linear(768, 768)
+        self.fc_main = nn.Linear(768, self.num_main_segment)
+        self.fc_sub = nn.Linear(768, self.num_sub_segment)
+
+    def forward(self, input_ids: torch.tensor, attention_masks: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        outputs = self.bert(input_ids=input_ids,
+                            attention_mask=attention_masks)
+        last_hidden_state_cls = outputs[0][:, 0, :]
+        return self.fc_main(last_hidden_state_cls), self.fc_sub(last_hidden_state_cls)

config.json

@@ -3,6 +3,9 @@
     "BertClassificationModel"
   ],
   "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModel": "bert_classification_model.BertClassificationModel"
+  },
   "classifier_dropout": null,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c8640f9aa43de09fc7f3aacdb97faf1ef0095821627f70315bd3a986d16221fe
+oid sha256:18b607e5e17d75afbb4b65e8c03fda9e28d87a1c9259a4c9cc5f291696317003
 size 671857125