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])