{
        "negative": ["bad"],
        "positive": ["good", "wonderful", "great"],
}

def aggregate(label_words_logits: torch.Tensor, label_words_mask: torch.Tensor) -> torch.Tensor:
    r"""Use weight to aggregate the logits of label words.

    Args:
        label_words_logits(:obj:`torch.Tensor`): The logits of the label words.

    Returns:
        :obj:`torch.Tensor`: The aggregated logits from the label words.
    """
    label_words_logits = (label_words_logits * label_words_mask).sum(-1) / label_words_mask.sum(-1)
    return label_words_logits

output = torch.rand(2, 28996)
# [tokenizer.vocab.get(i) for i in ['bad', 'good', 'wonderful', 'great']]
# [2213, 1363, 7310, 1632]
verbalizer = torch.tensor([[2213,    0,    0],
                    [1363, 7310, 1632]])
mask = torch.tensor([[1, 0, 0],
                    [1, 1, 1]])
logits = output[:, verbalizer]
# tensor([[[0.0508, 0.0000, 0.0000],
#          [0.1770, 0.9680, 0.0091]],
#         [[0.2045, 0.0000, 0.0000],
#          [0.9504, 0.2127, 0.1417]]])
print(aggreate(logits, mask))
#  negative， positive
# tensor([[0.0508, 0.3847],
#         [0.2045, 0.4349]])

print(aggreate(logits, mask))
# tensor([[0.0508, 0.3847],
#         [0.2045, 0.4349]])

# -*- coding: utf8 -*-
#
# step1:
import torch

from openprompt.data_utils import InputExample

classes = [  # There are two classes in Sentiment Analysis, one for negative and one for positive
    "negative",
    "positive"
]
dataset = [  # For simplicity, there's only two examples
    # text_a is the input text of the data, some other datasets may have multiple input sentences in one example.
    InputExample(
        guid=0,
        text_a="Albert Einstein was one of the greatest intellects of his time.",
    ),
    InputExample(
        guid=1,
        text_a="The film was badly made.",
    ),
]
# step 2
from openprompt.plms import load_plm

plm, tokenizer, model_config, WrapperClass = load_plm("bert", "bert-base-cased")
# step3
from openprompt.prompts import ManualTemplate

promptTemplate = ManualTemplate(
    text='{"placeholder":"text_a"} It was {"mask"}',
    tokenizer=tokenizer,
)
# step4
from openprompt.prompts import ManualVerbalizer

promptVerbalizer = ManualVerbalizer(
    classes=classes,
    label_words={
        "negative": ["bad"],
        "positive": ["good", "wonderful", "great"],
    },
    tokenizer=tokenizer,
)
# step 5
from openprompt import PromptForClassification

promptModel = PromptForClassification(
    template=promptTemplate,
    plm=plm,
    verbalizer=promptVerbalizer,
)
# step 6
from openprompt import PromptDataLoader

data_loader = PromptDataLoader(
    dataset=dataset,
    tokenizer=tokenizer,
    template=promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
)
# step 7
# making zero-shot inference using pretrained MLM with prompt
promptModel.eval()
with torch.no_grad():
    for batch in data_loader:
        logits = promptModel(batch)
        preds = torch.argmax(logits, dim=-1)
        print(classes[preds])
# predictions would be 1, 0 for classes 'positive', 'negative'

from paddlenlp import Taskflow


# 跨任务抽取
schema = ['法院', {'原告': '委托代理人'}, {'被告': '委托代理人'}]
ie = Taskflow('information_extraction', schema=schema)
ie.set_schema(schema)
print(ie("北京市海淀区人民法院\n民事判决书\n(199x)建初字第xxx号\n原告：张三。\n委托代理人李四，北京市 A律师事务所律师。\n被告：B公司，法定代表人王五，开发公司总经理。\n委托代理人赵六，北京市 C律师事务所律师。")) # Better print results using pprint

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import paddle
import paddle.nn as nn
from paddlenlp.transformers import ErniePretrainedModel


