bert 文本分类 tensorflow

        本文主要是参考了网上的文本分类例子，但网上的例子不够完善，只实现了训练的步骤，在此基础上，增加了模型数据保存，及如何调用模型。废话少说，上代码：（其中训练数据请自行下载，头条新闻数据下载链接：

链接：https://pan.baidu.com/s/1smvf5IzOMh4-lSK0kyPWNQ 
提取码：aaaa

预训练模型用的是“chinese_roberta_wwm_ext_L-12_H-768_A-12”，请自行下载：

链接：https://pan.baidu.com/s/1iUplG3al92X1qDX4vABX5w 
提取码：aaaa）

# 模型训练及模型保存代码文件(bert_train.py)

import pickle
from keras_bert import load_trained_model_from_checkpoint, Tokenizer
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
from sklearn.preprocessing import LabelEncoder
from sklearn.utils import shuffle
from keras.utils.vis_utils import plot_model
import codecs, gc
import keras.backend as K
import os
import pandas as pd

# 文件主路径定义
mainPath = '/home/dev/pythonproject/pros/bert_classifier/'

# 从文件中读取数据，获取训练集和验证集
rc = pd.read_csv(mainPath + 'data/tnews/toutiao_news_dataset.txt', delimiter="_!_", names=['labels', 'text'],
                 header=None, encoding='utf-8')

rc = shuffle(rc)  # shuffle数据，打乱

# 把类别转换为数字
# 一共15个类别："教育","科技","军事","旅游","国际","证券股票","农业","电竞游戏"，
# "民生故事","文化","娱乐","体育","财经","房产","汽车"
class_le = LabelEncoder()
rc.iloc[:, 0] = class_le.fit_transform(rc.iloc[:, 0].values)

# 保存标签文件
output_label2id_file = os.path.join(mainPath, "model/keras_class/label2id.pkl")
if not os.path.exists(output_label2id_file):
    with open(output_label2id_file, 'wb') as w:
        pickle.dump(class_le.classes_, w)

# 构建全部所需数据集
data_list = []
for d in rc.iloc[:].itertuples():
    data_list.append((d.text, d.labels))

# 取一部分数据做训练和验证
train_data = data_list[0:20000]
valid_data = data_list[20000:22000]

maxlen = 100  # 设置序列长度为100，要保证序列长度不超过512

# 设置预训练模型
configPath = mainPath + 'chinese_roberta_wwm_ext_L-12_H-768_A-12/bert_config.json'
ckpPath = mainPath + 'chinese_roberta_wwm_ext_L-12_H-768_A-12/bert_model.ckpt'
vocabPath = mainPath + 'chinese_roberta_wwm_ext_L-12_H-768_A-12/vocab.txt'

# 将词表中的词编号转换为字典
tokenDict = {}
with codecs.open(vocabPath, 'r', encoding='utf-8') as reader:
    for line in reader:
        token = line.strip()
        tokenDict[token] = len(tokenDict)


# 重写tokenizer
class OurTokenizer(Tokenizer):
    def _tokenize(self, content):
        reList = []
        for t in content:
            if t in self._token_dict:
                reList.append(t)
            elif self._is_space(t):

                # 用[unused1]来表示空格类字符
                reList.append('[unused1]')
            else:
                # 不在列表的字符用[UNK]表示
                reList.append('[UNK]')
        return reList


tokenizer = OurTokenizer(tokenDict)


def seqPadding(X, padding=0):
    L = [len(x) for x in X]
    ML = max(L)
    return np.array([np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X])


class data_generator:
    def __init__(self, data, batch_size=32, shuffle=True):
        self.data = data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.steps = len(self.data) // self.batch_size
        if len(self.data) % self.batch_size != 0:
            self.steps += 1

    def __len__(self):
        return self.steps

    def __iter__(self):
        while True:
            idxs = list(range(len(self.data)))

            if self.shuffle:
                np.random.shuffle(idxs)

            X1, X2, Y = [], [], []
            for i in idxs:
                d = self.data[i]
                text = d[0][:maxlen]
                x1, x2 = tokenizer.encode(first=text)
                y = d[1]
                X1.append(x1)
                X2.append(x2)
                Y.append([y])
                if len(X1) == self.batch_size or i == idxs[-1]:
                    X1 = seqPadding(X1)
                    X2 = seqPadding(X2)
                    Y = seqPadding(Y)
                    yield [X1, X2], Y
                    [X1, X2, Y] = [], [], []


# bert模型设置
bert_model = load_trained_model_from_checkpoint(configPath, ckpPath, seq_len=None)  # 加载预训练模型

for l in bert_model.layers:
    l.trainable = True

x1_in = Input(shape=(None,))
x2_in = Input(shape=(None,))

x = bert_model([x1_in, x2_in])

