role_prompt_roberta.py

import os

os.environ['CUDA_VISIBLE_DEVICES'] = '3'

from transformers import RobertaTokenizer, RobertaModel, BertTokenizer, BertModel
from torch.utils.data import DataLoader, TensorDataset
import torch
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report
import argparse
import json
import re
from collections import Counter, defaultdict
from constraints import Constraints
import scoring_utils as util
import numpy as np

max_padding_len = 512
pretrain_model_path = './roberta-large'
de_bert_dem = 236
myF1 = 0.0
temF1 = 0.0
data_rate = 1

parser = argparse.ArgumentParser()
parser.add_argument('--num_epoch', type=int, default=20)
parser.add_argument('--lr', type=float, default=3e-5)
parser.add_argument('--canshu', type=str, default='canshu.pt')
parser.add_argument('--max_acc', type=float, default=0.3)
parser.add_argument('-g', '--gold_file', type=str,
                    help='Gold file path')
parser.add_argument('-p', '--pred_file', type=str, default=None,
                    help='Predictions file path')
parser.add_argument('--reuse_gold_format', dest='reuse_gold_format',
                    default=False, action='store_true',
                    help="Reuse gold file format for pred file.")
parser.add_argument('-t', '--ontology_file', type=str, default=None,
                    help='Path to ontology file')
parser.add_argument('-cd', '--type_constrained_decoding', dest="cd",
                    default=False, action='store_true',
                    help="Use type constrained decoding" +
                         '(only possible when ontology file is given')
parser.add_argument('--do_all', dest='do_all', default=False,
                    action='store_true', help="Do everything.")
parser.add_argument('--metrics', dest='metrics', default=False,
                    action='store_true',
                    help="Compute overall p, r, f1.")
parser.add_argument('--distance', dest='distance', default=False,
                    action='store_true',
                    help="Compute p, r, f1 by distance.")
parser.add_argument('--role_table', dest='role_table', default=False,
                    action='store_true',
                    help="Compute p, r, f1 per role.")
parser.add_argument('--confusion', dest='confusion', default=False,
                    action='store_true',
                    help="Compute an error confusion matrix.")
args = parser.parse_args()
num_epoch = args.num_epoch
lr = args.lr
canshu = args.canshu
max_acc = args.max_acc

with open('./event_role_multiplicities.txt', encoding='utf-8') as f:
    event_in = [line.strip().split() for line in f.readlines()]

m_event_role = {}
for event_ in event_in:
    m_event_role[event_[0]] = event_[1:]

# ================================================= 抽样 =================================================
with open('./RAMS_1.0/scorer/event_role_multiplicities.txt', encoding='utf-8') as f:
    label = [line.strip().split()[0].split('.') for line in f.readlines()]
label_2 = []
for tem in label:
    label_2.append(tem[0] + '.' + tem[1])
set_label = set(label_2)
# 字典化
type_to_ix = {word: i for i, word in enumerate(set_label)}
data_list = []
for _ in range(0, 38):
    data_list.append([])
with open('./RAMS_1.0/data/train.jsonlines', encoding='utf-8') as f:
    train_seq_in_tem = [line for line in f.readlines()]

# utf-8
del train_seq_in_tem[3781]
del train_seq_in_tem[6246]
# 其他
del train_seq_in_tem[6937]
del train_seq_in_tem[6847]
del train_seq_in_tem[5477]
del train_seq_in_tem[4476]
del train_seq_in_tem[3603]
del train_seq_in_tem[1876]
del train_seq_in_tem[1067]

# 随机打乱
import random

for seq_in_ in train_seq_in_tem:
    m = eval(seq_in_)
    # [[69, 69, [["life.die.deathcausedbyviolentevents", 1.0]]]]
    event_type = m["evt_triggers"]
    tem_split = event_type[0][2][0][0].split('.')
    tem_event_type = tem_split[0] + '.' + tem_split[1]
    data_list[type_to_ix[tem_event_type]].append(seq_in_)
    assert type_to_ix[tem_event_type] < 38

len_ = 0
for tem in data_list:
    len_ += len(tem)
print('总数据：')
print(len_)

for i in range(0, 38):
    data_list[i] = data_list[i][:int(len(data_list[i]) * data_rate) + 1]

len_ = 0
for tem in data_list:
    len_ += len(tem)
print('按比例采样数据:')
print(len_)
train_seq_in_1 = []
for i in range(0, 38):
    train_seq_in_1 += data_list[i]

# ================================================= 数据读取预处理 =================================================
with open('./RAMS_1.0/data/dev.jsonlines', encoding='utf-8') as f:
    dev_seq_in_1 = [line.strip() for line in f.readlines()]
with open('./RAMS_1.0/data/dev.jsonlines', encoding='utf-8') as f:
    test_seq_in_1 = [line.strip() for line in f.readlines()]


