/
model.py
58 lines (51 loc) · 2.4 KB
/
model.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
from transformers.models.bert.modeling_bert import *
from torch.nn.utils.rnn import pad_sequence
from torchcrf import CRF
class BertNER(BertPreTrainedModel):
def __init__(self, config):
super(BertNER, self).__init__(config)
self.num_labels = config.num_labels
self.bert = BertModel(config)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.bilstm = nn.LSTM(
input_size=config.lstm_embedding_size, # 1024 768
hidden_size=config.hidden_size // 2, # 1024
batch_first=True,
num_layers=2,
dropout=config.lstm_dropout_prob, # 0.5
bidirectional=True
)
#emisson matrix
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
self.softmax = nn.Softmax(dim=1)
self.crf = CRF(config.num_labels, batch_first=True)
self.init_weights()
def forward(self, input_data, token_type_ids=None, attention_mask=None, labels=None,
position_ids=None, inputs_embeds=None, head_mask=None):
input_ids, input_token_starts = input_data
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)
sequence_output = outputs[0]
# 去除[CLS]标签等位置,获得与label对齐的pre_label表示
origin_sequence_output = [layer[starts.nonzero().squeeze(1)]
for layer, starts in zip(sequence_output, input_token_starts)]
# 将sequence_output的pred_label维度padding到最大长度
padded_sequence_output = pad_sequence(origin_sequence_output, batch_first=True)
# dropout pred_label的一部分feature
padded_sequence_output = self.dropout(padded_sequence_output)
lstm_output, _ = self.bilstm(padded_sequence_output)
# 得到判别值
logits = self.classifier(lstm_output)
#score=self.softmax(logits)
outputs = (logits,)
#outputs = (score,)
if labels is not None:
loss_mask = labels.gt(-1)
loss = self.crf(logits, labels, loss_mask) * (-1)
outputs = (loss,) + outputs
# contain: (loss), scores
return outputs