 View on TensorFlow.org
|
 Run in Google Colab
|
 View source on GitHub
|
 Download notebook
|
This tutorial demonstrates how to load BERT, ALBERT and ELECTRA pretrained checkpoints and use them for downstream tasks.
Model Garden contains a collection of state-of-the-art models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.
Install TF Model Garden package
pip install -U -q "tf-models-official"Import necessary libraries
import os
import yaml
import json
import tensorflow as tf
import tensorflow_models as tfm
from official.core import exp_factory
/tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tensorflow_hub/__init__.py:61: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81. from pkg_resources import parse_version
Load BERT model pretrained checkpoints
Select required BERT model
# @title Download Checkpoint of the Selected Model { display-mode: "form", run: "auto" }
model_display_name = 'BERT-base cased English' # @param ['BERT-base uncased English','BERT-base cased English','BERT-large uncased English', 'BERT-large cased English', 'BERT-large, Uncased (Whole Word Masking)', 'BERT-large, Cased (Whole Word Masking)', 'BERT-base MultiLingual','BERT-base Chinese']
if model_display_name == 'BERT-base uncased English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-12_H-768_A-12.tar.gz"
!tar -xvf "uncased_L-12_H-768_A-12.tar.gz"
elif model_display_name == 'BERT-base cased English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-12_H-768_A-12.tar.gz"
!tar -xvf "cased_L-12_H-768_A-12.tar.gz"
elif model_display_name == "BERT-large uncased English":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-24_H-1024_A-16.tar.gz"
!tar -xvf "uncased_L-24_H-1024_A-16.tar.gz"
elif model_display_name == "BERT-large cased English":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-24_H-1024_A-16.tar.gz"
!tar -xvf "cased_L-24_H-1024_A-16.tar.gz"
elif model_display_name == "BERT-large, Uncased (Whole Word Masking)":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/wwm_uncased_L-24_H-1024_A-16.tar.gz"
!tar -xvf "wwm_uncased_L-24_H-1024_A-16.tar.gz"
elif model_display_name == "BERT-large, Cased (Whole Word Masking)":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/wwm_cased_L-24_H-1024_A-16.tar.gz"
!tar -xvf "wwm_cased_L-24_H-1024_A-16.tar.gz"
elif model_display_name == "BERT-base MultiLingual":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/multi_cased_L-12_H-768_A-12.tar.gz"
!tar -xvf "multi_cased_L-12_H-768_A-12.tar.gz"
elif model_display_name == "BERT-base Chinese":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/chinese_L-12_H-768_A-12.tar.gz"
!tar -xvf "chinese_L-12_H-768_A-12.tar.gz"
--2026-10-31 11:29:35-- https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-12_H-768_A-12.tar.gz
Resolving storage.googleapis.com (storage.googleapis.com)... 173.194.64.207, 142.251.189.207, 142.250.152.207, ...
Connecting to storage.googleapis.com (storage.googleapis.com)|173.194.64.207|:443... connected.
HTTP request sent, awaiting response...
200 OK
Length: 401886728 (383M) [application/octet-stream]
Saving to: ‘cased_L-12_H-768_A-12.tar.gz’
cased_L-1 0%[ ] 0 --.-KB/s
cased_L-12 2%[ ] 11.11M 55.6MB/s
cased_L-12_ 6%[> ] 25.35M 63.4MB/s
cased_L-12_H 12%[=> ] 48.01M 78.4MB/s
cased_L-12_H- 18%[==> ] 72.01M 79.7MB/s
cased_L-12_H-7 21%[===> ] 81.36M 73.7MB/s
cased_L-12_H-76 27%[====> ] 104.01M 79.2MB/s
cased_L-12_H-768 33%[=====> ] 130.27M 86.1MB/s
cased_L-12_H-768_ 37%[======> ] 144.01M 83.7MB/s
cased_L-12_H-768_A 43%[=======> ] 168.01M 86.9MB/s
cased_L-12_H-768_A- 52%[=========> ] 200.01M 93.3MB/s
ased_L-12_H-768_A-1 59%[==========> ] 226.85M 96.8MB/s
sed_L-12_H-768_A-12 65%[============> ] 251.66M 98.9MB/s
ed_L-12_H-768_A-12. 73%[=============> ] 280.01M 102MB/s
d_L-12_H-768_A-12.t 79%[==============> ] 304.57M 103MB/s
_L-12_H-768_A-12.ta 85%[================> ] 328.01M 104MB/s eta 1s
L-12_H-768_A-12.tar 91%[=================> ] 352.01M 104MB/s eta 1s
cased_L-12_H-768_A- 100%[===================>] 383.27M 107MB/s in 3.6s
2026-10-31 11:29:39 (107 MB/s) - ‘cased_L-12_H-768_A-12.tar.gz’ saved [401886728/401886728]
cased_L-12_H-768_A-12/
cased_L-12_H-768_A-12/vocab.txt
cased_L-12_H-768_A-12/bert_model.ckpt.index
cased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001
cased_L-12_H-768_A-12/params.yaml
cased_L-12_H-768_A-12/bert_config.json
# Lookup table of the directory name corresponding to each model checkpoint
folder_bert_dict = {
'BERT-base uncased English': 'uncased_L-12_H-768_A-12',
'BERT-base cased English': 'cased_L-12_H-768_A-12',
'BERT-large uncased English': 'uncased_L-24_H-1024_A-16',
'BERT-large cased English': 'cased_L-24_H-1024_A-16',
'BERT-large, Uncased (Whole Word Masking)': 'wwm_uncased_L-24_H-1024_A-16',
'BERT-large, Cased (Whole Word Masking)': 'wwm_cased_L-24_H-1024_A-16',
'BERT-base MultiLingual': 'multi_cased_L-12_H-768_A-1',
'BERT-base Chinese': 'chinese_L-12_H-768_A-12'
}
folder_bert = folder_bert_dict.get(model_display_name)
folder_bert
'cased_L-12_H-768_A-12'
Construct BERT Model Using the New params.yaml
params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here.
config_file = os.path.join(folder_bert, "params.yaml")
config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())
config_dict
{'task': {'model': {'encoder': {'bert': {'attention_dropout_rate': 0.1,
'dropout_rate': 0.1,
'hidden_activation': 'gelu',
'hidden_size': 768,
'initializer_range': 0.02,
'intermediate_size': 3072,
'max_position_embeddings': 512,
'num_attention_heads': 12,
'num_layers': 12,
'type_vocab_size': 2,
'vocab_size': 28996},
'type': 'bert'} } } }
# Method 1: pass encoder config dict into EncoderConfig
encoder_config = tfm.nlp.encoders.EncoderConfig(config_dict["task"]["model"]["encoder"])
encoder_config.get().as_dict()
{'vocab_size': 28996,
'hidden_size': 768,
'num_layers': 12,
'num_attention_heads': 12,
'hidden_activation': 'gelu',
'intermediate_size': 3072,
'dropout_rate': 0.1,
'attention_dropout_rate': 0.1,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02,
'embedding_size': None,
'output_range': None,
'return_all_encoder_outputs': False,
'return_attention_scores': False,
'return_word_embeddings': False,
'norm_first': False}
# Method 2: use override_params_dict function to override default Encoder params
encoder_config = tfm.nlp.encoders.EncoderConfig()
tfm.hyperparams.override_params_dict(encoder_config, config_dict["task"]["model"]["encoder"], is_strict=True)
encoder_config.get().as_dict()
{'vocab_size': 28996,
'hidden_size': 768,
'num_layers': 12,
'num_attention_heads': 12,
'hidden_activation': 'gelu',
'intermediate_size': 3072,
'dropout_rate': 0.1,
'attention_dropout_rate': 0.1,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02,
'embedding_size': None,
'output_range': None,
'return_all_encoder_outputs': False,
'return_attention_scores': False,
'return_word_embeddings': False,
'norm_first': False}
Construct BERT Model Using the Old bert_config.json
bert_config_file = os.path.join(folder_bert, "bert_config.json")
config_dict = json.loads(tf.io.gfile.GFile(bert_config_file).read())
config_dict
{'hidden_size': 768,
'initializer_range': 0.02,
'intermediate_size': 3072,
'max_position_embeddings': 512,
'num_attention_heads': 12,
'num_layers': 12,
'type_vocab_size': 2,
'vocab_size': 28996,
'hidden_activation': 'gelu',
'dropout_rate': 0.1,
'attention_dropout_rate': 0.1}
encoder_config = tfm.nlp.encoders.EncoderConfig({
'type':'bert',
'bert': config_dict
})
encoder_config.get().as_dict()
{'vocab_size': 28996,
'hidden_size': 768,
'num_layers': 12,
'num_attention_heads': 12,
'hidden_activation': 'gelu',
'intermediate_size': 3072,
'dropout_rate': 0.1,
'attention_dropout_rate': 0.1,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02,
'embedding_size': None,
'output_range': None,
'return_all_encoder_outputs': False,
'return_attention_scores': False,
'return_word_embeddings': False,
'norm_first': False}
Construct a classifier with encoder_config
Here, we construct a new BERT Classifier with 2 classes and plot its model architecture. A BERT Classifier consists of a BERT encoder using the selected encoder config, a Dropout layer and a MLP classification head.
bert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)
bert_classifier = tfm.nlp.models.BertClassifier(network=bert_encoder, num_classes=2)
tf.keras.utils.plot_model(bert_classifier)
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1793446185.022988 31558 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13638 MB memory: -> device: 0, name: Tesla T4, pci bus id: 0000:00:05.0, compute capability: 7.5 I0000 00:00:1793446185.025291 31558 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 13756 MB memory: -> device: 1, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5 I0000 00:00:1793446185.027437 31558 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 13756 MB memory: -> device: 2, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5 I0000 00:00:1793446185.029659 31558 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 13756 MB memory: -> device: 3, name: Tesla T4, pci bus id: 0000:00:08.0, compute capability: 7.5
Load Pretrained Weights into the BERT Classifier
The provided pretrained checkpoint only contains weights for the BERT Encoder within the BERT Classifier. Weights for the Classification Head is still randomly initialized.
checkpoint = tf.train.Checkpoint(encoder=bert_encoder)
checkpoint.read(
os.path.join(folder_bert, 'bert_model.ckpt')).expect_partial().assert_existing_objects_matched()
<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f314813cf90>
Load ALBERT model pretrained checkpoints
# @title Download Checkpoint of the Selected Model { display-mode: "form", run: "auto" }
albert_model_display_name = 'ALBERT-xxlarge English' # @param ['ALBERT-base English', 'ALBERT-large English', 'ALBERT-xlarge English', 'ALBERT-xxlarge English']
if albert_model_display_name == 'ALBERT-base English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_base.tar.gz"
!tar -xvf "albert_base.tar.gz"
elif albert_model_display_name == 'ALBERT-large English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_large.tar.gz"
!tar -xvf "albert_large.tar.gz"
elif albert_model_display_name == "ALBERT-xlarge English":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_xlarge.tar.gz"
!tar -xvf "albert_xlarge.tar.gz"
elif albert_model_display_name == "ALBERT-xxlarge English":
!wget "https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_xxlarge.tar.gz"
!tar -xvf "albert_xxlarge.tar.gz"
--2026-10-31 11:29:49-- https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_xxlarge.tar.gz Resolving storage.googleapis.com (storage.googleapis.com)... 34.3.0.27, 104.154.124.27, 34.144.170.27, ... Connecting to storage.googleapis.com (storage.googleapis.com)|34.3.0.27|:443... connected. HTTP request sent, awaiting response... 200 OK Length: 826059238 (788M) [application/octet-stream] Saving to: ‘albert_xxlarge.tar.gz’ albert_xxlarge.tar. 0%[ ] 0 --.-KB/s albert_xxlarge.tar. 2%[ ] 20.92M 105MB/s albert_xxlarge.tar. 7%[> ] 56.01M 105MB/s albert_xxlarge.tar. 11%[=> ] 88.01M 120MB/s albert_xxlarge.tar. 16%[==> ] 130.01M 139MB/s albert_xxlarge.tar. 17%[==> ] 136.01M 117MB/s albert_xxlarge.tar. 23%[===> ] 184.01M 128MB/s albert_xxlarge.tar. 27%[====> ] 219.64M 134MB/s albert_xxlarge.tar. 31%[=====> ] 245.94M 134MB/s albert_xxlarge.tar. 35%[======> ] 280.01M 131MB/s albert_xxlarge.tar. 40%[=======> ] 322.23M 138MB/s albert_xxlarge.tar. 44%[=======> ] 353.38M 140MB/s albert_xxlarge.tar. 48%[========> ] 385.14M 141MB/s albert_xxlarge.tar. 53%[=========> ] 421.75M 144MB/s albert_xxlarge.tar. 57%[==========> ] 455.30M 145MB/s eta 2s albert_xxlarge.tar. 61%[===========> ] 484.12M 145MB/s eta 2s albert_xxlarge.tar. 66%[============> ] 524.51M 148MB/s eta 2s albert_xxlarge.tar. 73%[=============> ] 580.19M 160MB/s eta 2s albert_xxlarge.tar. 80%[===============> ] 632.01M 162MB/s eta 2s albert_xxlarge.tar. 84%[===============> ] 669.20M 169MB/s eta 1s albert_xxlarge.tar. 89%[================> ] 702.95M 168MB/s eta 1s albert_xxlarge.tar. 94%[=================> ] 742.81M 177MB/s eta 1s albert_xxlarge.tar. 98%[==================> ] 776.73M 173MB/s eta 1s albert_xxlarge.tar. 100%[===================>] 787.79M 178MB/s in 4.8s 2026-10-31 11:29:54 (163 MB/s) - ‘albert_xxlarge.tar.gz’ saved [826059238/826059238] albert_xxlarge/ albert_xxlarge/bert_model.ckpt.index albert_xxlarge/30k-clean.model albert_xxlarge/30k-clean.vocab albert_xxlarge/bert_model.ckpt.data-00000-of-00001 albert_xxlarge/params.yaml albert_xxlarge/albert_config.json
# Lookup table of the directory name corresponding to each model checkpoint
folder_albert_dict = {
'ALBERT-base English': 'albert_base',
'ALBERT-large English': 'albert_large',
'ALBERT-xlarge English': 'albert_xlarge',
'ALBERT-xxlarge English': 'albert_xxlarge'
}
folder_albert = folder_albert_dict.get(albert_model_display_name)
folder_albert
'albert_xxlarge'
Construct ALBERT Model Using the New params.yaml
params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here.
config_file = os.path.join(folder_albert, "params.yaml")
config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())
config_dict
{'task': {'model': {'encoder': {'albert': {'attention_dropout_rate': 0.0,
'dropout_rate': 0.0,
'embedding_width': 128,
'hidden_activation': 'gelu',
'hidden_size': 4096,
'initializer_range': 0.02,
'intermediate_size': 16384,
'max_position_embeddings': 512,
'num_attention_heads': 64,
'num_layers': 12,
'type_vocab_size': 2,
'vocab_size': 30000},
'type': 'albert'} } } }
# Method 1: pass encoder config dict into EncoderConfig
encoder_config = tfm.nlp.encoders.EncoderConfig(config_dict["task"]["model"]["encoder"])
encoder_config.get().as_dict()
{'vocab_size': 30000,
'embedding_width': 128,
'hidden_size': 4096,
'num_layers': 12,
'num_attention_heads': 64,
'hidden_activation': 'gelu',
'intermediate_size': 16384,
'dropout_rate': 0.0,
'attention_dropout_rate': 0.0,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02}
# Method 2: use override_params_dict function to override default Encoder params
encoder_config = tfm.nlp.encoders.EncoderConfig()
tfm.hyperparams.override_params_dict(encoder_config, config_dict["task"]["model"]["encoder"], is_strict=True)
encoder_config.get().as_dict()
{'vocab_size': 30000,
'embedding_width': 128,
'hidden_size': 4096,
'num_layers': 12,
'num_attention_heads': 64,
'hidden_activation': 'gelu',
'intermediate_size': 16384,
'dropout_rate': 0.0,
'attention_dropout_rate': 0.0,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02}
Construct ALBERT Model Using the Old albert_config.json
albert_config_file = os.path.join(folder_albert, "albert_config.json")
config_dict = json.loads(tf.io.gfile.GFile(albert_config_file).read())
config_dict
{'hidden_size': 4096,
'initializer_range': 0.02,
'intermediate_size': 16384,
'max_position_embeddings': 512,
'num_attention_heads': 64,
'type_vocab_size': 2,
'vocab_size': 30000,
'embedding_width': 128,
'attention_dropout_rate': 0.0,
'dropout_rate': 0.0,
'num_layers': 12,
'hidden_activation': 'gelu'}
encoder_config = tfm.nlp.encoders.EncoderConfig({
'type':'albert',
'albert': config_dict
})
encoder_config.get().as_dict()
{'vocab_size': 30000,
'embedding_width': 128,
'hidden_size': 4096,
'num_layers': 12,
'num_attention_heads': 64,
'hidden_activation': 'gelu',
'intermediate_size': 16384,
'dropout_rate': 0.0,
'attention_dropout_rate': 0.0,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02}
Construct a Classifier with encoder_config
Here, we construct a new BERT Classifier with 2 classes and plot its model architecture. A BERT Classifier consists of a BERT encoder using the selected encoder config, a Dropout layer and a MLP classification head.
albert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)
albert_classifier = tfm.nlp.models.BertClassifier(network=albert_encoder, num_classes=2)
tf.keras.utils.plot_model(albert_classifier)
Load Pretrained Weights into the Classifier
The provided pretrained checkpoint only contains weights for the ALBERT Encoder within the ALBERT Classifier. Weights for the Classification Head is still randomly initialized.
checkpoint = tf.train.Checkpoint(encoder=albert_encoder)
checkpoint.read(
os.path.join(folder_albert, 'bert_model.ckpt')).expect_partial().assert_existing_objects_matched()
<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f313b376050>
Load ELECTRA model pretrained checkpoints
# @title Download Checkpoint of the Selected Model { display-mode: "form", run: "auto" }
electra_model_display_name = 'ELECTRA-small English' # @param ['ELECTRA-small English', 'ELECTRA-base English']
if electra_model_display_name == 'ELECTRA-small English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/electra/small.tar.gz"
!tar -xvf "small.tar.gz"
elif electra_model_display_name == 'ELECTRA-base English':
!wget "https://storage.googleapis.com/tf_model_garden/nlp/electra/base.tar.gz"
!tar -xvf "base.tar.gz"
--2026-10-31 11:30:04-- https://storage.googleapis.com/tf_model_garden/nlp/electra/small.tar.gz Resolving storage.googleapis.com (storage.googleapis.com)... 34.3.1.27, 104.154.125.27, 34.144.171.27, ... Connecting to storage.googleapis.com (storage.googleapis.com)|34.3.1.27|:443... connected. HTTP request sent, awaiting response... 200 OK Length: 157951922 (151M) [application/octet-stream] Saving to: ‘small.tar.gz’ small.tar.gz 0%[ ] 0 --.-KB/s small.tar.gz 5%[> ] 8.01M 35.6MB/s small.tar.gz 31%[=====> ] 48.01M 104MB/s small.tar.gz 51%[=========> ] 78.12M 118MB/s small.tar.gz 63%[===========> ] 96.01M 111MB/s small.tar.gz 79%[==============> ] 120.01M 108MB/s small.tar.gz 84%[===============> ] 128.01M 89.2MB/s small.tar.gz 100%[===================>] 150.63M 99.0MB/s in 1.5s 2026-10-31 11:30:06 (99.0 MB/s) - ‘small.tar.gz’ saved [157951922/157951922] small/ small/ckpt-1000000.data-00000-of-00001 small/params.yaml small/checkpoint small/ckpt-1000000.index
# Lookup table of the directory name corresponding to each model checkpoint
folder_electra_dict = {
'ELECTRA-small English': 'small',
'ELECTRA-base English': 'base'
}
folder_electra = folder_electra_dict.get(electra_model_display_name)
folder_electra
'small'
Construct BERT Model Using the params.yaml
params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here.
config_file = os.path.join(folder_electra, "params.yaml")
config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())
config_dict
{'model': {'cls_heads': [{'activation': 'tanh',
'cls_token_idx': 0,
'dropout_rate': 0.1,
'inner_dim': 64,
'name': 'next_sentence',
'num_classes': 2}],
'disallow_correct': False,
'discriminator_encoder': {'type': 'bert',
'bert': {'attention_dropout_rate': 0.1,
'dropout_rate': 0.1,
'embedding_size': 128,
'hidden_activation': 'gelu',
'hidden_size': 256,
'initializer_range': 0.02,
'intermediate_size': 1024,
'max_position_embeddings': 512,
'num_attention_heads': 4,
'num_layers': 12,
'type_vocab_size': 2,
'vocab_size': 30522} },
'discriminator_loss_weight': 50.0,
'generator_encoder': {'type': 'bert',
'bert': {'attention_dropout_rate': 0.1,
'dropout_rate': 0.1,
'embedding_size': 128,
'hidden_activation': 'gelu',
'hidden_size': 64,
'initializer_range': 0.02,
'intermediate_size': 256,
'max_position_embeddings': 512,
'num_attention_heads': 1,
'num_layers': 12,
'type_vocab_size': 2,
'vocab_size': 30522} },
'num_classes': 2,
'num_masked_tokens': 76,
'sequence_length': 512,
'tie_embeddings': True} }
disc_encoder_config = tfm.nlp.encoders.EncoderConfig(
config_dict['model']['discriminator_encoder']
)
disc_encoder_config.get().as_dict()
{'vocab_size': 30522,
'hidden_size': 256,
'num_layers': 12,
'num_attention_heads': 4,
'hidden_activation': 'gelu',
'intermediate_size': 1024,
'dropout_rate': 0.1,
'attention_dropout_rate': 0.1,
'max_position_embeddings': 512,
'type_vocab_size': 2,
'initializer_range': 0.02,
'embedding_size': 128,
'output_range': None,
'return_all_encoder_outputs': False,
'return_attention_scores': False,
'return_word_embeddings': False,
'norm_first': False}
Construct a Classifier with encoder_config
Here, we construct a Classifier with 2 classes and plot its model architecture. A Classifier consists of a ELECTRA discriminator encoder using the selected encoder config, a Dropout layer and a MLP classification head.
disc_encoder = tfm.nlp.encoders.build_encoder(disc_encoder_config)
elctra_dic_classifier = tfm.nlp.models.BertClassifier(network=disc_encoder, num_classes=2)
tf.keras.utils.plot_model(elctra_dic_classifier)
Load Pretrained Weights into the Classifier
The provided pretrained checkpoint contains weights for the entire ELECTRA model. We are only loading its discriminator (conveninently named as encoder) wights within the Classifier. Weights for the Classification Head is still randomly initialized.
checkpoint = tf.train.Checkpoint(encoder=disc_encoder)
checkpoint.read(
tf.train.latest_checkpoint(os.path.join(folder_electra))
).expect_partial().assert_existing_objects_matched()
<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f312c5f8410>