PrivacyNER_with_NoisyData_CCF/main.py at main · RacleRay/PrivacyNER_with_NoisyData_CCF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import os
import torch

from data.dataset import DatasetFactory
from model.trainer import Trainer
from model.mymodel import CascadeFLAT
from utils.postprocess.process import analyze_dev_results, generate_submit
from utils.tools import console
from config import Config


datamanager = DatasetFactory(config=Config)
trainer = Trainer()
model = CascadeFLAT().cuda()


def run_train_kflod():
    console.log(f" === RUNNING Train with {Config.fold_num} flod=== ")

    for fold_idx in range(Config.fold_num):
        weight_path = os.path.join(Config.WEIGHT_PATH,
                                    'model_{}.bin'.format(fold_idx))

        train_data = datamanager.fetch_dataset("train", fold_idx)
        dev_data = datamanager.fetch_dataset("dev", fold_idx)
        inputs = {
            'train_data': train_data,
            'dev_data': dev_data,
            'dev_res_file': os.path.join(Config.RESULT_PATH,'dev_result.json'),
            'epoch_start': 0
        }
        trainer.train(inputs, model)
        torch.save(model.state_dict(), weight_path)


        if Config.en_eval:
            inputs = {'data': dev_data, 'type_data': 'dev'}

            inputs['weight'] = weight_path
            inputs['outfile'] = os.path.join(
                Config.RESULT_PATH, 'dev_result_{}.json'.format(fold_idx))

            results = trainer.eval(inputs, model)
            console.log(results)

            if fold_idx == (Config.fold_num - 1):
                analyze_dev_results(Config.RESULT_PATH,
                                    os.path.join(Config.DATA_PATH, 'seg_train_data.json'),
                                    os.path.join(Config.RESULT_PATH, 'dev_result_all.json'),
                                    verbose=False)


def run_train():
    "Config.en_cross == False"

    console.log(" === RUNNING Train === ")

    weight_path = os.path.join(Config.WEIGHT_PATH, 'model.bin')

    train_data = datamanager.fetch_dataset("train")
    dev_data = datamanager.fetch_dataset("dev")

    inputs = {
        'train_data': train_data,
        'dev_data': dev_data,
        'dev_res_file': os.path.join(Config.RESULT_PATH, 'dev_result.json'),
        'epoch_start': 0
    }

    trainer.train(inputs, model)
    torch.save(model.state_dict(), weight_path)


    if Config.en_eval:
        inputs['weight'] = weight_path
        inputs['outfile'] = os.path.join(Config.RESULT_PATH, 'dev_result.json')

        results = trainer.eval(inputs, model)
        console.log(results)

        analyze_dev_results(Config.RESULT_PATH,
                            os.path.join(Config.DATA_PATH, 'seg_train_data.json'),
                            en_cross=False,
                            verbose=False)


def run_train_lossmask():
    from utils.postprocess.selfsupervise import generate_loss_mask
    from utils.preprocess.format import format_data

    # generate input tokens with mask
    # 需要先训练生成 dev_result_all.json 文件
    generate_loss_mask(os.path.join(Config.RESULT_PATH, 'dev_result_all.json'),
                       os.path.join(Config.DATA_PATH, 'seg_train_data.json'),
                       os.path.join(Config.DATA_PATH, 'seg_train_data_loss_mask.json'))

    format_data("train_loss_mask")

    console.log(" === Loss mask generated. Begin Training... === ")

    if Config.en_cross:
        run_train_kflod()
    else:
        run_train()


def run_test():

    console.log(" === RUNNING Test === ")

    test_data = datamanager.fetch_dataset("test")
    inputs = {'data': test_data, 'type_data': 'test'}

    if Config.en_cross:
        for fold_idx in range(Config.fold_num):
            inputs['weight'] = Config.WEIGHT_PATH + 'swa_model_{}.bin'.format(fold_idx)
            inputs['outfile'] = Config.RESULT_PATH + 'test_result_{}.json'.format(fold_idx)

            trainer.test(inputs, model)

        generate_submit([Config.RESULT_PATH + 'test_result_{}.json'.format(fold_idx) \
                         for fold_idx in range(Config.fold_num)],
                        Config.RESULT_PATH + 'predict_vote.csv',
                        cheat_with_clue=False)

    else:
        inputs['weight'] = Config.WEIGHT_PATH + 'swa_model.bin'
        inputs['outfile'] = Config.RESULT_PATH + 'test_result.json'

        trainer.test(inputs, model)

        generate_submit(Config.RESULT_PATH + 'test_result.json',
                        Config.RESULT_PATH + 'predict.csv',
                        cheat_with_clue=False)

        # generate_submit(Config.RESULT_PATH + 'test_result.json',
        #                 Config.RESULT_PATH + 'predict.csv',
        #                 cheat_with_clue=True,
        #                 clue_file=Config.ADDITIONAL_DATA_PATH + 'train_dev_data.json')


if __name__ == '__main__':
    if Config.en_cross:
        run_train_kflod()
    else:
        run_train()

    # run_test()

    # run_train_lossmask()

    del datamanager
    del trainer
    del model