def add_no_answer(tem_seq_in_1, tem_seq_in):
    # 向读取的json列表中，添加没有答案的。第一个是原始的，第二个是修改的。
    for i in range(0, len(tem_seq_in_1)):
        m = eval(tem_seq_in_1[i])
        result_list = m['gold_evt_links']
        pad_ = 'evt089arg02'
        trigger_begin = m['evt_triggers'][0][0]
        trigger_end = m['evt_triggers'][0][1]
        list_event = []
        for list_ in m['gold_evt_links']:  # 这里需要修改，引入一个额外的列表即可。第i 个test_seq_in 需要预测的数目
            # [[40, 40], [28, 28], "evt089arg02place"]
            list_event.append(list_[2][11:])
        c = m["evt_triggers"]  # [[35, 35, [['artifactexistence.damagedestroy.damage', 1.0]]]]
        c_s = c[0][2][0][0]  # 'artifactexistence.damagedestroy.damage'
        c_ss = m_event_role[c_s]  # 列表
        index_c_ss = 0
        while index_c_ss < len(c_ss):
            for _ in range(list_event.count(c_ss[index_c_ss]), int(c_ss[index_c_ss + 1])):
                tem_gold = [[trigger_begin, trigger_end], [0, -1], "evt089arg02" + c_ss[index_c_ss]]
                result_list.append(tem_gold)
            index_c_ss += 2
        m['gold_evt_links'] = result_list
        tem_seq_in.append(str(m))


train_seq_in = []
test_seq_in = []
dev_seq_in = []
add_no_answer(train_seq_in_1, train_seq_in)
add_no_answer(test_seq_in_1, test_seq_in)
add_no_answer(dev_seq_in_1, dev_seq_in)


# ================================================ 自定义F1 ============================================

class Scorer(object):
    def __init__(self, args):
        self.role_string_mapping = {}
        self.roles = set()
        self.gold = self.read_gold_file(args.gold_file)
        if args.reuse_gold_format:
            self.pred = self.read_gold_file(args.pred_file, confidence=False)
        else:
            self.pred = self.read_preds_file(args.pred_file)
        self.constraints = Constraints(args.ontology_file)

    def get_role_label(self, role):
        if role in self.role_string_mapping:
            return self.role_string_mapping[role]
        else:
            # Each role is of the form evt###arg##role, we only want role
            role_string = re.split(r'\d+', role)[-1]
            assert (role_string == role[11:])

            self.role_string_mapping[role] = role_string
            self.roles.add(role_string)
            return role_string

    def read_gold_file(self, file_path, confidence=False):
        """
        Returns dict mapping doc_key -> (pred, arg, role)
        """

        def process_example(json_blob):
            doc_key = json_blob["doc_key"]
            gold_evt = json_blob["gold_evt_links"]
            sents = json_blob["sentences"]
            sent_map = []
            for i, sent in enumerate(sents):
                for _ in sent:
                    sent_map.append(i)

            def span_to_sent(span):
                # assumes span does not cross boundaries
                sent_start = sent_map[span[0]]
                try:
                    sent_end = sent_map[span[1]]
                except:
                    sent_end = sent_map[span[0]]
                if sent_start != sent_end:
                    sent_end = sent_start
                assert (sent_start == sent_end)
                return sent_start

            # There should only be one predicate
            evt_triggers = json_blob["evt_triggers"]
            assert (len(evt_triggers) == 1)

            evt_trigger = evt_triggers[0]
            evt_trigger_span = util.list_to_span(evt_trigger[:2])
            evt_trigger_types = set([evt_trigger_type[0]
                                     for evt_trigger_type in evt_trigger[2]])

            gold_evt_links = [(util.list_to_span(arg[0]),
                               util.list_to_span(arg[1]),
                               self.get_role_label(arg[2])) for arg in gold_evt]
            if confidence:
                gold_evt_links = [(a, b, c, 0) for a, b, c in gold_evt_links]
            assert (all([arg[0] == evt_trigger_span
                         for arg in gold_evt_links]))
            return (doc_key, gold_evt_links, evt_trigger_types, span_to_sent)

        jsonlines = open(file_path, 'r').readlines()
        lines = [process_example(json.loads(line)) for line in jsonlines]
        file_dict = {doc_key: (evt_links, evt_trigger_types, span_to_sent)
                     for doc_key, evt_links, evt_trigger_types, span_to_sent
                     in lines}
        return file_dict

    def read_preds_file(self, file_path):
        """
        Ideally have only a single file reader
        Returns dict mapping doc_key -> (pred, arg, role)
        """

        def process_example(json_blob):
            doc_key = json_blob["doc_key"]
            pred_evt = json_blob["predictions"]
            # There should only be one predicate
            if len(pred_evt) == 0:
                return (doc_key, [], None)
            assert (len(pred_evt) == 1)
            pred_evt = pred_evt[0]
            # convention that the 0th one is the predicate span
            evt_span = util.list_to_span(pred_evt[0])
            evt_args = pred_evt[1:]
            pred_args = [(evt_span,
                          util.list_to_span(args[:2]),
                          args[2],
                          args[3])
                         for args in evt_args]
            return doc_key, pred_args, None

        jsonlines = open(file_path, 'r').readlines()
        lines = [process_example(json.loads(line)) for line in jsonlines]
        file_dict = {doc_key: (evt_links, evt_trigger_types)
                     for doc_key, evt_links, evt_trigger_types
                     in lines}
        return file_dict

    def create_role_table(self, correct, missing, overpred):
        role_table = {}
        for role in self.roles:
            c = float(correct[role])
            m = float(missing[role])
            o = float(overpred[role])
            p, r, f1 = util.compute_metrics(c, m, o)
            role_table[role] = {'CORRECT': c,
                                'MISSING': m,
                                'OVERPRED': o,
                                'PRECISION': p,
                                'RECALL': r,
                                'F1': f1}
        total_c = sum(correct.values())
        total_m = sum(missing.values())
        total_o = sum(overpred.values())
        total_p, total_r, total_f1 = util.compute_metrics(total_c,
                                                          total_m,
                                                          total_o)
        totals = {'CORRECT': total_c,
                  'MISSING': total_m,
                  'OVERPRED': total_o,
                  'PRECISION': total_p,
                  'RECALL': total_r,
                  'F1': total_f1}
        return (role_table, totals)

    def evaluate(self, constrained_decoding=True):
        self.metrics = None
        self.distance_metrics = None
        self.role_table = None
        self.confusion = None
        # Also computes confusion counters
        global_confusion = defaultdict(Counter)
        sentence_breakdowns = [{
            "correct": Counter(),
            "missing": Counter(),
            "overpred": Counter()
        } for i in range(5)]
        total_lost = 0

        global_correct = Counter()
        global_missing = Counter()
        global_overpred = Counter()
        for doc_key, (gold_structure, evt_type, span_to_sent) in self.gold.items():
            pred_structure = self.pred.get(doc_key, ([], None))[0]
            pred_structure, lost = self.constraints.filter_preds(
                pred_structure,
                evt_type,
                constrained_decoding)

            total_lost += lost
            pred_set = Counter(pred_structure)
            gold_set = Counter(gold_structure)
            assert (sum(pred_set.values()) == len(pred_structure))
            assert (sum(gold_set.values()) == len(gold_structure))
            intersection = gold_set & pred_set
            missing = gold_set - pred_set
            overpred = pred_set - gold_set
            # Update confusion and counters
            util.compute_confusion(global_confusion, intersection,
                                   missing, overpred)
            util.update(intersection, global_correct)
            util.update(missing, global_missing)
            util.update(overpred, global_overpred)
            util.update_sentence_breakdowns(intersection, missing, overpred,
                                            sentence_breakdowns, span_to_sent)
        precision, recall, f1, _ = util.compute_from_counters(global_correct,
                                                              global_missing,
                                                              global_overpred)
        distance_metrics = []
        for i in range(5):
            i_p, i_r, i_f1, counts = util.compute_from_counters(
                sentence_breakdowns[i]["correct"],
                sentence_breakdowns[i]["missing"],
                sentence_breakdowns[i]["overpred"]
            )
            distance_metrics.append((i, (i_p, i_r, i_f1), counts))
        self.metrics = {'precision': precision,
                        'recall': recall,
                        'f1': f1}
        self.distance_metrics = distance_metrics
        self.role_table = self.create_role_table(global_correct,
                                                 global_missing,
                                                 global_overpred)
        return {"role_table": self.role_table,
                "confusion": global_confusion,
                "metrics": self.metrics,
                "distance_metrics": self.distance_metrics}


