BIDINLTK/dev/bertCRF/ElectraCRF.py

import torch
import torch.nn as nn
import torch.nn.functional as F


import codecs
import re
import numpy as np
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class CRF(nn.Module):

    def __init__(self,num_tags,use_mask=False,id_to_tag=None):
        super(CRF,self).__init__()

        self.start_transition = nn.Parameter(torch.empty(num_tags))
        self.transition = nn.Parameter(torch.empty(num_tags,num_tags))

        self.reset_parameters()
        self.decode_transition = None

        self.use_mask = use_mask
        if use_mask:
            self.transition_mask = torch.as_tensor(self.init_transition_mask(id_to_tag,(num_tags,num_tags),_type="BMESO"),dtype=torch.float32).to(device)

    def get_transition(self,):
        if self.use_mask:
            trans = torch.minimum(self.transition,self.transition_mask)
            return trans
        return self.transition

    def init_transition_mask(self,id2tag,_shape,_type,_split="_"):

        list_legal_trans = []
        if _type=="BMESO":
            list_legal_trans.append(("B%s"%(_split),"M%s"%(_split)))
            list_legal_trans.append(("B%s"%(_split),"E%s"%(_split)))
            list_legal_trans.append(("M%s"%(_split),"M%s"%(_split)))
            list_legal_trans.append(("M%s"%(_split),"E%s"%(_split)))
            list_legal_trans.append(("E","S"))
            list_legal_trans.append(("E","O"))
            list_legal_trans.append(("E","B"))
            list_legal_trans.append(("S","O"))
            list_legal_trans.append(("S","B"))
            list_legal_trans.append(("O","B"))
            list_legal_trans.append(("O","S"))
            list_legal_trans.append(("O","O"))
        if _type=="BISO":
            list_legal_trans.append(("B%s"%(_split),"I%s"%(_split)))
            list_legal_trans.append(("B","S"))
            list_legal_trans.append(("B","O"))
            list_legal_trans.append(("I%s"%_split,"S%s"%_split))
            list_legal_trans.append(("I","S"))
            list_legal_trans.append(("I","O"))
            list_legal_trans.append(("I","B"))
            list_legal_trans.append(("S","O"))
            list_legal_trans.append(("S","B"))
            list_legal_trans.append(("O","B"))
            list_legal_trans.append(("O","S"))
            list_legal_trans.append(("O","O"))
        mask = np.ones(_shape)
        for col_index,col_tag in id2tag.items():
            for row_index,row_tag in id2tag.items():
                is_legal = False

                for legal_a,legal_b in list_legal_trans:
                    if col_tag.startswith(legal_a) and row_tag.startswith(legal_b):
                        if legal_a[-1]!=_split or legal_b[-1]!=_split:
                            is_legal = True
                            break
                        else:
                            if len(col_tag.split(_split))>1 and len(row_tag.split(_split))>1:
                                if col_tag.split(_split)[-1]==row_tag.split(_split)[-1]:
                                    is_legal = True
                                    break
                if not is_legal:
                    mask[col_index,row_index] = -1.0
        mask *= 100
        return mask


    def reset_parameters(self):
        '''
        要将初始化的概率限制在一个小的范围，否则loss将会很大导致nan
        :return:
        '''
        nn.init.xavier_normal_(self.transition)
        nn.init.uniform_(self.start_transition,-0.1,0.1)

    def forward(self,logits,labels,seq_length):

        return self.forward_path(logits,labels,seq_length)
        #按照batch和path计算loss，可以理解内部计算，但速度较慢
        # return self.forward_batch_path(logits,labels,seq_length)

    def forward_batch_path(self,logits,labels,seq_length):
        '''
        :param logits:the observation probabilities shape[batch,max_length,num_tags]
        :param labels:the labels shape[batch,max_length
        :param seq_length: the seq_length shape[batch]
        :return:the Z(x) loss
        '''
        # generate mask
        batch,max_length,num_tags = logits.shape
        mask = (torch.arange(max_length)<=(seq_length-1).unsqueeze(1)).to(torch.float).unsqueeze(2)

        logits = logits*mask
        whole_loss = 0
        for i in range(batch):
            _loss = self._compute_path_loss(logits[i,...],labels[i,...],seq_length[i])
            whole_loss += _loss

        return whole_loss

    def _compute_path_loss(self,logits,labels,seq_length):
        '''
        compute the path loss
        :param logits:[max_length,num_tags]
        :param labels:[max_length
        :return:
        '''
        max_length,num_tags = logits.shape

        mask = (torch.arange(max_length)<=(seq_length-1))
        #compute the Z(x)
        prob = logits[0] +self.start_transition
        for _i in range(1,max_length):
            prob = torch.where(mask[_i],torch.logsumexp(prob+self.get_transition(),dim=1)+logits[_i],prob)
            # 以下形式是错误的，0不是概率为0
            # prob = torch.logsumexp(prob+self.get_transition()*mask[_i],dim=1)+logits[_i]*mask[_i]
            # 写成以下形式是一样的
            # prob = torch.where(mask[_i],torch.logsumexp(prob+self.get_transition()+logits[_i],dim=1),prob)
        prob = torch.logsumexp(prob,dim=0)

        #generate label [max_length,num_tags]
        path_prob = logits[0,labels[0]]+self.start_transition[labels[0]]
        last_i = labels[0]
        for _i in range(1,max_length):
            path_prob += (self.get_transition()[last_i,labels[_i]]+logits[_i,labels[_i]])*mask[_i]
        # path_prob = torch.clamp_max(path_prob,1000)
        return prob-path_prob

    def forward_path(self,logits,labels,seq_length):
        """计算给定的标签序列tags的负对数似然

        Args:
            emissions (torch.Tensor):  发射分数P 形状 (seq_len, batch_size, num_tags), 代表序列中每个单词产生每个标签的得分
            tags (torch.LongTensor): 标签序列 如果batch_first=False 形状 (seq_len, batch_size) ，否则 形状为 (batch_size, seq_len)
            mask (torch.ByteTensor | None, optional): 表明哪些元素为填充符，和tags的形状一致。  如果batch_first=False  形状 (seq_len, batch_size) ，否则 形状为 (batch_size, seq_len)
                默认为None，表示没有填充符。
            reduction (str): 汇聚函数： none|sum|mean|token_mean 。 none：不应用汇聚函数； 默认为sum。

        Returns:
            torch.Tensor: 输入tags的负对数似然
        """

        max_length = labels.size(1)
        mask = (torch.arange(max_length).to(device)<=(seq_length-1).unsqueeze(1)).to(torch.float)


        # 转换为seq维度在前的形式
        emissions = logits.transpose(0, 1) # (seq_len, batch_size, num_tags)
        tags = labels.transpose(0, 1) #  (seq_len, batch_size)
        mask = mask.transpose(0, 1) #  (seq_len, batch_size)




        # 计算标签序列tags的得分
        score = self._compute_score(emissions, tags, mask)
        # 计算配分函数 partition Z(x)
        partition = self._compute_partition(emissions, mask)

        return torch.mean(partition - score)



    def _compute_score(
            self,
            emissions: torch.Tensor,
            tags: torch.LongTensor,
            mask: torch.ByteTensor,
    ) -> torch.Tensor:
        """计算标签序列tags的得分


        Args:
            emissions (torch.Tensor): 发射分数P 形状 (seq_len, batch_size, num_tags)
            tags (torch.LongTensor): 标签序列 形状 (seq_len, batch_size)
            mask (torch.ByteTensor): 表明哪些元素为填充符 形状 (seq_len, batch_size)

        Returns:
            torch.Tensor: 批次内标签tags的得分， 形状(batch_size,)
        """
        seq_len, batch_size = tags.shape
        # first_tags (batch_size,)
        first_tags = tags[0]

        # 由start标签转移到批次内所有标签序列第一个标签的得分
        score = self.start_transition[first_tags]
        # 加上 批次内第一个(index=0)发射得分，即批次内第0个输入产生批次内对应第0个标签的得分

        score += emissions[0, torch.arange(batch_size), first_tags]

        mask = mask.type_as(emissions)  # 类型保持一致
        # 这里的index从1开始，也就是第二个时间步开始
        for i in range(1, seq_len):
            # 第i-1个标签转移到第i个标签的得分 + 第i个单词产生第i个标签的得分
            # 乘以mask[i]不需要计算填充单词的得分
            score += (
                             self.get_transition()[tags[i - 1], tags[i]]
                             + emissions[i, torch.arange(batch_size), tags[i]]
                     ) * mask[i]

        # last_tags = tags[-1] × 这是错误的！，因为可能包含填充单词
        valid_last_idx = mask.long().sum(dim=0) - 1  # 有效的最后一个索引
        last_tags = tags[valid_last_idx, torch.arange(batch_size)]

        # 最后加上最后一个标签转移到end标签的转移得分
        return score



    def _compute_partition(
            self, emissions: torch.Tensor, mask: torch.ByteTensor
    ) -> torch.Tensor:
        """利用CRF的前向算法计算partition的分数

        Args:
            emissions (torch.Tensor): 发射分数P 形状 (seq_len, batch_size, num_tags)
            mask (torch.ByteTensor): 表明哪些元素为填充符  (seq_len, batch_size)

        Returns:
            torch.Tensor: 批次内的partition分数 形状(batch_size,)
        """
        seq_len = emissions.shape[0]
        # score (batch_size, num_tags) 对于每个批次来说，第i个元素保存到目前为止以i结尾的所有可能序列的得分
        score = self.start_transition.unsqueeze(0) + emissions[0]

        for i in range(1, seq_len):
            # broadcast_score: (batch_size, num_tags, 1) = (batch_size, pre_tag, current_tag)
            # 所有可能的当前标签current_tag广播
            broadcast_score = score.unsqueeze(2)
            # 广播成 (batch_size, 1, num_tags)
            # shape: (batch_size, 1, num_tags)
            broadcast_emissions = emissions[i].unsqueeze(1)
            # (batch_size, num_tags, num_tags) = (batch_size, num_tags, 1) + (num_tags, num_tags) + (batch_size, 1, num_tags)
            current_score = broadcast_score + self.get_transition() + broadcast_emissions
            # 在前一时间步标签上求和  -> (batch_size, num_tags)
            # 对于每个批次来说，第i个元素保存到目前为止以i结尾的所有可能标签序列的得分
            current_score = torch.logsumexp(current_score, dim=1)
            # mask[i].unsqueeze(1) -> (batch_size, 1)
            # 只有mask[i]是有效标签的current_score才将值设置到score中，否则保持原来的score
            score = torch.where(mask[i].bool().unsqueeze(1), current_score, score)

        # 加上到end标签的转移得分 end_transitions本身保存的是所有的标签到end标签的得分
        # score (batch_size, num_tags)
        # 在所有的标签上求(logsumexp)和
        # return (batch_size,)
        return torch.logsumexp(score, dim=1)

    def viterbi_decode(self,score, transition_params):
        """Decode the highest scoring sequence of tags outside of TensorFlow.

        This should only be used at test time.

        Args:
          score: A [seq_len, num_tags] matrix of unary potentials.
          transition_params: A [num_tags, num_tags] matrix of binary potentials.

        Returns:
          viterbi: A [seq_len] list of integers containing the highest scoring tag
              indices.
          viterbi_score: A float containing the score for the Viterbi sequence.
        """
        trellis = np.zeros_like(score)
        backpointers = np.zeros_like(score, dtype=np.int32)
        trellis[0] = score[0]

        for t in range(1, score.shape[0]):
            v = np.expand_dims(trellis[t - 1], 1) + transition_params
            trellis[t] = score[t] + np.max(v, 0)
            backpointers[t] = np.argmax(v, 0)

        viterbi = [np.argmax(trellis[-1])]
        for bp in reversed(backpointers[1:]):
            viterbi.append(bp[viterbi[-1]])
        viterbi.reverse()

        viterbi_score = np.max(trellis[-1])
        return viterbi, viterbi_score

    def get_decode_transition(self):
        if self.decode_transition is None:
            transition = self.get_transition().detach().cpu().numpy()
            start_transition = self.start_transition.detach().cpu().numpy()
            num_tags = transition.shape[0]
            trans = np.ones((num_tags+1,num_tags+1))
            trans *= -1000
            trans[:num_tags,:num_tags] = transition
            trans[num_tags,:num_tags] = start_transition
            self.decode_transition = trans
        return self.decode_transition

    def decode(self,logits,seq_length):
        '''
        :param logits:the observation probabilities shape[batch,max_length,num_tags]
        :param seq_length:the sequence length [batch,]
        :return:the best path and its probability
        '''
        batch_size,max_length,num_tags = logits.shape
        batch_path = []

        logits = logits.detach().cpu().numpy()
        small = -1000
        for _i in range(batch_size):

            seq_len = seq_length[_i].item()
            scores = logits[_i,:seq_len,...]

            start = np.asarray([[small]*num_tags+[0]])
            pad = small*np.ones((seq_len,1))
            scores = np.concatenate([scores,pad],1)
            scores = np.concatenate([start,scores],0)


            path,path_score = self.viterbi_decode(scores,self.get_decode_transition())
            batch_path.append(path[1:])
        return batch_path

    def getAcc(self,y_batch,logits,lengths):
        index = 0

        preds = []
        true_tags = []
        for path in self.decode(logits,lengths):
            preds.append(path)
            index += 1

        for y, length in zip(y_batch, lengths):
            y = y.tolist()
            true_tags.append(y[: length])
        acc_char = 0
        acc_seq = 0
        char_len = 0
        for pred,y in zip(preds,true_tags):
            is_same = True
            for a,b in zip(pred,y):
                char_len += 1
                if a!=b:
                    is_same = False
                else:
                    acc_char += 1
            if is_same:
                acc_seq += 1
        return acc_char/char_len,acc_seq/len(preds)