class UIE(ErniePretrainedModel):

    def __init__(self, encoding_model):
        super(UIE, self).__init__()
        self.encoder = encoding_model
        hidden_size = self.encoder.config["hidden_size"]
        self.linear_start = paddle.nn.Linear(hidden_size, 1)
        self.linear_end = paddle.nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_ids, token_type_ids, pos_ids, att_mask):
        sequence_output, pooled_output = self.encoder(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            position_ids=pos_ids,
            attention_mask=att_mask)
        start_logits = self.linear_start(sequence_output)
        start_logits = paddle.squeeze(start_logits, -1)
        start_prob = self.sigmoid(start_logits)
        end_logits = self.linear_end(sequence_output)
        end_logits = paddle.squeeze(end_logits, -1)
        end_prob = self.sigmoid(end_logits)
        return start_prob, end_prob

# -*- coding: utf8 -*-
#
from paddlenlp import Taskflow

schema = {'地震触发词': ['地震强度', '时间', '震中位置', '震源深度']}
ie = Taskflow('information_extraction', schema=schema)
ie.set_schema(schema)
output = ie('中国地震台网正式测定：5月16日06时08分在云南临沧市凤庆县(北纬24.34度，东经99.98度)发生3.5级地震，震源深度10千米。')
print(output)

class NewConditionalLayerNorm(nn.Module):
    def __init__(self, normalized_shape1, normalized_shape2):
        super().__init__()
        # 触发词的
        self.trigger_feature_weight = ...
        self.trigger_feature_bias = ...

        self.trigger_feature_weight_dense = ...
        self.trigger_feature_bias_dense = ...
        # 客体的
        self.object_feature_weight = ...
        self.object_feature_bias = ...
        self.object_feature_weight_dense = ...
        self.object_feature_bias_dense = ...

    def forward(self, bert_output, trigger_feature, object_feature):
        # 1. 先计算bert_output的layer norm
        # 2. 计算trigger的weight和bias
        # 3. 计算object的weight和bias
        # 4. 三者相乘        

# 假设我们为主语
input_tokens = ['我', '们', '喜', '欢', '晴', '天', '。']
# token_type_ids
token_type_ids = [0, 0, 1, 1, 1, 1, 0]
# label
label = [[1, 0], [0, 1], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]

args = TrainArgs().get_parser()
args.gpu_ids = '0'
args.mode = "train"
args.raw_data_dir = './data/final/raw_data'
args.mid_data_dir = './data/final/mid_data'
args.aux_data_dir = "./data/final/preliminary_clean"
args.bert_dir = '/home/yuzhang/PycharmProjects/xf_event_extraction2020Top1/bert/torch_roberta_wwm'
args.output_dir = './out'
args.bert_type = 'roberta_wwm'
args.task_type = 'attribution'
args.max_seq_len = 320
args.train_epochs = 6
args.train_batch_size = 3
args.lr = 2e-5
args.other_lr = 2e-4
args.attack_train = "pgd"
args.swa_start = 4
args.eval_model = True
args.enhance_data = True
args.use_trigger_distance = True
args.use_distant_trigger = True

# 假设爱是trigger word。
input = ['我', '爱', '你', '北', '京', '。', ..., '你', '好', '呀']
mask =  [1, 0, 1, 1, 1, 1, ..., 0, 0,0] # 超过window_size就设置为0
# 即除这些mask外的都设置为-math.inf
bert_output.fill(~mask, -math.inf)

# 下面三步就是这面示例
seq_out = torch.transpose(seq_out, -1, -2)  # (bs, hidden, seq_len)
pooling_masks = torch.unsqueeze(pooling_masks, 1)
seq_out = seq_out + (1 - pooling_masks) * (-1e7)  # mask 无关区域

# 这里拿到max,这里看备忘那里的AdaptiveMaxPool1d
pooled_out = self.pooling_layer(seq_out).squeeze(-1)

logits = torch.cat([pooled_out, trigger_label_feature], dim=-1)

bert_outputs = self.bert_module(
    input_ids=token_ids,
    attention_mask=attention_masks,
    token_type_ids=token_type_ids
)

seq_out = bert_outputs[0]
# 拿到trigger feature
trigger_label_feature = self._batch_gather(seq_out, trigger_index)

trigger_label_feature = trigger_label_feature.view([trigger_label_feature.size()[0], -1])

# 做动态池化
seq_out = torch.transpose(seq_out, -1, -2)  # (bs, hidden, seq_len)
pooling_masks = torch.unsqueeze(pooling_masks, 1)
seq_out = seq_out + (1 - pooling_masks) * (-1e7)  # mask 无关区域

pooled_out = self.pooling_layer(seq_out).squeeze(-1)

# 合并两个feature
logits = torch.cat([pooled_out, trigger_label_feature], dim=-1)