def run_evaluation(args):
    """This is a separate wrapper around args so that other programs
    can call evaluation without resorting to an os-level call
    """
    scorer = Scorer(args)
    return_dict = scorer.evaluate(constrained_decoding=args.cd)
    if args.confusion or args.do_all:
        pass
        # util.print_confusion(return_dict['confusion'])
    if args.role_table or args.do_all:
        pass
        # util.print_table(*return_dict['role_table'])
    if args.distance or args.do_all:
        for (i, (p, r, f1), (gold, pred)) in return_dict['distance_metrics']:
            print(" {} & {} & {:.1f} & {:.1f} & {:.1f} \\\\ [p r f1 {} gold/{} pred. ]".format(
                i - 2, pred, p, r, f1, gold, pred))
    if args.metrics or args.do_all:
        print("Precision: {:.4f} Recall: {:.4f} F1: {:.4f}".format(
            return_dict['metrics']['precision'],
            return_dict['metrics']['recall'],
            return_dict['metrics']['f1']))
        global temF1
        temF1 = return_dict['metrics']['f1']

    return return_dict

def sentence_add_trigger_specal_token(trigger_begin, trigger_end, s_list):
    # 为句子列表添加特殊的token。
    # 如 I really love you , my baby !    ->  I really <t> love you </t> , my baby !
    tem_s_list = s_list.copy()
    tem_s_list.insert(trigger_begin, '<t>')
    tem_s_list.insert(trigger_end + 2, '</t>')
    return tem_s_list

def label_change_trigger_specal_token(trigger_begin, trigger_end,  label_begin, label_end):
    # 修正添加<t>和</t>之后，
    begin_add = 0
    end_add = 0
    if label_begin >= trigger_begin:
        begin_add += 1
    if label_begin > trigger_end:
        begin_add += 1
    if label_end >= trigger_begin:
        end_add += 1
    if label_end > trigger_end:
        end_add += 1

    return label_begin + begin_add, label_end + end_add


