依存句法分析

import torch# 假设768是mlp出来的hidden_size.# batch_size, sequence_length, hidden_size = 32, 128,768class Biaffine(object):    def __init__(self, n_in=768, n_out=2, bias_x=True, bias_y=True):        self.n_in = n_in        self.n_out = n_out        self.bias_x = bias_x        self.bias_y = bias_y        self.weight = nn.Parameter(torch.Tensor(n_out, n_in + bias_x, n_in + bias_y))    def forward(self, x, y):        if self.bias_x:            x = torch.cat((x, torch.ones_like(x[..., :1])), -1)        if self.bias_y:            y = torch.cat((y, torch.ones_like(y[..., :1])), -1)                b = x.shape[0] # 32        o = self.weight.shape[0] # 2        x = x.unsqueeze(1).expand(-1, o, -1, -1) # torch.Size([32, 2, 128, 769])        weight = self.weight.unsqueeze(0).expand(b, -1, -1, -1)  # torch.Size([32, 2, 769, 769])        y = y.unsqueeze(1).expand(-1, o, -1, -1)  # torch.Size([32, 2, 128, 769])        # torch.matmul(x, weight): torch.Size([32, 2, 128, 769])        # y.permute((0, 1, 3, 2)).shape: torch.Size([32, 2, 769, 128])        s = torch.matmul(torch.matmul(x, weight), y.permute((0, 1, 3, 2)))        if s.shape[1] == 1:            s = s.squeeze(dim=1)        return s # torch.Size([32, 2, 128, 128])if __name__ == '__main__':    biaffine = Biaffine()    x = torch.rand(32, 128, 768)    y = torch.rand(32, 128, 768)    print(biaffine.forward(x, y).shape)

def forward(self, input_ids, token_type_ids=None, attention_mask=None):    """    Args:        input_ids: bert input tokens, tensor of shape [seq_len]        token_type_ids: 0 for query, 1 for context, tensor of shape [seq_len]        attention_mask: attention mask, tensor of shape [seq_len]    Returns:        start_logits: start/non-start probs of shape [seq_len]        end_logits: end/non-end probs of shape [seq_len]        match_logits: start-end-match probs of shape [seq_len, 1]    """    bert_outputs = self.bert(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)    sequence_heatmap = bert_outputs[0]  # [batch, seq_len, hidden]    batch_size, seq_len, hid_size = sequence_heatmap.size()    start_logits = self.start_outputs(sequence_heatmap).squeeze(-1)  # [batch, seq_len, 1]    end_logits = self.end_outputs(sequence_heatmap).squeeze(-1)  # [batch, seq_len, 1]    # for every position $i$ in sequence, should concate $j$ to    # predict if $i$ and $j$ are start_pos and end_pos for an entity.    # [batch, seq_len, seq_len, hidden]    start_extend = sequence_heatmap.unsqueeze(2).expand(-1, -1, seq_len, -1)    # [batch, seq_len, seq_len, hidden]    end_extend = sequence_heatmap.unsqueeze(1).expand(-1, seq_len, -1, -1)    # [batch, seq_len, seq_len, hidden*2]    span_matrix = torch.cat([start_extend, end_extend], 3)    # [batch, seq_len, seq_len]    span_logits = self.span_embedding(span_matrix).squeeze(-1)    return start_logits, end_logits, span_logits