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,