# ================================================ 数据读取 ============================================
# train_seq_in = train_seq_in[:1]
def generate_input(tem_seq_in, tem_all_sentence_1_role, tem_label_1_role, tem_segment_embedding_1_role,
                         tem_all_role_span_for_one_sentence, tem_event_type):
    # 产生 输入 结构的数据。以及列表中需要的所有东西。
    for seq_in_ in tem_seq_in:
        m = eval(seq_in_)
        seq_in_list = []  # 列表结构的句子表示
        for list_ in m['sentences']:
            seq_in_list += list_

        # 获取到trigger    在句向量中。 role 和 trigger 都应该标注为 1.
        seq_seg = []
        # 产生和原始句子一样长的，
        for i in range(0, len(seq_in_list)):
            seq_seg.append(0)
        # trigger相应位置变成1
        for i in range(m['evt_triggers'][0][0], m['evt_triggers'][0][1] + 1):
            seq_seg[i] = 1
        # cls + role + sep +句子 + sep  对seq_seg的后处理。
        seq_seg = [0] + [1] + [0] + [0] + seq_seg + [0] + [0]

        tem_tem_all_role_span_for_one_sentence = []

        for event in m['gold_evt_links']:
            seq_label = []
            for i in range(0, len(seq_in_list)):
                seq_label.append(0)
            # [[31, 31], [27, 27], "evt043arg01communicator"]  event的结构
            begin_ = event[1][0]
            end_ = event[1][1]
            event_ = event[2][11:]
            seq_label_ = seq_label
            begin_2, end_2 = label_change_trigger_specal_token(m['evt_triggers'][0][0], m['evt_triggers'][0][1], begin_,
                                                               end_)
            seq_label_.append(0)
            seq_label_.append(0)
            # 由于trigger的特殊字符的添加，修改位置。
            for i in range(begin_2, end_2 + 1):
                seq_label_[i] = 1

            bert_seq_in_ = '<s> ' + event_ + ' </s> ' + (' '.join(
                sentence_add_trigger_specal_token(m['evt_triggers'][0][0], m['evt_triggers'][0][1],
                                                  seq_in_list))) + ' </s>'

            role_ = event_
            span_ = seq_in_list[begin_: end_ + 1]

            tem_all_role_span_for_one_sentence_ = [role_, span_]
            tem_tem_all_role_span_for_one_sentence.append(tem_all_role_span_for_one_sentence_)

            if begin_ == 0 and end_ == -1:
                tem_all_sentence_1_role.append(bert_seq_in_)
                tem_label_1_role.append([0] + [0] + [1] + seq_label_ + [0])
                tem_segment_embedding_1_role.append(seq_seg)
            else:
                tem_all_sentence_1_role.append(bert_seq_in_)
                tem_label_1_role.append([0] + [0] + [0] + seq_label_ + [0])
                tem_segment_embedding_1_role.append(seq_seg)

        for _ in m['gold_evt_links']:
            tem_all_role_span_for_one_sentence.append(tem_tem_all_role_span_for_one_sentence)
            tem_event_type.append(m['evt_triggers'][0][2][0][0])

    assert len(tem_all_role_span_for_one_sentence) == len(tem_all_sentence_1_role)

#
train_all_sentence_1_role = []
train_label_1_role = []
train_segment_embedding_1_role = []
test_all_sentence_1_role = []
test_label_1_role = []
test_segment_embedding_1_role = []
dev_all_sentence_1_role = []
dev_label_1_role = []
dev_segment_embedding_1_role = []

train_all_role_span_for_one_sentence = []
test_all_role_span_for_one_sentence = []
dev_all_role_span_for_one_sentence = []
# 注意每一个句子产生多组 用于训练的句子。
# 存放事件类型
train_event_type = []
test_event_type = []
dev_event_type = []

generate_input(train_seq_in, train_all_sentence_1_role, train_label_1_role, train_segment_embedding_1_role,
                         train_all_role_span_for_one_sentence, train_event_type)
generate_input(test_seq_in, test_all_sentence_1_role, test_label_1_role, test_segment_embedding_1_role,
                         test_all_role_span_for_one_sentence, test_event_type)
generate_input(dev_seq_in, dev_all_sentence_1_role, dev_label_1_role, dev_segment_embedding_1_role,
                         dev_all_role_span_for_one_sentence, dev_event_type)

print('未引入前  样本数量：')
print(len(train_all_sentence_1_role))
print('')


def all_role_to_list(list, role, template, mrole):
    # print(list)
    # print(role)
    # print(template)
    # [['communicator', ['Bill', 'Clinton']], ['recipient', ['Congress']], ['place', ['Congress']]]
    # communicator
    # <communicator> communicated to <recipient> at <place> place

    tem_role_ = []
    for tem_ in list:
        tem_role_.append(tem_[0])
    # print(tem_role_)
    # violator
    # <violator> violated an agreement with <otherparticipant> in <place> place
    c = 0
    for tem_ in list:
        # ['vehicle', ['RNoAF', 'P-3B']]
        tem_role = tem_[0]
        tem_role_spe = ' '.join(tem_[1])
        if tem_role_spe == '':
            tem_role_spe = 'null'
        if tem_role == role and tem_role_.count(role) == 1:
            continue
        if tem_role == role and tem_role_.count(role) == 2:
            # print('22222222222')
            if mrole == 1 and c == 0:
                c += 1
                continue
            elif mrole == 1 and c == 1:
                template = template.replace('<' + tem_role + '>', '[' + tem_role + ']', 1).replace('<' + tem_role + '>',
                                                                                                   tem_role_spe,
                                                                                                   1).replace(
                    '[' + tem_role + ']', '<' + tem_role + '>', 1)
                continue
            elif mrole == -1 and c == 0:
                template = template.replace('<' + tem_role + '>', tem_role_spe, 1)
                c += 1
                continue
            elif mrole == -1 and c == 1:
                continue
        template = template.replace('<' + tem_role + '>', tem_role_spe, 1)

    # print(template)
    # print('-------------------------')
    if tem_role_.count(role) == 2:
        mrole = -1 * mrole
    # print(template.split())
    # ['<communicator>', 'communicated', 'to', 'Congress', 'at', 'Congress', 'place']
    return template.split(), mrole


import csv

list_csv = []
with open('./aida_ontology_cleaned.csv', 'r') as f:
    reader = csv.reader(f)
    for line in reader:
        list_csv.append(line)

list_csv = list_csv[1:]

for i in range(0, len(list_csv)):
    list_csv[i][0] = list_csv[i][0].replace('unspecified', 'n/a')
    # print(list_csv[i])

