# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """CoKE model.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import six import json import logging import numpy as np import paddle.fluid as fluid from model.transformer_encoder import encoder, pre_process_layer logging.basicConfig( format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO) logger = logging.getLogger(__name__) class CoKEModel(object): def __init__(self, src_ids, position_ids, input_mask, config, soft_label=0.9, weight_sharing=True, use_fp16=False): self._emb_size = config['hidden_size'] self._n_layer = config['num_hidden_layers'] self._n_head = config['num_attention_heads'] self._voc_size = config['vocab_size'] self._n_relation = config['num_relations'] self._max_position_seq_len = config['max_position_embeddings'] self._hidden_act = config['hidden_act'] self._prepostprocess_dropout = config['hidden_dropout_prob'] self._attention_dropout = config['attention_probs_dropout_prob'] self._intermediate_size = config['intermediate_size'] self._soft_label = soft_label self._weight_sharing = weight_sharing self._word_emb_name = "word_embedding" self._pos_emb_name = "pos_embedding" self._dtype = "float16" if use_fp16 else "float32" # Initialize all weigths by truncated normal initializer, and all biases # will be initialized by constant zero by default. self._param_initializer = fluid.initializer.TruncatedNormal( scale=config['initializer_range']) self._build_model(src_ids, position_ids, input_mask) def _build_model(self, src_ids, position_ids, input_mask): # padding id in vocabulary must be set to 0 emb_out = fluid.layers.embedding( input=src_ids, size=[self._voc_size, self._emb_size], dtype=self._dtype, param_attr=fluid.ParamAttr( name=self._word_emb_name, initializer=self._param_initializer), is_sparse=False) position_emb_out = fluid.layers.embedding( input=position_ids, size=[self._max_position_seq_len, self._emb_size], dtype=self._dtype, param_attr=fluid.ParamAttr( name=self._pos_emb_name, initializer=self._param_initializer)) emb_out = emb_out + position_emb_out emb_out = pre_process_layer( emb_out, 'nd', self._prepostprocess_dropout, name='pre_encoder') if self._dtype == "float16": input_mask = fluid.layers.cast(x=input_mask, dtype=self._dtype) self_attn_mask = fluid.layers.matmul( x=input_mask, y=input_mask, transpose_y=True) self_attn_mask = fluid.layers.scale( x=self_attn_mask, scale=10000.0, bias=-1.0, bias_after_scale=False) n_head_self_attn_mask = fluid.layers.stack( x=[self_attn_mask] * self._n_head, axis=1) n_head_self_attn_mask.stop_gradient = True self._enc_out = encoder( enc_input=emb_out, attn_bias=n_head_self_attn_mask, n_layer=self._n_layer, n_head=self._n_head, d_key=self._emb_size // self._n_head, d_value=self._emb_size // self._n_head, d_model=self._emb_size, d_inner_hid=self._intermediate_size, prepostprocess_dropout=self._prepostprocess_dropout, attention_dropout=self._attention_dropout, relu_dropout=0, hidden_act=self._hidden_act, preprocess_cmd="", postprocess_cmd="dan", param_initializer=self._param_initializer, name='encoder') #def get_sequence_output(self): # return self._enc_out def get_pretraining_output(self, mask_label, mask_pos): """Get the loss & fc_out for training""" mask_pos = fluid.layers.cast(x=mask_pos, dtype='int32') reshaped_emb_out = fluid.layers.reshape( x=self._enc_out, shape=[-1, self._emb_size]) # extract masked tokens' feature mask_feat = fluid.layers.gather(input=reshaped_emb_out, index=mask_pos) # transform: fc mask_trans_feat = fluid.layers.fc( input=mask_feat, size=self._emb_size, act=self._hidden_act, param_attr=fluid.ParamAttr( name='mask_lm_trans_fc.w_0', initializer=self._param_initializer), bias_attr=fluid.ParamAttr(name='mask_lm_trans_fc.b_0')) # transform: layer norm mask_trans_feat = pre_process_layer( mask_trans_feat, 'n', name='mask_lm_trans') mask_lm_out_bias_attr = fluid.ParamAttr( name="mask_lm_out_fc.b_0", initializer=fluid.initializer.Constant(value=0.0)) if self._weight_sharing: fc_out = fluid.layers.matmul( x=mask_trans_feat, y=fluid.default_main_program().global_block().var( self._word_emb_name), transpose_y=True) fc_out += fluid.layers.create_parameter( shape=[self._voc_size], dtype=self._dtype, attr=mask_lm_out_bias_attr, is_bias=True) else: fc_out = fluid.layers.fc(input=mask_trans_feat, size=self._voc_size, param_attr=fluid.ParamAttr( name="mask_lm_out_fc.w_0", initializer=self._param_initializer), bias_attr=mask_lm_out_bias_attr) #generate soft labels for loss cross entropy loss one_hot_labels = fluid.layers.one_hot( input=mask_label, depth=self._voc_size) entity_indicator = fluid.layers.fill_constant_batch_size_like( input=mask_label, shape=[-1, (self._voc_size - self._n_relation)], dtype='int64', value=0) relation_indicator = fluid.layers.fill_constant_batch_size_like( input=mask_label, shape=[-1, self._n_relation], dtype='int64', value=1) is_relation = fluid.layers.concat( input=[entity_indicator, relation_indicator], axis=-1) soft_labels = one_hot_labels * self._soft_label \ + (1.0 - one_hot_labels - is_relation) \ * ((1.0 - self._soft_label) / (self._voc_size - 1 - self._n_relation)) soft_labels.stop_gradient = True mask_lm_loss = fluid.layers.softmax_with_cross_entropy( logits=fc_out, label=soft_labels, soft_label=True) mean_mask_lm_loss = fluid.layers.mean(mask_lm_loss) return mean_mask_lm_loss, fc_out