class BiLSTMCRF(nn.Module):

    def __init__(self,):
        super(BiLSTMCRF,self).__init__()
        config = {'lstm_dim':100,
                  'num_chars':6591,
                  'num_tags':25,
                  'char_dim':100,
                  'lr':0.00002,
                  'input_dropout_keep':1.0,
                  'optimizer':'adam',
                  'clip':5}

        # (self.char_to_id,embedding,self.id_to_tag,self.tag_to_id) = load_fool_map()
        # config["num_chars"] = len(self.char_to_id.keys())
        config["lstm_dim"] = 128
        # config["num_tags"] = len(self.id_to_tag.keys())
        config["char_dim"] = 128
        config["lr"] = 1e-2

        self.embedding = nn.Embedding(num_embeddings=config["num_chars"],embedding_dim=config["char_dim"])
        self.bilstm = nn.LSTM(input_size=config["char_dim"],hidden_size=config["lstm_dim"],num_layers=2, bidirectional=True)
        self.linear = nn.Linear(config["lstm_dim"]*2,config["lstm_dim"])
        self.out = nn.Linear(config["lstm_dim"],config["num_tags"])
        self.crf = CRF(config["num_tags"])


    def forward(self,text_ids):

        B,N = text_ids.shape
        seq_length = (text_ids!=0).sum(1).squeeze()
        x = self.embedding(text_ids)

        #pack为lstm所需的变长类型
        _seq_length = seq_length.cpu()
        x = torch.nn.utils.rnn.pack_padded_sequence(x,_seq_length,batch_first=True,enforce_sorted=False)
        seq_length.to(device)
        x,(hidden,cell) = self.bilstm(x)
        #pad初始形状
        x,output_length = torch.nn.utils.rnn.pad_packed_sequence(x,batch_first=True)

        # x = self.norm(x.transpose(1,2)).transpose(1,2)
        x = self.linear(x)
        # x = torch.relu(self.norm(x.transpose(1,2)).transpose(1,2))
        x = self.out(x)
        return x

