docs/nlp/load_lm_ckpts.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "30155835fc9f"
},
"source": [
"##### Copyright 2022 The TensorFlow Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "906e07f6e562"
},
"outputs": [],
"source": [
"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5hrbPTziJK15"
},
"source": [
"# Load LM Checkpoints using Model Garden"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-PYqCW1II75I"
},
"source": [
"\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
" \u003ctd\u003e\n",
" \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/tfmodels/nlp/load_lm_ckpts\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/models/blob/master/docs/nlp/load_lm_ckpts.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/models/blob/master/docs/nlp/load_lm_ckpts.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/models/docs/nlp/load_lm_ckpts.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
" \u003c/td\u003e\n",
"\u003c/table\u003e"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yyyk1KMlJdWd"
},
"source": [
"This tutorial demonstrates how to load BERT, ALBERT and ELECTRA pretrained checkpoints and use them for downstream tasks.\n",
"\n",
"[Model Garden](https://www.tensorflow.org/tfmodels) 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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uEG4RYHolQij"
},
"source": [
"## Install TF Model Garden package"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "kPfC1NJZnJq1"
},
"outputs": [],
"source": [
"!pip install -U -q \"tf-models-official\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Op9R3zy3lUk8"
},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6_y4Rfq23wK-"
},
"outputs": [],
"source": [
"import os\n",
"import yaml\n",
"import json\n",
"\n",
"import tensorflow as tf"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xjgv3gllzbYQ"
},
"outputs": [],
"source": [
"import tensorflow_models as tfm\n",
"\n",
"from official.core import exp_factory"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "J-t2mo6VQNfY"
},
"source": [
"## Load BERT model pretrained checkpoints"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hdBsFnI20LDE"
},
"source": [
"### Select required BERT model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "apn3VgxUlr5G"
},
"outputs": [],
"source": [
"# @title Download Checkpoint of the Selected Model { display-mode: \"form\", run: \"auto\" }\n",
"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']\n",
"\n",
"if model_display_name == 'BERT-base uncased English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-12_H-768_A-12.tar.gz\"\n",
" !tar -xvf \"uncased_L-12_H-768_A-12.tar.gz\"\n",
"elif model_display_name == 'BERT-base cased English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-12_H-768_A-12.tar.gz\"\n",
" !tar -xvf \"cased_L-12_H-768_A-12.tar.gz\"\n",
"elif model_display_name == \"BERT-large uncased English\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-24_H-1024_A-16.tar.gz\"\n",
" !tar -xvf \"uncased_L-24_H-1024_A-16.tar.gz\"\n",
"elif model_display_name == \"BERT-large cased English\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-24_H-1024_A-16.tar.gz\"\n",
" !tar -xvf \"cased_L-24_H-1024_A-16.tar.gz\"\n",
"elif model_display_name == \"BERT-large, Uncased (Whole Word Masking)\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/wwm_uncased_L-24_H-1024_A-16.tar.gz\"\n",
" !tar -xvf \"wwm_uncased_L-24_H-1024_A-16.tar.gz\"\n",
"elif model_display_name == \"BERT-large, Cased (Whole Word Masking)\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/wwm_cased_L-24_H-1024_A-16.tar.gz\"\n",
" !tar -xvf \"wwm_cased_L-24_H-1024_A-16.tar.gz\"\n",
"elif model_display_name == \"BERT-base MultiLingual\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/multi_cased_L-12_H-768_A-12.tar.gz\"\n",
" !tar -xvf \"multi_cased_L-12_H-768_A-12.tar.gz\"\n",
"elif model_display_name == \"BERT-base Chinese\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/chinese_L-12_H-768_A-12.tar.gz\"\n",
" !tar -xvf \"chinese_L-12_H-768_A-12.tar.gz\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jzxyziRuaC95"
},
"outputs": [],
"source": [
"# Lookup table of the directory name corresponding to each model checkpoint\n",
"folder_bert_dict = {\n",
" 'BERT-base uncased English': 'uncased_L-12_H-768_A-12',\n",
" 'BERT-base cased English': 'cased_L-12_H-768_A-12',\n",
" 'BERT-large uncased English': 'uncased_L-24_H-1024_A-16',\n",
" 'BERT-large cased English': 'cased_L-24_H-1024_A-16',\n",
" 'BERT-large, Uncased (Whole Word Masking)': 'wwm_uncased_L-24_H-1024_A-16',\n",
" 'BERT-large, Cased (Whole Word Masking)': 'wwm_cased_L-24_H-1024_A-16',\n",
" 'BERT-base MultiLingual': 'multi_cased_L-12_H-768_A-1',\n",
" 'BERT-base Chinese': 'chinese_L-12_H-768_A-12'\n",
"}\n",
"\n",
"folder_bert = folder_bert_dict.get(model_display_name)\n",
"folder_bert"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "q1WrYswpZPlc"
},
"source": [
"### Construct BERT Model Using the New `params.yaml`\n",
"\n",
"params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "quu1s8Hi2szo"
},
"outputs": [],
"source": [
"config_file = os.path.join(folder_bert, \"params.yaml\")\n",
"config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())\n",
"config_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3t8o0iG9v8ac"
},
"outputs": [],
"source": [
"# Method 1: pass encoder config dict into EncoderConfig\n",
"encoder_config = tfm.nlp.encoders.EncoderConfig(config_dict[\"task\"][\"model\"][\"encoder\"])\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2I5PetB6wPvb"
},
"outputs": [],
"source": [
"# Method 2: use override_params_dict function to override default Encoder params\n",
"encoder_config = tfm.nlp.encoders.EncoderConfig()\n",
"tfm.hyperparams.override_params_dict(encoder_config, config_dict[\"task\"][\"model\"][\"encoder\"], is_strict=True)\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5yHiG_9oS3Uw"
},
"source": [
"### Construct BERT Model Using the Old `bert_config.json`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WEyaqLcW3nne"
},
"outputs": [],
"source": [
"bert_config_file = os.path.join(folder_bert, \"bert_config.json\")\n",
"config_dict = json.loads(tf.io.gfile.GFile(bert_config_file).read())\n",
"config_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xSIcaW9tdrl4"
},
"outputs": [],
"source": [
"encoder_config = tfm.nlp.encoders.EncoderConfig({\n",
" 'type':'bert',\n",
" 'bert': config_dict\n",
"})\n",
"\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yZznAP--TDLe"
},
"source": [
"### Construct a classifier with `encoder_config`\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ny962I8nqs4n"
},
"outputs": [],
"source": [
"bert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)\n",
"bert_classifier = tfm.nlp.models.BertClassifier(network=bert_encoder, num_classes=2)\n",
"\n",
"tf.keras.utils.plot_model(bert_classifier)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IStKfxXkTJMu"
},
"source": [
"### Load Pretrained Weights into the BERT Classifier\n",
"\n",
"The provided pretrained checkpoint only contains weights for the BERT Encoder within the BERT Classifier. Weights for the Classification Head is still randomly initialized."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "G9_XCBpOEo4y"
},
"outputs": [],
"source": [
"checkpoint = tf.train.Checkpoint(encoder=bert_encoder)\n",
"checkpoint.read(\n",
" os.path.join(folder_bert, 'bert_model.ckpt')).expect_partial().assert_existing_objects_matched()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "E6Hu1FFgQWUU"
},
"source": [
"## Load ALBERT model pretrained checkpoints"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TWUtFeWxQn0V"
},
"outputs": [],
"source": [
"# @title Download Checkpoint of the Selected Model { display-mode: \"form\", run: \"auto\" }\n",
"albert_model_display_name = 'ALBERT-xxlarge English' # @param ['ALBERT-base English', 'ALBERT-large English', 'ALBERT-xlarge English', 'ALBERT-xxlarge English']\n",
"\n",
"if albert_model_display_name == 'ALBERT-base English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_base.tar.gz\"\n",
" !tar -xvf \"albert_base.tar.gz\"\n",
"elif albert_model_display_name == 'ALBERT-large English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_large.tar.gz\"\n",
" !tar -xvf \"albert_large.tar.gz\"\n",
"elif albert_model_display_name == \"ALBERT-xlarge English\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_xlarge.tar.gz\"\n",
" !tar -xvf \"albert_xlarge.tar.gz\"\n",
"elif albert_model_display_name == \"ALBERT-xxlarge English\":\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/albert/albert_xxlarge.tar.gz\"\n",
" !tar -xvf \"albert_xxlarge.tar.gz\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5lZDWD7zUAAO"
},
"outputs": [],
"source": [
"# Lookup table of the directory name corresponding to each model checkpoint\n",
"folder_albert_dict = {\n",
" 'ALBERT-base English': 'albert_base',\n",
" 'ALBERT-large English': 'albert_large',\n",
" 'ALBERT-xlarge English': 'albert_xlarge',\n",
" 'ALBERT-xxlarge English': 'albert_xxlarge'\n",
"}\n",
"\n",
"folder_albert = folder_albert_dict.get(albert_model_display_name)\n",
"folder_albert"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ftXwmObdU2fS"
},
"source": [
"### Construct ALBERT Model Using the New `params.yaml`\n",
"\n",
"params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VXn20q2oU1UJ"
},
"outputs": [],
"source": [
"config_file = os.path.join(folder_albert, \"params.yaml\")\n",
"config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())\n",
"config_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Uo_TSMSvWOX_"
},
"outputs": [],
"source": [
"# Method 1: pass encoder config dict into EncoderConfig\n",
"encoder_config = tfm.nlp.encoders.EncoderConfig(config_dict[\"task\"][\"model\"][\"encoder\"])\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "u7oJe93uWcy0"
},
"outputs": [],
"source": [
"# Method 2: use override_params_dict function to override default Encoder params\n",
"encoder_config = tfm.nlp.encoders.EncoderConfig()\n",
"tfm.hyperparams.override_params_dict(encoder_config, config_dict[\"task\"][\"model\"][\"encoder\"], is_strict=True)\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "abpQFw80Wx6c"
},
"source": [
"### Construct ALBERT Model Using the Old `albert_config.json`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Xb99qms6WuPa"
},
"outputs": [],
"source": [
"albert_config_file = os.path.join(folder_albert, \"albert_config.json\")\n",
"config_dict = json.loads(tf.io.gfile.GFile(albert_config_file).read())\n",
"config_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mCW0RJHcEtVV"
},
"outputs": [],
"source": [
"encoder_config = tfm.nlp.encoders.EncoderConfig({\n",
" 'type':'albert',\n",
" 'albert': config_dict\n",
"})\n",
"\n",
"encoder_config.get().as_dict()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EIAMaOxdZw5u"
},
"source": [
"### Construct a Classifier with `encoder_config`\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xTkUisEEFEey"
},
"outputs": [],
"source": [
"albert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)\n",
"albert_classifier = tfm.nlp.models.BertClassifier(network=albert_encoder, num_classes=2)\n",
"\n",
"tf.keras.utils.plot_model(albert_classifier)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "m6EG_7CaZ2rI"
},
"source": [
"### Load Pretrained Weights into the Classifier\n",
"\n",
"The provided pretrained checkpoint only contains weights for the ALBERT Encoder within the ALBERT Classifier. Weights for the Classification Head is still randomly initialized."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7dOG3agXZ9Dx"
},
"outputs": [],
"source": [
"checkpoint = tf.train.Checkpoint(encoder=albert_encoder)\n",
"checkpoint.read(\n",
" os.path.join(folder_albert, 'bert_model.ckpt')).expect_partial().assert_existing_objects_matched()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6xsbeS-EcCqu"
},
"source": [
"## Load ELECTRA model pretrained checkpoints"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VpwIrAR4cIBF"
},
"outputs": [],
"source": [
"# @title Download Checkpoint of the Selected Model { display-mode: \"form\", run: \"auto\" }\n",
"electra_model_display_name = 'ELECTRA-small English' # @param ['ELECTRA-small English', 'ELECTRA-base English']\n",
"\n",
"if electra_model_display_name == 'ELECTRA-small English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/electra/small.tar.gz\"\n",
" !tar -xvf \"small.tar.gz\"\n",
"elif electra_model_display_name == 'ELECTRA-base English':\n",
" !wget \"https://storage.googleapis.com/tf_model_garden/nlp/electra/base.tar.gz\"\n",
" !tar -xvf \"base.tar.gz\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fy4FmsNOhlNa"
},
"outputs": [],
"source": [
"# Lookup table of the directory name corresponding to each model checkpoint\n",
"folder_electra_dict = {\n",
" 'ELECTRA-small English': 'small',\n",
" 'ELECTRA-base English': 'base'\n",
"}\n",
"\n",
"folder_electra = folder_electra_dict.get(electra_model_display_name)\n",
"folder_electra"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rgAcf-Fl3RTG"
},
"source": [
"### Construct BERT Model Using the `params.yaml`\n",
"\n",
"params.yaml can be used for training with the bundled trainer in addition to constructing the BERT encoder here."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZNBg5xzqh0Gr"
},
"outputs": [],
"source": [
"config_file = os.path.join(folder_electra, \"params.yaml\")\n",
"config_dict = yaml.safe_load(tf.io.gfile.GFile(config_file).read())\n",
"config_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "i-yX-KgJyduv"
},
"outputs": [],
"source": [
"disc_encoder_config = tfm.nlp.encoders.EncoderConfig(\n",
" config_dict['model']['discriminator_encoder']\n",
")\n",
"\n",
"disc_encoder_config.get().as_dict()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1AdrMkH73VYz"
},
"source": [
"### Construct a Classifier with `encoder_config`\n",
"\n",
"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.\n",
"\n",
"**Note**: The generator is discarded and the discriminator is used for downstream tasks"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "98Pt-SxszAvN"
},
"outputs": [],
"source": [
"disc_encoder = tfm.nlp.encoders.build_encoder(disc_encoder_config)\n",
"elctra_dic_classifier = tfm.nlp.models.BertClassifier(network=disc_encoder, num_classes=2)\n",
"tf.keras.utils.plot_model(elctra_dic_classifier)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aWQ2FKj64X5U"
},
"source": [
"### Load Pretrained Weights into the Classifier\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "99pznFJszQfV"
},
"outputs": [],
"source": [
"checkpoint = tf.train.Checkpoint(encoder=disc_encoder)\n",
"checkpoint.read(\n",
" tf.train.latest_checkpoint(os.path.join(folder_electra))\n",
" ).expect_partial().assert_existing_objects_matched()"
]
}
],
"metadata": {
"colab": {
"name": "load_lm_ckpts.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}