Keras 實現對bert編碼的封裝實現

bert編碼突襲了NLP各大榜單，爲了讓大家更方便的在自己的數據集合上測試bert的威力，在此分享一下，keras版本的nert編碼層。直接上代碼了。

class b_embeding_layer_b(keras.layers.Layer):
    """自定義層"""

    def __init__(self, max_seq_len=50,model_dir = r"F:\glove.6B\chinese_L-12_H-768_A-12",mode="word", **kwargs):
        self.max_seq_len = max_seq_len
        self.model_dir = model_dir
        self.mode = mode

        super(b_embeding_layer_b, self).__init__(**kwargs)


    def build(self, input_shape):
        #構建訓練參數
        pass

    def call(self, x, mask=None):
        base_location = 'F:/glove.6B/chinese_L-12_H-768_A-12/'
        bert_config = BertConfig.from_json_file(base_location + 'bert_config.json')
        init_checkpoint = base_location + 'bert_model.ckpt'
        model = BertModel(
            config=bert_config,
            is_training=True,
            input_ids=tf.cast(x[0], tf.int32),
            input_mask=tf.cast(x[1], tf.int32),
            token_type_ids=tf.cast(x[2], tf.int32),
            use_one_hot_embeddings=False)
        tvars = tf.trainable_variables()
        scaffold_fn = None
        (assignment_map, _) = get_assignment_map_from_checkpoint(
            tvars, init_checkpoint)
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
        # 獲取對應的embedding 輸入數據[batch_size, seq_length, embedding_size]
        embedding = None
        if self.mode == "sent":  #如此判斷會有問題 --- 從jsom加載時不會有初始化時候的傳值，只有默認值
            embedding = model.get_pooled_output()
        if self.mode == "word":
            embedding = model.get_sequence_output()

        return embedding

    def compute_output_shape(self, input_shape):
        if self.mode == "word":
            return input_shape[0][0], self.max_seq_len,768
        if self.mode == "sent":
            return input_shape[0][0],  768

    def get_config(self):
        config = {'mode': self.mode,
                  "max_seq_len":self.max_seq_len
                  }
        base_config = super(b_embeding_layer_b, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

調用方式，在前面分享的NER模型中有相應內容，在這裏不在贅述了。