# artifactexistence.artifactfailure.mechanicalfailure
# 把 <arg1> mechanical artifact failed due to <arg2> instrument at <arg3> place  中的role替换为后面的。
map_aida = {}
for line in list_csv:
    template_ = line[1]
    # print(template_)
    num_i = 0
    for tem in line[2:]:
        num_i += 1
        if tem == '':
            continue
        tem = tem[11:]
        tem_replace = '<arg' + str(num_i) + '>'
        # print(tem, tem_replace)
        template_ = template_.replace(tem_replace, '<' + tem.replace(' ', '') + '>')
    if '<arg' in template_:
        print('wrong!!!')
        assert 1 == 2
    map_aida[line[0]] = template_
    # print(template_)




def add_role_meaning(tem_segment_embedding_1_role, tem_all_sentence_1_role, tem_label_1_role,
                     tem_role, tem_all_sentence, tem_len, tem_label, tem_segment_embedding, tem_segment_embedding_222,
                     tem_all_role_span_for_one_sentence, tem_event_type):

    mrole = 1
    # seg_embedding222 是用来训练的。
    for i in range(0, len(tem_segment_embedding_1_role)):
        tem_sentence = tem_all_sentence_1_role[i]
        tem_tem_label = tem_label_1_role[i]
        tem_segment = tem_segment_embedding_1_role[i]
        # role_meaning = map_role_meaning[tem_sentence.split()[1]].lower().split()
        # role_meaning = [tem_sentence.split()[1]]
        # role_meaning = [trigger_, tem_sentence.split()[1]]
        temKB = tem_all_role_span_for_one_sentence[i]
        tem_tem_role = tem_sentence.split()[1]

        role_KB = []
        role_meaning_, mrole = all_role_to_list(temKB, tem_tem_role, map_aida[tem_event_type[i]], mrole)
        
        role_i = -1
        for ii in range(0, len(role_meaning_)):
            tem = role_meaning_[ii]
            if tem[0] == '<' and (tem[-1] == '>' or tem[-2] == '>'):
                role_i = ii
                break
        assert role_i != -1

        # role_meaning = []
        # for tem in role_meaning_:
        #     if tem[0] == '<' and (tem[-1] == '>' or tem[-2] == '>'):
        #         role_meaning.append('<')
        #         if tem[-1] == '>':
        #             role_meaning.append(tem[1:-1])
        #         else:
        #             role_meaning.append(tem[1:-2])
        #         role_meaning.append('>')
        #     else:
        #         role_meaning.append(tem)
        # role_i += 1
        role_meaning = [tem_sentence.split()[1]]
        role_meaning.append('</s>')
        tem_role.append(tem_sentence.split()[1])
        tem_all_sentence.append(' '.join([tem_sentence.split()[0]] + role_meaning + tem_sentence.split()[2:]))
        tem_len.append(len(role_meaning) + 2)
        tem_label.append([tem_tem_label[0]] + len(role_meaning) * [0] + tem_tem_label[2:])
        # 用来标识位置的。
        tem_segment_embedding.append([tem_segment[0]] + len(role_meaning) * [1] + tem_segment[2:])
        # 用于训练的。
        tem__ = [tem_segment[0]] + (len(role_meaning)) * [0] + tem_segment[2:]
        tem_segment_embedding_222.append(tem__)


# train
train_all_sentence = []
train_label = []
train_segment_embedding = []

# test
test_all_sentence = []
test_label = []
test_segment_embedding = []
test_segment_embedding_222 = []
test_len = []  # role解释的长度 + 【CLS】 + 第一个【SEP】
test_role = []

# 对于train和dev没用，仅仅为了封装。
train_segment_embedding_222 = []
train_len = []  # role解释的长度 + 【CLS】 + 第一个【SEP】
train_role = []
dev_all_sentence = []
dev_label = []
dev_segment_embedding = []
dev_segment_embedding_222 = []
dev_len = []  # role解释的长度 + 【CLS】 + 第一个【SEP】
dev_role = []


add_role_meaning(train_segment_embedding_1_role, train_all_sentence_1_role, train_label_1_role,
                     train_role, train_all_sentence, train_len, train_label, train_segment_embedding_222, train_segment_embedding
                 , train_all_role_span_for_one_sentence, train_event_type)
add_role_meaning(test_segment_embedding_1_role, test_all_sentence_1_role, test_label_1_role,
                     test_role, test_all_sentence, test_len, test_label, test_segment_embedding, test_segment_embedding_222,
                 test_all_role_span_for_one_sentence, test_event_type)

add_role_meaning(dev_segment_embedding_1_role, dev_all_sentence_1_role, dev_label_1_role,
                     dev_role, dev_all_sentence, dev_len, dev_label, dev_segment_embedding_222, dev_segment_embedding,
                 dev_all_role_span_for_one_sentence, dev_event_type)

# ============================================ 截取 填充 tokenizer ============================================
# [CLS] Role [SEP] Sentence [SEP]
# dev_all_sentence          字符串格式
# dev_label                 列表格式
# dev_segment_embedding     列表格式

def replace_sequence_for_t(lst):
    for i in range(len(lst) - 2):
        if lst[i:i+3] == ['Ġ<','t','>']:
            lst[i:i+3] = ['<t>']
    for i in range(len(lst) - 2):
        if lst[i:i+3] == ['Ġ</','t','>']:
            lst[i:i+3] = ['</t>']
    return lst

