from fastNLP.models.biaffine_parser import BiaffineParser from fastNLP.models.biaffine_parser import ArcBiaffine, LabelBilinear import numpy as np import torch from torch import nn from torch.nn import functional as F from fastNLP.modules.dropout import TimestepDropout from fastNLP.modules.encoder.variational_rnn import VarLSTM from fastNLP import seq_len_to_mask from fastNLP.embeddings import Embedding def drop_input_independent(word_embeddings, dropout_emb): batch_size, seq_length, _ = word_embeddings.size() word_masks = word_embeddings.new(batch_size, seq_length).fill_(1 - dropout_emb) word_masks = torch.bernoulli(word_masks) word_masks = word_masks.unsqueeze(dim=2) word_embeddings = word_embeddings * word_masks return word_embeddings class CharBiaffineParser(BiaffineParser): def __init__(self, char_vocab_size, emb_dim, bigram_vocab_size, trigram_vocab_size, num_label, rnn_layers=3, rnn_hidden_size=800, #单向的数量 arc_mlp_size=500, label_mlp_size=100, dropout=0.3, encoder='lstm', use_greedy_infer=False, app_index = 0, pre_chars_embed=None, pre_bigrams_embed=None, pre_trigrams_embed=None): super(BiaffineParser, self).__init__() rnn_out_size = 2 * rnn_hidden_size self.char_embed = Embedding((char_vocab_size, emb_dim)) self.bigram_embed = Embedding((bigram_vocab_size, emb_dim)) self.trigram_embed = Embedding((trigram_vocab_size, emb_dim)) if pre_chars_embed: self.pre_char_embed = Embedding(pre_chars_embed) self.pre_char_embed.requires_grad = False if pre_bigrams_embed: self.pre_bigram_embed = Embedding(pre_bigrams_embed) self.pre_bigram_embed.requires_grad = False if pre_trigrams_embed: self.pre_trigram_embed = Embedding(pre_trigrams_embed) self.pre_trigram_embed.requires_grad = False self.timestep_drop = TimestepDropout(dropout) self.encoder_name = encoder if encoder == 'var-lstm': self.encoder = VarLSTM(input_size=emb_dim*3, hidden_size=rnn_hidden_size, num_layers=rnn_layers, bias=True, batch_first=True, input_dropout=dropout, hidden_dropout=dropout, bidirectional=True) elif encoder == 'lstm': self.encoder = nn.LSTM(input_size=emb_dim*3, hidden_size=rnn_hidden_size, num_layers=rnn_layers, bias=True, batch_first=True, dropout=dropout, bidirectional=True) else: raise ValueError('unsupported encoder type: {}'.format(encoder)) self.mlp = nn.Sequential(nn.Linear(rnn_out_size, arc_mlp_size * 2 + label_mlp_size * 2), nn.LeakyReLU(0.1), TimestepDropout(p=dropout),) self.arc_mlp_size = arc_mlp_size self.label_mlp_size = label_mlp_size self.arc_predictor = ArcBiaffine(arc_mlp_size, bias=True) self.label_predictor = LabelBilinear(label_mlp_size, label_mlp_size, num_label, bias=True) self.use_greedy_infer = use_greedy_infer self.reset_parameters() self.dropout = dropout self.app_index = app_index self.num_label = num_label if self.app_index != 0: raise ValueError("现在app_index必须等于0") def reset_parameters(self): for name, m in self.named_modules(): if 'embed' in name: pass elif hasattr(m, 'reset_parameters') or hasattr(m, 'init_param'): pass else: for p in m.parameters(): if len(p.size())>1: nn.init.xavier_normal_(p, gain=0.1) else: nn.init.uniform_(p, -0.1, 0.1) def forward(self, chars, bigrams, trigrams, seq_lens, gold_heads=None, pre_chars=None, pre_bigrams=None, pre_trigrams=None): """ max_len是包含root的 :param chars: batch_size x max_len :param ngrams: batch_size x max_len*ngram_per_char :param seq_lens: batch_size :param gold_heads: batch_size x max_len :param pre_chars: batch_size x max_len :param pre_ngrams: batch_size x max_len*ngram_per_char :return dict: parsing results arc_pred: [batch_size, seq_len, seq_len] label_pred: [batch_size, seq_len, seq_len] mask: [batch_size, seq_len] head_pred: [batch_size, seq_len] if gold_heads is not provided, predicting the heads """ # prepare embeddings batch_size, seq_len = chars.shape # print('forward {} {}'.format(batch_size, seq_len)) # get sequence mask mask = seq_len_to_mask(seq_lens).long() chars = self.char_embed(chars) # [N,L] -> [N,L,C_0] bigrams = self.bigram_embed(bigrams) # [N,L] -> [N,L,C_1] trigrams = self.trigram_embed(trigrams) if pre_chars is not None: pre_chars = self.pre_char_embed(pre_chars) # pre_chars = self.pre_char_fc(pre_chars) chars = pre_chars + chars if pre_bigrams is not None: pre_bigrams = self.pre_bigram_embed(pre_bigrams) # pre_bigrams = self.pre_bigram_fc(pre_bigrams) bigrams = bigrams + pre_bigrams if pre_trigrams is not None: pre_trigrams = self.pre_trigram_embed(pre_trigrams) # pre_trigrams = self.pre_trigram_fc(pre_trigrams) trigrams = trigrams + pre_trigrams x = torch.cat([chars, bigrams, trigrams], dim=2) # -> [N,L,C] # encoder, extract features if self.training: x = drop_input_independent(x, self.dropout) sort_lens, sort_idx = torch.sort(seq_lens, dim=0, descending=True) x = x[sort_idx] x = nn.utils.rnn.pack_padded_sequence(x, sort_lens, batch_first=True) feat, _ = self.encoder(x) # -> [N,L,C] feat, _ = nn.utils.rnn.pad_packed_sequence(feat, batch_first=True) _, unsort_idx = torch.sort(sort_idx, dim=0, descending=False) feat = feat[unsort_idx] feat = self.timestep_drop(feat) # for arc biaffine # mlp, reduce dim feat = self.mlp(feat) arc_sz, label_sz = self.arc_mlp_size, self.label_mlp_size arc_dep, arc_head = feat[:,:,:arc_sz], feat[:,:,arc_sz:2*arc_sz] label_dep, label_head = feat[:,:,2*arc_sz:2*arc_sz+label_sz], feat[:,:,2*arc_sz+label_sz:] # biaffine arc classifier arc_pred = self.arc_predictor(arc_head, arc_dep) # [N, L, L] # use gold or predicted arc to predict label if gold_heads is None or not self.training: # use greedy decoding in training if self.training or self.use_greedy_infer: heads = self.greedy_decoder(arc_pred, mask) else: heads = self.mst_decoder(arc_pred, mask) head_pred = heads else: assert self.training # must be training mode if gold_heads is None: heads = self.greedy_decoder(arc_pred, mask) head_pred = heads else: head_pred = None heads = gold_heads # heads: batch_size x max_len batch_range = torch.arange(start=0, end=batch_size, dtype=torch.long, device=chars.device).unsqueeze(1) label_head = label_head[batch_range, heads].contiguous() label_pred = self.label_predictor(label_head, label_dep) # [N, max_len, num_label] # 这里限制一下，只有当head为下一个时，才能预测app这个label arange_index = torch.arange(1, seq_len+1, dtype=torch.long, device=chars.device).unsqueeze(0)\ .repeat(batch_size, 1) # batch_size x max_len app_masks = heads.ne(arange_index) # batch_size x max_len, 为1的位置不可以预测app app_masks = app_masks.unsqueeze(2).repeat(1, 1, self.num_label) app_masks[:, :, 1:] = 0 label_pred = label_pred.masked_fill(app_masks, -np.inf) res_dict = {'arc_pred': arc_pred, 'label_pred': label_pred, 'mask': mask} if head_pred is not None: res_dict['head_pred'] = head_pred return res_dict @staticmethod def loss(arc_pred, label_pred, arc_true, label_true, mask): """ Compute loss. :param arc_pred: [batch_size, seq_len, seq_len] :param label_pred: [batch_size, seq_len, n_tags] :param arc_true: [batch_size, seq_len] :param label_true: [batch_size, seq_len] :param mask: [batch_size, seq_len] :return: loss value """ batch_size, seq_len, _ = arc_pred.shape flip_mask = (mask.eq(False)) # _arc_pred = arc_pred.clone() _arc_pred = arc_pred.masked_fill(flip_mask.unsqueeze(1), -float('inf')) arc_true.data[:, 0].fill_(-1) label_true.data[:, 0].fill_(-1) arc_nll = F.cross_entropy(_arc_pred.view(-1, seq_len), arc_true.view(-1), ignore_index=-1) label_nll = F.cross_entropy(label_pred.view(-1, label_pred.size(-1)), label_true.view(-1), ignore_index=-1) return arc_nll + label_nll def predict(self, chars, bigrams, trigrams, seq_lens, pre_chars, pre_bigrams, pre_trigrams): """ max_len是包含root的 :param chars: batch_size x max_len :param ngrams: batch_size x max_len*ngram_per_char :param seq_lens: batch_size :param pre_chars: batch_size x max_len :param pre_ngrams: batch_size x max_len*ngram_per_cha :return: """ res = self(chars, bigrams, trigrams, seq_lens, pre_chars=pre_chars, pre_bigrams=pre_bigrams, pre_trigrams=pre_trigrams, gold_heads=None) output = {} output['arc_pred'] = res.pop('head_pred') _, label_pred = res.pop('label_pred').max(2) output['label_pred'] = label_pred return output class CharParser(nn.Module): def __init__(self, char_vocab_size, emb_dim, bigram_vocab_size, trigram_vocab_size, num_label, rnn_layers=3, rnn_hidden_size=400, #单向的数量 arc_mlp_size=500, label_mlp_size=100, dropout=0.3, encoder='var-lstm', use_greedy_infer=False, app_index = 0, pre_chars_embed=None, pre_bigrams_embed=None, pre_trigrams_embed=None): super().__init__() self.parser = CharBiaffineParser(char_vocab_size, emb_dim, bigram_vocab_size, trigram_vocab_size, num_label, rnn_layers, rnn_hidden_size, #单向的数量 arc_mlp_size, label_mlp_size, dropout, encoder, use_greedy_infer, app_index, pre_chars_embed=pre_chars_embed, pre_bigrams_embed=pre_bigrams_embed, pre_trigrams_embed=pre_trigrams_embed) def forward(self, chars, bigrams, trigrams, seq_lens, char_heads, char_labels, pre_chars=None, pre_bigrams=None, pre_trigrams=None): res_dict = self.parser(chars, bigrams, trigrams, seq_lens, gold_heads=char_heads, pre_chars=pre_chars, pre_bigrams=pre_bigrams, pre_trigrams=pre_trigrams) arc_pred = res_dict['arc_pred'] label_pred = res_dict['label_pred'] masks = res_dict['mask'] loss = self.parser.loss(arc_pred, label_pred, char_heads, char_labels, masks) return {'loss': loss} def predict(self, chars, bigrams, trigrams, seq_lens, pre_chars=None, pre_bigrams=None, pre_trigrams=None): res = self.parser(chars, bigrams, trigrams, seq_lens, gold_heads=None, pre_chars=pre_chars, pre_bigrams=pre_bigrams, pre_trigrams=pre_trigrams) output = {} output['head_preds'] = res.pop('head_pred') _, label_pred = res.pop('label_pred').max(2) output['label_preds'] = label_pred return output