forked from ojas-sethi/enhancing-bert-using-syntax
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmodified_multi_label_classifier.py
65 lines (48 loc) · 1.96 KB
/
modified_multi_label_classifier.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import torch
import torch.nn as nn
from bert_model import BertPooler, BertPreTrainedModel, BertModel
class ModifiedBertForMultiLabelClassification(BertPreTrainedModel):
"""
First Version of Modified Go Emotions Bert Model (Model 2 in our writeup)
Output of first layer and final layer is given to the classifier.
"""
def __init__(self, config):
super().__init__(config)
self.num_labels = config.num_labels
self.bert = BertModel(config)
self.hidden_pooler = BertPooler(config)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.classifier = nn.Linear(config.hidden_size * 2, self.config.num_labels)
self.loss_fct = nn.BCEWithLogitsLoss()
self.init_weights()
def forward(
self,
input_ids=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
labels=None,
):
outputs = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
output_hidden_states=True,
return_dict=True,
)
first_hidden_state_result = outputs.hidden_states[0]
first_hidden_state_result = self.hidden_pooler(first_hidden_state_result)
pooled_output = outputs.pooler_output
pooled_output = torch.cat((pooled_output, first_hidden_state_result), 1)
pooled_output = self.dropout(pooled_output)
logits = self.classifier(pooled_output)
outputs = (logits,) + outputs[2:] # add hidden states and attention if they are here
if labels is not None:
loss = self.loss_fct(logits, labels)
outputs = (loss,) + outputs
return outputs # (loss), logits, (hidden_states), (attentions)