|
|
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 modules.token_embedders.bert_encoder import BertLayerNorm |
|
|
from transformers import BertModel |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
class RelDecoder(nn.Module): |
|
|
def __init__(self, cfg, vocab, ent_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.num_labels = 3 |
|
|
self.vocab = vocab |
|
|
self.max_span_length = cfg.max_span_length |
|
|
self.device = cfg.device |
|
|
|
|
|
|
|
|
self.embedding_model = BertEmbedModel(cfg, vocab, True) |
|
|
|
|
|
|
|
|
self.encoder_output_size = self.embedding_model.get_hidden_size() |
|
|
|
|
|
self.layer_norm = BertLayerNorm(self.encoder_output_size * 2) |
|
|
self.classifier = nn.Linear(self.encoder_output_size * 2, self.num_labels) |
|
|
self.classifier.weight.data.normal_(mean=0.0, std=0.02) |
|
|
self.classifier.bias.data.zero_() |
|
|
|
|
|
if cfg.logit_dropout > 0: |
|
|
self.dropout = nn.Dropout(p=cfg.logit_dropout) |
|
|
else: |
|
|
self.dropout = lambda x: x |
|
|
|
|
|
self.none_idx = self.vocab.get_token_index('None', 'ent_rel_id') |
|
|
|
|
|
self.rel_label = torch.LongTensor(ent_rel_file["relation"]) |
|
|
if self.device > -1: |
|
|
self.rel_label = self.rel_label.cuda(device=self.device, non_blocking=True) |
|
|
|
|
|
self.element_loss = nn.CrossEntropyLoss() |
|
|
|
|
|
def forward(self, batch_inputs): |
|
|
"""forward |
|
|
|
|
|
Arguments: |
|
|
batch_inputs {dict} -- batch input data |
|
|
|
|
|
Returns: |
|
|
dict -- results: ent_loss, ent_pred |
|
|
""" |
|
|
|
|
|
results = {} |
|
|
|
|
|
self.embedding_model(batch_inputs) |
|
|
batch_seq_tokens_encoder_repr = batch_inputs['seq_encoder_reprs'] |
|
|
relation_tokens = batch_seq_tokens_encoder_repr[torch.arange(batch_seq_tokens_encoder_repr.shape[0]).unsqueeze(-1), |
|
|
batch_inputs["relation_ids"]] |
|
|
argument_tokens = batch_seq_tokens_encoder_repr[torch.arange(batch_seq_tokens_encoder_repr.shape[0]).unsqueeze(-1), |
|
|
batch_inputs["argument_ids"]] |
|
|
batch_input_rep = torch.cat((relation_tokens, argument_tokens), dim=-1) |
|
|
batch_input_rep = self.layer_norm(batch_input_rep) |
|
|
batch_input_rep = self.dropout(batch_input_rep) |
|
|
batch_logits = self.classifier(batch_input_rep) |
|
|
|
|
|
if not self.training: |
|
|
results['label_preds'] = torch.argmax(batch_logits, dim=-1) * batch_inputs['label_ids_mask'] |
|
|
return results |
|
|
|
|
|
results['loss'] = self.element_loss( |
|
|
batch_logits[batch_inputs['label_ids_mask']], |
|
|
batch_inputs['label_ids'][batch_inputs['label_ids_mask']]) |
|
|
|
|
|
return results |
|
|
|
