Skip to content

一行代码使用BERT生成句向量,BERT做文本分类、文本相似度计算

License

Notifications You must be signed in to change notification settings

frostjsy/bert-utils

 
 

Repository files navigation

bert-utils

本文基于Google开源的BERT代码进行了进一步的简化,方便生成句向量与做文本分类


***** New July 1st, 2019 *****

  • 修改句向量graph文件的生成方式,提升句向量启动速度。不再每次以临时文件的方式生成,首次执行extract_feature.py时会创建tmp/result/graph, 再次执行时直接读取该文件,如果args.py文件内容有修改,需要删除tmp/result/graph文件
  • 修复同时启动两个进程生成句向量时代码报错的bug
  • 修改文本匹配数据集为QA_corpus,该份数据相比于蚂蚁金服的数据更有权威性

1、下载BERT中文模型

下载地址: https://storage.googleapis.com/bert_models/2018_11_03/chinese_L-12_H-768_A-12.zip

2、把下载好的模型添加到当前目录下

3、句向量生成

生成句向量不需要做fine tune,使用预先训练好的模型即可,可参考extract_feature.pymain方法,注意参数必须是一个list。

首次生成句向量时需要加载graph,并在output_dir路径下生成一个新的graph文件,因此速度比较慢,再次调用速度会很快

from bert.extrac_feature import BertVector
bv = BertVector()
bv.encode(['今天天气不错'])

4、文本分类

文本分类需要做fine tune,首先把数据准备好存放在data目录下,训练集的名字必须为train.csv,验证集的名字必须为dev.csv,测试集的名字必须为test.csv, 必须先调用set_mode方法,可参考similarity.pymain方法,

训练:

from similarity import BertSim
import tensorflow as tf

bs = BertSim()
bs.set_mode(tf.estimator.ModeKeys.TRAIN)
bs.train()

验证:

from similarity import BertSim
import tensorflow as tf

bs = BertSim()
bs.set_mode(tf.estimator.ModeKeys.EVAL)
bs.eval()

测试:

from similarity import BertSim
import tensorflow as tf

bs = BertSim()
bs.set_mode(tf.estimator.ModeKeys.PREDICT)
bs.test()

5、DEMO中自带了QA_corpus数据集,这里给出地址, 该份数据的生成方式请参阅附件中的论文The BQ Corpus.pdf

About

一行代码使用BERT生成句向量,BERT做文本分类、文本相似度计算

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 100.0%