bert得到向量

modeling,tokenization为bert在github上的代码，链接: https://github.com/google-research/bert.

chinese_L-12_H-768_A-12是中文语料训练的模型，链接: https://storage.googleapis.com/bert_models/2018_11_03/chinese_L-12_H-768_A-12.zip.

from bert_demo import modeling
from bert_demo import tokenization
import numpy as np
import tensorflow as tf

class bert_vec():
def __init__(self):
# graph
self.input_ids = tf.placeholder(tf.int32, shape=[None, None], name='input_ids')
self.input_mask = tf.placeholder(tf.int32, shape=[None, None], name='input_masks')
self.segment_ids = tf.placeholder(tf.int32, shape=[None, None], name='segment_ids')

bert_config = modeling.BertConfig.from_json_file('chinese_L-12_H-768_A-12/bert_config.json')
# 初始化BERT
self.model = modeling.BertModel(
config=bert_config,
is_training=False,
input_ids=self.input_ids,
input_mask=self.input_mask,
token_type_ids=self.segment_ids,
use_one_hot_embeddings=False
)
# bert模型地址
init_checkpoint = "chinese_L-12_H-768_A-12/bert_model.ckpt"
# 模型的训练参数
tvars = tf.trainable_variables()
# 加载模型
(assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars,init_checkpoint)
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())

def get_embedding(self,char_lists,mask_lists,seg_lists):
# 获取最后一层和倒数第二层
encoder_last_layer = self.model.get_sequence_output()
# encoder_last2_layer = model.all_encoder_layers[-2]

feed_data = {self.input_ids: np.asarray(char_lists), self.input_mask: np.asarray(mask_lists), self.segment_ids: np.asarray(seg_lists)}
embedding = self.sess.run(encoder_last_layer, feed_dict=feed_data)
return embedding

if __name__ == '__main__':

#数据处理
string = '设置一个随机种子'
char_list = ['[CLS]'] + list(string) +['[SEP]']
#不做masked处理
mask_list = [1] * (len(string)+2)
#不做分词处理
seg_list = [0] * (len(string)+2)

# 根据bert的词表做一个char_to_id的操作
# 未登录词会报错,更改报错代码使未登录词时为'[UNK]'
# 也可以自己实现
token = tokenization.FullTokenizer(vocab_file='chinese_L-12_H-768_A-12/vocab.txt')
char_list = token.convert_tokens_to_ids(char_list)
bertVec = bert_vec()

#得到bert的embedding
embedding = bertVec.get_embedding([char_list], [mask_list], [seg_list])