models/language_modeling/kpplm.py

# -*- coding: utf-8 -*-
# @Time    : 2022/3/15 21:26
# @Author  : ruihan.wjn
# @File    : pk-plm.py

"""
This code is implemented for the paper ""Knowledge Prompting in Pre-trained Langauge Models for Natural Langauge Understanding""
"""

from time import time
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn import CrossEntropyLoss
from collections import OrderedDict
from transformers.models.bert import BertPreTrainedModel, BertModel
from transformers.models.roberta import RobertaModel, RobertaPreTrainedModel, RobertaTokenizer, RobertaForMaskedLM
from transformers.models.deberta import DebertaModel, DebertaPreTrainedModel, DebertaTokenizer, DebertaForMaskedLM
from transformers.models.bert.modeling_bert import BertOnlyMLMHead, BertPreTrainingHeads
from transformers.models.roberta.modeling_roberta import RobertaModel, RobertaLMHead
from transformers.models.deberta.modeling_deberta import DebertaModel, DebertaLMPredictionHead

"""
kg enhanced corpus structure example:
{
    "token_ids": [20, 46098, 3277, 680, 10, 4066, 278, 9, 11129, 4063, 877, 579, 8, 8750, 14720, 8, 22498, 548,
    19231, 46098, 3277, 6, 25, 157, 25, 130, 3753, 46098, 3277, 4, 3684, 19809, 10960, 9, 5, 30731, 2788, 914, 5,
    1675, 8151, 35], "entity_pos": [[8, 11], [13, 15], [26, 27]],
    "entity_qid": ["Q17582", "Q231978", "Q427013"],
    "relation_pos": null,
    "relation_pid": null
}
"""


from enum import Enum
class SiameseDistanceMetric(Enum):
    """
    The metric for the contrastive loss
    """
    EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2)
    MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1)
    COSINE_DISTANCE = lambda x, y: 1-F.cosine_similarity(x, y)


class ContrastiveLoss(nn.Module):
    """
    Contrastive loss. Expects as input two texts and a label of either 0 or 1. If the label == 1, then the distance between the
    two embeddings is reduced. If the label == 0, then the distance between the embeddings is increased.
    Further information: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
    :param model: SentenceTransformer model
    :param distance_metric: Function that returns a distance between two emeddings. The class SiameseDistanceMetric contains pre-defined metrices that can be used
    :param margin: Negative samples (label == 0) should have a distance of at least the margin value.
    :param size_average: Average by the size of the mini-batch.
    Example::
        from sentence_transformers import SentenceTransformer,  SentencesDataset, LoggingHandler, losses
        from sentence_transformers.readers import InputExample
        model = SentenceTransformer("distilbert-base-nli-mean-tokens")
        train_examples = [InputExample(texts=["This is a positive pair", "Where the distance will be minimized"], label=1),
            InputExample(texts=["This is a negative pair", "Their distance will be increased"], label=0)]
        train_dataset = SentencesDataset(train_examples, model)
        train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=train_batch_size)
        train_loss = losses.ContrastiveLoss(model=model)
    """

    def __init__(self, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE, margin: float = 0.5, size_average:bool = True):
        super(ContrastiveLoss, self).__init__()
        self.distance_metric = distance_metric
        self.margin = margin
        self.size_average = size_average

    def forward(self, sent_embs1, sent_embs2, labels: torch.Tensor):
        rep_anchor, rep_other = sent_embs1, sent_embs2
        distances = self.distance_metric(rep_anchor, rep_other)
        losses = 0.5 * (labels.float() * distances.pow(2) + (1 - labels).float() * F.relu(self.margin - distances).pow(2))
        return losses.mean() if self.size_average else losses.sum()



class NSPHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, pooled_output):
        seq_relationship_score = self.seq_relationship(pooled_output)
        return seq_relationship_score



