{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "58MWWgq75lMh"
},
"source": [
"##### Copyright 2020 The TensorFlow Hub Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "jM3hCI1UUzar"
},
"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": "4_NEJlxKKjyI"
},
"source": [
"\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
" \u003ctd\u003e\n",
" \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/text/tutorials/bert_glue\"\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/text/blob/master/docs/tutorials/bert_glue.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/text/blob/master/docs/tutorials/bert_glue.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView on GitHub\u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/text/docs/tutorials/bert_glue.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca href=\"https://tfhub.dev/google/collections/bert/1\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" /\u003eSee TF Hub model\u003c/a\u003e\n",
" \u003c/td\u003e\n",
"\u003c/table\u003e"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "U5POcTVNB_dv"
},
"source": [
"# Solve GLUE tasks using BERT on TPU\n",
"\n",
"BERT can be used to solve many problems in natural language processing. You will learn how to fine-tune BERT for many tasks from the [GLUE benchmark](https://gluebenchmark.com/):\n",
"\n",
"1. [CoLA](https://nyu-mll.github.io/CoLA/) (Corpus of Linguistic Acceptability): Is the sentence grammatically correct?\n",
"\n",
"1. [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank): The task is to predict the sentiment of a given sentence.\n",
"\n",
"1. [MRPC](https://www.microsoft.com/en-us/download/details.aspx?id=52398) (Microsoft Research Paraphrase Corpus): Determine whether a pair of sentences are semantically equivalent.\n",
"\n",
"1. [QQP](https://data.quora.com/First-Quora-Dataset-Release-Question-Pairs) (Quora Question Pairs2): Determine whether a pair of questions are semantically equivalent.\n",
"\n",
"1. [MNLI](http://www.nyu.edu/projects/bowman/multinli/) (Multi-Genre Natural Language Inference): Given a premise sentence and a hypothesis sentence, the task is to predict whether the premise entails the hypothesis (entailment), contradicts the hypothesis (contradiction), or neither (neutral).\n",
"\n",
"1. [QNLI](https://rajpurkar.github.io/SQuAD-explorer/)(Question-answering Natural Language Inference): The task is to determine whether the context sentence contains the answer to the question.\n",
"\n",
"1. [RTE](https://aclweb.org/aclwiki/Recognizing_Textual_Entailment)(Recognizing Textual Entailment): Determine if a sentence entails a given hypothesis or not.\n",
"\n",
"1. [WNLI](https://cs.nyu.edu/faculty/davise/papers/WinogradSchemas/WS.html)(Winograd Natural Language Inference): The task is to predict if the sentence with the pronoun substituted is entailed by the original sentence.\n",
"\n",
"This tutorial contains complete end-to-end code to train these models on a TPU. You can also run this notebook on a GPU, by changing one line (described below).\n",
"\n",
"In this notebook, you will:\n",
"\n",
"- Load a BERT model from TensorFlow Hub\n",
"- Choose one of GLUE tasks and download the dataset\n",
"- Preprocess the text\n",
"- Fine-tune BERT (examples are given for single-sentence and multi-sentence datasets)\n",
"- Save the trained model and use it\n",
"\n",
"Key point: The model you develop will be end-to-end. The preprocessing logic will be included in the model itself, making it capable of accepting raw strings as input.\n",
"\n",
"Note: This notebook should be run using a TPU. In Colab, choose **Runtime -\u003e Change runtime type** and verify that a **TPU** is selected.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SCjmX4zTCkRK"
},
"source": [
"## Setup\n",
"\n",
"You will use a separate model to preprocess text before using it to fine-tune BERT. This model depends on [tensorflow/text](https://github.com/tensorflow/text), which you will install below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rE2dA1XO9clA"
},
"outputs": [],
"source": [
"!pip install -q -U tensorflow-text"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "WMaudPO1a2Hx"
},
"source": [
"You will use the AdamW optimizer from [tensorflow/models](https://github.com/tensorflow/models) to fine-tune BERT, which you will install as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5zwJyopqa3uH"
},
"outputs": [],
"source": [
"!pip install -q -U tf-models-official"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Kx_Acxvo1nje"
},
"outputs": [],
"source": [
"!pip install -U tfds-nightly"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_XgTpm9ZxoN9"
},
"outputs": [],
"source": [
"import os\n",
"import tensorflow as tf\n",
"import tensorflow_hub as hub\n",
"import tensorflow_datasets as tfds\n",
"import tensorflow_text as text # A dependency of the preprocessing model\n",
"import tensorflow_addons as tfa\n",
"from official.nlp import optimization\n",
"import numpy as np\n",
"\n",
"tf.get_logger().setLevel('ERROR')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sv7A19G32Kfw"
},
"source": [
"Next, configure TFHub to read checkpoints directly from TFHub's Cloud Storage buckets. This is only recommended when running TFHub models on TPU.\n",
"\n",
"Without this setting TFHub would download the compressed file and extract the checkpoint locally. Attempting to load from these local files will fail with the following error:\n",
"\n",
"```\n",
"InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented\n",
"```\n",
"\n",
"This is because the [TPU can only read directly from Cloud Storage buckets](https://cloud.google.com/tpu/docs/troubleshooting#cannot_use_local_filesystem).\n",
"\n",
"Note: This setting is automatic in Colab."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Sz6P5pK3ldxQ"
},
"outputs": [],
"source": [
"os.environ[\"TFHUB_MODEL_LOAD_FORMAT\"]=\"UNCOMPRESSED\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "__Uqe2vNETAu"
},
"source": [
"### Connect to the TPU worker\n",
"\n",
"The following code connects to the TPU worker and changes TensorFlow's default device to the CPU device on the TPU worker. It also defines a TPU distribution strategy that you will use to distribute model training onto the 8 separate TPU cores available on this one TPU worker. See TensorFlow's [TPU guide](https://www.tensorflow.org/guide/tpu) for more information."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cpHWNs1nV0Zn"
},
"outputs": [],
"source": [
"import os\n",
"\n",
"if os.environ['COLAB_TPU_ADDR']:\n",
" cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n",
" tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
" tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
" strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
" print('Using TPU')\n",
"elif tf.config.list_physical_devices('GPU'):\n",
" strategy = tf.distribute.MirroredStrategy()\n",
" print('Using GPU')\n",
"else:\n",
" raise ValueError('Running on CPU is not recommended.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UVtEyxDFpKE1"
},
"source": [
"## Loading models from TensorFlow Hub\n",
"\n",
"Here you can choose which BERT model you will load from TensorFlow Hub and fine-tune.\n",
"There are multiple BERT models available to choose from.\n",
"\n",
" - [BERT-Base](https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3), [Uncased](https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3) and [seven more models](https://tfhub.dev/google/collections/bert/1) with trained weights released by the original BERT authors.\n",
" - [Small BERTs](https://tfhub.dev/google/collections/bert/1) have the same general architecture but fewer and/or smaller Transformer blocks, which lets you explore tradeoffs between speed, size and quality.\n",
" - [ALBERT](https://tfhub.dev/google/collections/albert/1): four different sizes of \"A Lite BERT\" that reduces model size (but not computation time) by sharing parameters between layers.\n",
" - [BERT Experts](https://tfhub.dev/google/collections/experts/bert/1): eight models that all have the BERT-base architecture but offer a choice between different pre-training domains, to align more closely with the target task.\n",
" - [Electra](https://tfhub.dev/google/collections/electra/1) has the same architecture as BERT (in three different sizes), but gets pre-trained as a discriminator in a set-up that resembles a Generative Adversarial Network (GAN).\n",
" - BERT with Talking-Heads Attention and Gated GELU [[base](https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_base/1), [large](https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_large/1)] has two improvements to the core of the Transformer architecture.\n",
"\n",
"See the model documentation linked above for more details.\n",
"\n",
"In this tutorial, you will start with BERT-base. You can use larger and more recent models for higher accuracy, or smaller models for faster training times. To change the model, you only need to switch a single line of code (shown below). All the differences are encapsulated in the SavedModel you will download from TensorFlow Hub."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "y8_ctG55-uTX"
},
"outputs": [],
"source": [
"#@title Choose a BERT model to fine-tune\n",
"\n",
"bert_model_name = 'bert_en_uncased_L-12_H-768_A-12' #@param [\"bert_en_uncased_L-12_H-768_A-12\", \"bert_en_uncased_L-24_H-1024_A-16\", \"bert_en_wwm_uncased_L-24_H-1024_A-16\", \"bert_en_cased_L-12_H-768_A-12\", \"bert_en_cased_L-24_H-1024_A-16\", \"bert_en_wwm_cased_L-24_H-1024_A-16\", \"bert_multi_cased_L-12_H-768_A-12\", \"small_bert/bert_en_uncased_L-2_H-128_A-2\", \"small_bert/bert_en_uncased_L-2_H-256_A-4\", \"small_bert/bert_en_uncased_L-2_H-512_A-8\", \"small_bert/bert_en_uncased_L-2_H-768_A-12\", \"small_bert/bert_en_uncased_L-4_H-128_A-2\", \"small_bert/bert_en_uncased_L-4_H-256_A-4\", \"small_bert/bert_en_uncased_L-4_H-512_A-8\", \"small_bert/bert_en_uncased_L-4_H-768_A-12\", \"small_bert/bert_en_uncased_L-6_H-128_A-2\", \"small_bert/bert_en_uncased_L-6_H-256_A-4\", \"small_bert/bert_en_uncased_L-6_H-512_A-8\", \"small_bert/bert_en_uncased_L-6_H-768_A-12\", \"small_bert/bert_en_uncased_L-8_H-128_A-2\", \"small_bert/bert_en_uncased_L-8_H-256_A-4\", \"small_bert/bert_en_uncased_L-8_H-512_A-8\", \"small_bert/bert_en_uncased_L-8_H-768_A-12\", \"small_bert/bert_en_uncased_L-10_H-128_A-2\", \"small_bert/bert_en_uncased_L-10_H-256_A-4\", \"small_bert/bert_en_uncased_L-10_H-512_A-8\", \"small_bert/bert_en_uncased_L-10_H-768_A-12\", \"small_bert/bert_en_uncased_L-12_H-128_A-2\", \"small_bert/bert_en_uncased_L-12_H-256_A-4\", \"small_bert/bert_en_uncased_L-12_H-512_A-8\", \"small_bert/bert_en_uncased_L-12_H-768_A-12\", \"albert_en_base\", \"albert_en_large\", \"albert_en_xlarge\", \"albert_en_xxlarge\", \"electra_small\", \"electra_base\", \"experts_pubmed\", \"experts_wiki_books\", \"talking-heads_base\", \"talking-heads_large\"]\n",
"\n",
"map_name_to_handle = {\n",
" 'bert_en_uncased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3',\n",
" 'bert_en_uncased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_L-24_H-1024_A-16/3',\n",
" 'bert_en_wwm_uncased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_wwm_uncased_L-24_H-1024_A-16/3',\n",
" 'bert_en_cased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_cased_L-12_H-768_A-12/3',\n",
" 'bert_en_cased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_cased_L-24_H-1024_A-16/3',\n",
" 'bert_en_wwm_cased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_wwm_cased_L-24_H-1024_A-16/3',\n",
" 'bert_multi_cased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_multi_cased_L-12_H-768_A-12/3',\n",
" 'small_bert/bert_en_uncased_L-2_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-2_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-2_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-2_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-768_A-12/1',\n",
" 'small_bert/bert_en_uncased_L-4_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-4_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-4_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-4_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-768_A-12/1',\n",
" 'small_bert/bert_en_uncased_L-6_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-6_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-6_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-6_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-768_A-12/1',\n",
" 'small_bert/bert_en_uncased_L-8_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-8_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-8_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-8_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-8_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-8_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-8_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-8_H-768_A-12/1',\n",
" 'small_bert/bert_en_uncased_L-10_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-10_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-10_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-10_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-10_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-10_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-10_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-10_H-768_A-12/1',\n",
" 'small_bert/bert_en_uncased_L-12_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-12_H-128_A-2/1',\n",
" 'small_bert/bert_en_uncased_L-12_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-12_H-256_A-4/1',\n",
" 'small_bert/bert_en_uncased_L-12_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-12_H-512_A-8/1',\n",
" 'small_bert/bert_en_uncased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-12_H-768_A-12/1',\n",
" 'albert_en_base':\n",
" 'https://tfhub.dev/tensorflow/albert_en_base/2',\n",
" 'albert_en_large':\n",
" 'https://tfhub.dev/tensorflow/albert_en_large/2',\n",
" 'albert_en_xlarge':\n",
" 'https://tfhub.dev/tensorflow/albert_en_xlarge/2',\n",
" 'albert_en_xxlarge':\n",
" 'https://tfhub.dev/tensorflow/albert_en_xxlarge/2',\n",
" 'electra_small':\n",
" 'https://tfhub.dev/google/electra_small/2',\n",
" 'electra_base':\n",
" 'https://tfhub.dev/google/electra_base/2',\n",
" 'experts_pubmed':\n",
" 'https://tfhub.dev/google/experts/bert/pubmed/2',\n",
" 'experts_wiki_books':\n",
" 'https://tfhub.dev/google/experts/bert/wiki_books/2',\n",
" 'talking-heads_base':\n",
" 'https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_base/1',\n",
" 'talking-heads_large':\n",
" 'https://tfhub.dev/tensorflow/talkheads_ggelu_bert_en_large/1',\n",
"}\n",
"\n",
"map_model_to_preprocess = {\n",
" 'bert_en_uncased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'bert_en_uncased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'bert_en_wwm_cased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_cased_preprocess/3',\n",
" 'bert_en_cased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_cased_preprocess/3',\n",
" 'bert_en_cased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_cased_preprocess/3',\n",
" 'bert_en_wwm_uncased_L-24_H-1024_A-16':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-2_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-2_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-2_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-2_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-4_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-4_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-4_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-4_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-6_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-6_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-6_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-6_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-8_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-8_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-8_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-8_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-10_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-10_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-10_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-10_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-12_H-128_A-2':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-12_H-256_A-4':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-12_H-512_A-8':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'small_bert/bert_en_uncased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'bert_multi_cased_L-12_H-768_A-12':\n",
" 'https://tfhub.dev/tensorflow/bert_multi_cased_preprocess/3',\n",
" 'albert_en_base':\n",
" 'https://tfhub.dev/tensorflow/albert_en_preprocess/3',\n",
" 'albert_en_large':\n",
" 'https://tfhub.dev/tensorflow/albert_en_preprocess/3',\n",
" 'albert_en_xlarge':\n",
" 'https://tfhub.dev/tensorflow/albert_en_preprocess/3',\n",
" 'albert_en_xxlarge':\n",
" 'https://tfhub.dev/tensorflow/albert_en_preprocess/3',\n",
" 'electra_small':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'electra_base':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'experts_pubmed':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'experts_wiki_books':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'talking-heads_base':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
" 'talking-heads_large':\n",
" 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3',\n",
"}\n",
"\n",
"tfhub_handle_encoder = map_name_to_handle[bert_model_name]\n",
"tfhub_handle_preprocess = map_model_to_preprocess[bert_model_name]\n",
"\n",
"print('BERT model selected :', tfhub_handle_encoder)\n",
"print('Preprocessing model auto-selected:', tfhub_handle_preprocess)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7WrcxxTRDdHi"
},
"source": [
"## Preprocess the text\n",
"\n",
"On the [Classify text with BERT colab](https://www.tensorflow.org/text/tutorials/classify_text_with_bert) the preprocessing model is used directly embedded with the BERT encoder.\n",
"\n",
"This tutorial demonstrates how to do preprocessing as part of your input pipeline for training, using Dataset.map, and then merge it into the model that gets exported for inference. That way, both training and inference can work from raw text inputs, although the TPU itself requires numeric inputs.\n",
"\n",
"TPU requirements aside, it can help performance have preprocessing done asynchronously in an input pipeline (you can learn more in the [tf.data performance guide](https://www.tensorflow.org/guide/data_performance)).\n",
"\n",
"This tutorial also demonstrates how to build multi-input models, and how to adjust the sequence length of the inputs to BERT.\n",
"\n",
"Let's demonstrate the preprocessing model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-ePjboKOPmv4"
},
"outputs": [],
"source": [
"bert_preprocess = hub.load(tfhub_handle_preprocess)\n",
"tok = bert_preprocess.tokenize(tf.constant(['Hello TensorFlow!']))\n",
"print(tok)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tRMCuruaQb5X"
},
"source": [
"Each preprocessing model also provides a method, `.bert_pack_inputs(tensors, seq_length)`, which takes a list of tokens (like `tok` above) and a sequence length argument. This packs the inputs to create a dictionary of tensors in the format expected by the BERT model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lraoc4csP0lY"
},
"outputs": [],
"source": [
"text_preprocessed = bert_preprocess.bert_pack_inputs([tok, tok], tf.constant(20))\n",
"\n",
"print('Shape Word Ids : ', text_preprocessed['input_word_ids'].shape)\n",
"print('Word Ids : ', text_preprocessed['input_word_ids'][0, :16])\n",
"print('Shape Mask : ', text_preprocessed['input_mask'].shape)\n",
"print('Input Mask : ', text_preprocessed['input_mask'][0, :16])\n",
"print('Shape Type Ids : ', text_preprocessed['input_type_ids'].shape)\n",
"print('Type Ids : ', text_preprocessed['input_type_ids'][0, :16])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KyBsEhoa0_7r"
},
"source": [
"Here are some details to pay attention to:\n",
"- `input_mask` The mask allows the model to cleanly differentiate between the content and the padding. The mask has the same shape as the `input_word_ids`, and contains a 1 anywhere the `input_word_ids` is not padding.\n",
"- `input_type_ids` has the same shape as `input_mask`, but inside the non-padded region, contains a 0 or a 1 indicating which sentence the token is a part of."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "H63KFuKcRwjO"
},
"source": [
"Next, you will create a preprocessing model that encapsulates all this logic. Your model will take strings as input, and return appropriately formatted objects which can be passed to BERT.\n",
"\n",
"Each BERT model has a specific preprocessing model, make sure to use the proper one described on the BERT's model documentation.\n",
"\n",
"Note: BERT adds a \"position embedding\" to the token embedding of each input, and these come from a fixed-size lookup table. That imposes a max seq length of 512 (which is also a practical limit, due to the quadratic growth of attention computation). For this Colab 128 is good enough."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "KeHEYKXGqjAZ"
},
"outputs": [],
"source": [
"def make_bert_preprocess_model(sentence_features, seq_length=128):\n",
" \"\"\"Returns Model mapping string features to BERT inputs.\n",
"\n",
" Args:\n",
" sentence_features: a list with the names of string-valued features.\n",
" seq_length: an integer that defines the sequence length of BERT inputs.\n",
"\n",
" Returns:\n",
" A Keras Model that can be called on a list or dict of string Tensors\n",
" (with the order or names, resp., given by sentence_features) and\n",
" returns a dict of tensors for input to BERT.\n",
" \"\"\"\n",
"\n",
" input_segments = [\n",
" tf.keras.layers.Input(shape=(), dtype=tf.string, name=ft)\n",
" for ft in sentence_features]\n",
"\n",
" # Tokenize the text to word pieces.\n",
" bert_preprocess = hub.load(tfhub_handle_preprocess)\n",
" tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')\n",
" segments = [tokenizer(s) for s in input_segments]\n",
"\n",
" # Optional: Trim segments in a smart way to fit seq_length.\n",
" # Simple cases (like this example) can skip this step and let\n",
" # the next step apply a default truncation to approximately equal lengths.\n",
" truncated_segments = segments\n",
"\n",
" # Pack inputs. The details (start/end token ids, dict of output tensors)\n",
" # are model-dependent, so this gets loaded from the SavedModel.\n",
" packer = hub.KerasLayer(bert_preprocess.bert_pack_inputs,\n",
" arguments=dict(seq_length=seq_length),\n",
" name='packer')\n",
" model_inputs = packer(truncated_segments)\n",
" return tf.keras.Model(input_segments, model_inputs)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kk5SS1bStmfP"
},
"source": [
"Let's demonstrate the preprocessing model. You will create a test with two sentences input (input1 and input2). The output is what a BERT model would expect as input: `input_word_ids`, `input_masks` and `input_type_ids`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BehJu3wLtAg-"
},
"outputs": [],
"source": [
"test_preprocess_model = make_bert_preprocess_model(['my_input1', 'my_input2'])\n",
"test_text = [np.array(['some random test sentence']),\n",
" np.array(['another sentence'])]\n",
"text_preprocessed = test_preprocess_model(test_text)\n",
"\n",
"print('Keys : ', list(text_preprocessed.keys()))\n",
"print('Shape Word Ids : ', text_preprocessed['input_word_ids'].shape)\n",
"print('Word Ids : ', text_preprocessed['input_word_ids'][0, :16])\n",
"print('Shape Mask : ', text_preprocessed['input_mask'].shape)\n",
"print('Input Mask : ', text_preprocessed['input_mask'][0, :16])\n",
"print('Shape Type Ids : ', text_preprocessed['input_type_ids'].shape)\n",
"print('Type Ids : ', text_preprocessed['input_type_ids'][0, :16])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qXU6bQWmNfhp"
},
"source": [
"Let's take a look at the model's structure, paying attention to the two inputs you just defined."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "a2_XrcVPFiz_"
},
"outputs": [],
"source": [
"tf.keras.utils.plot_model(test_preprocess_model, show_shapes=True, show_dtype=True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GRVVol5G9i0b"
},
"source": [
"To apply the preprocessing in all the inputs from the dataset, you will use the `map` function from the dataset. The result is then cached for [performance](https://www.tensorflow.org/guide/data_performance#top_of_page)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1zhR-SVwx4_J"
},
"outputs": [],
"source": [
"AUTOTUNE = tf.data.AUTOTUNE\n",
"\n",
"\n",
"def load_dataset_from_tfds(in_memory_ds, info, split, batch_size,\n",
" bert_preprocess_model):\n",
" is_training = split.startswith('train')\n",
" dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[split])\n",
" num_examples = info.splits[split].num_examples\n",
"\n",
" if is_training:\n",
" dataset = dataset.shuffle(num_examples)\n",
" dataset = dataset.repeat()\n",
" dataset = dataset.batch(batch_size)\n",
" dataset = dataset.map(lambda ex: (bert_preprocess_model(ex), ex['label']))\n",
" dataset = dataset.cache().prefetch(buffer_size=AUTOTUNE)\n",
" return dataset, num_examples"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pDNKfAXbDnJH"
},
"source": [
"## Define your model\n",
"\n",
"You are now ready to define your model for sentence or sentence pair classification by feeding the preprocessed inputs through the BERT encoder and putting a linear classifier on top (or other arrangement of layers as you prefer), and using dropout for regularization."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aksj743St9ga"
},
"outputs": [],
"source": [
"def build_classifier_model(num_classes):\n",
"\n",
" class Classifier(tf.keras.Model):\n",
" def __init__(self, num_classes):\n",
" super(Classifier, self).__init__(name=\"prediction\")\n",
" self.encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True)\n",
" self.dropout = tf.keras.layers.Dropout(0.1)\n",
" self.dense = tf.keras.layers.Dense(num_classes)\n",
"\n",
" def call(self, preprocessed_text):\n",
" encoder_outputs = self.encoder(preprocessed_text)\n",
" pooled_output = encoder_outputs[\"pooled_output\"]\n",
" x = self.dropout(pooled_output)\n",
" x = self.dense(x)\n",
" return x\n",
"\n",
" model = Classifier(num_classes)\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TTa5VZssizDm"
},
"source": [
"Let's try running the model on some preprocessed inputs."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e6mJ3WXhwUa8"
},
"outputs": [],
"source": [
"test_classifier_model = build_classifier_model(2)\n",
"bert_raw_result = test_classifier_model(text_preprocessed)\n",
"print(tf.sigmoid(bert_raw_result))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1s0xpHS-XQcP"
},
"source": [
"## Choose a task from GLUE\n",
"\n",
"You are going to use a TensorFlow DataSet from the [GLUE](https://www.tensorflow.org/datasets/catalog/glue) benchmark suite.\n",
"\n",
"Colab lets you download these small datasets to the local filesystem, and the code below reads them entirely into memory, because the separate TPU worker host cannot access the local filesystem of the colab runtime.\n",
"\n",
"For bigger datasets, you'll need to create your own [Google Cloud Storage](https://cloud.google.com/storage) bucket and have the TPU worker read the data from there. You can learn more in the [TPU guide](https://www.tensorflow.org/guide/tpu#input_datasets).\n",
"\n",
"It's recommended to start with the CoLa dataset (for single sentence) or MRPC (for multi sentence) since these are small and don't take long to fine tune."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "RhL__V2mwRNH"
},
"outputs": [],
"source": [
"tfds_name = 'glue/cola' #@param ['glue/cola', 'glue/sst2', 'glue/mrpc', 'glue/qqp', 'glue/mnli', 'glue/qnli', 'glue/rte', 'glue/wnli']\n",
"\n",
"tfds_info = tfds.builder(tfds_name).info\n",
"\n",
"sentence_features = list(tfds_info.features.keys())\n",
"sentence_features.remove('idx')\n",
"sentence_features.remove('label')\n",
"\n",
"available_splits = list(tfds_info.splits.keys())\n",
"train_split = 'train'\n",
"validation_split = 'validation'\n",
"test_split = 'test'\n",
"if tfds_name == 'glue/mnli':\n",
" validation_split = 'validation_matched'\n",
" test_split = 'test_matched'\n",
"\n",
"num_classes = tfds_info.features['label'].num_classes\n",
"num_examples = tfds_info.splits.total_num_examples\n",
"\n",
"print(f'Using {tfds_name} from TFDS')\n",
"print(f'This dataset has {num_examples} examples')\n",
"print(f'Number of classes: {num_classes}')\n",
"print(f'Features {sentence_features}')\n",
"print(f'Splits {available_splits}')\n",
"\n",
"with tf.device('/job:localhost'):\n",
" # batch_size=-1 is a way to load the dataset into memory\n",
" in_memory_ds = tfds.load(tfds_name, batch_size=-1, shuffle_files=True)\n",
"\n",
"# The code below is just to show some samples from the selected dataset\n",
"print(f'Here are some sample rows from {tfds_name} dataset')\n",
"sample_dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[train_split])\n",
"\n",
"labels_names = tfds_info.features['label'].names\n",
"print(labels_names)\n",
"print()\n",
"\n",
"sample_i = 1\n",
"for sample_row in sample_dataset.take(5):\n",
" samples = [sample_row[feature] for feature in sentence_features]\n",
" print(f'sample row {sample_i}')\n",
" for sample in samples:\n",
" print(sample.numpy())\n",
" sample_label = sample_row['label']\n",
"\n",
" print(f'label: {sample_label} ({labels_names[sample_label]})')\n",
" print()\n",
" sample_i += 1"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lFhjoYtsoVNF"
},
"source": [
"The dataset also determines the problem type (classification or regression) and the appropriate loss function for training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "OWPOZE-L3AgE"
},
"outputs": [],
"source": [
"def get_configuration(glue_task):\n",
"\n",
" loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
"\n",
" if glue_task == 'glue/cola':\n",
" metrics = tfa.metrics.MatthewsCorrelationCoefficient(num_classes=2)\n",
" else:\n",
" metrics = tf.keras.metrics.SparseCategoricalAccuracy(\n",
" 'accuracy', dtype=tf.float32)\n",
"\n",
" return metrics, loss"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EgJoTvo2DsWq"
},
"source": [
"## Train your model\n",
"\n",
"Finally, you can train the model end-to-end on the dataset you chose.\n",
"\n",
"### Distribution\n",
"\n",
"Recall the set-up code at the top, which has connected the colab runtime to\n",
"a TPU worker with multiple TPU devices. To distribute training onto them, you will create and compile your main Keras model within the scope of the TPU distribution strategy. (For details, see [Distributed training with Keras](https://www.tensorflow.org/tutorials/distribute/keras).)\n",
"\n",
"Preprocessing, on the other hand, runs on the CPU of the worker host, not the TPUs, so the Keras model for preprocessing as well as the training and validation datasets mapped with it are built outside the distribution strategy scope. The call to `Model.fit()` will take care of distributing the passed-in dataset to the model replicas.\n",
"\n",
"Note: The single TPU worker host already has the resource objects (think: a lookup table) needed for tokenization. Scaling up to multiple workers requires use of `Strategy.experimental_distribute_datasets_from_function` with a function that loads the preprocessing model separately onto each worker.\n",
"\n",
"### Optimizer\n",
"\n",
"Fine-tuning follows the optimizer set-up from BERT pre-training (as in [Classify text with BERT](https://www.tensorflow.org/text/tutorials/classify_text_with_bert)): It uses the AdamW optimizer with a linear decay of a notional initial learning rate, prefixed with a linear warm-up phase over the first 10% of training steps (`num_warmup_steps`). In line with the BERT paper, the initial learning rate is smaller for fine-tuning (best of 5e-5, 3e-5, 2e-5)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AiU5_ioh_fEr"
},
"outputs": [],
"source": [
"epochs = 3\n",
"batch_size = 32\n",
"init_lr = 2e-5\n",
"\n",
"print(f'Fine tuning {tfhub_handle_encoder} model')\n",
"bert_preprocess_model = make_bert_preprocess_model(sentence_features)\n",
"\n",
"with strategy.scope():\n",
"\n",
" # metric have to be created inside the strategy scope\n",
" metrics, loss = get_configuration(tfds_name)\n",
"\n",
" train_dataset, train_data_size = load_dataset_from_tfds(\n",
" in_memory_ds, tfds_info, train_split, batch_size, bert_preprocess_model)\n",
" steps_per_epoch = train_data_size // batch_size\n",
" num_train_steps = steps_per_epoch * epochs\n",
" num_warmup_steps = num_train_steps // 10\n",
"\n",
" validation_dataset, validation_data_size = load_dataset_from_tfds(\n",
" in_memory_ds, tfds_info, validation_split, batch_size,\n",
" bert_preprocess_model)\n",
" validation_steps = validation_data_size // batch_size\n",
"\n",
" classifier_model = build_classifier_model(num_classes)\n",
"\n",
" optimizer = optimization.create_optimizer(\n",
" init_lr=init_lr,\n",
" num_train_steps=num_train_steps,\n",
" num_warmup_steps=num_warmup_steps,\n",
" optimizer_type='adamw')\n",
"\n",
" classifier_model.compile(optimizer=optimizer, loss=loss, metrics=[metrics])\n",
"\n",
" classifier_model.fit(\n",
" x=train_dataset,\n",
" validation_data=validation_dataset,\n",
" steps_per_epoch=steps_per_epoch,\n",
" epochs=epochs,\n",
" validation_steps=validation_steps)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Rtn7jewb6dg4"
},
"source": [
"## Export for inference\n",
"\n",
"You will create a final model that has the preprocessing part and the fine-tuned BERT we've just created.\n",
"\n",
"At inference time, preprocessing needs to be part of the model (because there is no longer a separate input queue as for training data that does it). Preprocessing is not just computation; it has its own resources (the vocab table) that must be attached to the Keras Model that is saved for export.\n",
"This final assembly is what will be saved.\n",
"\n",
"You are going to save the model on colab and later you can download to keep it for the future (**View -\u003e Table of contents -\u003e Files**).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ShcvqJAgVera"
},
"outputs": [],
"source": [
"main_save_path = './my_models'\n",
"bert_type = tfhub_handle_encoder.split('/')[-2]\n",
"saved_model_name = f'{tfds_name.replace(\"/\", \"_\")}_{bert_type}'\n",
"\n",
"saved_model_path = os.path.join(main_save_path, saved_model_name)\n",
"\n",
"preprocess_inputs = bert_preprocess_model.inputs\n",
"bert_encoder_inputs = bert_preprocess_model(preprocess_inputs)\n",
"bert_outputs = classifier_model(bert_encoder_inputs)\n",
"model_for_export = tf.keras.Model(preprocess_inputs, bert_outputs)\n",
"\n",
"print('Saving', saved_model_path)\n",
"\n",
"# Save everything on the Colab host (even the variables from TPU memory)\n",
"save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')\n",
"model_for_export.save(saved_model_path, include_optimizer=False,\n",
" options=save_options)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "C2qyM9Q9z12v"
},
"source": [
"## Test the model\n",
"\n",
"The final step is testing the results of your exported model.\n",
"\n",
"Just to make some comparison, let's reload the model and test it using some inputs from the test split from the dataset.\n",
"\n",
"Note: The test is done on the colab host, not the TPU worker that it has connected to, so it appears below with explicit device placements. You can omit those when loading the SavedModel elsewhere."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BhI0_W0kbXji"
},
"outputs": [],
"source": [
"with tf.device('/job:localhost'):\n",
" reloaded_model = tf.saved_model.load(saved_model_path)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "4yl-CEcDDXzX"
},
"outputs": [],
"source": [
"#@title Utility methods\n",
"\n",
"def prepare(record):\n",
" model_inputs = [[record[ft]] for ft in sentence_features]\n",
" return model_inputs\n",
"\n",
"\n",
"def prepare_serving(record):\n",
" model_inputs = {ft: record[ft] for ft in sentence_features}\n",
" return model_inputs\n",
"\n",
"\n",
"def print_bert_results(test, bert_result, dataset_name):\n",
"\n",
" bert_result_class = tf.argmax(bert_result, axis=1)[0]\n",
"\n",
" if dataset_name == 'glue/cola':\n",
" print('sentence:', test[0].numpy())\n",
" if bert_result_class == 1:\n",
" print('This sentence is acceptable')\n",
" else:\n",
" print('This sentence is unacceptable')\n",
"\n",
" elif dataset_name == 'glue/sst2':\n",
" print('sentence:', test[0])\n",
" if bert_result_class == 1:\n",
" print('This sentence has POSITIVE sentiment')\n",
" else:\n",
" print('This sentence has NEGATIVE sentiment')\n",
"\n",
" elif dataset_name == 'glue/mrpc':\n",
" print('sentence1:', test[0])\n",
" print('sentence2:', test[1])\n",
" if bert_result_class == 1:\n",
" print('Are a paraphrase')\n",
" else:\n",
" print('Are NOT a paraphrase')\n",
"\n",
" elif dataset_name == 'glue/qqp':\n",
" print('question1:', test[0])\n",
" print('question2:', test[1])\n",
" if bert_result_class == 1:\n",
" print('Questions are similar')\n",
" else:\n",
" print('Questions are NOT similar')\n",
"\n",
" elif dataset_name == 'glue/mnli':\n",
" print('premise :', test[0])\n",
" print('hypothesis:', test[1])\n",
" if bert_result_class == 1:\n",
" print('This premise is NEUTRAL to the hypothesis')\n",
" elif bert_result_class == 2:\n",
" print('This premise CONTRADICTS the hypothesis')\n",
" else:\n",
" print('This premise ENTAILS the hypothesis')\n",
"\n",
" elif dataset_name == 'glue/qnli':\n",
" print('question:', test[0])\n",
" print('sentence:', test[1])\n",
" if bert_result_class == 1:\n",
" print('The question is NOT answerable by the sentence')\n",
" else:\n",
" print('The question is answerable by the sentence')\n",
"\n",
" elif dataset_name == 'glue/rte':\n",
" print('sentence1:', test[0])\n",
" print('sentence2:', test[1])\n",
" if bert_result_class == 1:\n",
" print('Sentence1 DOES NOT entails sentence2')\n",
" else:\n",
" print('Sentence1 entails sentence2')\n",
"\n",
" elif dataset_name == 'glue/wnli':\n",
" print('sentence1:', test[0])\n",
" print('sentence2:', test[1])\n",
" if bert_result_class == 1:\n",
" print('Sentence1 DOES NOT entails sentence2')\n",
" else:\n",
" print('Sentence1 entails sentence2')\n",
"\n",
" print('BERT raw results:', bert_result[0])\n",
" print()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "12VA4BcKuR7n"
},
"source": [
"### Test"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "dt-O94gcwbIi"
},
"outputs": [],
"source": [
"with tf.device('/job:localhost'):\n",
" test_dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[test_split])\n",
" for test_row in test_dataset.shuffle(1000).map(prepare).take(5):\n",
" if len(sentence_features) == 1:\n",
" result = reloaded_model(test_row[0])\n",
" else:\n",
" result = reloaded_model(list(test_row))\n",
"\n",
" print_bert_results(test_row, result, tfds_name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3cOmih754Y_M"
},
"source": [
"If you want to use your model on [TF Serving](https://www.tensorflow.org/tfx/guide/serving), remember that it will call your SavedModel through one of its named signatures. Notice there are some small differences in the input. In Python, you can test them as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "b0vTQAXKN_K0"
},
"outputs": [],
"source": [
"with tf.device('/job:localhost'):\n",
" serving_model = reloaded_model.signatures['serving_default']\n",
" for test_row in test_dataset.shuffle(1000).map(prepare_serving).take(5):\n",
" result = serving_model(**test_row)\n",
" # The 'prediction' key is the classifier's defined model name.\n",
" print_bert_results(list(test_row.values()), result['prediction'], tfds_name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GOA5bX2g3wCW"
},
"source": [
"You did it! Your saved model could be used for serving or simple inference in a process, with a simpler api with less code and easier to maintain.\n",
"\n",
"## Next Steps\n",
"\n",
"Now that you've tried one of the base BERT models, you can try other ones to achieve more accuracy or maybe with smaller model versions.\n",
"\n",
"You can also try in other datasets."
]
}
],
"metadata": {
"accelerator": "TPU",
"colab": {
"collapsed_sections": [],
"name": "bert_glue.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