print('loading tokenizer...')
# print(train_all_sentence[0])
# assert False
tokenizer = RobertaTokenizer.from_pretrained(pretrain_model_path)
train_feature = [
    tokenizer.tokenize(line) for line
    in train_all_sentence]
test_feature = [
    tokenizer.tokenize(line) for line
    in test_all_sentence]
dev_feature = [
    tokenizer.tokenize(line) for line
    in dev_all_sentence]


# 和句子的长度，原始标签的长度保持一致 的  列表，每一个位置表示tokenizer后对应的词的数量。
train_tokenizer_num = []
test_tokenizer_num = []
dev_tokenizer_num = []

def deal_tokenizer_num(tem_all_sentence, tem_tokenizer_num):
    # 产生XX_tokenizer_num, 也就是每个词对应的tokenizer的个数。
    for line in tem_all_sentence:
        s0 = line
        s3 = tokenizer.tokenize(line)
        s3 = replace_sequence_for_t(s3)
        seq_token_len = [1,]
        for i in range(1, len(s0.split())):
            if s0.split()[i]=='</s>':
                seq_token_len.append(2)
            elif s0.split()[i]=='</t>' or s0.split()[i]=='<t>':
                seq_token_len.append(1)
            else:
                seq_token_len.append(len(tokenizer(s0.split()[i], add_prefix_space=True)['input_ids']) - 2)

        assert sum(seq_token_len) == len(s3)
        assert len(seq_token_len) == len(s0.split())

        tem_tokenizer_num.append(seq_token_len)

deal_tokenizer_num(test_all_sentence, test_tokenizer_num)
deal_tokenizer_num(train_all_sentence, train_tokenizer_num)
deal_tokenizer_num(dev_all_sentence, dev_tokenizer_num)

# 根据tokenizer，对segment embedding 和 label 等的 数据进行修正。
# 进行相应的修改。因为此时必须包含trigger，如果不包含，则处理截取的窗口。
# 需要留存截取初始位置前面元素的个数，

tokenizer.add_tokens(['<t>', '</t>'])
train_feature_id = [tokenizer.convert_tokens_to_ids(replace_sequence_for_t(line)) for line in train_feature]
test_feature_id = [tokenizer.convert_tokens_to_ids(replace_sequence_for_t(line)) for line in test_feature]
dev_feature_id = [tokenizer.convert_tokens_to_ids(replace_sequence_for_t(line)) for line in dev_feature]

def deal_tokenizer_convert_to_tokenizer(tem_list, tem_tokenizer_num, tem_target_list):
    # 根据 tokenizer 处理，label和segment。
    for i in range(0, len(tem_list)):
        # train_segment_embedding[i] 是第i个句子的列表。
        tem_tem_list = []
        for j in range(0, len(tem_list[i])):
            # train_segment_embedding[i][j] 表示原始的位置元素。
            # train_tokenizer_num[i][j] 表示存放的数目。
            for k in range(0, tem_tokenizer_num[i][j]):
                tem_tem_list.append(tem_list[i][j])
        tem_target_list.append(tem_tem_list)

train_segment_embedding_1 = []
test_segment_embedding_1 = []
test_segment_embedding_1_222 = []
dev_segment_embedding_1 = []

deal_tokenizer_convert_to_tokenizer(train_segment_embedding, train_tokenizer_num, train_segment_embedding_1)
deal_tokenizer_convert_to_tokenizer(dev_segment_embedding, dev_tokenizer_num, dev_segment_embedding_1)
deal_tokenizer_convert_to_tokenizer(test_segment_embedding, test_tokenizer_num, test_segment_embedding_1)
deal_tokenizer_convert_to_tokenizer(test_segment_embedding_222, test_tokenizer_num, test_segment_embedding_1_222)



train_label_1 = []
test_label_1 = []
dev_label_1 = []
deal_tokenizer_convert_to_tokenizer(train_label, train_tokenizer_num, train_label_1)
deal_tokenizer_convert_to_tokenizer(test_label, test_tokenizer_num, test_label_1)
deal_tokenizer_convert_to_tokenizer(dev_label, dev_tokenizer_num, dev_label_1)




def padding_to_maxlength(max_length, padding_token_id, tem_paddinglist):
    # 进行最大长度的填充
    for j in range(len(tem_paddinglist)):
        # 将样本数据填充至长度为 max_padding_len
        i = tem_paddinglist[j]
        if len(i) < max_length:
            tem_paddinglist[j].extend([padding_token_id] * (max_length - len(i)))
        else:
            tem_paddinglist[j] = tem_paddinglist[j][0:max_length - 1] + [tem_paddinglist[j][-1]]

# feature  句子的截取 和 填充
# ====== 不同的与训练模型的padding当然只需要修改token_id的padding啦。 =========
padding_to_maxlength(max_padding_len, 0, train_feature_id)
padding_to_maxlength(max_padding_len, 0, test_feature_id)
padding_to_maxlength(max_padding_len, 0, dev_feature_id)
# seg
padding_to_maxlength(max_padding_len, 0, train_segment_embedding_1)
padding_to_maxlength(max_padding_len, 0, dev_segment_embedding_1)
padding_to_maxlength(max_padding_len, 0, test_segment_embedding_1)
padding_to_maxlength(max_padding_len, 0, test_segment_embedding_1_222)

# tokenizer_num   全部按照512.
padding_to_maxlength(512, 0, train_tokenizer_num)
padding_to_maxlength(512, 0, dev_tokenizer_num)
padding_to_maxlength(512, 0, test_tokenizer_num)