class RoBertaKPPLMForProcessedWikiKGPLM(RobertaForMaskedLM):

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        # self.roberta = RobertaModel(config)
        try:
            classifier_dropout = (
                config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
            )
        except:
            classifier_dropout = (config.hidden_dropout_prob)
        self.dropout = nn.Dropout(classifier_dropout)
        # self.cls = BertOnlyMLMHead(config)
        # self.lm_head = RobertaLMHead(config)  # Masked Language Modeling head
        self.detector = NSPHead(config)  # Knowledge Noise Detection head
        self.entity_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        self.relation_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        # self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, config.num_ner_labels) for _ in range(config.entity_type_num)])

        self.contrastive_loss_fn = ContrastiveLoss()
        self.post_init()

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            labels=None,
            # entity_label=None,
            entity_candidate=None,
            # relation_label=None,
            relation_candidate=None,
            noise_detect_label=None,
            task_id=None,
            mask_id=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
    ):
        # start_time = time()
        mlm_labels = labels

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        # print("attention_mask.shape=", attention_mask.shape)
        # print("input_ids[0]=", input_ids[0])
        # print("token_type_ids[0]=", token_type_ids[0])
        # attention_mask = None

        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)  # mlm head
        # noise_detect_scores = self.detector(pooled_output)  # knowledge noise detector use pool output
        noise_detect_scores = self.detector(sequence_output[:, 0, :])  # knowledge noise detector use cls embedding

        # ner
        # sequence_output = self.dropout(sequence_output)
        # ner_logits = torch.stack([classifier(sequence_output) for classifier in self.classifiers]).movedim(1, 0)

        # mlm
        masked_lm_loss, noise_detect_loss, entity_loss, total_loss = None, None, None, None
        total_loss = list()
        if mlm_labels is not None:
            loss_fct = CrossEntropyLoss()  # -100 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), mlm_labels.view(-1))
            total_loss.append(masked_lm_loss)

        # if noise_detect_label is not None:
        #     noise_detect_scores = noise_detect_scores[task_id == 1]
        #     noise_detect_label = noise_detect_label[task_id == 1]
        #
        #     if len(noise_detect_label) > 0:
        #         loss_fct = CrossEntropyLoss()
        #         noise_detect_loss = loss_fct(noise_detect_scores.view(-1, 2), noise_detect_label.view(-1))
        #         total_loss.append(noise_detect_loss)

        entity_candidate = entity_candidate[task_id == 2]
        if len(entity_candidate) > 0:
            batch_size = entity_candidate.shape[0]
            candidate_num = entity_candidate.shape[1]
            # print("negative_num=", negative_num)
            # 获取被mask实体的embedding
            batch_entity_query_embedding = list()
            for ei, input_id in enumerate(input_ids[task_id == 2]):
                batch_entity_query_embedding.append(
                    torch.mean(sequence_output[task_id == 2][ei][input_id == mask_id[task_id == 2][ei]], 0))  # [hidden_dim]
            batch_entity_query_embedding = torch.stack(batch_entity_query_embedding)  # [bz, dim]
            # print("batch_entity_query_embedding.shape=", batch_entity_query_embedding.shape)
            batch_entity_query_embedding = self.entity_mlp(batch_entity_query_embedding)  # [bz, dim]
            batch_entity_query_embedding = batch_entity_query_embedding.unsqueeze(1).repeat((1, candidate_num, 1))  # [bz, 11, dim]
            batch_entity_query_embedding = batch_entity_query_embedding.view(-1, batch_entity_query_embedding.shape[-1])  # [bz * 11, dim]
            # print("batch_entity_query_embedding.shape=", batch_entity_query_embedding.shape)

            # 获得positive和negative的BERT表示
            # entity_candidiate: [bz, 11, len]
            entity_candidate = entity_candidate.view(-1, entity_candidate.shape[-1])  # [bz * 11, len]
            entity_candidate_embedding = self.roberta.embeddings(input_ids=entity_candidate)  # [bz * 11, len, dim]
            entity_candidate_embedding = self.entity_mlp(torch.mean(entity_candidate_embedding, 1))  # [bz * 11, dim]

            contrastive_entity_label = torch.Tensor([0] * (candidate_num - 1) + [1]).float().cuda()
            contrastive_entity_label = contrastive_entity_label.unsqueeze(0).repeat([batch_size, 1]).view(-1)  # [bz * 11]

            entity_loss = self.contrastive_loss_fn(
                batch_entity_query_embedding, entity_candidate_embedding, contrastive_entity_label
            )
            total_loss.append(entity_loss)

        relation_candidate = relation_candidate[task_id == 3]
        if len(relation_candidate) > 0:
            batch_size = relation_candidate.shape[0]
            candidate_num = relation_candidate.shape[1]
            # print("negative_num=", negative_num)
            # 获取被mask relation的embedding
            batch_relation_query_embedding = list()
            for ei, input_id in enumerate(input_ids[task_id == 3]):
                batch_relation_query_embedding.append(
                    torch.mean(sequence_output[task_id == 3][ei][input_id == mask_id[task_id == 3][ei]], 0))  # [hidden_dim]
            batch_relation_query_embedding = torch.stack(batch_relation_query_embedding)  # [bz, dim]
            # print("batch_relation_query_embedding.shape=", batch_relation_query_embedding.shape)
            batch_relation_query_embedding = self.relation_mlp(batch_relation_query_embedding)  # [bz, dim]
            batch_relation_query_embedding = batch_relation_query_embedding.unsqueeze(1).repeat(
                (1, candidate_num, 1))  # [bz, 11, dim]
            batch_relation_query_embedding = batch_relation_query_embedding.view(-1, batch_relation_query_embedding.shape[-1])  # [bz * 11, dim]
            # print("batch_relation_query_embedding.shape=", batch_relation_query_embedding.shape)

            # 获得positive和negative的BERT表示
            # entity_candidiate: [bz, 11, len]
            relation_candidate = relation_candidate.view(-1, relation_candidate.shape[-1])  # [bz * 11, len]
            relation_candidate_embedding = self.roberta.embeddings(input_ids=relation_candidate)  # [bz * 11, len, dim]
            relation_candidate_embedding = self.relation_mlp(torch.mean(relation_candidate_embedding, 1))  # [bz * 11, dim]

            contrastive_relation_label = torch.Tensor([0] * (candidate_num - 1) + [1]).float().cuda()
            contrastive_relation_label = contrastive_relation_label.unsqueeze(0).repeat([batch_size, 1]).view(-1)  # [bz * 11]

            relation_loss = self.contrastive_loss_fn(
                batch_relation_query_embedding, relation_candidate_embedding, contrastive_relation_label
            )
            total_loss.append(relation_loss)

        total_loss = torch.sum(torch.stack(total_loss), -1)

        # end_time = time()
        # print("neural_mode_time: {}".format(end_time - start_time))
        # print("masked_lm_loss.unsqueeze(0)=", masked_lm_loss.unsqueeze(0))
        # print("masked_lm_loss.unsqueeze(0).shape=", masked_lm_loss.unsqueeze(0).shape)
        # print("logits=", prediction_scores.argmax(2))
        # print("logits.shape=", prediction_scores.argmax(2).shape)


        return OrderedDict([
            ("loss", total_loss),
            ("mlm_loss", masked_lm_loss.unsqueeze(0)),
            # ("noise_detect_loss", noise_detect_loss.unsqueeze(0) if noise_detect_loss is not None else None),
            # ("entity_loss", entity_loss.unsqueeze(0) if entity_loss is not None else None),
            # ("relation_loss", relation_loss.unsqueeze(0) if relation_loss is not None else None),
            ("logits", prediction_scores.argmax(2)),
            # ("noise_detect_logits", noise_detect_scores.argmax(-1) if noise_detect_scores is not None and len(noise_detect_scores) > 0 else None),
        ])


