import json
from tensorflow import keras
import numpy as np
import tensorflow as tf
from .bert import get_model
__all__ = [
'build_model_from_config',
'load_model_weights_from_checkpoint',
'load_trained_model_from_checkpoint',
]
def checkpoint_loader(checkpoint_file):
def _loader(name):
return tf.train.load_variable(checkpoint_file, name)
return _loader
def build_model_from_config(config_file,
training=False,
trainable=None,
seq_len=None):
"""Build the model from config file.
:param config_file: The path to the JSON configuration file.
:param training: If training, the whole model will be returned.
:param trainable: Whether the model is trainable.
:param seq_len: If it is not None and it is shorter than the value in the config file, the weights in
position embeddings will be sliced to fit the new length.
:return: model and config
"""
with open(config_file, 'r') as reader:
config = json.loads(reader.read())
if seq_len is not None:
config['max_position_embeddings'] = min(
seq_len, config['max_position_embeddings'])
if trainable is None:
trainable = training
model = get_model(
token_num=config['vocab_size'],
pos_num=config['max_position_embeddings'],
seq_len=config['max_position_embeddings'],
embed_dim=config['hidden_size'],
transformer_num=config['num_hidden_layers'],
head_num=config['num_attention_heads'],
feed_forward_dim=config['intermediate_size'],
training=training,
trainable=trainable,
)
model.build(input_shape=[(None, None), (None, None), (None, None)])
return model, config
def load_model_weights_from_checkpoint(model,
config,
checkpoint_file,
training=False):
"""Load trained official model from checkpoint.
:param model: Built keras model.
:param config: Loaded configuration file.
:param checkpoint_file: The path to the checkpoint files, should end with '.ckpt'.
:param training: If training, the whole model will be returned.
Otherwise, the MLM and NSP parts will be ignored.
"""
loader = checkpoint_loader(checkpoint_file)
model.get_layer(name='Embedding-Token').set_weights([
loader('bert/embeddings/word_embeddings'),
])
model.get_layer(name='Embedding-Position').set_weights([
loader(
'bert/embeddings/position_embeddings')[:config['max_position_embeddings'], :],
])
model.get_layer(name='Embedding-Segment').set_weights([
loader('bert/embeddings/token_type_embeddings'),
])
model.get_layer(name='Embedding-Norm').set_weights([
loader('bert/embeddings/LayerNorm/gamma'),
loader('bert/embeddings/LayerNorm/beta'),
])
for i in range(config['num_hidden_layers']):
model.get_layer(name='Encoder-%d-MultiHeadSelfAttention' % (i + 1)).set_weights([
loader('bert/encoder/layer_%d/attention/self/query/kernel' % i),
loader('bert/encoder/layer_%d/attention/self/query/bias' % i),
loader('bert/encoder/layer_%d/attention/self/key/kernel' % i),
loader('bert/encoder/layer_%d/attention/self/key/bias' % i),
loader('bert/encoder/layer_%d/attention/self/value/kernel' % i),
loader('bert/encoder/layer_%d/attention/self/value/bias' % i),
loader('bert/encoder/layer_%d/attention/output/dense/kernel' % i),
loader('bert/encoder/layer_%d/attention/output/dense/bias' % i),
])
model.get_layer(name='Encoder-%d-MultiHeadSelfAttention-Norm' % (i + 1)).set_weights([
loader('bert/encoder/layer_%d/attention/output/LayerNorm/gamma' % i),
loader('bert/encoder/layer_%d/attention/output/LayerNorm/beta' % i),
])
model.get_layer(name='Encoder-%d-MultiHeadSelfAttention-Norm' % (i + 1)).set_weights([
loader('bert/encoder/layer_%d/attention/output/LayerNorm/gamma' % i),
loader('bert/encoder/layer_%d/attention/output/LayerNorm/beta' % i),
])
model.get_layer(name='Encoder-%d-FeedForward' % (i + 1)).set_weights([
loader('bert/encoder/layer_%d/intermediate/dense/kernel' % i),
loader('bert/encoder/layer_%d/intermediate/dense/bias' % i),
loader('bert/encoder/layer_%d/output/dense/kernel' % i),
loader('bert/encoder/layer_%d/output/dense/bias' % i),
])
model.get_layer(name='Encoder-%d-FeedForward-Norm' % (i + 1)).set_weights([
loader('bert/encoder/layer_%d/output/LayerNorm/gamma' % i),
loader('bert/encoder/layer_%d/output/LayerNorm/beta' % i),
])
if training:
model.get_layer(name='MLM-Dense').set_weights([
loader('cls/predictions/transform/dense/kernel'),
loader('cls/predictions/transform/dense/bias'),
])
model.get_layer(name='MLM-Norm').set_weights([
loader('cls/predictions/transform/LayerNorm/gamma'),
loader('cls/predictions/transform/LayerNorm/beta'),
])
model.get_layer(name='MLM-Sim').set_weights([
loader('cls/predictions/output_bias'),
])
model.get_layer(name='NSP-Dense').set_weights([
loader('bert/pooler/dense/kernel'),
loader('bert/pooler/dense/bias'),
])
model.get_layer(name='NSP').set_weights([
np.transpose(loader('cls/seq_relationship/output_weights')),
loader('cls/seq_relationship/output_bias'),
])
def load_trained_model_from_checkpoint(config_file,
checkpoint_file,
training=False,
trainable=None,
seq_len=None):
"""Load trained official model from checkpoint.
:param config_file: The path to the JSON configuration file.
:param checkpoint_file: The path to the checkpoint files, should end with '.ckpt'.
:param training: If training, the whole model will be returned.
Otherwise, the MLM and NSP parts will be ignored.
:param trainable: Whether the model is trainable. The default value is the same with `training`.
:param seq_len: If it is not None and it is shorter than the value in the config file, the weights in
position embeddings will be sliced to fit the new length.
:return: model
"""
model, config = build_model_from_config(
config_file, training=training, trainable=trainable, seq_len=seq_len)
load_model_weights_from_checkpoint(
model, config, checkpoint_file, training=training)
return model
