| | import copy |
| | import logging |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import numpy as np |
| |
|
| | from models.embedding_models.bert_embedding_model import BertEmbedModel |
| | from models.embedding_models.pretrained_embedding_model import PretrainedEmbedModel |
| | from modules.token_embedders.bert_encoder import BertLinear |
| | from collections import defaultdict |
| | from transformers import AutoTokenizer |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class EntRelJointDecoder(nn.Module): |
| | Argument_START_NER = '<START=Argument>'.lower() |
| | Argument_END_NER = '<END=Argument>'.lower() |
| | Relation_START_NER = '<START=Relation>'.lower() |
| | Relation_END_NER = '<END=Relation>'.lower() |
| |
|
| | def __init__(self, cfg, vocab, ent_rel_file, rel_file): |
| | """__init__ constructs `EntRelJointDecoder` components and |
| | sets `EntRelJointDecoder` parameters. This class adopts a joint |
| | decoding algorithm for entity relation joint decoing and facilitates |
| | the interaction between entity and relation. |
| | |
| | Args: |
| | cfg (dict): config parameters for constructing multiple models |
| | vocab (Vocabulary): vocabulary |
| | ent_rel_file (dict): entity and relation file (joint id, entity id, relation id, symmetric id, asymmetric id) |
| | """ |
| |
|
| | super().__init__() |
| | self.auto_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") |
| | self.cls = self.auto_tokenizer.cls_token |
| | self.sep = self.auto_tokenizer.sep_token |
| | self.rel_file = rel_file |
| | self.add_marker_tokens() |
| | self.vocab = vocab |
| | self.max_span_length = cfg.max_span_length |
| | self.activation = nn.GELU() |
| | self.device = cfg.device |
| | self.separate_threshold = cfg.separate_threshold |
| |
|
| | if cfg.embedding_model == 'bert': |
| | self.embedding_model = BertEmbedModel(cfg, vocab) |
| | elif cfg.embedding_model == 'pretrained': |
| | self.embedding_model = PretrainedEmbedModel(cfg, vocab) |
| | self.encoder_output_size = self.embedding_model.get_hidden_size() |
| |
|
| | self.head_mlp = BertLinear(input_size=self.encoder_output_size, |
| | output_size=cfg.mlp_hidden_size, |
| | activation=self.activation, |
| | dropout=cfg.dropout) |
| | self.tail_mlp = BertLinear(input_size=self.encoder_output_size, |
| | output_size=cfg.mlp_hidden_size, |
| | activation=self.activation, |
| | dropout=cfg.dropout) |
| |
|
| | self.U = nn.Parameter( |
| | torch.FloatTensor(self.vocab.get_vocab_size('ent_rel_id'), cfg.mlp_hidden_size + 1, |
| | cfg.mlp_hidden_size + 1)) |
| | self.U.data.zero_() |
| |
|
| | if cfg.logit_dropout > 0: |
| | self.logit_dropout = nn.Dropout(p=cfg.logit_dropout) |
| | else: |
| | self.logit_dropout = lambda x: x |
| |
|
| | self.none_idx = self.vocab.get_token_index('None', 'ent_rel_id') |
| |
|
| | self.symmetric_label = torch.LongTensor(ent_rel_file["symmetric"]) |
| | self.asymmetric_label = torch.LongTensor(ent_rel_file["asymmetric"]) |
| | self.ent_label = torch.LongTensor(ent_rel_file["entity"]) |
| | self.rel_label = torch.LongTensor(ent_rel_file["relation"]) |
| | |
| | if self.device > -1: |
| | self.symmetric_label = self.symmetric_label.cuda(device=self.device, non_blocking=True) |
| | self.asymmetric_label = self.asymmetric_label.cuda(device=self.device, non_blocking=True) |
| | self.ent_label = self.ent_label.cuda(device=self.device, non_blocking=True) |
| | self.rel_label = self.rel_label.cuda(device=self.device, non_blocking=True) |
| |
|
| | self.element_loss = nn.CrossEntropyLoss() |
| |
|
| | def add_marker_tokens(self): |
| | new_tokens = ['<START>', '<END>'] |
| | for label in self.rel_file["entity_text"]: |
| | new_tokens.append('<START=%s>' % label) |
| | new_tokens.append('<END=%s>' % label) |
| | self.auto_tokenizer.add_tokens(new_tokens) |
| | |
| |
|
| | def forward(self, batch_inputs, rel_model, dataset_vocab): |
| | """forward |
| | |
| | Arguments: |
| | batch_inputs {dict} -- batch input data |
| | |
| | Returns: |
| | dict -- results: ent_loss, ent_pred |
| | """ |
| |
|
| | results = {} |
| | batch_seq_tokens_lens = batch_inputs['tokens_lens'] |
| | batch_tokens = batch_inputs['tokens'] |
| |
|
| | self.embedding_model(batch_inputs) |
| | batch_seq_tokens_encoder_repr = batch_inputs['seq_encoder_reprs'] |
| |
|
| | batch_seq_tokens_head_repr = self.head_mlp(batch_seq_tokens_encoder_repr) |
| | batch_seq_tokens_head_repr = torch.cat( |
| | [batch_seq_tokens_head_repr, |
| | torch.ones_like(batch_seq_tokens_head_repr[..., :1])], dim=-1) |
| | batch_seq_tokens_tail_repr = self.tail_mlp(batch_seq_tokens_encoder_repr) |
| | batch_seq_tokens_tail_repr = torch.cat( |
| | [batch_seq_tokens_tail_repr, |
| | torch.ones_like(batch_seq_tokens_tail_repr[..., :1])], dim=-1) |
| |
|
| | batch_joint_score = torch.einsum('bxi, oij, byj -> boxy', batch_seq_tokens_head_repr, self.U, |
| | batch_seq_tokens_tail_repr).permute(0, 2, 3, 1) |
| |
|
| | batch_normalized_joint_score = torch.softmax( |
| | batch_joint_score, dim=-1) * batch_inputs['joint_label_matrix_mask'].unsqueeze(-1).float() |
| |
|
| | if not self.training: |
| | results['entity_label_preds'] = torch.argmax(batch_normalized_joint_score, dim=-1) |
| |
|
| | separate_position_preds, ent_preds, rel_preds = self.soft_joint_decoding( |
| | batch_normalized_joint_score, rel_model, batch_tokens, batch_seq_tokens_lens, dataset_vocab) |
| |
|
| | results['all_separate_position_preds'] = separate_position_preds |
| | results['all_ent_preds'] = ent_preds |
| | results['all_rel_preds'] = rel_preds |
| |
|
| | return results |
| |
|
| | results['element_loss'] = self.element_loss( |
| | self.logit_dropout(batch_joint_score[batch_inputs['joint_label_matrix_mask']]), |
| | batch_inputs['joint_label_matrix'][batch_inputs['joint_label_matrix_mask']]) |
| |
|
| | batch_symmetric_normalized_joint_score = batch_normalized_joint_score[..., self.symmetric_label] |
| | results['symmetric_loss'] = torch.abs(batch_symmetric_normalized_joint_score - |
| | batch_symmetric_normalized_joint_score.transpose(1, 2)).sum( |
| | dim=-1)[batch_inputs['joint_label_matrix_mask']].mean() |
| |
|
| | batch_rel_normalized_joint_score = torch.max(batch_normalized_joint_score[..., self.rel_label], dim=-1).values |
| | batch_diag_ent_normalized_joint_score = torch.max( |
| | batch_normalized_joint_score[..., self.ent_label].diagonal(0, 1, 2), |
| | dim=1).values.unsqueeze(-1).expand_as(batch_rel_normalized_joint_score) |
| |
|
| | results['implication_loss'] = ( |
| | torch.relu(batch_rel_normalized_joint_score - batch_diag_ent_normalized_joint_score).sum(dim=2) + |
| | torch.relu( |
| | batch_rel_normalized_joint_score.transpose(1, 2) - batch_diag_ent_normalized_joint_score).sum( |
| | dim=2))[batch_inputs['joint_label_matrix_mask'][..., 0]].mean() |
| |
|
| | relation_entity_mask = batch_inputs['joint_label_matrix'].diagonal(0, 1, 2) |
| | relation_entity_mask = torch.eq(relation_entity_mask, self.ent_label[1]) |
| |
|
| | batch_row_subject_normalized_joint_score = torch.max(batch_normalized_joint_score[..., self.rel_label[0]], |
| | dim=-1).values |
| | batch_column_subject_normalized_joint_score = torch.max( |
| | batch_normalized_joint_score.transpose(1, 2)[..., self.rel_label[0]], dim=-1).values |
| | batch_row_object_normalized_joint_score = torch.max(batch_normalized_joint_score[..., self.rel_label[1]], |
| | dim=-1).values |
| | batch_column_object_normalized_joint_score = torch.max( |
| | batch_normalized_joint_score.transpose(1, 2)[..., self.rel_label[1]], dim=-1).values |
| |
|
| | results['triple_loss'] = ( |
| | (torch.relu(batch_row_object_normalized_joint_score - batch_row_subject_normalized_joint_score) + |
| | torch.relu( |
| | batch_column_object_normalized_joint_score - batch_column_subject_normalized_joint_score)) / 2 |
| | )[relation_entity_mask].mean() |
| |
|
| | return results |
| |
|
| | def soft_joint_decoding(self, batch_normalized_entity_score, rel_model, batch_tokens, batch_seq_tokens_lens, |
| | dataset_vocab): |
| | separate_position_preds = [] |
| | ent_preds = [] |
| | rel_preds = [] |
| |
|
| | batch_normalized_entity_score = batch_normalized_entity_score.cpu().numpy() |
| | ent_label = self.ent_label.cpu().numpy() |
| | rel_label = self.rel_label.cpu().numpy() |
| |
|
| | for idx, seq_len in enumerate(batch_seq_tokens_lens): |
| | |
| | tokens = [dataset_vocab.get_token_from_index(token.item(), 'tokens') for token in |
| | batch_tokens[idx][:seq_len]] |
| |
|
| | ent_pred = {} |
| | rel_pred = {} |
| | entity_score = batch_normalized_entity_score[idx][:seq_len, :seq_len, :] |
| | entity_score = (entity_score + entity_score.transpose((1, 0, 2))) / 2 |
| |
|
| | entity_score_feature = entity_score.reshape(seq_len, -1) |
| | transposed_entity_score_feature = entity_score.transpose((1, 0, 2)).reshape(seq_len, -1) |
| | separate_pos = ( |
| | (np.linalg.norm(entity_score_feature[0:seq_len - 1] - entity_score_feature[1:seq_len], axis=1) + |
| | np.linalg.norm( |
| | transposed_entity_score_feature[0:seq_len - 1] - transposed_entity_score_feature[1:seq_len], |
| | axis=1)) |
| | * 0.5 > self.separate_threshold).nonzero()[0] |
| | separate_position_preds.append([pos.item() for pos in separate_pos]) |
| | if len(separate_pos) > 0: |
| | spans = [(0, separate_pos[0].item() + 1)] + [ |
| | (separate_pos[idx].item() + 1, separate_pos[idx + 1].item() + 1) |
| | for idx in range(len(separate_pos) - 1)] + [(separate_pos[-1].item() + 1, seq_len)] |
| | else: |
| | spans = [(0, seq_len)] |
| |
|
| | merged_spans = [(span,) for span in spans] |
| | ents = [] |
| | relations = [] |
| | arguments = [] |
| | index2span = {} |
| | for span in merged_spans: |
| | target_indices = [] |
| | for sp in span: |
| | target_indices += [idx for idx in range(sp[0], sp[1])] |
| | score = np.mean(entity_score[target_indices, :, :][:, target_indices, :], axis=(0, 1)) |
| | if not (np.max(score[ent_label]) < score[self.none_idx]): |
| | pred = ent_label[np.argmax(score[ent_label])].item() |
| | pred_label = self.vocab.get_token_from_index(pred, 'ent_rel_id') |
| | if pred_label == "Relation": |
| | relations.append(target_indices) |
| | else: |
| | arguments.append(target_indices) |
| | ents.append(target_indices) |
| | index2span[tuple(target_indices)] = span |
| | ent_pred[span] = pred_label |
| |
|
| | |
| | for rel in relations: |
| | subj_found = False |
| | obj_found = False |
| | |
| | sorted_arguments = sorted(arguments, key=lambda a: abs(a[0] - rel[0])) |
| | sorted_indices = [arguments.index(arg) for arg in sorted_arguments] |
| | argument_start_ids = [arg[0] for arg in sorted_arguments] |
| | argument_end_ids = [arg[-1] for arg in sorted_arguments] |
| | relation_indices = [] |
| | argument_indices = [] |
| | wordpiece_tokens = [self.cls] |
| | for i, token in enumerate(tokens): |
| | if i == rel[0]: |
| | relation_indices.append(len(wordpiece_tokens)) |
| | wordpiece_tokens.append(self.Relation_START_NER) |
| | if i in argument_start_ids: |
| | argument_indices.append(len(wordpiece_tokens)) |
| | wordpiece_tokens.append(self.Argument_START_NER) |
| |
|
| | tokenized_token = list(self.auto_tokenizer.tokenize(token)) |
| | wordpiece_tokens.extend(tokenized_token) |
| | if i == rel[-1]: |
| | wordpiece_tokens.append(self.Relation_END_NER) |
| | if i in argument_end_ids: |
| | wordpiece_tokens.append(self.Argument_END_NER) |
| |
|
| | wordpiece_tokens.append(self.sep) |
| | wordpiece_segment_ids = [1] * (len(wordpiece_tokens)) |
| | wordpiece_tokens = [rel_model.vocab.get_token_index(token, 'wordpiece') for token in wordpiece_tokens] |
| | rel_input = { |
| | "wordpiece_tokens": torch.LongTensor([wordpiece_tokens]), |
| | "relation_ids": torch.LongTensor([relation_indices * len(argument_indices)]), |
| | "argument_ids": torch.LongTensor([argument_indices]), |
| | "label_ids_mask": torch.LongTensor([[1] * len(argument_indices)]), |
| | "wordpiece_segment_ids": torch.LongTensor([wordpiece_segment_ids]) |
| | } |
| | output = rel_model(rel_input) |
| | output = output['label_preds'][0].cpu().numpy() |
| | sorted_output_labels = [output[i] for i in sorted_indices] |
| | assert len(argument_start_ids) == len(output) |
| | prev_subj = 0 |
| | prev_obj = 0 |
| | for idx, label_id in enumerate(sorted_output_labels): |
| | if label_id == 0 and subj_found and obj_found: |
| | break |
| |
|
| | pred_label = "None" |
| | pred_t_label = "None" |
| | score = np.mean(entity_score[rel, :, :][:, sorted_arguments[idx], :], axis=(0, 1)) |
| | score_t = np.mean(entity_score[sorted_arguments[idx], :, :][:, rel, :], axis=(0, 1)) |
| | if not (np.max(score[self.rel_label]) < score[self.none_idx]) or \ |
| | not (np.max(score_t[self.rel_label]) < score_t[self.none_idx]): |
| | pred = rel_label[np.argmax(score[self.rel_label])].item() |
| | pred_label = self.vocab.get_token_from_index(pred, 'ent_rel_id') |
| |
|
| | pred = rel_label[np.argmax(score_t[self.rel_label])].item() |
| | pred_t_label = self.vocab.get_token_from_index(pred, 'ent_rel_id') |
| |
|
| | |
| | if label_id == 1 and sorted_arguments[idx][0] > rel[-1]: |
| | obj_found = True |
| | if (pred_label == "Object" or pred_t_label == "Object") and \ |
| | (not obj_found or (prev_obj != 0 and prev_obj + 1 <= sorted_arguments[idx][0] <= prev_obj + 3)): |
| | rel_pred[(index2span[tuple(rel)], index2span[tuple(sorted_arguments[idx])])] = "Object" |
| | prev_obj = sorted_arguments[idx][-1] |
| | continue |
| |
|
| | |
| | if (label_id == 2 and sorted_arguments[idx][0] < rel[0]): |
| | if (pred_label == "Subject" or pred_t_label == "Subject") and \ |
| | (not subj_found or (prev_subj != 0 and prev_subj - 1 == sorted_arguments[idx][-1])): |
| | rel_pred[(index2span[tuple(rel)], index2span[tuple(sorted_arguments[idx])])] = "Subject" |
| | subj_found = True |
| | prev_subj = sorted_arguments[idx][0] |
| | continue |
| |
|
| | if label_id == 1 and (not subj_found or ( |
| | prev_subj != 0 and sorted_arguments[idx][-1] == prev_subj - 1)): |
| | rel_pred[(index2span[tuple(rel)], index2span[tuple(sorted_arguments[idx])])] = "Subject" |
| | subj_found = True |
| | prev_subj = sorted_arguments[idx][0] |
| |
|
| | elif label_id == 2 and (not obj_found or (prev_obj != 0 and prev_obj + 1 == sorted_arguments[idx][0])): |
| | rel_pred[(index2span[tuple(rel)], index2span[tuple(sorted_arguments[idx])])] = "Object" |
| | obj_found = True |
| | prev_obj = sorted_arguments[idx][-1] |
| |
|
| | ent_preds.append(ent_pred) |
| | rel_preds.append(rel_pred) |
| |
|
| | return separate_position_preds, ent_preds, rel_preds |
| |
|
