BERT预训练模型字向量提取工具预训练模型字向量提取工具–使用使用BERT编码句子编码句子
本文将介绍两个使用BERT编码句子（从BERT中提取向量）的例子。
（1）BERT预训练模型字向量提取工具
本工具直接读取BERT预训练模型，从中提取样本文件中所有使用到字向量，保存成向量文件，为后续模型提供embdding。
本工具直接读取预训练模型，不需要其它的依赖，同时把样本中所有 出现的字符对应的字向量全部提取，后续的模型可以非
常快速进行embdding
github完整源码
#!/usr/bin/env python
# coding: utf-8
__author__ = 'xmxoxo'
'''
BERT预训练模型字向量提取工具
版本： v 0.3.2
更新: 2020/3/25 11:11
git: https://github.com/xmxoxo/BERT-Vector/
'''
import argparse
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
import numpy as np
import os
import sys
import traceback
import pickle
gblVersion = '0.3.2'
# 如果模型的文件名不同，可修改此处
model_name = 'bert_model.ckpt'
vocab_name = 'vocab.txt'
# BERT embdding提取类
class bert_embdding():
def __init__(self, model_path='', fmt='pkl'):
# 模型和词表的文件名
ckpt_path = os.path.join(model_path, model_name)
vocab_file = os.path.join(model_path, vocab_name)
if not os.path.isfile(vocab_file):
print('词表文件不存在，请检查...')
#sys.exit()
return
# 从模型读出指定层
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_path)
#param_dict = reader.get_variable_to_shape_map()
self.emb = reader.get_tensor("bert/embeddings/word_embeddings")
self.vocab = open(vocab_file,'r', encoding='utf-8').read().split("")
print('embeddings size: %s' % str(self.emb.shape))
print('词表大小：%d' % len(self.vocab))
# 兼容不同格式
self.fmt=fmt
# 取出指定字符的embdding,返回向量
def get_embdding (self, char):
if char in self.vocab:
index = self.vocab.index(char)
return self.emb[index,:] else:
return None
# 根据字符串提取向量并保存到文件
def export (self, txt_all, out_file=''):
# 过滤重复，形成字典
txt_lst = sorted(list(set(txt_all)))
资源评论
