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)