polarity_logits = self.polarity_classifier(self.dropout_layer(logits))
tense_logits = self.tense_classifier(self.dropout_layer(logits))

out = (torch.softmax(polarity_logits, dim=-1), torch.softmax(tense_logits, dim=-1),)

if labels is not None:
    labels = labels.long()

    tense_loss = self.criterion(tense_logits, labels[:, 0])
    polarity_loss = self.criterion(polarity_logits, labels[:, 1])

    loss = polarity_loss + tense_loss

    out = (loss,) + out

return out

a = torch.arange(24, dtype=torch.float32).view(2,3,4)
a
Out[11]: 
tensor([[[ 0.,  1.,  2.,  3.],
         [ 4.,  5.,  6.,  7.],
         [ 8.,  9., 10., 11.]],
        [[12., 13., 14., 15.],
         [16., 17., 18., 19.],
         [20., 21., 22., 23.]]])

torch.functional.F.max_pool1d(a,kernel_size=4)
Out[12]: 
tensor([[[ 3.],
         [ 7.],
         [11.]],
        [[15.],
         [19.],
         [23.]]])
torch.nn.AdaptiveMaxPool1d(1)(a)
Out[13]: 
tensor([[[ 3.],
         [ 7.],
         [11.]],
        [[15.],
         [19.],
         [23.]]])

args = TrainArgs().get_parser()
args.gpu_ids = '0'
args.mode = "train"
args.raw_data_dir = './data/final/raw_data'
args.mid_data_dir = './data/final/mid_data'
args.aux_data_dir = "./data/final/preliminary_clean"
args.bert_dir = '/home/yuzhang/PycharmProjects/xf_event_extraction2020Top1/bert/torch_roberta_wwm'
args.output_dir = './out'
args.bert_type = 'roberta_wwm'
args.task_type = 'role1' # 改动这里
args.max_seq_len = 320
args.train_epochs = 6
args.train_batch_size = 3
args.lr = 2e-5
args.other_lr = 2e-4
args.attack_train = "pgd"
args.swa_start = 4
args.eval_model = True
args.enhance_data = True
args.use_trigger_distance = True
args.use_distant_trigger = True

label = [
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [1, 0, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 0],

]

distince_feature = [3,2,1,0,0,1,2,3]

class ConditionalLayerNorm(nn.Module):
    def __init__(self,
                 normalized_shape,
                 eps=1e-12):
        super().__init__()

        self.eps = eps

        self.weight = nn.Parameter(torch.Tensor(normalized_shape))
        self.bias = nn.Parameter(torch.Tensor(normalized_shape))

        self.weight_dense = nn.Linear(normalized_shape * 2, normalized_shape, bias=False)
        self.bias_dense = nn.Linear(normalized_shape * 2, normalized_shape, bias=False)

        self.reset_weight_and_bias()

    def reset_weight_and_bias(self):
        """
        此处初始化的作用是在训练开始阶段不让 conditional layer norm 起作用
        """
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)

        nn.init.zeros_(self.weight_dense.weight)
        nn.init.zeros_(self.bias_dense.weight)

    def forward(self, inputs, cond=None):
        assert cond is not None, 'Conditional tensor need to input when use conditional layer norm'
        cond = torch.unsqueeze(cond, 1)  # (b, 1, h*2)

        weight = self.weight_dense(cond) + self.weight  # (b, 1, h)
        bias = self.bias_dense(cond) + self.bias  # (b, 1, h)

        mean = torch.mean(inputs, dim=-1, keepdim=True)  # （b, s, 1）
        outputs = inputs - mean  # (b, s, h)

        variance = torch.mean(outputs ** 2, dim=-1, keepdim=True)
        std = torch.sqrt(variance + self.eps)  # (b, s, 1)

        outputs = outputs / std  # (b, s, h)
        # 这里做了一个交互
        outputs = outputs * weight + bias

        return outputs

if self.use_trigger_distance:
    assert trigger_distance is not None, \
        'When using trigger distance features, trigger distance should be implemented'

    trigger_distance_feature = self.trigger_distance_embedding(trigger_distance)
    seq_out = torch.cat([seq_out, trigger_distance_feature], dim=-1)
    seq_out = self.layer_norm(seq_out)

