def get_encoder(
model_name,
args,
trainable: bool = True,
prefix=None,
):
MODEL_PATH = model_name
if prefix is not None:
MODEL_PATH = os.path.join(prefix, model_name)
if model_name in {'bert-base-uncased', 'NlpHUST/vibert4news-base-cased'}:
encoder = TFBertModel.from_pretrained(MODEL_PATH, trainable=trainable)
elif model_name.find('bigbird') > -1:
encoder = modeling.BertModel({
"attention_probs_dropout_prob": 0.1,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.1,
"hidden_size": 768,
"initializer_range": 0.02,
"intermediate_size": 3072,
"max_position_embeddings": 4096,
"max_encoder_length": args.max_context_length,
"num_attention_heads": 12,
"num_hidden_layers": 12,
"type_vocab_size": 2,
"use_bias": True,
"rescale_embedding": False,
"use_gradient_checkpointing": False,
"scope": "bert",
"attention_type": "block_sparse",
"norm_type": "postnorm",
"block_size": 16,
"num_rand_blocks": 3,
"vocab_size": 50358
})
checkpoint_path = 'gs://bigbird-transformer/pretrain/bigbr_base/model.ckpt-0'
checkpoint_reader = tf.compat.v1.train.NewCheckpointReader(
checkpoint_path)
encoder.set_weights([
checkpoint_reader.get_tensor(v.name[:-2])
for v in encoder.trainable_weights
])
encoder.trainable = True
else:
raise Exception("Model {} not supported".format(model_name))
if not args.use_pooler:
if model_name in {
'bert-base-uncased', 'NlpHUST/vibert4news-base-cased'
}:
encoder.bert.pooler.trainable = False
elif model_name == 'bigbird':
encoder.pooler.trainable = False
return encoder
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
model = modeling.BertModel(bert_config)
masked_lm = MaskedLMLayer(bert_config["hidden_size"],
bert_config["vocab_size"],
model.embeder,
initializer=utils.create_initializer(
bert_config["initializer_range"]),
activation_fn=utils.get_activation(
bert_config["hidden_act"]))
next_sentence = NSPLayer(bert_config["hidden_size"],
initializer=utils.create_initializer(
bert_config["initializer_range"]))
sequence_output, pooled_output = model(
features["input_ids"],
training=is_training,
token_type_ids=features.get("segment_ids"))
masked_lm_loss, masked_lm_log_probs = masked_lm(
sequence_output,
label_ids=features.get("masked_lm_ids"),
label_weights=features.get("masked_lm_weights"),
masked_lm_positions=features.get("masked_lm_positions"))
next_sentence_loss, next_sentence_log_probs = next_sentence(
pooled_output, features.get("next_sentence_labels"))
total_loss = masked_lm_loss
if bert_config["use_nsp"]:
total_loss += next_sentence_loss
tvars = tf.compat.v1.trainable_variables()
utils.log_variables(tvars, bert_config["ckpt_var_list"])
output_spec = None
if mode == tf.estimator.ModeKeys.TRAIN:
learning_rate = optimization.get_linear_warmup_linear_decay_lr(
init_lr=bert_config["learning_rate"],
num_train_steps=bert_config["num_train_steps"],
num_warmup_steps=bert_config["num_warmup_steps"])
optimizer = optimization.get_optimizer(bert_config, learning_rate)
global_step = tf.compat.v1.train.get_global_step()
gradients = optimizer.compute_gradients(total_loss, tvars)
train_op = optimizer.apply_gradients(gradients,
global_step=global_step)
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
host_call=utils.add_scalars_to_summary(
bert_config["output_dir"],
{"learning_rate": learning_rate}))
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(masked_lm_loss_value, masked_lm_log_probs,
masked_lm_ids, masked_lm_weights,
next_sentence_loss_value, next_sentence_log_probs,
next_sentence_labels):
"""Computes the loss and accuracy of the model."""
masked_lm_predictions = tf.argmax(masked_lm_log_probs,
axis=-1,
output_type=tf.int32)
masked_lm_accuracy = tf.compat.v1.metrics.accuracy(
labels=masked_lm_ids,
predictions=masked_lm_predictions,
weights=masked_lm_weights)
masked_lm_mean_loss = tf.compat.v1.metrics.mean(
values=masked_lm_loss_value)
next_sentence_predictions = tf.argmax(next_sentence_log_probs,
axis=-1,
output_type=tf.int32)
next_sentence_accuracy = tf.compat.v1.metrics.accuracy(
labels=next_sentence_labels,
predictions=next_sentence_predictions)
next_sentence_mean_loss = tf.compat.v1.metrics.mean(
values=next_sentence_loss_value)
return {
"masked_lm_accuracy": masked_lm_accuracy,
"masked_lm_loss": masked_lm_mean_loss,
"next_sentence_accuracy": next_sentence_accuracy,
"next_sentence_loss": next_sentence_mean_loss,
}
eval_metrics = (metric_fn, [
masked_lm_loss, masked_lm_log_probs, features["masked_lm_ids"],
features["masked_lm_weights"], next_sentence_loss,
next_sentence_log_probs, features["next_sentence_labels"]
])
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode, loss=total_loss, eval_metrics=eval_metrics)
else:
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
predictions={
"log-probabilities": masked_lm_log_probs,
"seq-embeddings": sequence_output
})
return output_spec
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
if isinstance(features, dict):
if not labels and "labels" in features:
labels = features["labels"]
features = features["input_ids"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
model = modeling.BertModel(bert_config)
headl = ClassifierLossLayer(
bert_config["num_labels"], bert_config["hidden_dropout_prob"],
utils.create_initializer(bert_config["initializer_range"]),
name=bert_config["scope"]+"/classifier")
_, pooled_output = model(features, training=is_training)
total_loss, log_probs = headl(pooled_output, labels, is_training)
tvars = tf.compat.v1.trainable_variables()
utils.log_variables(tvars, bert_config["ckpt_var_list"])
output_spec = None
if mode == tf.estimator.ModeKeys.TRAIN:
learning_rate = optimization.get_linear_warmup_linear_decay_lr(
init_lr=bert_config["learning_rate"],
num_train_steps=bert_config["num_train_steps"],
num_warmup_steps=bert_config["num_warmup_steps"])
optimizer = optimization.get_optimizer(bert_config, learning_rate)
global_step = tf.compat.v1.train.get_or_create_global_step()
gradients = optimizer.compute_gradients(total_loss, tvars)
train_op = optimizer.apply_gradients(gradients, global_step=global_step)
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
host_call=utils.add_scalars_to_summary(
bert_config["output_dir"], {"learning_rate": learning_rate}))
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(loss_value, label_ids, log_probs):
loss = tf.compat.v1.metrics.mean(values=loss_value)
predictions = tf.argmax(log_probs, axis=-1, output_type=tf.int32)
accuracy = tf.compat.v1.metrics.accuracy(
labels=label_ids, predictions=predictions)
p1, p1_op = tf.compat.v1.metrics.precision_at_k(
labels=tf.cast(label_ids, tf.int64), predictions=log_probs, k=1)
r1, r1_op = tf.compat.v1.metrics.recall_at_k(
labels=tf.cast(label_ids, tf.int64), predictions=log_probs, k=1)
f11 = tf.math.divide_no_nan(2*p1*r1, p1+r1)
metric_dict = {
"P@1": (p1, p1_op),
"R@1": (r1, r1_op),
"f1@1": (f11, tf.no_op()),
"classification_accuracy": accuracy,
"classification_loss": loss,
}
return metric_dict
eval_metrics = (metric_fn,
[tf.expand_dims(total_loss, 0), labels, log_probs])
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics)
else:
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
predictions={"log-probabilities": log_probs})
return output_spec
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
# BigBird model 정의
model = modeling.BertModel(bert_config,
features["input_ids"],
training=is_training,
token_type_ids=features.get("segment_ids"))
# attention feature와 cls token에 대한 pooling feature를 가져옴
sequence_output, pooled_output = model.get_output_feature()
masked_lm = MaskedLMLayer( # masked language output 계산 모델 정의
bert_config["hidden_size"],
bert_config["vocab_size"],
model.embeder,
input_tensor=sequence_output,
label_ids=features.get("masked_lm_ids"),
label_weights=features.get("masked_lm_weights"),
masked_lm_positions=features.get("masked_lm_positions"),
initializer=utils.create_initializer(
bert_config["initializer_range"]),
activation_fn=utils.get_activation(bert_config["hidden_act"]))
masked_lm_loss, masked_lm_log_probs = masked_lm.get_mlm_loss()
total_loss = masked_lm_loss
tvars = tf.compat.v1.trainable_variables()
utils.LogVariable(tvars, bert_config["ckpt_var_list"])
output_spec = None
if mode == tf.estimator.ModeKeys.TRAIN:
# optimize 계산
opt_model = optimization.LinearWarmupLinearDecay( # optimize model 불러옴
init_lr=bert_config["learning_rate"],
num_train_steps=bert_config["num_train_steps"],
num_warmup_steps=bert_config["num_warmup_steps"])
learning_rate = opt_model.get_learning_rate() # laernin rate 가져옴
optimizer = optimization.Optimizer(bert_config, learning_rate)
optimizer = optimizer.get_optimizer()
global_step = tf.compat.v1.train.get_global_step()
gradients = optimizer.compute_gradients(total_loss, tvars)
train_op = optimizer.apply_gradients(gradients,
global_step=global_step)
logging_hook = [
tf.compat.v1.train.LoggingTensorHook(
{"loss is -> ": total_loss}, every_n_iter=256),
tf.compat.v1.train.LoggingTensorHook(
{"global step -> ": global_step}, every_n_iter=256),
tf.compat.v1.train.LoggingTensorHook(
{"learning rate -> ": learning_rate}, every_n_iter=256)
]
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
training_hooks=logging_hook,
host_call=utils.add_scalars_to_summary(
bert_config["output_dir"],
{"learning_rate": learning_rate}))
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(masked_lm_loss_value, masked_lm_log_probs,
masked_lm_ids, masked_lm_weights):
masked_lm_predictions = tf.argmax(masked_lm_log_probs,
axis=-1,
output_type=tf.int32)
masked_lm_accuracy = tf.compat.v1.metrics.accuracy(
labels=masked_lm_ids,
predictions=masked_lm_predictions,
weights=masked_lm_weights)
masked_lm_mean_loss = tf.compat.v1.metrics.mean(
values=masked_lm_loss_value)
return {
"masked_lm_accuracy": masked_lm_accuracy,
"masked_lm_loss": masked_lm_mean_loss,
}
eval_metrics = (metric_fn, [
masked_lm_loss, masked_lm_log_probs, features["masked_lm_ids"],
features["masked_lm_weights"]
])
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode, loss=total_loss, eval_metrics=eval_metrics)
else:
output_spec = tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=mode,
predictions={
"log-probabilities": masked_lm_log_probs,
"seq-embeddings": sequence_output
})
return output_spec