Skip to main content

extract vector from BERT pre-train model

Project description

BERT预训练模型字向量提取工具

BERT预训练模型字向量提取工具

版本: v 0.3.7

更新: 2020/4/20 16:39

升级日志

v0.3.7

  • 把测试程序加入到包中,可直接在命令行中使用 BERTVector_test 运行测试程序;

v0.3.6

  • 发布到pypi中,可直接在命令行使用;

v0.3.3

  • 增加了测试的样本及使用示例:短句相似度,词向量分布图等;

v0.3.2

  • 同时兼容txt,pkl两种格式输出;

v0.3

  • 使用pkl格式保存输出文件;
  • 去掉了txt的格式输出;
  • 参数--in_file可指定文件或者目录,自动识别处理;

v0.2

  • 可对单个数据文件进行处理;

工具说明

本工具直接读取BERT预训练模型,从中提取样本文件中所有使用到字向量,保存成向量文件,为后续模型提供字向量。

本工具直接读取预训练模型,不需要其它的依赖,同时把样本中所有出现的字符对应的字向量全部提取, 后续的模型可以非常快速进行索引,生成自己的句向量,不再需要庞大的预训练模型或者bert-as-service服务了。

安装:

pip install BERTVector

使用:

BERTVector -h

测试:

BERTVector_test

输出文件格式

支持txt和pkl两种文件格式,可自由选择,默认为pkl格式。 (>v0.3.2版本)

txt格式为:

一行一个字符向量,中间使用空格分隔;

格式为:字符 768大小向量

例如:

! 0.06241071 -0.012108288 0.035899848....

文件格式为pkl,数据格式为dic字典格式,更方便使用。(v0.3版本)

key: 字符

value: 对应的向量,768大小

参数说明

参数说明:

--model_path    BERT预训练模型存放的目录,必填;

--in_file       待提取的样本文件名或者目录名,必填;

--out_file      提取后保存的向量文件,可选项,默认为'./bert_embedding.pkl'

--ext           当指定为目录时,可选项,要处理的文件后缀,默认为['csv','txt']

--fmt           指定输出文件的格式,可设置为'txt'或者'pkl', 默认为pkl

调用示例

命令行示例:

示例一:

处理单个文件./data/train_interger.csv,保存到./data/need_bertembedding.pkl

BERTVector \
    --model_path=/mnt/sda1/models/chinese_L-12_H-768_A-12/ \
    --in_file=./data/train_interger.csv \
    --out_file=./data/need_bertembedding.pkl

示例二:

处理目录下的所有tsv,txt文件,默认保存为:./need_bertembedding.pkl

BERTVector \
    --model_path=/mnt/sda1/models/chinese_L-12_H-768_A-12/ \
    --in_file=./data/ \
    --ext tsv txt

示例三: 处理文件:dat.csv,保存为txt格式,输出文件名为:./need_bertembedding

BERTVector \
    --model_path=/mnt/sda1/models/chinese_L-12_H-768_A-12/ \
    --in_file=./dat.csv \
    --out_file=./need_bertembedding \
    --fmt=txt

测试数据及示例

直接运行以下命令即可运行测试程序:

BERTVector_test

示例文件跟随项目安装在python的目录下: \Lib\site-packages\BERTVector\test

可使用以下命令生成测试的向量字典:

BERTVector --model_path=d:\\model\chinese_L-12_H-768_A-12 --in_file=test/test.txt --out_file=test/test.pkl

其中d:\\model\chinese_L-12_H-768_A-12是BERT预训练模型的目录

生成的字向量字典可以方便地为文本进行编码。 示例文件及数据在test目录下,可运行以下命令查看

cd test
python BERTVector_test.py

运行结果:

运行截图

代码输出如下:

17:35:27.08|X:>load_vector.py
-----------------字向量加载------------------
字典大小:221
向量维度:768

字典首个元素:
字符:(
向量(仅输出前20维):
[-0.0701572  -0.02538603  0.01186333  0.03037252 -0.01405589  0.04111631
 -0.00601521  0.08385976 -0.06300574  0.06307469 -0.08003023  0.01581223
 -0.0285087  -0.04477277  0.0110429   0.05696263  0.00558841  0.05059543
  0.01589187  0.01623219]
-----------------句向量示例------------------
待编码句子:
今日数据趣谈:阿杜比肩魔术师 热火中锋另类纪录新浪体育讯北京时间4月28日

本示例中简单把各字的向量相加作为句向量...
句向量(仅输出前20维):
[-0.31147036 -0.04251522  0.16138141 -0.39930761 -1.07815197  0.06680285
 -0.1597157  -0.37606054 -0.42281876  0.34424412 -0.05901627  0.34751215
 -1.04583078 -0.50036957 -0.17383758  0.54849081  0.13344521 -0.02983405
  0.10145024  0.55525893]

补充使用示例 v0.3.3

生成了句向量,计算句子之间的相似度;

----------------余弦相似度计算-----------------
[基金投资的热度近期有所回升] 与 [新浪体育纪录另类热火中锋] 的相似度: 0.4444
[基金投资的热度近期有所回升] 与 [基金投资的热度近期有所回升] 的相似度: 1.0000
[基金投资的热度近期有所回升] 与 [封闭式基金折价率近期有所上涨] 的相似度: 0.7164
[基金投资的热度近期有所回升] 与 [今日数据趣谈:阿杜比肩魔术师 热火中锋另类纪录新浪体育讯北京时间4月28日] 的相似度: 0.510
4
[基金投资的热度近期有所回升] 与 [以下是今日比赛中诞生的一些有趣数据] 的相似度: 0.5744
[基金投资的热度近期有所回升] 与 [最年轻纪录属于“魔术师”约翰逊,他1980年总决赛对阵76人得到42分,] 的相似度: 0.4471
[基金投资的热度近期有所回升] 与 [首次有三人组合在季后赛做到这一点] 的相似度: 0.5647
[基金投资的热度近期有所回升] 与 [安东尼此役打了38分59秒没有任何运动战出手] 的相似度: 0.4799
[基金投资的热度近期有所回升] 与 [本周受权重股业绩超市场预期] 的相似度: 0.5688
[基金投资的热度近期有所回升] 与 [券商研究机构在本周密集发布二季度的基金投资策略报告] 的相似度: 0.6885

生成词向量,将词向量显示到二维和三维分布图中展示;

运行截图

词语二维图 词语二维图-放大

词语三维图

请运行例子自行生成:

注意:每次生成的图有可能不一样。

应用案例

使用BERT预训练模型的向量进行编码,可再搭建LSTM,textCNN等模型进行文本分类。

实际案例: BERT字向量+LSTM文本分类模型 https://github.com/mathCrazyy/text_classify

注:该案例使用参数--fmt=txt生成输出文件。

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

BERTVector-0.3.7.tar.gz (1.8 MB view hashes)

Uploaded source

Built Distribution

BERTVector-0.3.7-py3-none-any.whl (1.8 MB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page