def maskedBert(
input_ids_blob,
input_mask_blob,
token_type_ids_blob,
masked_lm_positions_blob,
# masked_lm_positions_blob,
# masked_lm_ids_blob,
vocab_size,
seq_length=512,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
hidden_act="gelu",
hidden_dropout_prob=0.1,
attention_probs_dropout_prob=0.1,
max_position_embeddings=512,
type_vocab_size=16,
# max_predictions_per_seq=20,
initializer_range=0.02,
):
backbone = bert_util.BertBackbone(
input_ids_blob=input_ids_blob,
input_mask_blob=input_mask_blob,
token_type_ids_blob=token_type_ids_blob,
vocab_size=vocab_size,
seq_length=seq_length,
hidden_size=hidden_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
hidden_act=hidden_act,
hidden_dropout_prob=hidden_dropout_prob,
attention_probs_dropout_prob=attention_probs_dropout_prob,
max_position_embeddings=max_position_embeddings,
type_vocab_size=type_vocab_size,
initializer_range=initializer_range,
)
predictions = _AddMaskedLanguageModel(
input_blob=backbone.sequence_output(),
output_weights_blob=backbone.embedding_table(),
positions_blob=masked_lm_positions_blob,
seq_length=seq_length,
hidden_size=hidden_size,
vocab_size=vocab_size,
hidden_act=bert_util.GetActivation(hidden_act),
initializer_range=initializer_range,
)
pooled_output = PooledOutput(backbone.sequence_output(), hidden_size,
initializer_range)
return predictions
def PreTrain(
input_ids_blob,
input_mask_blob,
token_type_ids_blob,
masked_lm_positions_blob,
masked_lm_ids_blob,
masked_lm_weights_blob,
next_sentence_label_blob,
vocab_size,
seq_length=512,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
hidden_act="gelu",
hidden_dropout_prob=0.1,
attention_probs_dropout_prob=0.1,
max_position_embeddings=512,
type_vocab_size=16,
max_predictions_per_seq=20,
initializer_range=0.02,
use_fp16=False,
):
backbone = bert_util.BertBackbone(
input_ids_blob=input_ids_blob,
input_mask_blob=input_mask_blob,
token_type_ids_blob=token_type_ids_blob,
vocab_size=vocab_size,
seq_length=seq_length,
hidden_size=hidden_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
hidden_act=hidden_act,
hidden_dropout_prob=hidden_dropout_prob,
attention_probs_dropout_prob=attention_probs_dropout_prob,
max_position_embeddings=max_position_embeddings,
type_vocab_size=type_vocab_size,
initializer_range=initializer_range,
)
(lm_loss, _, _) = _AddMaskedLanguageModelLoss(
input_blob=backbone.sequence_output(),
output_weights_blob=backbone.embedding_table(),
positions_blob=masked_lm_positions_blob,
label_id_blob=masked_lm_ids_blob,
label_weight_blob=masked_lm_weights_blob,
seq_length=seq_length,
hidden_size=hidden_size,
vocab_size=vocab_size,
max_predictions_per_seq=max_predictions_per_seq,
hidden_act=bert_util.GetActivation(hidden_act),
initializer_range=initializer_range,
)
pooled_output = PooledOutput(backbone.sequence_output(), hidden_size,
initializer_range)
(ns_loss, _, _) = _AddNextSentenceOutput(
input_blob=pooled_output,
label_blob=next_sentence_label_blob,
hidden_size=hidden_size,
initializer_range=initializer_range,
)
with flow.scope.namespace("cls-loss"):
lm_loss = flow.math.reduce_mean(lm_loss)
ns_loss = flow.math.reduce_mean(ns_loss)
total_loss = lm_loss + ns_loss
return total_loss, lm_loss, ns_loss