class DeBertaKPPLMForProcessedWikiKGPLM(DebertaForMaskedLM):

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        # self.roberta = RobertaModel(config)
        try:
            classifier_dropout = (
                config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
            )
        except:
            classifier_dropout = (config.hidden_dropout_prob)
        self.dropout = nn.Dropout(classifier_dropout)
        # self.cls = BertOnlyMLMHead(config)
        # self.lm_head = RobertaLMHead(config)  # Masked Language Modeling head
        self.detector = NSPHead(config)  # Knowledge Noise Detection head
        self.entity_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        self.relation_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        # self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, config.num_ner_labels) for _ in range(config.entity_type_num)])

        self.contrastive_loss_fn = ContrastiveLoss()
        self.post_init()

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            labels=None,
            # entity_label=None,
            entity_candidate=None,
            # relation_label=None,
            relation_candidate=None,
            noise_detect_label=None,
            task_id=None,
            mask_id=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
    ):
        # start_time = time()
        mlm_labels = labels

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        # print("attention_mask.shape=", attention_mask.shape)
        # print("input_ids[0]=", input_ids[0])
        # print("token_type_ids[0]=", token_type_ids[0])
        # attention_mask = None

        outputs = self.deberta(
            input_ids,
            # attention_mask=attention_mask,
            attention_mask=None,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)  # mlm head
        # noise_detect_scores = self.detector(pooled_output)  # knowledge noise detector use pool output
        noise_detect_scores = self.detector(sequence_output[:, 0, :])  # knowledge noise detector use cls embedding

        # ner
        # sequence_output = self.dropout(sequence_output)
        # ner_logits = torch.stack([classifier(sequence_output) for classifier in self.classifiers]).movedim(1, 0)

        # mlm
        masked_lm_loss, noise_detect_loss, entity_loss, total_loss = None, None, None, None
        total_loss = list()
        if mlm_labels is not None:
            loss_fct = CrossEntropyLoss()  # -100 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), mlm_labels.view(-1))
            total_loss.append(masked_lm_loss)

        # if noise_detect_label is not None:
        #     noise_detect_scores = noise_detect_scores[task_id == 1]
        #     noise_detect_label = noise_detect_label[task_id == 1]
        #
        #     if len(noise_detect_label) > 0:
        #         loss_fct = CrossEntropyLoss()
        #         noise_detect_loss = loss_fct(noise_detect_scores.view(-1, 2), noise_detect_label.view(-1))
        #         total_loss.append(noise_detect_loss)

        entity_candidate = entity_candidate[task_id == 2]
        if len(entity_candidate) > 0:
            batch_size = entity_candidate.shape[0]
            candidate_num = entity_candidate.shape[1]
            # print("negative_num=", negative_num)
            # 获取被mask实体的embedding
            batch_entity_query_embedding = list()
            for ei, input_id in enumerate(input_ids[task_id == 2]):
                batch_entity_query_embedding.append(
                    torch.mean(sequence_output[task_id == 2][ei][input_id == mask_id[task_id == 2][ei]], 0))  # [hidden_dim]
            batch_entity_query_embedding = torch.stack(batch_entity_query_embedding)  # [bz, dim]
            # print("batch_entity_query_embedding.shape=", batch_entity_query_embedding.shape)
            batch_entity_query_embedding = self.entity_mlp(batch_entity_query_embedding)  # [bz, dim]
            batch_entity_query_embedding = batch_entity_query_embedding.unsqueeze(1).repeat((1, candidate_num, 1))  # [bz, 11, dim]
            batch_entity_query_embedding = batch_entity_query_embedding.view(-1, batch_entity_query_embedding.shape[-1])  # [bz * 11, dim]
            # print("batch_entity_query_embedding.shape=", batch_entity_query_embedding.shape)

            # 获得positive和negative的BERT表示
            # entity_candidiate: [bz, 11, len]
            entity_candidate = entity_candidate.view(-1, entity_candidate.shape[-1])  # [bz * 11, len]
            entity_candidate_embedding = self.deberta.embeddings(input_ids=entity_candidate)  # [bz * 11, len, dim]
            entity_candidate_embedding = self.entity_mlp(torch.mean(entity_candidate_embedding, 1))  # [bz * 11, dim]

            contrastive_entity_label = torch.Tensor([0] * (candidate_num - 1) + [1]).float().cuda()
            contrastive_entity_label = contrastive_entity_label.unsqueeze(0).repeat([batch_size, 1]).view(-1)  # [bz * 11]

            entity_loss = self.contrastive_loss_fn(
                batch_entity_query_embedding, entity_candidate_embedding, contrastive_entity_label
            )
            total_loss.append(entity_loss)

        relation_candidate = relation_candidate[task_id == 3]
        if len(relation_candidate) > 0:
            batch_size = relation_candidate.shape[0]
            candidate_num = relation_candidate.shape[1]
            # print("negative_num=", negative_num)
            # 获取被mask relation的embedding
            batch_relation_query_embedding = list()
            for ei, input_id in enumerate(input_ids[task_id == 3]):
                batch_relation_query_embedding.append(
                    torch.mean(sequence_output[task_id == 3][ei][input_id == mask_id[task_id == 3][ei]], 0))  # [hidden_dim]
            batch_relation_query_embedding = torch.stack(batch_relation_query_embedding)  # [bz, dim]
            # print("batch_relation_query_embedding.shape=", batch_relation_query_embedding.shape)
            batch_relation_query_embedding = self.relation_mlp(batch_relation_query_embedding)  # [bz, dim]
            batch_relation_query_embedding = batch_relation_query_embedding.unsqueeze(1).repeat(
                (1, candidate_num, 1))  # [bz, 11, dim]
            batch_relation_query_embedding = batch_relation_query_embedding.view(-1, batch_relation_query_embedding.shape[-1])  # [bz * 11, dim]
            # print("batch_relation_query_embedding.shape=", batch_relation_query_embedding.shape)

            # 获得positive和negative的BERT表示
            # entity_candidiate: [bz, 11, len]
            relation_candidate = relation_candidate.view(-1, relation_candidate.shape[-1])  # [bz * 11, len]
            relation_candidate_embedding = self.deberta.embeddings(input_ids=relation_candidate)  # [bz * 11, len, dim]
            relation_candidate_embedding = self.relation_mlp(torch.mean(relation_candidate_embedding, 1))  # [bz * 11, dim]

            contrastive_relation_label = torch.Tensor([0] * (candidate_num - 1) + [1]).float().cuda()
            contrastive_relation_label = contrastive_relation_label.unsqueeze(0).repeat([batch_size, 1]).view(-1)  # [bz * 11]

            relation_loss = self.contrastive_loss_fn(
                batch_relation_query_embedding, relation_candidate_embedding, contrastive_relation_label
            )
            total_loss.append(relation_loss)

        total_loss = torch.sum(torch.stack(total_loss), -1)

        # end_time = time()
        # print("neural_mode_time: {}".format(end_time - start_time))
        # print("masked_lm_loss.unsqueeze(0)=", masked_lm_loss.unsqueeze(0))
        # print("masked_lm_loss.unsqueeze(0).shape=", masked_lm_loss.unsqueeze(0).shape)
        # print("logits=", prediction_scores.argmax(2))
        # print("logits.shape=", prediction_scores.argmax(2).shape)


        return OrderedDict([
            ("loss", total_loss),
            ("mlm_loss", masked_lm_loss.unsqueeze(0)),
            # ("noise_detect_loss", noise_detect_loss.unsqueeze(0) if noise_detect_loss is not None else None),
            # ("entity_loss", entity_loss.unsqueeze(0) if entity_loss is not None else None),
            # ("relation_loss", relation_loss.unsqueeze(0) if relation_loss is not None else None),
            ("logits", prediction_scores.argmax(2)),
            # ("noise_detect_logits", noise_detect_scores.argmax(-1) if noise_detect_scores is not None and len(noise_detect_scores) > 0 else None),
        ])