train_label = train_label_1
test_label = test_label_1
dev_label = dev_label_1

padding_to_maxlength(max_padding_len, 0, train_label)
padding_to_maxlength(max_padding_len, 0, test_label)
padding_to_maxlength(max_padding_len, 0, dev_label)




train_set = TensorDataset(torch.LongTensor(train_feature_id), torch.LongTensor(train_label),
                          torch.LongTensor(train_segment_embedding_1), torch.LongTensor(train_tokenizer_num))
train_loader = DataLoader(dataset=train_set, batch_size=16, shuffle=True)

test_set = TensorDataset(torch.LongTensor(test_feature_id), torch.LongTensor(test_label),
                         torch.LongTensor(test_segment_embedding_1), torch.LongTensor(test_tokenizer_num),
                         torch.LongTensor(test_segment_embedding_1_222))
test_loader = DataLoader(dataset=test_set, batch_size=16, shuffle=False)

dev_set = TensorDataset(torch.LongTensor(dev_feature_id), torch.LongTensor(dev_label),
                        torch.LongTensor(dev_segment_embedding_1), torch.LongTensor(dev_tokenizer_num))
dev_loader = DataLoader(dataset=dev_set, batch_size=16, shuffle=False)

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

class Bert(torch.nn.Module):
    def __init__(self):
        super(Bert, self).__init__()
        self.model = RobertaModel.from_pretrained(pretrain_model_path).cuda()  # , config=modelConfig
        new_embeddings = self.model.resize_token_embeddings(len(tokenizer))
        self.model.shared = new_embeddings

        embedding_dim = self.model.config.hidden_size
        self.dropout = torch.nn.Dropout(0.5)
        self.linear_0 = torch.nn.Linear(embedding_dim, de_bert_dem)
        self.linear_1 = torch.nn.Linear(embedding_dim, 1)
        self.linear_2 = torch.nn.Linear(embedding_dim, 1)

    def forward(self, tokens, seg_embedding, attention_mask):
        output = self.model(tokens, attention_mask=attention_mask)
        output = output[0]
        output = self.dropout(output)
        # output = self.linear_0(output)
        # for i in range(len(self.reasoner)):
        #     output = self.reasoner[i](output)
        output_1 = self.linear_1(output)
        output_2 = self.linear_2(output)
        return output_1.squeeze(-1), output_2.squeeze(-1)


loss_func = torch.nn.CrossEntropyLoss()
model = Bert()
model = torch.nn.DataParallel(model, device_ids=[0])
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
max_acc = args.max_acc
print('start trainning....')


# model.load_state_dict(torch.load('./model_2/static_dict_0.pkl'))


def test(model, dev_dataloader):
    test_loss, test_f1, n = 0.0, 0.0, 0
    all_label = []
    all_prediction = []
    all_prediction_list_begin = []
    all_prediction_list_end = []
    model.eval()
    with torch.no_grad():
        for data, label, seg_embedding, tokenizer_num, seg_embedding222 in dev_dataloader:
            out_1, out_2 = model(data.cuda(), seg_embedding222.cuda(), attention_mask=(data > 0).cuda())
            # [batch_size, len]
            n += 1
            label_begin = []
            label_end = []
            tokenizer_id = 0
            for label_ in label:
                label_begin_ = tokenizer_num[tokenizer_id][0]
                label_end_ = tokenizer_num[tokenizer_id][0]
                for tem_seg_index in range(1, len(seg_embedding[tokenizer_id])):
                    # print(seg_embedding[tokenizer_id][tem_seg_index])
                    if seg_embedding[tokenizer_id][tem_seg_index] == 0:
                        label_begin_ = tem_seg_index
                        label_end_ = tem_seg_index
                        break
                tokenizer_id += 1
                for i in range(0, max_padding_len):
                    if label_[i] == 1:
                        label_begin_ = i
                        break
                for i in range(0, max_padding_len):
                    if label_[max_padding_len - 1 - i] == 1:
                        label_end_ = max_padding_len - 1 - i
                        break
                label_begin.append(label_begin_)
                label_end.append(label_end_)
            label_begin = torch.LongTensor(label_begin)
            label_end = torch.LongTensor(label_end)
            loss = loss_func(out_1.cuda(), label_begin.cuda()) + loss_func(out_2.cuda(), label_end.cuda())
            prediction_begin = out_1.argmax(dim=1).view(-1).data.cpu().numpy().tolist()
            prediction_end = out_2.argmax(dim=1).view(-1).data.cpu().numpy().tolist()
            all_prediction_list_begin.extend(prediction_begin)
            all_prediction_list_end.extend(prediction_end)
            prediction = []
            for i in range(0, len(prediction_begin)):
                prediction_tem = []
                for _ in range(0, max_padding_len):
                    prediction_tem.append(0)
                for j in range(prediction_begin[i], prediction_end[i] + 1):
                    prediction_tem[j] = 1
                prediction.extend(prediction_tem)
            label = label.view(-1).squeeze().data.cpu().numpy().tolist()
            test_loss += loss.item()
            all_label.extend(label)
            all_prediction.extend(prediction)
    test_f1 = f1_score(all_label, all_prediction, average='macro')

    # all_prediction_list_begin 和 all_prediction_list_end 中存放的  分别是预测的开头的结尾。
    # 而且是下标，而不是第几个位置处。
    # 对 test_segment_embedding 进行修正。一方面是前面的 cls 事件
    # 另一方面是 子词的切分。
    for i in range(0, len(test_tokenizer_num)):
        # begin 进行修正。
        all = 0
        c = 0
        for n_ in test_tokenizer_num[i]:
            all += n_
            c = c + n_ - 1
            if all >= all_prediction_list_begin[i] + 1:
                break
        all_prediction_list_begin[i] = all - 1 - c - test_len[i]
        # end 进行修正
        all = 0
        c = 0
        for n_ in test_tokenizer_num[i]:
            all += n_
            c = c + n_ - 1
            if all >= all_prediction_list_end[i] + 1:
                break
        all_prediction_list_end[i] = all - 1 - c - test_len[i]
    # 但是前面 两个额外的，所以-2
    test_index = 0
    f_1 = open('1.json', mode='w', encoding='UTF-8')
    for i in range(0, len(test_seq_in)):
        m = eval(test_seq_in[i])
        trigger_begin = m['evt_triggers'][0][0]
        trigger_end = m['evt_triggers'][0][1]
        result_list = []
        pad_ = 'evt089arg02'
        for _ in range(0, len(m['gold_evt_links'])):  # 这里需要修改，引入一个额外的列表即可。第i 个test_seq_in 需要预测的数目
            # [[40, 40], [28, 28], "evt089arg02place"]
            if len(m['gold_evt_links']) <= 0:
                break
            ss = [[m['gold_evt_links'][0][0][0], m['gold_evt_links'][0][0][1]], [28, 28], ""]
            ss[1][0] = all_prediction_list_begin[test_index]
            ss[1][1] = all_prediction_list_end[test_index]
            tem_add_1 = 0
            tem_add_2 = 0
            if ss[1][0] >= trigger_begin + 1:
                tem_add_1 += 1
            if ss[1][0] > trigger_end + 1:
                tem_add_1 += 1
            if ss[1][1] >= trigger_begin + 1:
                tem_add_2 += 1
            if ss[1][1] > trigger_end + 1:
                tem_add_2 += 1
            ss[1][0] = ss[1][0] - tem_add_1
            ss[1][1] = ss[1][1] - tem_add_2

            # ss[2] = pad_ + (test_all_sentence[test_index].split())[1]
            ss[2] = pad_ + test_role[test_index]
            if ss[1][0] <= ss[1][1] and ss[1][0] > 0 and ss[1][1] < len(test_all_sentence[test_index].split()) - 2 - \
                    test_len[test_index]:
                result_list.append(ss)
            test_index += 1
        m['gold_evt_links'] = result_list
        json.dump(m, f_1, ensure_ascii=False)
        f_1.write('\n')
    f_1.close()

    return_dict = run_evaluation(args)

    return test_loss / n, test_f1


