-
Notifications
You must be signed in to change notification settings - Fork 1
/
bert.py
34 lines (25 loc) · 899 Bytes
/
bert.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
import torch
import torch.nn as nn
class mBERT(nn.Module):
def __init__(self, bert, max_length, slot_classes, num_classes, hidden_layers, drop=0.1):
super(mBERT, self).__init__()
self.bert = bert
#self.relu = nn.ReLU()
#self.dropout = nn.Dropout(drop)
self.dropout = nn.Dropout(p=drop, inplace=False)
self.fcSl = nn.Linear(hidden_layers,slot_classes)
#define forward pass
def forward(self, sent_id, mask, labels=None):
if labels is None:
out = self.bert(sent_id,
token_type_ids=None,
attention_mask=mask)
else:
out = self.bert(sent_id,
token_type_ids=None,
attention_mask=mask,
labels=labels)
slots_vec = out.hidden_states[-1]
slots = self.dropout(slots_vec)
slots = self.fcSl(slots)
return out[0], slots