class RoBertaForWikiKGPLM(RobertaPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        self.roberta = RobertaModel(config)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        # self.cls = BertOnlyMLMHead(config)
        self.lm_head = RobertaLMHead(config) # Masked Language Modeling head
        self.detector = NSPHead(config) # Knowledge Noise Detection head
        self.entity_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        self.relation_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        # self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, config.num_ner_labels) for _ in range(config.entity_type_num)])

        self.contrastive_loss_fn = ContrastiveLoss()
        self.post_init()

        self.tokenizer = RobertaTokenizer.from_pretrained(config.name_or_path)

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            mlm_labels=None,
            entity_label=None,
            entity_negative=None,
            relation_label=None,
            relation_negative=None,
            noise_detect_label=None,
            task_id=None,
            mask_id=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
    ):
        # start_time = time()

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        # print("attention_mask.shape=", attention_mask.shape)
        # print("input_ids[0]=", input_ids[0])
        # print("token_type_ids[0]=", token_type_ids[0])
        # attention_mask = None


        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output, pooled_output = outputs[:2]
        prediction_scores = self.lm_head(sequence_output) # mlm head
        noise_detect_scores = self.detector(pooled_output) # knowledge noise detector


        # ner
        # sequence_output = self.dropout(sequence_output)
        # ner_logits = torch.stack([classifier(sequence_output) for classifier in self.classifiers]).movedim(1, 0)

        # mlm
        masked_lm_loss, noise_detect_loss, entity_loss, total_loss = None, None, None, None
        if mlm_labels is not None:
            loss_fct = CrossEntropyLoss()  # -100 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), mlm_labels.view(-1))

        if noise_detect_label is not None:
            loss_fct = CrossEntropyLoss()
            noise_detect_loss = loss_fct(noise_detect_scores.view(-1, 2), noise_detect_label.view(-1))
            total_loss = masked_lm_loss + noise_detect_loss

        if entity_label is not None and entity_negative is not None:
            batch_size = input_ids.shape[0]
            negative_num = entity_negative.shape[1]
            # print("negative_num=", negative_num)
            # 获取被mask实体的embedding
            batch_query_embedding = list()
            for ei, input_id in enumerate(input_ids):
                batch_query_embedding.append(torch.mean(sequence_output[ei][input_id == mask_id[ei]], 0)) # [hidden_dim]
            batch_query_embedding = torch.stack(batch_query_embedding) # [bz, dim]
            # print("batch_query_embedding.shape=", batch_query_embedding.shape)
            batch_query_embedding = self.entity_mlp(batch_query_embedding) # [bz, dim]
            batch_query_embedding = batch_query_embedding.unsqueeze(1).repeat((1, negative_num + 1, 1)) # [bz, 11, dim]
            batch_query_embedding = batch_query_embedding.view(-1, batch_query_embedding.shape[-1]) # [bz * 11, dim]
            # print("batch_query_embedding.shape=", batch_query_embedding.shape)

            # 获得positive和negative的BERT表示
            # entity_label: [bz, len], entity_negative: [bz, 10, len]
            entity_negative = entity_negative.view(-1, entity_negative.shape[-1]) # [bz * 10, len]
            entity_label_embedding = self.roberta.embeddings(input_ids=entity_label) # [bz, len, dim]
            entity_label_embedding = self.entity_mlp(torch.mean(entity_label_embedding, 1)) # [bz, dim]
            entity_label_embedding = entity_label_embedding.unsqueeze(1) # [bz, 1, dim]

            entity_negative_embedding = self.roberta.embeddings(input_ids=entity_negative) # [bz * 10, len, dim]
            entity_negative_embedding = self.entity_mlp(torch.mean(entity_negative_embedding, 1)) # [bz * 10, dim]
            entity_negative_embedding = entity_negative_embedding \
                .view(input_ids.shape[0], -1, entity_negative_embedding.shape[-1]) # [bz, 10, dim]

            contrastive_label = torch.Tensor([0] * negative_num + [1]).float().cuda()
            contrastive_label = contrastive_label.unsqueeze(0).repeat([batch_size, 1]).view(-1) # [bz * 11]
            # print("entity_negative_embedding.shape=", entity_negative_embedding.shape)
            # print("entity_label_embedding.shape=", entity_label_embedding.shape)
            candidate_embedding = torch.cat([entity_negative_embedding, entity_label_embedding], 1) # [bz, 11, dim]
            candidate_embedding = candidate_embedding.view(-1, candidate_embedding.shape[-1]) # [bz * 11, dim]
            # print("candidate_embedding.shape=", candidate_embedding.shape)

            entity_loss = self.contrastive_loss_fn(batch_query_embedding, candidate_embedding, contrastive_label)
            total_loss = masked_lm_loss + entity_loss


        # if ner_labels is not None:
        #     loss_fct = CrossEntropyLoss()
        #     # Only keep active parts of the loss
        #
        #     active_loss = attention_mask.repeat(self.config.entity_type_num, 1, 1).view(-1) == 1
        #     active_logits = ner_logits.reshape(-1, self.config.num_ner_labels)
        #     active_labels = torch.where(
        #         active_loss, ner_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(ner_labels)
        #     )
        #     ner_loss = loss_fct(active_logits, active_labels)
        #
        # if masked_lm_loss:
        #     total_loss = masked_lm_loss + ner_loss * 4
        # print("total_loss=", total_loss)
        # print("mlm_loss=", masked_lm_loss)


        # end_time = time()
        # print("neural_mode_time: {}".format(end_time - start_time))

        return OrderedDict([
            ("loss", total_loss),
            ("mlm_loss", masked_lm_loss.unsqueeze(0)),
            ("noise_detect_loss", noise_detect_loss.unsqueeze(0) if noise_detect_loss is not None else None),
            ("entity_loss", entity_loss.unsqueeze(0) if entity_label is not None else None),
            ("logits", prediction_scores.argmax(2)),
            ("noise_detect_logits", noise_detect_scores.argmax(-1) if noise_detect_scores is not None else None),
        ])
        # MaskedLMOutput(
        #     loss=total_loss,
        #     logits=prediction_scores.argmax(2),
        #     ner_l
        #     hidden_states=outputs.hidden_states,
        #     attentions=outputs.attentions,
        # )




