使用Bert预训练模型进行中文文本分类（基于pytorch）

基于BERT预训练模型的中文文本分类任务的实现

17不想当码农

10643人浏览 · 2022-02-15 18:18:00

17不想当码农 · 2022-02-15 18:18:00 发布

前言

最近在做一个关于图书系统的项目，需要先对图书进行分类，想到Bert模型是有中文文本分类功能的，于是打算使用Bert模型进行预训练和实现下游文本分类任务

数据预处理

2.1 输入介绍

在选择数据集时，我们首先要知道的是模型到底应该接受什么样的输入，只有让模型接收正确的输入，才能构建出正确的数据形式。Bert模型的输入如图：

图 1 BERT模型输入图

在Segment embeddings里面，中文模型的处理与原始BERT的处理有一些不一样，由于对于文本分类这个场景输入只有一个序列，所以不需要区分0或1，只需要构造全0的输入即可，其他地方与原始BERT模型无较大差别，只需要构造原始文本对应的token序列，并在首尾分别加上[CLS]和[SEP]符作为输入即可

2.2 语料介绍

这里我是用的语料是THUCNews中的20万条新闻标题，文本长度在20到30之间。一共10个类别，每类2万条。数据以字为单位输入模型。

类别：财经、房产、股票、教育、科技、社会、时政、体育、游戏、娱乐

如果你想用不同的数据集来训练，可以参考本数据集的格式来格式化自己想要训练的数据集

我的数据集分配比例如下：

数据集	数据量
训练集	18万
验证集	1万
测试集	1万

2.3 数据集构建流程

第1步：将原始的数据样本进行分字（tokenize）处理

体现在代码中是 utils.py的line 20

token = config.tokenizer.tokenize(content)

第2步：根据第1步tokenize后的结果构造一个字典，但是由于我们使用的BERT

预训练模型，所以这个字典不需要我们自己来构造，只需要引入谷歌开源的 vocab.txt文件构造字典即可，因为只有vocab.txt中每个字的索引顺序才与开源模型中每个字的Embedding向量一一对应。

体现在代码中是 utils.py的line 24

token_ids = config.tokenizer.convert_tokens_to_ids(token)

第3步：根据字典将tokenize后的文本序列转换为token 序列，同时在token序列的首尾分别加上[CLS]和[SEP]符号，并进行padding（填充）

体现在代码中是 utils.py的line 26-33

if pad_size:
    if len(token) < pad_size:
        mask = [1] * len(token_ids) + [0] * (pad_size - len(token))
        token_ids += ([0] * (pad_size - len(token)))
    else:
        mask = [1] * pad_size
        token_ids = token_ids[:pad_size]
        seq_len = pad_size

第4步则是根据第3步处理后的结果生成对应的Padding Mask向量。

最后在模型训练的时候，只需要将第3步和第4步处理后的结果一起喂给模型就行了。

代码

源码链接：https://github.com/649453932/Bert-Chinese-Text-Classification-Pytorch

python 版本：3.7

pytorch版本：1.1

run.py

import time
import torch
import numpy as np
from train_eval import train, init_network
from importlib import import_module
import argparse
from utils import build_dataset, build_iterator, get_time_dif

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True, help='choose a model: Bert, ERNIE')
args = parser.parse_args()


if __name__ == '__main__':
    dataset = 'THUCNews'  # 数据集

    model_name = args.model  # bert
    x = import_module('models.' + model_name)
    config = x.Config(dataset)
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True  # 保证每次结果一样

    start_time = time.time()
    print("Loading data...")
    train_data, dev_data, test_data = build_dataset(config)
    train_iter = build_iterator(train_data, config)
    dev_iter = build_iterator(dev_data, config)
    test_iter = build_iterator(test_data, config)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

    # train
    model = x.Model(config).to(config.device)
    train(config, model, train_iter, dev_iter, test_iter)