def test_google_weights(self):
albert_model_name = "albert_base"
albert_dir = bert.fetch_tfhub_albert_model(albert_model_name,
".models")
albert_params = bert.albert_params(albert_model_name)
l_bert = bert.BertModelLayer.from_params(albert_params, name="albert")
l_input_ids = keras.layers.Input(shape=(128, ),
dtype='int32',
name="input_ids")
l_token_type_ids = keras.layers.Input(shape=(128, ),
dtype='int32',
name="token_type_ids")
output = l_bert([l_input_ids, l_token_type_ids])
output = keras.layers.Lambda(lambda x: x[:, 0, :])(output)
output = keras.layers.Dense(2)(output)
model = keras.Model(inputs=[l_input_ids, l_token_type_ids],
outputs=output)
model.build(input_shape=(None, 128))
model.compile(
optimizer=keras.optimizers.Adam(),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")])
for weight in l_bert.weights:
print(weight.name)
bert.load_albert_weights(l_bert, albert_dir)
model.summary()
def test_albert_load_base_google_weights(self): # for coverage mainly
albert_model_name = "albert_base"
albert_dir = bert.fetch_tfhub_albert_model(albert_model_name,
".models")
model_params = bert.albert_params(albert_model_name)
l_bert = bert.BertModelLayer.from_params(model_params, name="albert")
model = keras.models.Sequential([
keras.layers.InputLayer(input_shape=(8, ),
dtype=tf.int32,
name="input_ids"),
l_bert,
keras.layers.Lambda(lambda x: x[:, 0, :]),
keras.layers.Dense(2),
])
model.build(input_shape=(None, 8))
model.compile(
optimizer=keras.optimizers.Adam(),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")])
bert.load_albert_weights(l_bert, albert_dir)
model.summary()
def test_albert_params(self):
albert_model_name = "albert_base"
albert_dir = bert.fetch_tfhub_albert_model(albert_model_name,
".models")
dir_params = bert.albert_params(albert_dir)
dir_params.attention_dropout = 0.1 # diff between README and assets/albert_config.json
dir_params.hidden_dropout = 0.1
name_params = bert.albert_params(albert_model_name)
self.assertEqual(name_params, dir_params)
# coverage
model_params = dir_params
model_params.vocab_size = model_params.vocab_size + 2
model_params.adapter_size = 1
l_bert = bert.BertModelLayer.from_params(model_params, name="albert")
l_bert(tf.zeros((1, 128)))
bert.load_albert_weights(l_bert, albert_dir)
def test_albert_google_weights(self):
albert_model_name = "albert_base"
albert_dir = bert.fetch_tfhub_albert_model(albert_model_name, ".models")
albert_params = bert.albert_params(albert_model_name)
model, l_bert = self.build_model(albert_params)
skipped_weight_value_tuples = bert.load_albert_weights(l_bert, albert_dir)
self.assertEqual(0, len(skipped_weight_value_tuples))
model.summary()
def test_albert_google_weights_non_tfhub(self):
albert_model_name = "albert_base_v2"
albert_dir = bert.fetch_google_albert_model(albert_model_name, ".models")
model_ckpt = os.path.join(albert_dir, "model.ckpt-best")
albert_params = bert.albert_params(albert_dir)
model, l_bert = self.build_model(albert_params)
skipped_weight_value_tuples = bert.load_albert_weights(l_bert, model_ckpt)
self.assertEqual(0, len(skipped_weight_value_tuples))
model.summary()
def create_model(
model_dir, model_type, max_seq_len, n_classes, load_pretrained_weights=True, summary=False,
):
"""Creates keras model with pretrained BERT/ALBERT layer.
Args:
model_dir: String. Path to model.
model_type: String. Expects either "albert" or "bert"
max_seq_len: Int. Maximum length of a classificaton example.
n_classes: Int. Number of training classes.
load_pretrained_weights: Boolean. Load pretrained model weights.
summary: Boolean. Print model summary.
Returns:
Keras model
"""
if model_type == "albert":
model_ckpt = os.path.join(model_dir, "model.ckpt-best")
model_params = bert.albert_params(model_dir)
elif model_type == "bert":
model_ckpt = os.path.join(model_dir, "bert_model.ckpt")
model_params = bert.params_from_pretrained_ckpt(model_dir)
layer_bert = bert.BertModelLayer.from_params(model_params, name=model_type)
input_ids = keras.layers.Input(shape=(max_seq_len,), dtype="int32", name="input_ids")
output = layer_bert(input_ids)
cls_out = keras.layers.Lambda(lambda seq: seq[:, 0, :])(output)
cls_out = keras.layers.Dropout(0.5)(cls_out)
logits = keras.layers.Dense(units=model_params["hidden_size"], activation="relu")(cls_out)
logits = keras.layers.Dropout(0.5)(logits)
logits = keras.layers.Dense(units=n_classes, activation="softmax")(logits)
model = keras.Model(inputs=input_ids, outputs=logits)
model.build(input_shape=(None, max_seq_len))
if load_pretrained_weights:
if model_type == "albert":
bert.load_albert_weights(layer_bert, model_ckpt)
elif model_type == "bert":
bert.load_bert_weights(layer_bert, model_ckpt)
model.compile(
optimizer=keras.optimizers.Adam(),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
)
if summary:
model.summary()
return model
def fetch_bert_layer():
"""
Function to return ALBERT layer and weights
Returns:
l_bert (bert.model.BertModelLayer): BERT layer
model_ckpt (str): path to best model checkpoint
"""
model_name = "albert_base_v2"
model_dir = bert.fetch_google_albert_model(model_name, ".models")
model_ckpt = os.path.join(model_dir, "model.ckpt-best")
model_params = bert.albert_params(model_dir)
l_bert = bert.BertModelLayer.from_params(model_params, name="albert")
return l_bert, model_ckpt
def Albert_model(max_seq_len):
model_name = "albert_large"
model_dir = bert.fetch_tfhub_albert_model(model_name, ".models")
model_params = bert.albert_params(model_name)
model_params.shared_layer = True
model_params.embedding_size = 1024
l_bert = bert.BertModelLayer.from_params(model_params, name="albert")
l_input_ids = keras.layers.Input(shape=(max_seq_len, ), dtype='int32')
# using the default token_type/segment id 0
output = l_bert(l_input_ids) # output: [batch_size, max_seq_len, hidden_size]
output = keras.layers.GlobalAveragePooling1D()(output)
model = keras.Model(inputs=l_input_ids, outputs=output)
model.build(input_shape=(None, max_seq_len))
# use in a Keras Model here, and call model.build()
bert.load_albert_weights(l_bert, model_dir) # should be called after model.build()
return model, model_dir
def load_bert_model(name_model, max_seq_len, trainable=False):
"""
models name supported, same as tf-2.0-bert
"""
model_name = name_model
model_dir = bert.fetch_tfhub_albert_model(model_name, ".models")
model_params = bert.albert_params(model_name)
l_bert = bert.BertModelLayer.from_params(model_params, name=name_model)
l_input_ids = tf.keras.layers.Input(shape=(max_seq_len, ), dtype='int32')
output = l_bert(
l_input_ids) # output: [batch_size, max_seq_len, hidden_size]
model = tf.keras.Model(inputs=l_input_ids, outputs=output)
model.build(input_shape=(None, max_seq_len))
# load google albert original weights after the build
bert.load_albert_weights(l_bert, model_dir)
model.trainable = trainable
return model
if isinstance(root_layer, keras.layers.Layer):
yield root_layer
for layer in root_layer._layers:
yield from flatten_layers(layer)
def freeze_layers(root_layer, exclude=None):
exclude = [] if exclude is None else exclude
root_layer.trainable = False
for layer in flatten_layers(root_layer):
if layer.name in exclude:
layer.trainable = True
# Create ALBERT layer
model_params = bert.albert_params(MODEL_NAME)
model_params.adapter_size = ADAPTER_SIZE
bert_layer = bert.BertModelLayer.from_params(model_params, name="albert")
# Define model architecture
input_ids = keras.layers.Input(shape=(MAX_SEQ_LEN,), dtype='int32')
# NOTE: Following line not required if using default token_type/segment id 0
# token_type_ids = keras.layers.Input(shape=(MAX_SEQ_LEN,), dtype='int32')
output = bert_layer(input_ids) # output:[batch_size, MAX_SEQ_LEN, hidden_size]
# NOTE: The following is an alternative for classification taken from
# https://github.com/kpe/bert-for-tf2/blob/master/examples/gpu_movie_reviews.ipynb
# The Lambda layer just takes one output from the sequence
cls_out = keras.layers.Lambda(lambda seq: seq[:, 0, :])(output)
# TODO: Try with more regularisation
# cls_out = keras.layers.Dropout(0.5)(cls_out)