class BertForWikiKGPLM(BertPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        self.bert = BertModel(config)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        # self.cls = BertOnlyMLMHead(config)
        self.cls = BertPreTrainedModel(config)
        self.entity_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        self.relation_mlp = nn.Linear(config.hidden_size, config.hidden_size)
        # self.classifiers = nn.ModuleList([nn.Linear(config.hidden_size, config.num_ner_labels) for _ in range(config.entity_type_num)])

        self.contrastive_loss_fn = ContrastiveLoss()
        self.post_init()

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            mlm_labels=None,
            entity_label=None,
            entity_negative=None,
            relation_label=None,
            relation_negative=None,
            noise_detect_label=None,
            task_id=None,
            mask_id=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
    ):

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        print("attention_mask.shape=", attention_mask.shape)
        print("input_ids[0]=", input_ids[0])
        print("token_type_ids[0]=", token_type_ids[0])
        attention_mask = None
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output, pooled_output = outputs[:2]
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

        # ner
        # sequence_output = self.dropout(sequence_output)
        # ner_logits = torch.stack([classifier(sequence_output) for classifier in self.classifiers]).movedim(1, 0)

        # mlm
        masked_lm_loss, noise_detect_loss, entity_loss, total_loss = None, None, None, None

        if mlm_labels is not None:
            loss_fct = CrossEntropyLoss()  # -100 index = padding token
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), mlm_labels.view(-1))

        if noise_detect_label is not None:
            loss_fct = CrossEntropyLoss()
            noise_detect_loss = loss_fct(seq_relationship_score.view(-1, 2), noise_detect_label.view(-1))
            total_loss = masked_lm_loss + noise_detect_loss

        if entity_label is not None and entity_negative is not None:
            negative_num = entity_negative.shape[1]
            # 获取被mask实体的embedding
            batch_query_embedding = list()
            for ei, input_id in enumerate(input_ids):
                batch_query_embedding.append(torch.mean(sequence_output[ei][input_id == mask_id[ei]], 0)) # [hidden_dim]
            batch_query_embedding = torch.stack(batch_query_embedding) # [bz, dim]
            batch_query_embedding = self.entity_mlp(batch_query_embedding) # [bz, dim]
            batch_query_embedding = batch_query_embedding.repeat((1, negative_num + 1, 1)) # [bz, 11, dim]

            # 获得positive和negative的BERT表示
            # entity_label: [bz, len], entity_negative: [bz, 10, len]
            entity_negative = entity_negative.view(-1, entity_negative.shape[-1]) # [bz * 10, len]
            entity_label_embedding = self.bert.embeddings(input_id=entity_label) # [bz, len, dim]
            entity_label_embedding = self.entity_mlp(torch.mean(entity_label_embedding, 1)) # [bz, dim]
            entity_label_embedding = entity_label_embedding.unsqueeze(1) # [bz, 1, dim]

            entity_negative_embedding = self.bert.embeddings(input_id=entity_negative) # [bz * 10, len, dim]
            entity_negative_embedding = self.entity_mlp(torch.mean(entity_negative_embedding, 1)) # [bz * 10, dim]
            entity_negative_embedding = entity_negative_embedding \
                .view(input_ids.shape[0], -1, entity_negative_embedding.shape[-1]) # [bz, 10, dim]

            contrastive_label = torch.Tensor([0] * negative_num + [1]).float().cuda()
            candidate_embedding = torch.cat([entity_negative_embedding, entity_label_embedding], 1) # [bz, 11, dim]

            entity_loss = self.contrastive_loss_fn(batch_query_embedding, candidate_embedding, contrastive_label)
            total_loss = masked_lm_loss + entity_loss


        # if ner_labels is not None:
        #     loss_fct = CrossEntropyLoss()
        #     # Only keep active parts of the loss
        #
        #     active_loss = attention_mask.repeat(self.config.entity_type_num, 1, 1).view(-1) == 1
        #     active_logits = ner_logits.reshape(-1, self.config.num_ner_labels)
        #     active_labels = torch.where(
        #         active_loss, ner_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(ner_labels)
        #     )
        #     ner_loss = loss_fct(active_logits, active_labels)
        #
        # if masked_lm_loss:
        #     total_loss = masked_lm_loss + ner_loss * 4

        return OrderedDict([
            ("loss", total_loss),
            ("mlm_loss", masked_lm_loss.unsqueeze(0)),
            ("noise_detect_loss", noise_detect_loss.unsqueeze(0)),
            ("entity_loss", entity_loss.unsqueeze(0)),
            ("logits", prediction_scores.argmax(2)),
            ("noise_detect_logits", seq_relationship_score.argmax(3)),
            ()
        ])
        # MaskedLMOutput(
        #     loss=total_loss,
        #     logits=prediction_scores.argmax(2),
        #     ner_l
        #     hidden_states=outputs.hidden_states,
        #     attentions=outputs.attentions,
        # )