| 5 | |
| 6 | |
| 7 | class MultiInferBert(torch.nn.Module): |
| 8 | def __init__(self, args): |
| 9 | super(MultiInferBert, self).__init__() |
| 10 | |
| 11 | self.args = args |
| 12 | self.bert = BertModel.from_pretrained(args.bert_model_path) |
| 13 | self.tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer_path) |
| 14 | |
| 15 | self.cls_linear = torch.nn.Linear(args.bert_feature_dim*2, args.class_num) |
| 16 | self.feature_linear = torch.nn.Linear(args.bert_feature_dim*2 + args.class_num*3, args.bert_feature_dim*2) |
| 17 | self.dropout_output = torch.nn.Dropout(0.1) |
| 18 | |
| 19 | def multi_hops(self, features, mask, k): |
| 20 | '''generate mask''' |
| 21 | max_length = features.shape[1] |
| 22 | mask = mask[:, :max_length] |
| 23 | mask_a = mask.unsqueeze(1).expand([-1, max_length, -1]) |
| 24 | mask_b = mask.unsqueeze(2).expand([-1, -1, max_length]) |
| 25 | mask = mask_a * mask_b |
| 26 | mask = torch.triu(mask).unsqueeze(3).expand([-1, -1, -1, self.args.class_num]) |
| 27 | |
| 28 | '''save all logits''' |
| 29 | logits_list = [] |
| 30 | logits = self.cls_linear(features) |
| 31 | logits_list.append(logits) |
| 32 | |
| 33 | for i in range(k): |
| 34 | #probs = torch.softmax(logits, dim=3) |
| 35 | probs = logits |
| 36 | logits = probs * mask |
| 37 | |
| 38 | logits_a = torch.max(logits, dim=1)[0] |
| 39 | logits_b = torch.max(logits, dim=2)[0] |
| 40 | logits = torch.cat([logits_a.unsqueeze(3), logits_b.unsqueeze(3)], dim=3) |
| 41 | logits = torch.max(logits, dim=3)[0] |
| 42 | |
| 43 | logits = logits.unsqueeze(2).expand([-1,-1, max_length, -1]) |
| 44 | logits_T = logits.transpose(1, 2) |
| 45 | logits = torch.cat([logits, logits_T], dim=3) |
| 46 | |
| 47 | new_features = torch.cat([features, logits, probs], dim=3) |
| 48 | features = self.feature_linear(new_features) |
| 49 | logits = self.cls_linear(features) |
| 50 | logits_list.append(logits) |
| 51 | return logits_list |
| 52 | |
| 53 | def forward(self, tokens, masks): |
| 54 | bert_feature, _ = self.bert(tokens, masks) |
| 55 | bert_feature = self.dropout_output(bert_feature) |
| 56 | |
| 57 | bert_feature = bert_feature.unsqueeze(2).expand([-1, -1, self.args.max_sequence_len, -1]) |
| 58 | bert_feature_T = bert_feature.transpose(1, 2) |
| 59 | features = torch.cat([bert_feature, bert_feature_T], dim=3) |
| 60 | logits = self.multi_hops(features, masks, self.args.nhops) |
| 61 | |
| 62 | return logits[-1] |