add model

model.py

@@ -29,6 +29,7 @@ class BertMesh(PreTrainedModel):
         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.id2label = self.config.id2label
 
         self.bert = AutoModel.from_pretrained(self.pretrained_model)  # 768
         self.multilabel_attention_layer = MultiLabelAttention(
@@ -39,7 +40,7 @@ class BertMesh(PreTrainedModel):
         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, **kwargs):
+    def forward(self, input_ids, return_labels=False, **kwargs):
         if type(input_ids) is list:
             # coming from tokenizer
             input_ids = torch.tensor(input_ids)
@@ -55,4 +56,7 @@ class BertMesh(PreTrainedModel):
             outs = torch.nn.functional.relu(self.linear_1(cls))
             outs = self.dropout_layer(outs)
             outs = torch.sigmoid(self.linear_out(outs))
+        if return_labels:
+            # TODO Vectorize
+           outs =  [[self.id2label[label_id] for label_id, label_prob in enumerate(out) if label_prob > 0.5] for out in outs]
         return outs