def body(self, features, mode):
"""Body of the model, aka Bert
Arguments:
features {dict} -- feature dict,
keys: input_ids, input_mask, segment_ids
mode {mode} -- mode
Returns:
dict -- features extracted from bert.
keys: 'seq', 'pooled', 'all', 'embed'
seq:
tensor, [batch_size, seq_length, hidden_size]
pooled:
tensor, [batch_size, hidden_size]
all:
list of tensor, num_hidden_layers * [batch_size, seq_length, hidden_size]
embed:
tensor, [batch_size, seq_length, hidden_size]
"""
config = self.config
input_ids = features["input_ids"]
input_mask = features["input_mask"]
segment_ids = features["segment_ids"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
model = BertModel(config=config.bert_config,
is_training=is_training,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=config.use_one_hot_embeddings)
feature_dict = {}
for logit_type in ['seq', 'pooled', 'all', 'embed', 'embed_table']:
if logit_type == 'seq':
# tensor, [batch_size, seq_length, hidden_size]
feature_dict[logit_type] = model.get_sequence_output()
elif logit_type == 'pooled':
# tensor, [batch_size, hidden_size]
feature_dict[logit_type] = model.get_pooled_output()
elif logit_type == 'all':
# list, num_hidden_layers * [batch_size, seq_length, hidden_size]
feature_dict[logit_type] = model.get_all_encoder_layers()
elif logit_type == 'embed':
# for res connection
feature_dict[logit_type] = model.get_embedding_output()
elif logit_type == 'embed_table':
feature_dict[logit_type] = model.get_embedding_table()
return feature_dict
class BertEncoder(object):
def __init__(self,
config,
is_training,
input_ids,
input_mask=None,
token_type_ids=None):
self.model = BertModel(config=config,
is_training=is_training,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=token_type_ids)
self.embeddings_table = self.model.get_embedding_table()
def encode(self):
#encoded is => sequence_output` shape = [batch_size, seq_length, hidden_size].
output = self.model.get_sequence_output()
states = ()
for layer in self.model.get_all_encoder_layers():
states += (tf.reduce_mean(layer, axis=1), )
return output, states,