import json
import tensorflow as tf
from tensorflow import keras
import tensorflow.keras.backend as K
from keras_bert.keras_pos_embd import PositionEmbedding
from keras_bert.layers import get_inputs, get_embedding, TokenEmbedding, EmbeddingSimilarity, Masked, Extract
from keras_bert.keras_layer_normalization import LayerNormalization
from keras_bert.keras_multi_head import MultiHeadAttention
from keras_bert.keras_position_wise_feed_forward import FeedForward
def gelu(x):
return 0.5 * x * (1.0 + tf.math.erf(x / tf.sqrt(2.0)))
class Bert(keras.Model):
def __init__(
self,
token_num,
pos_num=512,
seq_len=512,
embed_dim=768,
transformer_num=12,
head_num=12,
feed_forward_dim=3072,
dropout_rate=0.1,
attention_activation=None,
feed_forward_activation=gelu,
custom_layers=None,
training=True,
trainable=None,
lr=1e-4,
name='Bert'):
super().__init__(name=name)
self.token_num = token_num
self.pos_num = pos_num
self.seq_len = seq_len
self.embed_dim = embed_dim
self.transformer_num = transformer_num
self.head_num = head_num
self.feed_forward_dim = feed_forward_dim
self.dropout_rate = dropout_rate
self.attention_activation = attention_activation
self.feed_forward_activation = feed_forward_activation
self.custom_layers = custom_layers
self.training = training
self.trainable = trainable
self.lr = lr
# build layers
# embedding
self.token_embedding_layer = TokenEmbedding(
input_dim=token_num,
output_dim=embed_dim,
mask_zero=True,
trainable=trainable,
name='Embedding-Token',
)
self.segment_embedding_layer = keras.layers.Embedding(
input_dim=2,
output_dim=embed_dim,
trainable=trainable,
name='Embedding-Segment',
)
self.position_embedding_layer = PositionEmbedding(
input_dim=pos_num,
output_dim=embed_dim,
mode=PositionEmbedding.MODE_ADD,
trainable=trainable,
name='Embedding-Position',
)
self.embedding_layer_norm = LayerNormalization(
trainable=trainable,
name='Embedding-Norm',
)
self.encoder_multihead_layers = []
self.encoder_ffn_layers = []
self.encoder_attention_norm = []
self.encoder_ffn_norm = []
# attention layers
for i in range(transformer_num):
base_name = 'Encoder-%d' % (i + 1)
attention_name = '%s-MultiHeadSelfAttention' % base_name
feed_forward_name = '%s-FeedForward' % base_name
self.encoder_multihead_layers.append(MultiHeadAttention(
head_num=head_num,
activation=attention_activation,
history_only=False,
trainable=trainable,
name=attention_name,
))
self.encoder_ffn_layers.append(FeedForward(
units=feed_forward_dim,
activation=feed_forward_activation,
trainable=trainable,
name=feed_forward_name,
))
self.encoder_attention_norm.append(LayerNormalization(
trainable=trainable,
name='%s-Norm' % attention_name,
))
self.encoder_ffn_norm.append(LayerNormalization(
trainable=trainable,
name='%s-Norm' % feed_forward_name,
))
def call(self, inputs):
embeddings = [
self.token_embedding_layer(inputs[0]),
self.segment_embedding_layer(inputs[1])
]
embeddings[0], embed_weights = embeddings[0]
embed_layer = keras.layers.Add(
name='Embedding-Token-Segment')(embeddings)
embed_layer = self.position_embedding_layer(embed_layer)
if self.dropout_rate > 0.0:
dropout_layer = keras.layers.Dropout(
rate=self.dropout_rate,
name='Embedding-Dropout',
)(embed_layer)
else:
dropout_layer = embed_layer
embedding_output = self.embedding_layer_norm(dropout_layer)
def _wrap_layer(name, input_layer, build_func, norm_layer, dropout_rate=0.0, trainable=True):
"""Wrap layers with residual, normalization and dropout.
:param name: Prefix of names for internal layers.
:param input_layer: Input layer.
:param build_func: A callable that takes the input tensor and generates the output tensor.
:param dropout_rate: Dropout rate.
:param trainable: Whether the layers are trainable.
:return: Output layer.
"""
build_output = build_func(input_layer)
if dropout_rate > 0.0:
dropout_layer = keras.layers.Dropout(
rate=dropout_rate,
name='%s-Dropout' % name,
)(build_output)
else:
dropout_layer = build_output
if isinstance(input_layer, list):
input_layer = input_layer[0]
add_layer = keras.layers.Add(name='%s-Add' %
name)([input_layer, dropout_layer])
normal_layer = norm_layer(add_layer)
return normal_layer
last_layer = embedding_output
output_tensor_list = [last_layer]
# self attention
for i in range(self.transformer_num):
base_name = 'Encoder-%d' % (i + 1)
attention_name = '%s-MultiHeadSelfAttention' % base_name
feed_forward_name = '%s-FeedForward' % base_name
self_attention_output = _wrap_layer(
name=attention_name,
input_layer=last_layer,
build_func=self.encoder_multihead_layers[i],
norm_layer=self.encoder_attention_norm[i],
dropout_rate=self.dropout_rate,
trainable=self.trainable)
last_layer = _wrap_layer(
name=attention_name,
input_layer=self_attention_output,
build_func=self.encoder_ffn_layers[i],
norm_layer=self.encoder_ffn_norm[i],
dropout_rate=self.dropout_rate,
trainable=self.trainable)
output_tensor_list.append(last_layer)
return output_tensor_list
def build_model_from_config(config_file,
training=False,
trainable=None,
seq_len=None,
build=True):
"""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 = Bert(
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,
)
if build:
model.build(input_shape=[(None, None), (None, None), (None, None)])
return model, config
