TextMatch/train_model/train_bert.py at master · MachineLP/TextMatch · 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
# -*- coding:utf-8 -*-
'''
-------------------------------------------------
   Description :  bert train
   Author :       machinelp
   Date :         2020-06-03
-------------------------------------------------

'''
import numpy as np
import pandas as pd
from keras.layers import *
from bert4keras.backend import keras, set_gelu
from bert4keras.bert import build_bert_model
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding, DataGenerator
from bert4keras.tokenizer import Tokenizer
from textmatch.config.constant import Constant as const
from textmatch.models.text_embedding.bert_embedding import BertEmbedding

set_gelu('tanh')  # 切换gelu版本

maxlen = 32
batch_size = 16
num_classes = 2
epochs = 20
learning_rate = 2e-5


# sim roeberta_zh
# 【百度网盘】链接:https://pan.baidu.com/s/1RVAHqL1CfLGltPWpoTyThw  密码:bp71
config_path = 'publish/bert_config.json'
checkpoint_path = 'publish/bert_model.ckpt'
dict_path = 'publish/vocab.txt'


def load_data(filename):
    D = []
    data = pd.read_csv(filename)
    data.dropna(axis=0, how='any', inplace=True)
    data = data.values.tolist()
    for per_data in data:
        D.append( (per_data[0],per_data[1],int(per_data[2])) )
    return D


# 加载数据集
train_val_data = load_data('./data/train_data.csv')
# test_data = load_data('dev.csv')
# 查看一下数据
print ( 'train>>>>', train_val_data[0] )
print ( '训练验证集数量:', len(train_val_data) )

random_order = range(len(train_val_data))
np.random.shuffle(list(random_order))
train_data = [train_val_data[j] for i, j in enumerate(random_order) if i % 5 != 1 ]
valid_data = [train_val_data[j] for i, j in enumerate(random_order) if i % 5 == 1 ]
test_data = valid_data
print ( '训练集数量:', len(train_data) )
print ( '验证集数量:', len(valid_data) )
print ( '测试集数量:', len(test_data) )


# 建立分词器
tokenizer = Tokenizer(dict_path, do_lower_case=True)
class data_generator(DataGenerator):
    """数据生成器
    """

    def __iter__(self, random=False):
        idxs = list(range(len(self.data)))
        if random:
            np.random.shuffle(idxs)
        batch_token_ids, batch_segment_ids, batch_labels = [], [], []
        for i in idxs:
            text1, text2, label = self.data[i]
#             print(text1, text2, label)
            token_ids, segment_ids = tokenizer.encode(text1, text2, max_length=maxlen)
            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append([label])
            if len(batch_token_ids) == self.batch_size or i == idxs[-1]:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids, batch_segment_ids], batch_labels
                batch_token_ids, batch_segment_ids, batch_labels = [], [], []

const.BERT_CONFIG_PATH = config_path
const.BERT_CHECKPOINT_PATH = checkpoint_path
const.BERT_DICT_PATH = dict_path

bert_embedding = BertEmbedding(const.BERT_CONFIG_PATH, const.BERT_CHECKPOINT_PATH, const.BERT_DICT_PATH, train_mode=True)
bert = bert_embedding.bert
output = Dropout(rate=0.1)(bert.model.output)
output = Dense(units=num_classes,
               activation='softmax',
               kernel_initializer=bert.initializer)(output)

model = keras.models.Model(bert.model.input, output)
model.summary()

model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=Adam(learning_rate),  # 用足够小的学习率
    # optimizer=PiecewiseLinearLearningRate(Adam(5e-5), {10000: 1, 30000: 0.1}),
    metrics=['accuracy'],
)

# 转换数据集
train_generator = data_generator(train_data, batch_size)
valid_generator = data_generator(valid_data, batch_size)
test_generator = data_generator(test_data, batch_size)


def evaluate(data):
    total, right = 0., 0.
    for x_true, y_true in data:
        y_pred = model.predict(x_true).argmax(axis=1)
        y_true = y_true[:, 0]
        total += len(y_true)
        right += (y_true == y_pred).sum()
    return right / total


class Evaluator(keras.callbacks.Callback):
    def __init__(self):
        self.best_val_acc = 0.

    def on_epoch_end(self, epoch, logs=None):
        val_acc = evaluate(valid_generator)
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            model.save_weights('best_bert_model.weights')
        test_acc = evaluate(test_generator)
        print(u'val_acc: %.5f, best_val_acc: %.5f, test_acc: %.5f\n'
              % (val_acc, self.best_val_acc, test_acc))

evaluator = Evaluator()
model.fit_generator(train_generator.forfit(),
                    steps_per_epoch=len(train_generator),
                    epochs=epochs,
                    callbacks=[evaluator])

model.load_weights('best_bert_model.weights')
print(u'final test acc: %05f\n' % (evaluate(test_generator)))