bert_embeddings_with_tensorflow_2_0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2IjSWx7-O8yY"
   },
   "source": [
    "# BERT Embeddings with TensorFlow 2.0\n",
    "With the new release of TensorFlow, this Notebook aims to show a simple use of the BERT model.\n",
    "- See BERT on paper: https://arxiv.org/pdf/1810.04805.pdf\n",
    "- See BERT on GitHub: https://github.com/google-research/bert\n",
    "- See BERT on TensorHub: https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1\n",
    "- See 'old' use of BERT for comparison: https://colab.research.google.com/github/google-research/bert/blob/master/predicting_movie_reviews_with_bert_on_tf_hub.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SQktrOSAPq_n"
   },
   "source": [
    "## Install Libraries\n",
    "\n",
    "We need Tensorflow 2.0 and TensorHub 0.7 for this notebook\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_Iwew0KP8vRM"
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow\n",
    "!pip install tensorflow_hub\n",
    "!pip install bert-for-tf2\n",
    "!pip install sentencepiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 50
    },
    "colab_type": "code",
    "id": "NV-yr9ulP_E-",
    "outputId": "c37c1b21-1aa5-456e-ddc4-446636805189"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TF version:  2.0.0\n",
      "Hub version:  0.8.0.dev\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "print(\"TF version: \", tf.__version__)\n",
    "print(\"Hub version: \", hub.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4svE1x7iP3af"
   },
   "source": [
    "If TensorFlow Hub is not 0.7 yet on release, use dev:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "id": "pUo1k6rd9iaP"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tf-hub-nightly in /home/ayesha/anaconda3/lib/python3.7/site-packages (0.8.0.dev201912310004)\r\n",
      "Requirement already satisfied: numpy>=1.12.0 in /home/ayesha/anaconda3/lib/python3.7/site-packages (from tf-hub-nightly) (1.16.2)\r\n",
      "Requirement already satisfied: six>=1.12.0 in /home/ayesha/anaconda3/lib/python3.7/site-packages (from tf-hub-nightly) (1.12.0)\r\n",
      "Requirement already satisfied: protobuf>=3.8.0 in /home/ayesha/anaconda3/lib/python3.7/site-packages (from tf-hub-nightly) (3.10.1)\r\n",
      "Requirement already satisfied: setuptools in /home/ayesha/anaconda3/lib/python3.7/site-packages (from protobuf>=3.8.0->tf-hub-nightly) (41.6.0.post20191030)\r\n"
     ]
    }
   ],
   "source": [
    "!pip install tf-hub-nightly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "av1oBm8m-Ajz"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.8.0.dev'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hub.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cMeXU54uQUew"
   },
   "source": [
    "## Import modules\n",
    "\n",
    "First we will import the modules to be use in the following code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KBfktvAc-CNb"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import Model \n",
    "import bert\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now these TensorFlow and bert libraries are imported, now its time to import the BERT's tokenizer. The methodology on which BERT was trained using the WordPiece tokenization. It means that a word can be broken down into more than one sub-words.\n",
    "\n",
    "For example, of we have the input sentence like this\n",
    "\n",
    "input_sentence= 'Hi we are using BERT'\n",
    "\n",
    "### OUTPUT\n",
    "['hi', 'we' , 'are', 'using', 'bert']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FullTokenizer = bert.bert_tokenization.FullTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tU2OpvYrRFNf"
   },
   "source": [
    "Building model using tf.keras and hub. from sentences to embeddings.\n",
    "\n",
    "#### max_seq_length = 128 \n",
    "\n",
    "BERT has a constraint on the maximum length of a sequence after tokenizing. For any BERT model, the maximum sequence length after tokenization is 512. But we can set any sequence length equal to or below this value.\n",
    "\n",
    "\n",
    "### Inputs:\n",
    " - input token ids (tokenizer converts tokens using vocab file)\n",
    " - input masks (1 for useful tokens, 0 for padding)\n",
    " - segment ids (for 2 text training: 0 for the first one, 1 for the second one)\n",
    "\n",
    "### Outputs:\n",
    " - pooled_output of shape `[batch_size, 768]` with representations for the entire input sequences\n",
    " - sequence_output of shape `[batch_size, max_seq_length, 768]` with representations for each input token (in context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IW6V3afD-q1K"
   },
   "outputs": [],
   "source": [
    "max_seq_length = 128  # Your choice here.\n",
    "input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,\n",
    "                                       name=\"input_word_ids\")\n",
    "input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,\n",
    "                                   name=\"input_mask\")\n",
    "segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,\n",
    "                                    name=\"segment_ids\")\n",
    "bert_layer = hub.KerasLayer(\"https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1\",\n",
    "                            trainable=True)\n",
    "pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gmR3jHYE_y3X"
   },
   "outputs": [],
   "source": [
    "model = Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=[pooled_output, sequence_output])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BFDpzy1-STOh"
   },
   "source": [
    "Generating segments and masks based on the original BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "e4Y_r3lmFO1E"
   },
   "outputs": [],
   "source": [
    "# See BERT paper: https://arxiv.org/pdf/1810.04805.pdf\n",
    "# And BERT implementation convert_single_example() at https://github.com/google-research/bert/blob/master/run_classifier.py\n",
    "\n",
    "def get_masks(tokens, max_seq_length):\n",
    "    \"\"\"Mask for padding\"\"\"\n",
    "    if len(tokens)>max_seq_length:\n",
    "        raise IndexError(\"Token length more than max seq length!\")\n",
    "    return [1]*len(tokens) + [0] * (max_seq_length - len(tokens))\n",
    "\n",
    "\n",
    "def get_segments(tokens, max_seq_length):\n",
    "    \"\"\"Segments: 0 for the first sequence, 1 for the second\"\"\"\n",
    "    if len(tokens)>max_seq_length:\n",
    "        raise IndexError(\"Token length more than max seq length!\")\n",
    "    segments = []\n",
    "    current_segment_id = 0\n",
    "    for token in tokens:\n",
    "        segments.append(current_segment_id)\n",
    "        if token == \"[SEP]\":\n",
    "            current_segment_id = 1\n",
    "    return segments + [0] * (max_seq_length - len(tokens))\n",
    "\n",
    "\n",
    "def get_ids(tokens, tokenizer, max_seq_length):\n",
    "    \"\"\"Token ids from Tokenizer vocab\"\"\"\n",
    "    token_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "    input_ids = token_ids + [0] * (max_seq_length-len(token_ids))\n",
    "    return input_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "44u2pruZSbMX"
   },
   "source": [
    "Import tokenizer using the original vocab file, do lower case all the word pieces and then tokenize the sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sm3lGfQb-1J8"
   },
   "outputs": [],
   "source": [
    "vocab_file = bert_layer.resolved_object.vocab_file.asset_path.numpy()\n",
    "do_lower_case = bert_layer.resolved_object.do_lower_case.numpy()\n",
    "tokenizer = FullTokenizer(vocab_file, do_lower_case)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Y5ZV6MtqSmLa"
   },
   "source": [
    "## Test BERT embedding generator model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jVkTaR0lCcCM"
   },
   "outputs": [],
   "source": [
    "s = \"Hi we are using BERT\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AihvrFWcSzd6"
   },
   "source": [
    "Tokenizing the sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X798BKV_Co71"
   },
   "outputs": [],
   "source": [
    "stokens = tokenizer.tokenize(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PxroAvbjStsX"
   },
   "source": [
    "Adding separator tokens according to the paper\n",
    "\n",
    "[CLS] provided by BERT for sentence embeddings without any combination or processing from all the word vectors in the sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "znzQLHURDAfs"
   },
   "outputs": [],
   "source": [
    "stokens = [\"[CLS]\"] + stokens + [\"[SEP]\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kRwyPuauTc2z"
   },
   "source": [
    "Get the model inputs from the tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PyyjWY75Flns"
   },
   "outputs": [],
   "source": [
    "input_ids = get_ids(stokens, tokenizer, max_seq_length)\n",
    "input_masks = get_masks(stokens, max_seq_length)\n",
    "input_segments = get_segments(stokens, max_seq_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 104
    },
    "colab_type": "code",
    "id": "FgKzlBloMX_Q",
    "outputId": "0d6b41e9-6db5-42fb-e6e1-59199350cf6c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[CLS]', 'hi', 'we', 'are', 'using', 'bert', '[SEP]']\n",
      "[101, 7632, 2057, 2024, 2478, 14324, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "print(stokens)\n",
    "print(input_ids)\n",
    "print(input_masks)\n",
    "print(input_segments)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mi2mj4EUTi0X"
   },
   "source": [
    "Generate Embeddings using the pretrained model\n",
    "\n",
    "#### The output contains:\n",
    "\n",
    "##### pooled_output: \n",
    "pooled output of the entire sequence with shape [batch_size, hidden_size].\n",
    "##### sequence_output: \n",
    "representations of every token in the input sequence with shape [batch_size, max_sequence_length, hidden_size]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ik3xqHqXM_lN"
   },
   "outputs": [],
   "source": [
    "pool_embs, all_embs = model.predict([[input_ids],[input_masks],[input_segments]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References: \n",
    "\n",
    "https://github.com/google-research/bert/blob/master/run_classifier.py\n",
    "\n",
    "https://colab.research.google.com/drive/1hMLd5-r82FrnFnBub-B-fVW78Px4KPX1\n",
    "\n",
    "https://tfhub.dev/google/bert_cased_L-12_H-768_A-12/1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "BERT Embeddings with TensorFlow 2.0",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}