-
Notifications
You must be signed in to change notification settings - Fork 16
/
model.py
247 lines (218 loc) · 10.7 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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
# /usr/bin/env python
# coding=utf-8
"""model"""
from collections import Counter
import torch
import torch.nn as nn
from transformers import BertPreTrainedModel, BertModel
class MultiNonLinearClassifier(nn.Module):
def __init__(self, hidden_size, tag_size, dropout_rate):
super(MultiNonLinearClassifier, self).__init__()
self.tag_size = tag_size
self.linear = nn.Linear(hidden_size, int(hidden_size / 2))
self.hidden2tag = nn.Linear(int(hidden_size / 2), self.tag_size)
self.dropout = nn.Dropout(dropout_rate)
def forward(self, input_features):
features_tmp = self.linear(input_features)
features_tmp = nn.ReLU()(features_tmp)
features_tmp = self.dropout(features_tmp)
features_output = self.hidden2tag(features_tmp)
return features_output
class SequenceLabelForSO(nn.Module):
def __init__(self, hidden_size, tag_size, dropout_rate):
super(SequenceLabelForSO, self).__init__()
self.tag_size = tag_size
self.linear = nn.Linear(hidden_size, int(hidden_size / 2))
self.hidden2tag_sub = nn.Linear(int(hidden_size / 2), self.tag_size)
self.hidden2tag_obj = nn.Linear(int(hidden_size / 2), self.tag_size)
self.dropout = nn.Dropout(dropout_rate)
def forward(self, input_features):
"""
Args:
input_features: (bs, seq_len, h)
"""
features_tmp = self.linear(input_features)
features_tmp = nn.ReLU()(features_tmp)
features_tmp = self.dropout(features_tmp)
sub_output = self.hidden2tag_sub(features_tmp)
obj_output = self.hidden2tag_obj(features_tmp)
return sub_output, obj_output
class BertForRE(BertPreTrainedModel):
def __init__(self, config, params):
super().__init__(config)
self.max_seq_len = params.max_seq_length
self.seq_tag_size = params.seq_tag_size
self.rel_num = params.rel_num
# pretrain model
self.bert = BertModel(config)
# sequence tagging
self.sequence_tagging_sub = MultiNonLinearClassifier(config.hidden_size * 2, self.seq_tag_size, params.drop_prob)
self.sequence_tagging_obj = MultiNonLinearClassifier(config.hidden_size * 2, self.seq_tag_size, params.drop_prob)
self.sequence_tagging_sum = SequenceLabelForSO(config.hidden_size, self.seq_tag_size, params.drop_prob)
# global correspondence
self.global_corres = MultiNonLinearClassifier(config.hidden_size * 2, 1, params.drop_prob)
# relation judgement
self.rel_judgement = MultiNonLinearClassifier(config.hidden_size, params.rel_num, params.drop_prob)
self.rel_embedding = nn.Embedding(params.rel_num, config.hidden_size)
self.init_weights()
@staticmethod
def masked_avgpool(sent, mask):
mask_ = mask.masked_fill(mask == 0, -1e9).float()
score = torch.softmax(mask_, -1)
return torch.matmul(score.unsqueeze(1), sent).squeeze(1)
def forward(
self,
input_ids=None,
attention_mask=None,
seq_tags=None,
potential_rels=None,
corres_tags=None,
rel_tags=None,
ex_params=None
):
"""
Args:
input_ids: (batch_size, seq_len)
attention_mask: (batch_size, seq_len)
rel_tags: (bs, rel_num)
potential_rels: (bs,), only in train stage.
seq_tags: (bs, 2, seq_len)
corres_tags: (bs, seq_len, seq_len)
ex_params: experiment parameters
"""
# get params for experiments
corres_threshold, rel_threshold = ex_params.get('corres_threshold', 0.5), ex_params.get('rel_threshold', 0.1)
# ablation study
ensure_corres, ensure_rel = ex_params['ensure_corres'], ex_params['ensure_rel']
# pre-train model
outputs = self.bert(
input_ids,
attention_mask=attention_mask,
output_hidden_states=True
) # sequence_output, pooled_output, (hidden_states), (attentions)
sequence_output = outputs[0]
bs, seq_len, h = sequence_output.size()
if ensure_rel:
# (bs, h)
h_k_avg = self.masked_avgpool(sequence_output, attention_mask)
# (bs, rel_num)
rel_pred = self.rel_judgement(h_k_avg)
# before fuse relation representation
if ensure_corres:
# for every position $i$ in sequence, should concate $j$ to predict.
sub_extend = sequence_output.unsqueeze(2).expand(-1, -1, seq_len, -1) # (bs, s, s, h)
obj_extend = sequence_output.unsqueeze(1).expand(-1, seq_len, -1, -1) # (bs, s, s, h)
# batch x seq_len x seq_len x 2*hidden
corres_pred = torch.cat([sub_extend, obj_extend], 3)
# (bs, seq_len, seq_len)
corres_pred = self.global_corres(corres_pred).squeeze(-1)
mask_tmp1 = attention_mask.unsqueeze(-1)
mask_tmp2 = attention_mask.unsqueeze(1)
corres_mask = mask_tmp1 * mask_tmp2
# relation predict and data construction in inference stage
xi, pred_rels = None, None
if ensure_rel and seq_tags is None:
# (bs, rel_num)
rel_pred_onehot = torch.where(torch.sigmoid(rel_pred) > rel_threshold,
torch.ones(rel_pred.size(), device=rel_pred.device),
torch.zeros(rel_pred.size(), device=rel_pred.device))
# if potential relation is null
for idx, sample in enumerate(rel_pred_onehot):
if 1 not in sample:
# (rel_num,)
max_index = torch.argmax(rel_pred[idx])
sample[max_index] = 1
rel_pred_onehot[idx] = sample
# 2*(sum(x_i),)
bs_idxs, pred_rels = torch.nonzero(rel_pred_onehot, as_tuple=True)
# get x_i
xi_dict = Counter(bs_idxs.tolist())
xi = [xi_dict[idx] for idx in range(bs)]
pos_seq_output = []
pos_potential_rel = []
pos_attention_mask = []
for bs_idx, rel_idx in zip(bs_idxs, pred_rels):
# (seq_len, h)
pos_seq_output.append(sequence_output[bs_idx])
pos_attention_mask.append(attention_mask[bs_idx])
pos_potential_rel.append(rel_idx)
# (sum(x_i), seq_len, h)
sequence_output = torch.stack(pos_seq_output, dim=0)
# (sum(x_i), seq_len)
attention_mask = torch.stack(pos_attention_mask, dim=0)
# (sum(x_i),)
potential_rels = torch.stack(pos_potential_rel, dim=0)
# ablation of relation judgement
elif not ensure_rel and seq_tags is None:
# construct test data
sequence_output = sequence_output.repeat((1, self.rel_num, 1)).view(bs * self.rel_num, seq_len, h)
attention_mask = attention_mask.repeat((1, self.rel_num)).view(bs * self.rel_num, seq_len)
potential_rels = torch.arange(0, self.rel_num, device=input_ids.device).repeat(bs)
# (bs/sum(x_i), h)
rel_emb = self.rel_embedding(potential_rels)
# relation embedding vector fusion
rel_emb = rel_emb.unsqueeze(1).expand(-1, seq_len, h)
if ex_params['emb_fusion'] == 'concat':
# (bs/sum(x_i), seq_len, 2*h)
decode_input = torch.cat([sequence_output, rel_emb], dim=-1)
# (bs/sum(x_i), seq_len, tag_size)
output_sub = self.sequence_tagging_sub(decode_input)
output_obj = self.sequence_tagging_obj(decode_input)
elif ex_params['emb_fusion'] == 'sum':
# (bs/sum(x_i), seq_len, h)
decode_input = sequence_output + rel_emb
# (bs/sum(x_i), seq_len, tag_size)
output_sub, output_obj = self.sequence_tagging_sum(decode_input)
# train
if seq_tags is not None:
# calculate loss
attention_mask = attention_mask.view(-1)
# sequence label loss
loss_func = nn.CrossEntropyLoss(reduction='none')
loss_seq_sub = (loss_func(output_sub.view(-1, self.seq_tag_size),
seq_tags[:, 0, :].reshape(-1)) * attention_mask).sum() / attention_mask.sum()
loss_seq_obj = (loss_func(output_obj.view(-1, self.seq_tag_size),
seq_tags[:, 1, :].reshape(-1)) * attention_mask).sum() / attention_mask.sum()
loss_seq = (loss_seq_sub + loss_seq_obj) / 2
# init
loss_matrix, loss_rel = torch.tensor(0), torch.tensor(0)
if ensure_corres:
corres_pred = corres_pred.view(bs, -1)
corres_mask = corres_mask.view(bs, -1)
corres_tags = corres_tags.view(bs, -1)
loss_func = nn.BCEWithLogitsLoss(reduction='none')
loss_matrix = (loss_func(corres_pred,
corres_tags.float()) * corres_mask).sum() / corres_mask.sum()
if ensure_rel:
loss_func = nn.BCEWithLogitsLoss(reduction='mean')
loss_rel = loss_func(rel_pred, rel_tags.float())
loss = loss_seq + loss_matrix + loss_rel
return loss, loss_seq, loss_matrix, loss_rel
# inference
else:
# (sum(x_i), seq_len)
pred_seq_sub = torch.argmax(torch.softmax(output_sub, dim=-1), dim=-1)
pred_seq_obj = torch.argmax(torch.softmax(output_obj, dim=-1), dim=-1)
# (sum(x_i), 2, seq_len)
pred_seqs = torch.cat([pred_seq_sub.unsqueeze(1), pred_seq_obj.unsqueeze(1)], dim=1)
if ensure_corres:
corres_pred = torch.sigmoid(corres_pred) * corres_mask
# (bs, seq_len, seq_len)
pred_corres_onehot = torch.where(corres_pred > corres_threshold,
torch.ones(corres_pred.size(), device=corres_pred.device),
torch.zeros(corres_pred.size(), device=corres_pred.device))
return pred_seqs, pred_corres_onehot, xi, pred_rels
return pred_seqs, xi, pred_rels
if __name__ == '__main__':
from transformers import BertConfig
import utils
import os
params = utils.Params()
# Prepare model
bert_config = BertConfig.from_json_file(os.path.join(params.bert_model_dir, 'bert_config.json'))
model = BertForRE.from_pretrained(config=bert_config,
pretrained_model_name_or_path=params.bert_model_dir,
params=params)
model.to(params.device)
for n, _ in model.named_parameters():
print(n)