if labels is not None:
    masks = torch.unsqueeze(attention_masks, -1)

    labels = labels.float()
    obj_loss = self.criterion(obj_logits * masks, labels[:, :, :2])
    sub_loss = self.criterion(sub_logits * masks, labels[:, :, 2:])

    loss = obj_loss + sub_loss

# -*- coding: utf8 -*-
#

from torch import nn
import torch


class ConditionalLayerNorm(nn.Module):
    def __init__(self,
                 normalized_shape,
                 eps=1e-12):
        super().__init__()

        self.eps = eps

        self.weight = nn.Parameter(torch.Tensor(normalized_shape))
        self.bias = nn.Parameter(torch.Tensor(normalized_shape))

        self.weight_dense = nn.Linear(normalized_shape * 2, normalized_shape, bias=False)
        self.bias_dense = nn.Linear(normalized_shape * 2, normalized_shape, bias=False)

        self.reset_weight_and_bias()

    def reset_weight_and_bias(self):
        """
        此处初始化的作用是在训练开始阶段不让 conditional layer norm 起作用
        """
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)

        nn.init.zeros_(self.weight_dense.weight)
        nn.init.zeros_(self.bias_dense.weight)

    def forward(self, inputs, cond=None):
        assert cond is not None, 'Conditional tensor need to input when use conditional layer norm'
        cond = torch.unsqueeze(cond, 1)  # (b, 1, h*2)

        weight = self.weight_dense(cond) + self.weight  # (b, 1, h)
        bias = self.bias_dense(cond) + self.bias  # (b, 1, h)

        mean = torch.mean(inputs, dim=-1, keepdim=True)  # （b, s, 1）
        outputs = inputs - mean  # (b, s, h)

        variance = torch.mean(outputs ** 2, dim=-1, keepdim=True)
        std = torch.sqrt(variance + self.eps)  # (b, s, 1)

        outputs = outputs / std  # (b, s, h)
        # 这里做了一个交互
        outputs = outputs * weight + bias

        return outputs


if __name__ == '__main__':

    bert_output = torch.rand(32, 128, 768)
    # 当然，每个句子的trigger_index都是变化的
    # 此处假设这个batch的index span为(56, 58)
    trigger_feature = bert_output[:, 56:58, :]

    trigger_feature = trigger_feature.view(32, 768 * 2)
    cln = ConditionalLayerNorm(768)
    output = cln(bert_output, trigger_feature)
    print(output)

import torch
import torch.nn as nn


# FGM
class FGM:
    def __init__(self, model: nn.Module, eps=1.):
        self.model = (
            model.module if hasattr(model, "module") else model
        )
        self.eps = eps
        self.backup = {}

    # only attack word embedding
    def attack(self, emb_name='word_embeddings'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm and not torch.isnan(norm):
                    r_at = self.eps * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='word_embeddings'):
        for name, para in self.model.named_parameters():
            if para.requires_grad and emb_name in name:
                assert name in self.backup
                para.data = self.backup[name]

        self.backup = {}


# PGD
class PGD:
    def __init__(self, model, eps=1., alpha=0.3):
        self.model = (
            model.module if hasattr(model, "module") else model
        )
        self.eps = eps
        self.alpha = alpha
        self.emb_backup = {}
        self.grad_backup = {}

    def attack(self, emb_name='word_embeddings', is_first_attack=False):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                if is_first_attack:
                    self.emb_backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = self.alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data)

    def restore(self, emb_name='word_embeddings'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.emb_backup
                param.data = self.emb_backup[name]
        self.emb_backup = {}

    def project(self, param_name, param_data):
        r = param_data - self.emb_backup[param_name]
        if torch.norm(r) > self.eps:
            r = self.eps * r / torch.norm(r)
        return self.emb_backup[param_name] + r

    def backup_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None:
                self.grad_backup[name] = param.grad.clone()

    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None:
                param.grad = self.grad_backup[name]

# 初始化
fgm = FGM(model)
for batch_input, batch_label in data:
  # 正常训练
  loss = model(batch_input, batch_label)
  loss.backward() # 反向传播，得到正常的grad
  # 对抗训练
  fgm.attack() # embedding被修改了
  # optimizer.zero_grad() # 如果不想累加梯度，就把这里的注释取消
  loss_sum = model(batch_input, batch_label)
  loss_sum.backward() # 反向传播，在正常的grad基础上，累加对抗训练的梯度
  fgm.restore() # 恢复Embedding的参数
  # 梯度下降，更新参数
  optimizer.step()
  optimizer.zero_grad()