Is it possible to get the embedding table in tf_hub models?

user_17945966 0
Viewed 45

I'm having problems with finetuning a BERT model. I was previously using get_transformer_encoder() in official.nlp and MLM task in official.nlp to train a BERT. But this seems like tough, so I changed to finetuning a pretrained BERT model. But to also complete the MLM task, it seems like I need to get the embedding table from the tf_hub models. I think it's possible to extract the embedding table from the hub models, but I don't know whether there are some functions I can use to get them more quickly, Thanks.

pre-trained-model bert-language-model tensorflow nlp
1 Answers

I think this might could help finetune a pretrained model using tensorflow_hub with MLM task, we can use keras_nlp from official.nlp.

Here is a simple example:

bert_encoder_scratch = get_transformer_encoder(config)
bert_encoder_pretrain = hub.KerasLayer(hub.load(module_url), trainable=True)

if you are training from scratch:

mlm_model = keras_nlp.layers.MaskedLM(embedding_table=bert_encoder_scratch.get_embedding_table())

So in this way, what lacks for a hub model is to find the embedding table. Find the variable word_embedding by:

bert_encoder_pretrain.trainable_variables

or you can use tf.get_variable(). Here I simply write bert_encoder_pretrain.trainable_variables[0]

The mlm task would be:

mlm_model = keras_nlp.layers.MaskedLM(embedding_table=bert_encoder_pretrain.trainable_variables[0])

Then you get your seq_output from the model, and also masked_positions, here I define it as:

masked_positions = tf.broadcast_to(tf.range(0, seq_length), [batch_size, seq_length])

You will get the mlm_logits from the model by:

mlm_logits = mlm_model(seq_output, masked_positions)

I have tested this approach in my model.

user_179459660
Related