zejunwang1 / bert_text_classification

基于 BERT 模型的中文文本分类工具

bert_text_classification

基于 BERT 模型的中文文本分类工具。

环境

python 3.7

torch 1.7

tqdm

sklearn

transformers 4.8.1

中文数据集

从 THUCNews 中随机抽取20万条新闻标题，一共有10个类别：财经、房产、股票、教育、科技、社会、时政、体育、游戏、娱乐，每类2万条标题数据。数据集按如下划分：

• 训练集：18万条新闻标题，每个类别的标题数为18000
• 验证集：1万条新闻标题，每个类别的标题数为1000
• 测试集：1万条新闻标题，每个类别的标题数为1000

可以按照 data 文件夹中的数据格式来准备自己任务所需的数据，并调整 config.py 中的相关配置参数。

预训练 BERT 模型

从 huggingface 官网上下载 bert-base-chinese 模型权重、配置文件和词典到 pretrained_bert 文件夹中，下载地址：https://huggingface.co/bert-base-chinese/tree/main

模型训练

文本分类模型训练：

python main.py --mode train --data_dir ./data --pretrained_bert_dir ./pretrained_bert

训练中间日志如下：

模型在验证集上的效果如下：

Demo演示

文本分类 demo 展示：

python main.py --mode demo --data_dir ./data --pretrained_bert_dir ./pretrained_bert

模型预测

对 data 文件夹下的 input.txt 中的文本进行分类预测：

python main.py --mode predict --data_dir ./data --pretrained_bert_dir ./pretrained_bert --input_file ./data/input.txt

输出如下结果：

bert_text_classification's Issues

Need Help Train not running success?

When I run python main.py --mode train --data_dir ./data --pretrained_bert_dir ./pretrained_bert
loading data...
202752it [00:00, 1377265.38it/s]
18980it [00:00, 1325163.80it/s]
time usage: 0:00:00
It's train mode
Epoch [1/3]
Epoch [2/3]
Epoch [3/3]
then the shell exit，not print the line
Iter: 0, Batch Train Loss: 2.4, Batch Train Acc: 13.28%, Val Loss: 2.3, Val Acc: 13.82%, Time: 0:00:19 *
Iter: 100, Batch Train Loss: 0.31, Batch Train Acc: 90.62%, Val Loss: 0.35, Val Acc: 89.44%, Time: 0:01:42 *
Looking forward to your reply，Thanks