# 取出[CLS]对应的向量用来做分类
x = Lambda(lambda x: x[:, 0])(x)
p = Dense(15, activation='softmax')(x)

model = Model([x1_in, x2_in], p)
model.compile(loss='sparse_categorical_crossentropy', optimizer=Adam(1e-5), metrics=['accuracy'])
model.summary()

train_D = data_generator(train_data)
valid_D = data_generator(valid_data)

model.fit_generator(train_D.__iter__(), steps_per_epoch=len(train_D), epochs=5, validation_data=valid_D.__iter__(),
                    validation_steps=len(valid_D))

model.save(mainPath + 'model/keras_class/tnews.h5', True, True)

# 保存模型结构图
plot_model(model, to_file='model/keras_class/tnews.png', show_shapes=True)

del model

# 清理内存
gc.collect()

# clear_session就是清除一个session
K.clear_session()

#训练完后，在model/keras_class文件夹里，会存在一个H5文件、及一个标签文件，模型调用实例代码文件（bert_predict.py)，上代码：

from bert_base.client import BertClient
import os
from keras.models import load_model
from keras_bert import get_custom_objects
import codecs, gc
import keras.backend as K
import tensorflow as tf
import pickle
from keras_bert import load_trained_model_from_checkpoint, Tokenizer
import numpy as np
from keras.models import model_from_json

str1 = "我爱北京天安门,天安门上太阳升"
str2 = "普京总统会见了拜登总统"
str3 = "中国乒乓球队夺得奥运会5枚金牌,目前居金牌榜第一"

# 文件主路径定义
mainPath = '/home/dev/pythonproject/pros/bert_classifier/'

dict_path = mainPath + 'chinese_roberta_wwm_ext_L-12_H-768_A-12/vocab.txt'

# 设置序列长度为100，要保证序列长度不超过512
maxlen = 100


# 重写tokenizer
class OurTokenizer(Tokenizer):
    def _tokenize(self, content):
        reList = []
        for t in content:
            if t in self._token_dict:
                reList.append(t)
            elif self._is_space(t):

                # 用[unused1]来表示空格类字符
                reList.append('[unused1]')
            else:
                # 不在列表的字符用[UNK]表示
                reList.append('[UNK]')
        return reList


# 将词表中的词编号转换为字典
tokenDict = {}
with codecs.open(dict_path, 'r', encoding='utf-8') as reader:
    for line in reader:
        token = line.strip()
        tokenDict[token] = len(tokenDict)

tokenizer = OurTokenizer(tokenDict)


def seqPadding(X, padding=0):
    L = [len(x) for x in X]
    ML = max(L)
    return np.array([
        np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X
    ])


class data_generator:
    def __init__(self, data, batch_size=32, shuffle=True):
        self.data = data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.steps = len(self.data) // self.batch_size
        if len(self.data) % self.batch_size != 0:
            self.steps += 1

    def __len__(self):
        return self.steps

    def __iter__(self):
        while True:
            idxs = list(range(len(self.data)))

            if self.shuffle:
                np.random.shuffle(idxs)

            X1, X2, Y = [], [], []
            for i in idxs:
                d = self.data[i]
                text = d[0][:maxlen]
                x1, x2 = tokenizer.encode(first=text)
                y = d[1]
                X1.append(x1)
                X2.append(x2)
                Y.append([y])
                if len(X1) == self.batch_size or i == idxs[-1]:
                    X1 = seqPadding(X1)
                    X2 = seqPadding(X2)
                    Y = seqPadding(Y)
                    yield [X1, X2], Y
                    [X1, X2, Y] = [], [], []


predict_D = data_generator([(str1, 4), (str2, 5), (str3, 2)], shuffle=False)
labes = None

output_label2id_file = os.path.join(mainPath, "model/keras_class/label2id.pkl")
if os.path.exists(output_label2id_file):
    with open(output_label2id_file, 'rb') as w:
        labes = pickle.load(w)

custom_objects = get_custom_objects()


model = load_model(mainPath + 'model/keras_class/tnews.h5', custom_objects=custom_objects)
tmpData = predict_D.__iter__()
preds = model.predict_generator(tmpData, steps=len(predict_D), verbose=1)

# 求每行最大值得下标，其中，axis=1表示按行计算
index_maxs = np.argmax(preds, axis=1)
result = [(x, labes[x]) for x in index_maxs]

print(result)

del model

# 清理内存
gc.collect()

# #clear_session就是清除一个session
K.clear_session()

        上面代码测试了3句话，"我爱北京天安门,天安门上太阳升"、 "普京总统会见了拜登总统" 、 "中国乒乓球队夺得奥运会5枚金牌,目前居金牌榜第一"，对应的标签输出为：

[(7, '文化'), (3, '国际'), (0, '体育')]

到此，文本分类简单实例应用到此结束，可以根据自己的项目实际情况，进行改造，如改成接口调用方式等，仅供参考。