PaperText_Segmentation/PointerNet/model.py

import torch
import gc
import torch.nn as nn
from transformers import BertConfig, BertModel, BertForSequenceClassification


class UIEModel(nn.Module):
    def __init__(self, args):
        super(UIEModel, self).__init__()
        self.args = args
        
        self.tasks = args.tasks
        bert_dir = args.bert_dir
        self.bert_config = BertConfig.from_pretrained(bert_dir)
        self.bert_model = BertModel.from_pretrained(bert_dir)
        # self.bert_model.load_state_dict(torch.load(self.args.bert_pt_dir, map_location="cuda"))
        # self.bert_model = BertForSequenceClassification.from_pretrained(args.bert_pt_dir)

        if "ner" in args.tasks:
            self.ner_num_labels = args.ner_num_labels
            self.module_start_list = nn.ModuleList()
            self.module_end_list = nn.ModuleList()
            self.module_content_list = nn.ModuleList()  # 增加一层内容判断
            for i in range(args.ner_num_labels):
                self.module_start_list.append(nn.Linear(self.bert_config.hidden_size, 1))
                self.module_end_list.append(nn.Linear(self.bert_config.hidden_size, 1))
                self.module_content_list.append(nn.Linear(self.bert_config.hidden_size, 1))
            self.ner_criterion = nn.BCEWithLogitsLoss()
        self.dropout = nn.Dropout(self.bert_config.hidden_dropout_prob) 


    @staticmethod
    def build_dummpy_inputs():
        inputs = {}
        inputs['ner_input_ids'] = torch.LongTensor(
            torch.randint(low=1, high=10, size=(32, 56)))
        inputs['ner_attention_mask'] = torch.ones(size=(32, 56)).long()
        inputs['ner_token_type_ids'] = torch.zeros(size=(32, 56)).long()
        inputs['ner_start_labels'] = torch.zeros(size=(32, 8, 56)).float()
        inputs['ner_end_labels'] = torch.zeros(size=(32, 8, 56)).float()

        return inputs

    def get_pointer_loss(self,
                         start_logits,
                         end_logits,
                         attention_mask,
                         start_labels,
                         end_labels,
                         criterion):

        start_logits = start_logits.view(-1)
        end_logits = end_logits.view(-1)

        active_loss = attention_mask.view(-1) == 1

        active_start_logits = start_logits[active_loss]
        active_end_logits = end_logits[active_loss]

        active_start_labels = start_labels.view(-1)[active_loss]
        active_end_labels = end_labels.view(-1)[active_loss]

        start_loss = criterion(active_start_logits, active_start_labels)
        end_loss = criterion(active_end_logits, active_end_labels)
        loss = start_loss + end_loss

        return loss

    def ner_forward_1(self,
                    ner_input_ids,
                    ner_attention_mask,
                    ner_start_labels=None,
                    ner_end_labels=None):
        # 四个参数格式均为[tensor(), tensor(), ...]
        # 一次传入batch_size个样本，每个样本含多条句子
        # 编码还需要一个个样本进行，若每个样本句子太长，还需截断分批处理
    
        all_start_logits = []
        all_end_logits = []

        ner_loss = None
        for i in  range(len(ner_end_labels)):  # 有len(ner_end_labels)个样本/文档
            input_ids = ner_input_ids[i].to(self.args.device)
            attention_mask = ner_attention_mask[i].to(self.args.device)
            # start_labels = ner_start_labels[i].to(self.args.device)
            # end_labels = ner_end_labels[i].to(self.args.device)

            # 根据sent_num的大小分段进行编码（sent_num太大时，显存不够）
            max_encoder_len = self.args.max_encoder_sent_len
            batch_num = int(input_ids.size(0) / max_encoder_len)
            last_hidden_states = []
            if batch_num > 0:
                for i in range(batch_num):
                    truncated_input_ids = input_ids[i*max_encoder_len:(i+1)*max_encoder_len, :]
                    truncated_attention_mask = attention_mask[i*max_encoder_len:(i+1)*max_encoder_len, :]
                    truncated_outputs = self.bert_model(truncated_input_ids, attention_mask=truncated_attention_mask)
                    last_hidden_states.append(truncated_outputs.last_hidden_state)  # .detach().cpu()
            if input_ids.size(0) - batch_num * max_encoder_len > 0:
                truncated_input_ids = input_ids[batch_num*max_encoder_len:, :]
                truncated_attention_mask = attention_mask[batch_num*max_encoder_len:, :]
                truncated_outputs = self.bert_model(truncated_input_ids, attention_mask=truncated_attention_mask)
                last_hidden_states.append(truncated_outputs.last_hidden_state)  # .detach().cpu()
            if len(last_hidden_states) > 1:
                seq_bert_output = torch.cat(last_hidden_states, dim=0)
            else:
                seq_bert_output = last_hidden_states[0]   # [sent_num, seq_len, hidden_dim]
            # 忽略padding求均值的方法
            seq_bert_output = seq_bert_output[:, 1:-1, :] #.to(self.args.device)  # 忽略[CLS]和[SEP] ???
            expanded_attention_mask = attention_mask[:,1:-1].unsqueeze(-1).expand_as(seq_bert_output)  # [sent_num,seq_len,hidden_size] 
            sum_of_non_padded_output = (seq_bert_output * expanded_attention_mask).sum(dim=1)   # 仅对有效位置求和 
            mean_encoder_outputs = sum_of_non_padded_output / expanded_attention_mask.sum(dim=1)  # [sent_num, hidden_size] 
            # dropout
            pooled_output = self.dropout(mean_encoder_outputs)

            # 计算Pointer位置的loss
            # for i in range(self.ner_num_labels):   # 每个ner任务单独计算
            if self.ner_num_labels == 1:
                # 对每个pointer位置接个线性层
                start_logit = self.module_start_list[0](pooled_output).squeeze(1)  #[sent_num]
                end_logit = self.module_end_list[0](pooled_output).squeeze(1)  #[sent_num]
                # print(start_logit, start_logit.size())
                all_start_logits.append(start_logit)
                all_end_logits.append(end_logit)
        
        # 将批次数据合成一个loss值 
        concat_start_logits = torch.cat(all_start_logits, dim=0)
        concat_end_logits = torch.cat(all_end_logits, dim=0)
        all_start_labels = torch.cat(ner_start_labels, dim=0).to(self.args.device)
        all_end_labels = torch.cat(ner_end_labels, dim=0).to(self.args.device)

        start_loss = self.ner_criterion(concat_start_logits, all_start_labels)  # 起始位置loss值
        end_loss = self.ner_criterion(concat_end_logits, all_end_labels)  # 结束位置loss值
        if ner_loss is None:
            ner_loss = start_loss + end_loss
        else:
            ner_loss += (start_loss + end_loss)
        
        res = {
            "ner_start_logits": [a.detach().cpu() for a in all_start_logits],
            "ner_end_logits": [a.detach().cpu() for a in all_end_logits],
            "ner_loss": ner_loss,
        }
        return res

    def ner_bc_forward(self,
                       ner_input_ids,
                       ner_attention_mask,
                       ner_start_labels=None,
                       ner_end_labels=None,
                       ner_content_labels=None):
        """
        ner:topic识别任务；  bc:二分类任务
        这里将试题的开始、结束位置预测，与试题的判断(是否属于试题内容)任务合并在一起！！！
        原因：根据预测标签进行试题切分时，按start的位置划分错误最少，但会出现题型行也会被划分到试题中，故需要单独判断！
        # 四个参数格式均为[tensor(), tensor(), ...]
        # 一次传入batch_size=1个样本被截取的一部分，每个样本含多条句子
        # 编码还需要一个个样本进行，若每个样本句子太长，需截断分批处理
        """
        all_start_logits = []
        all_end_logits = []
        all_content_logits = []
        ner_loss = None

        outputs = self.bert_model(ner_input_ids, attention_mask=ner_attention_mask)
        # 取cls位置的表示
        cls_bert_output = outputs.pooler_output  # [sent_num, hidden_size] 
        # dropout
        pooled_output = self.dropout(cls_bert_output)
        # 尝试内存回收
        del outputs,cls_bert_output
        gc.collect()

        # 对每个pointer位置接个线性层
        # 对试题判断任务也接个线性层
        start_logit = self.module_start_list[0](pooled_output).squeeze(1)  #[sent_num]
        end_logit = self.module_end_list[0](pooled_output).squeeze(1)  #[sent_num]
        content_logit = self.module_content_list[0](pooled_output).squeeze(1)  #[sent_num]
        all_start_logits.append(start_logit)
        all_end_logits.append(end_logit)
        all_content_logits.append(content_logit)

        if ner_start_labels is not None and ner_end_labels is not None:   
            start_loss = self.ner_criterion(start_logit, ner_start_labels)  # 起始位置loss值
            end_loss = self.ner_criterion(end_logit, ner_end_labels)  # 结束位置loss值
            content_loss = self.ner_criterion(content_logit, ner_content_labels)

            if ner_loss is None:
                ner_loss = start_loss + end_loss + content_loss
            else:
                ner_loss += (start_loss + end_loss + content_loss)
        
        res = {
            "ner_start_logits": [a.detach().cpu() for a in all_start_logits],
            "ner_end_logits": [a.detach().cpu() for a in all_end_logits],
            "ner_content_logits": [a.detach().cpu() for a in all_content_logits],
            "ner_loss": ner_loss,
        }
        return res
        

    def ner_forward(self,
                    ner_input_ids,
                    ner_attention_mask,
                    ner_start_labels=None,
                    ner_end_labels=None):
        """
        # 四个参数格式均为[tensor(), tensor(), ...]
        # 一次传入batch_size=1个样本被截取的一部分，每个样本含多条句子
        # 编码还需要一个个样本进行，若每个样本句子太长，还需截断分批处理
        """
        all_start_logits = []
        all_end_logits = []
        ner_loss = None
        res = {
            "ner_start_logits": None,
            "ner_end_logits": None,
            "ner_loss": None
        }

        outputs = self.bert_model(ner_input_ids, attention_mask=ner_attention_mask)
        # 取每个token位置的表示再求平均
        # last_hidden_states = outputs.last_hidden_state  # [sent_num, seq_len, hidden_dim]
        # seq_bert_output = last_hidden_states[:, 1:-1, :]  # 忽略[CLS]和[SEP] ???
        # # 忽略padding求均值的方法
        # expanded_attention_mask = ner_attention_mask[:,1:-1].unsqueeze(-1).expand_as(seq_bert_output)  # [sent_num,seq_len,hidden_size] 
        # sum_of_non_padded_output = (seq_bert_output * expanded_attention_mask).sum(dim=1)   # 仅对有效位置求和 
        # mean_encoder_outputs = sum_of_non_padded_output / expanded_attention_mask.sum(dim=1)  # [sent_num, hidden_size] 
        # 取cls位置的表示
        cls_bert_output = outputs.pooler_output  # [sent_num, hidden_size] 
        # dropout
        pooled_output = self.dropout(cls_bert_output)
        
        # 计算Pointer位置的loss
        if self.ner_num_labels == 1:
            # 对每个pointer位置接个线性层
            # tensor.squeeze(1):移除大小为1的第二个维度
            start_logit = self.module_start_list[0](pooled_output).squeeze(1)  #[sent_num]
            end_logit = self.module_end_list[0](pooled_output).squeeze(1)  #[sent_num]
            all_start_logits.append(start_logit)
            all_end_logits.append(end_logit)

            if ner_start_labels is not None and ner_end_labels is not None:   
                start_loss = self.ner_criterion(start_logit, ner_start_labels)  # 起始位置loss值
                end_loss = self.ner_criterion(end_logit, ner_end_labels)  # 结束位置loss值

                if ner_loss is None:
                    ner_loss = start_loss + end_loss
                else:
                    ner_loss += (start_loss + end_loss)
            
            res = {
                "ner_start_logits": [a.detach().cpu() for a in all_start_logits],
                "ner_end_logits": [a.detach().cpu() for a in all_end_logits],
                "ner_loss": ner_loss,
            }
        return res


    def forward(self,
                ner_input_ids=None,
                # ner_token_type_ids=None,
                ner_attention_mask=None,
                ner_start_labels=None,
                ner_end_labels=None,
                ner_content_labels=None,
                ):

        res = {
            "ner_output": None,
            "re_output": None,
            "event_output": None
        }

        if "ner" in self.tasks:
            # ner_output = self.ner_forward(
            #     ner_input_ids,
            #     # ner_token_type_ids,
            #     ner_attention_mask,
            #     ner_start_labels,
            #     ner_end_labels,
            # )
            ner_output = self.ner_bc_forward(
                ner_input_ids,
                # ner_token_type_ids,
                ner_attention_mask,
                ner_start_labels,
                ner_end_labels,
                ner_content_labels,
            )
            res["ner_output"] = ner_output

        return res


if __name__ == '__main__':
    inputs = UIEModel.build_dummpy_inputs()


    class Args:
        bert_dir = "../chinese-bert-wwm-ext/"
        ner_num_labels = 8
        re_num_labels = 16
        tasks = ["re_rel"]


    args = Args()
    model = UIEModel(args)
    res = model(
        ner_input_ids=inputs["ner_input_ids"],
        ner_token_type_ids=inputs["ner_token_type_ids"],
        ner_attention_mask=inputs["ner_attention_mask"],
        ner_start_labels=inputs["ner_start_labels"],
        ner_end_labels=inputs["ner_end_labels"],
    )

    print(res)