for epoch in range(args.num_epoch):
    train_loss, train_f1, n = 0.0, 0.0, 0
    all_label = []
    all_prediction = []
    for data, label, seg_embedding, tokenizer_num in train_loader:
        model.train()
        out_1, out_2 = model(data.cuda(), seg_embedding.cuda(), attention_mask=(data > 0).cuda())
        # [batch_size, len]
        n += 1
        label_begin = []
        label_end = []
        tokenizer_id = 0
        for label_ in label:
            label_begin_ = tokenizer_num[tokenizer_id][0]
            label_end_ = tokenizer_num[tokenizer_id][0]
            for tem_seg_index in range(1, len(seg_embedding[tokenizer_id])):
                # print(seg_embedding[tokenizer_id][tem_seg_index])
                if seg_embedding[tokenizer_id][tem_seg_index] == 0:
                    break
                label_begin_ += tokenizer_num[tokenizer_id][tem_seg_index]
                label_end_ += tokenizer_num[tokenizer_id][tem_seg_index]

            tokenizer_id += 1
            for i in range(0, max_padding_len):
                if label_[i] == 1:
                    label_begin_ = i
                    break
            for i in range(0, max_padding_len):
                if label_[max_padding_len - 1 - i] == 1:
                    label_end_ = max_padding_len - 1 - i
                    break
            label_begin.append(label_begin_)
            label_end.append(label_end_)
        label_begin = torch.LongTensor(label_begin)
        label_end = torch.LongTensor(label_end)
        loss = loss_func(out_1.cuda(), label_begin.cuda()) + loss_func(out_2.cuda(), label_end.cuda())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        prediction_begin = out_1.argmax(dim=1).view(-1).data.cpu().numpy().tolist()
        prediction_end = out_2.argmax(dim=1).view(-1).data.cpu().numpy().tolist()
        prediction = []
        for i in range(0, len(prediction_begin)):
            prediction_tem = []
            for _ in range(0, max_padding_len):
                prediction_tem.append(0)
            for j in range(prediction_begin[i], prediction_end[i] + 1):
                prediction_tem[j] = 1
            prediction.extend(prediction_tem)
        label = label.view(-1).squeeze().data.cpu().numpy().tolist()
        train_loss += loss.item()
        all_label.extend(label)
        all_prediction.extend(prediction)

    train_f1 = f1_score(all_label, all_prediction, average='macro')
    train_loss = train_loss / n
    test_loss, test_f1 = test(model, test_loader)

    if temF1 > myF1:
        myF1 = temF1
        # torch.save(model.module.state_dict(), './model/static_dict' + str(epoch) + '_' + str(myF1) + '.pkl')
    myF1 = temF1

    print('epoch %d, train_loss %f, train_f1 %f, dev_loss %f, dev_f1 %f' %
          (epoch + 1, train_loss, train_f1, test_loss, test_f1))
    print('----------------------------')