from transformers import ElectraModel,ElectraTokenizer
class ElectraCRF(nn.Module):

    def __init__(self,hidden_dim= 256,num_tags = 26):
        super(ElectraCRF,self).__init__()
        model_path = "../bert/chinese-electra-180g-small-discriminator"

        self.tokenizers = ElectraTokenizer.from_pretrained(model_path)
        self.eletrac = ElectraModel.from_pretrained(model_path)

        self.pad_id = self.tokenizers.pad_token_id

        self.crf = CRF(num_tags)
        self.out = nn.Linear(hidden_dim,num_tags)

    def forward(self,inputs):
        '''
        :param inputs: #[B,N]
        :return:
        '''
        B,N = inputs.shape
        sequence_length = (inputs!=self.pad_id).sum(1).squeeze()
        mask = (torch.arange(N).unsqueeze(dim=0).to(device)<(sequence_length.unsqueeze(dim=1))).to(torch.float32).unsqueeze(dim=-1)
        x = self.eletrac(inputs,attention_mask=mask,output_hidden_states=True).__getitem__("hidden_states")[-1]
        out = self.out(x)
        return out

from dev.bert.BIDIBert import BIDIBert
class BIDIBertCRF(nn.Module):

    def __init__(self,n_src_vocab=6500,d_word_vec=128, n_layers=3, n_head=3, d_k=128, d_v=128,
                 d_model=128, d_inner=128,num_tags = 26):
        super(BIDIBertCRF,self).__init__()
        model_path = "../bert/chinese-electra-180g-small-discriminator"

        self.tokenizers = ElectraTokenizer.from_pretrained(model_path)

        self.pad_id = self.tokenizers.pad_token_id

        self.eletrac = BIDIBert(n_src_vocab=len(self.tokenizers.get_vocab()),d_word_vec=128, n_layers=4, n_head=2, d_k=128, d_v=128,
                                d_model=128, d_inner=128, pad_idx=self.pad_id)

        self.crf = CRF(num_tags)
        self.out = nn.Linear(d_model,num_tags)

    def forward(self,inputs):
        '''
        :param inputs: #[B,N]
        :return:
        '''
        B,N = inputs.shape
        sequence_length = (inputs!=self.pad_id).sum(1)
        mask = (torch.arange(N).unsqueeze(dim=0).to(device)<(sequence_length.unsqueeze(dim=1))).to(torch.float32).unsqueeze(dim=-1)
        x, = self.eletrac(inputs,attention_mask=mask)
        out = self.out(x)

        return out


def test_speed():
    # model = ElectraCRF()
    model = BIDIBertCRF()

    a = [[i for i in range(500)]]*10
    a = np.array(a,dtype=np.int32)
    print(a.shape)
    import time

    start_time = time.time()
    for _i in range(100):
        model(torch.from_numpy(a))
    print("cost",time.time()-start_time)

    model = BiLSTMCRF()
    start_time = time.time()
    for _i in range(100):
        model(torch.from_numpy(a))
    print("cost",time.time()-start_time)

if __name__ == '__main__':
    test_speed()