{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-1 fw=all -->\n",
    "\n",
    "# Bidirectional Encoder Representations from Transformers (BERT)\n",
    "\n",
    "We have introduced several word embedding models for natural language understanding.\n",
    "After pretraining, the output can be thought of as a matrix\n",
    "where each row is a vector that represents a word of a predefined vocabulary.\n",
    "In fact, these word embedding models are all *context-independent*.\n",
    "Let's begin by illustrating this property.\n",
    "\n",
    "\n",
    "## From Context-Independent to Context-Sensitive\n",
    "\n",
    "Recall the experiments in that section and that section.\n",
    "For instance, word2vec and GloVe both assign the same pretrained vector to the same word regardless of the context of the word (if any).\n",
    "Formally, a context-independent representation of any token $x$\n",
    "is a function $f(x)$ that only takes $x$ as its input.\n",
    "Given the abundance of polysemy and complex semantics in natural languages,\n",
    "context-independent representations have obvious limitations.\n",
    "For instance, the word \"crane\" in contexts\n",
    "\"a crane is flying\" and \"a crane driver came\" has completely different meanings;\n",
    "thus, the same word may be assigned different representations depending on contexts.\n",
    "\n",
    "This motivates the development of *context-sensitive* word representations,\n",
    "where representations of words depend on their contexts.\n",
    "Hence, a context-sensitive representation of token $x$ is a function $f(x, c(x))$\n",
    "depending on both $x$ and its context $c(x)$.\n",
    "Popular context-sensitive representations\n",
    "include TagLM (language-model-augmented sequence tagger) [@Peters.Ammar.Bhagavatula.ea.2017],\n",
    "CoVe (Context Vectors) [@McCann.Bradbury.Xiong.ea.2017],\n",
    "and ELMo (Embeddings from Language Models) [@Peters.Neumann.Iyyer.ea.2018].\n",
    "\n",
    "For example, by taking the entire sequence as input,\n",
    "ELMo is a function that assigns a representation to each word from the input sequence.\n",
    "Specifically, ELMo combines all the intermediate layer representations from pretrained bidirectional LSTM as the output representation.\n",
    "Then the ELMo representation will be added to a downstream task's existing supervised model\n",
    "as additional features, such as by concatenating ELMo representation and the original representation (e.g., GloVe) of tokens in the existing model.\n",
    "On the one hand,\n",
    "all the weights in the pretrained bidirectional LSTM model are frozen after ELMo representations are added.\n",
    "On the other hand,\n",
    "the existing supervised model is specifically customized for a given task.\n",
    "Leveraging different best models for different tasks at that time,\n",
    "adding ELMo improved the state of the art across six natural language processing tasks:\n",
    "sentiment analysis, natural language inference,\n",
    "semantic role labeling, coreference resolution,\n",
    "named entity recognition, and question answering.\n",
    "\n",
    "\n",
    "## From Task-Specific to Task-Agnostic\n",
    "\n",
    "Although ELMo has significantly improved solutions to a diverse set of natural language processing tasks,\n",
    "each solution still hinges on a *task-specific* architecture.\n",
    "However, it is practically non-trivial to craft a specific architecture for every natural language processing task.\n",
    "The GPT (Generative Pre-Training) model represents an effort in designing\n",
    "a general *task-agnostic* model for context-sensitive representations [@Radford.Narasimhan.Salimans.ea.2018].\n",
    "Built on a Transformer decoder,\n",
    "GPT pretrains a language model that will be used to represent text sequences.\n",
    "When applying GPT to a downstream task,\n",
    "the output of the language model will be fed into an added linear output layer\n",
    "to predict the label of the task.\n",
    "In sharp contrast to ELMo that freezes parameters of the pretrained model,\n",
    "GPT fine-tunes *all* the parameters in the pretrained Transformer decoder\n",
    "during supervised learning of the downstream task.\n",
    "GPT was evaluated on twelve tasks of natural language inference,\n",
    "question answering, sentence similarity, and classification,\n",
    "and improved the state of the art in nine of them with minimal changes\n",
    "to the model architecture.\n",
    "\n",
    "However, due to the autoregressive nature of language models,\n",
    "GPT only looks forward (left-to-right).\n",
    "In contexts \"i went to the bank to deposit cash\" and \"i went to the bank to sit down\",\n",
    "as \"bank\" is sensitive to the context to its left,\n",
    "GPT will return the same representation for \"bank\",\n",
    "though it has different meanings.\n",
    "\n",
    "\n",
    "## BERT: Combining the Best of Both Worlds\n",
    "\n",
    "As we have seen,\n",
    "ELMo encodes context bidirectionally but uses task-specific architectures;\n",
    "while GPT is task-agnostic but encodes context left-to-right.\n",
    "Combining the best of both worlds,\n",
    "BERT (Bidirectional Encoder Representations from Transformers)\n",
    "encodes context bidirectionally and requires minimal architecture changes\n",
    "for a wide range of natural language processing tasks [@Devlin.Chang.Lee.ea.2018].\n",
    "Using a pretrained Transformer encoder,\n",
    "BERT is able to represent any token based on its bidirectional context.\n",
    "During supervised learning of downstream tasks,\n",
    "BERT is similar to GPT in two aspects.\n",
    "First, BERT representations will be fed into an added output layer,\n",
    "with minimal changes to the model architecture depending on nature of tasks,\n",
    "such as predicting for every token vs. predicting for the entire sequence.\n",
    "Second,\n",
    "all the parameters of the pretrained Transformer encoder are fine-tuned,\n",
    "while the additional output layer will be trained from scratch.\n",
    "the figure depicts the differences among ELMo, GPT, and BERT.\n",
    "\n",
    "![A comparison of ELMo, GPT, and BERT.](../img/elmo-gpt-bert.svg)\n",
    "\n",
    "\n",
    "BERT further improved the state of the art on eleven natural language processing tasks\n",
    "under broad categories of (i) single text classification (e.g., sentiment analysis), (ii) text pair classification (e.g., natural language inference),\n",
    "(iii) question answering, (iv) text tagging (e.g., named entity recognition).\n",
    "All proposed in 2018,\n",
    "from context-sensitive ELMo to task-agnostic GPT and BERT,\n",
    "conceptually simple yet empirically powerful pretraining of deep representations for natural languages have revolutionized solutions to various natural language processing tasks.\n",
    "\n",
    "In the rest of this chapter,\n",
    "we will dive into the pretraining of BERT.\n",
    "When natural language processing applications are explained in that section,\n",
    "we will illustrate fine-tuning of BERT for downstream applications."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-bert-combining-the-best-of-both-worlds"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "import optax\n",
    "import numpy as np"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-2 fw=all -->\n",
    "\n",
    "## Input Representation\n",
    "\n",
    "In natural language processing,\n",
    "some tasks (e.g., sentiment analysis) take single text as input,\n",
    "while in some other tasks (e.g., natural language inference),\n",
    "the input is a pair of text sequences.\n",
    "The BERT input sequence unambiguously represents both single text and text pairs.\n",
    "In the former,\n",
    "the BERT input sequence is the concatenation of\n",
    "the special classification token “&lt;cls&gt;”,\n",
    "tokens of a text sequence,\n",
    "and the special separation token “&lt;sep&gt;”.\n",
    "In the latter,\n",
    "the BERT input sequence is the concatenation of\n",
    "“&lt;cls&gt;”, tokens of the first text sequence,\n",
    "“&lt;sep&gt;”, tokens of the second text sequence, and “&lt;sep&gt;”.\n",
    "We will consistently distinguish the terminology \"BERT input sequence\"\n",
    "from other types of \"sequences\".\n",
    "For instance, one *BERT input sequence* may include either one *text sequence* or two *text sequences*.\n",
    "\n",
    "To distinguish text pairs,\n",
    "the learned segment embeddings $\\mathbf{e}_A$ and $\\mathbf{e}_B$\n",
    "are added to the token embeddings of the first sequence and the second sequence, respectively.\n",
    "For single text inputs, only $\\mathbf{e}_A$ is used.\n",
    "\n",
    "The following `get_tokens_and_segments` takes either one sentence or two sentences\n",
    "as input, then returns tokens of the BERT input sequence\n",
    "and their corresponding segment IDs."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-input-representation-1"
   },
   "source": [
    "\n",
    "def get_tokens_and_segments(tokens_a, tokens_b=None):\n",
    "    \"\"\"Get tokens of the BERT input sequence and their segment IDs.\"\"\"\n",
    "    tokens = ['<cls>'] + tokens_a + ['<sep>']\n",
    "    # 0 and 1 are marking segment A and B, respectively\n",
    "    segments = [0] * (len(tokens_a) + 2)\n",
    "    if tokens_b is not None:\n",
    "        tokens += tokens_b + ['<sep>']\n",
    "        segments += [1] * (len(tokens_b) + 1)\n",
    "    return tokens, segments"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-3 fw=all -->\n",
    "\n",
    "BERT chooses the Transformer encoder as its bidirectional architecture.\n",
    "Common in the Transformer encoder,\n",
    "positional embeddings are added at every position of the BERT input sequence.\n",
    "However, different from the original Transformer encoder,\n",
    "BERT uses *learnable* positional embeddings.\n",
    "To sum up, the figure shows that\n",
    "the embeddings of the BERT input sequence are the sum\n",
    "of the token embeddings, segment embeddings, and positional embeddings.\n",
    "\n",
    "![The embeddings of the BERT input sequence are the sum\n",
    "of the token embeddings, segment embeddings, and positional embeddings.](../img/bert-input.svg)\n",
    "\n",
    "The following `BERTEncoder` class is similar to the `TransformerEncoder` class\n",
    "as implemented in that section.\n",
    "Different from `TransformerEncoder`, `BERTEncoder` uses\n",
    "segment embeddings and learnable positional embeddings."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-input-representation-2"
   },
   "source": [
    "\n",
    "class BERTEncoder(nnx.Module):\n",
    "    \"\"\"BERT encoder.\"\"\"\n",
    "    def __init__(self, vocab_size, num_hiddens, ffn_num_hiddens, num_heads,\n",
    "                 num_blks, dropout, max_len=1000, rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs\n",
    "        self.token_embedding = nnx.Embed(vocab_size, num_hiddens, rngs=rngs)\n",
    "        self.segment_embedding = nnx.Embed(2, num_hiddens, rngs=rngs)\n",
    "        self.blks = nnx.List([d2l.TransformerEncoderBlock(\n",
    "            num_hiddens, ffn_num_hiddens, num_heads, dropout, True, rngs=rngs)\n",
    "            for _ in range(num_blks)])\n",
    "        # In BERT, positional embeddings are learnable, thus we create a\n",
    "        # parameter of positional embeddings that are long enough\n",
    "        self.pos_embedding = nnx.Param(\n",
    "            jax.random.normal(rngs.params(), (1, max_len, num_hiddens)) * 0.02)\n",
    "\n",
    "    def __call__(self, tokens, segments, valid_lens):\n",
    "        # Shape of `X` remains unchanged in the following code snippet:\n",
    "        # (batch size, max sequence length, `num_hiddens`)\n",
    "        X = self.token_embedding(tokens) + self.segment_embedding(segments)\n",
    "        X = X + self.pos_embedding[:, :X.shape[1], :]\n",
    "        for blk in self.blks:\n",
    "            X, _ = blk(X, valid_lens)\n",
    "        return X"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-4 fw=all -->\n",
    "\n",
    "Suppose that the vocabulary size is 10000.\n",
    "To demonstrate forward inference of `BERTEncoder`,\n",
    "let's create an instance of it and initialize its parameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-input-representation-3"
   },
   "source": [
    "vocab_size, num_hiddens, ffn_num_hiddens, num_heads = 10000, 768, 1024, 4\n",
    "ffn_num_input, num_blks, dropout = 768, 2, 0.2\n",
    "encoder = BERTEncoder(vocab_size, num_hiddens, ffn_num_hiddens, num_heads,\n",
    "                      num_blks, dropout)\n",
    "tokens = jnp.ones((2, 8), dtype=jnp.int32)\n",
    "segments = jnp.array([[0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 1, 1, 1, 1]])"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-5 fw=all -->\n",
    "\n",
    "We define `tokens` to be 2 BERT input sequences of length 8,\n",
    "where each token is an index of the vocabulary.\n",
    "The forward inference of `BERTEncoder` with the input `tokens`\n",
    "returns the encoded result where each token is represented by a vector\n",
    "whose length is predefined by the hyperparameter `num_hiddens`.\n",
    "This hyperparameter is usually referred to as the *hidden size*\n",
    "(number of hidden units) of the Transformer encoder."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-input-representation-4"
   },
   "source": [
    "tokens = jax.random.randint(jax.random.PRNGKey(0), (2, 8), 0, vocab_size)\n",
    "segments = jnp.array([[0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 1, 1, 1, 1]])\n",
    "encoded_X = nnx.view(encoder, deterministic=True)(tokens, segments, None)\n",
    "encoded_X.shape"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 8, 768)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-6 fw=all -->\n",
    "\n",
    "## Pretraining Tasks\n",
    "\n",
    "The forward inference of `BERTEncoder` gives the BERT representation\n",
    "of each token of the input text and the inserted\n",
    "special tokens “&lt;cls&gt;” and “&lt;sep&gt;”.\n",
    "Next, we will use these representations to compute the loss function\n",
    "for pretraining BERT.\n",
    "The pretraining is composed of the following two tasks:\n",
    "masked language modeling and next sentence prediction.\n",
    "\n",
    "### Masked Language Modeling\n",
    "\n",
    "As illustrated in that section,\n",
    "a language model predicts a token using the context on its left.\n",
    "To encode context bidirectionally for representing each token,\n",
    "BERT randomly masks tokens and uses tokens from the bidirectional context to\n",
    "predict the masked tokens in a self-supervised fashion.\n",
    "This task is referred to as a *masked language model*.\n",
    "\n",
    "In this pretraining task,\n",
    "15% of tokens will be selected at random as the masked tokens for prediction.\n",
    "To predict a masked token without cheating by using the label,\n",
    "one straightforward approach is to always replace it with a special “&lt;mask&gt;” token in the BERT input sequence.\n",
    "However, the artificial special token “&lt;mask&gt;” will never appear\n",
    "in fine-tuning.\n",
    "To avoid such a mismatch between pretraining and fine-tuning,\n",
    "if a token is masked for prediction (e.g., \"great\" is selected to be masked and predicted in \"this movie is great\"),\n",
    "in the input it will be replaced with:\n",
    "\n",
    "* a special “&lt;mask&gt;” token for 80% of the time (e.g., \"this movie is great\" becomes \"this movie is &lt;mask&gt;\");\n",
    "* a random token for 10% of the time (e.g., \"this movie is great\" becomes \"this movie is drink\");\n",
    "* the unchanged label token for 10% of the time (e.g., \"this movie is great\" becomes \"this movie is great\").\n",
    "\n",
    "Note that for 10% of 15% time a random token is inserted.\n",
    "This occasional noise encourages BERT to be less biased towards the masked token (especially when the label token remains unchanged) in its bidirectional context encoding.\n",
    "\n",
    "We implement the following `MaskLM` class to predict masked tokens\n",
    "in the masked language model task of BERT pretraining.\n",
    "The prediction uses a one-hidden-layer MLP (`self.mlp`).\n",
    "In forward inference, it takes two inputs:\n",
    "the encoded result of `BERTEncoder` and the token positions for prediction.\n",
    "The output is the prediction results at these positions."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-masked-language-modeling-1"
   },
   "source": [
    "\n",
    "class MaskLM(nnx.Module):\n",
    "    \"\"\"The masked language model task of BERT.\"\"\"\n",
    "    def __init__(self, vocab_size, num_hiddens, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.dense1 = nnx.Linear(num_hiddens, num_hiddens, rngs=rngs)\n",
    "        self.layer_norm = nnx.LayerNorm(num_hiddens, rngs=rngs)\n",
    "        self.dense2 = nnx.Linear(num_hiddens, vocab_size, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X, pred_positions):\n",
    "        num_pred_positions = pred_positions.shape[1]\n",
    "        pred_positions = pred_positions.reshape(-1)\n",
    "        batch_size = X.shape[0]\n",
    "        batch_idx = jnp.arange(0, batch_size)\n",
    "        # Suppose that `batch_size` = 2, `num_pred_positions` = 3, then\n",
    "        # `batch_idx` is `jnp.array([0, 0, 0, 1, 1, 1])`\n",
    "        batch_idx = jnp.repeat(batch_idx, num_pred_positions)\n",
    "        masked_X = X[batch_idx, pred_positions]\n",
    "        masked_X = masked_X.reshape((batch_size, num_pred_positions, -1))\n",
    "        mlm_Y_hat = self.dense1(masked_X)\n",
    "        mlm_Y_hat = nnx.relu(mlm_Y_hat)\n",
    "        mlm_Y_hat = self.layer_norm(mlm_Y_hat)\n",
    "        mlm_Y_hat = self.dense2(mlm_Y_hat)\n",
    "        return mlm_Y_hat"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-7 fw=all -->\n",
    "\n",
    "To demonstrate the forward inference of `MaskLM`,\n",
    "we create its instance `mlm` and initialize it.\n",
    "Recall that `encoded_X` from the forward inference of `BERTEncoder`\n",
    "represents 2 BERT input sequences.\n",
    "We define `mlm_positions` as the 3 indices to predict in either BERT input sequence of `encoded_X`.\n",
    "The forward inference of `mlm` returns prediction results `mlm_Y_hat`\n",
    "at all the masked positions `mlm_positions` of `encoded_X`.\n",
    "For each prediction, the size of the result is equal to the vocabulary size."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-masked-language-modeling-2"
   },
   "source": [
    "mlm = MaskLM(vocab_size, num_hiddens)\n",
    "mlm_positions = jnp.array([[1, 5, 2], [6, 1, 5]])\n",
    "mlm_Y_hat = mlm(encoded_X, mlm_positions)\n",
    "mlm_Y_hat.shape"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 3, 10000)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-8 fw=all -->\n",
    "\n",
    "With the ground truth labels `mlm_Y` of the predicted tokens `mlm_Y_hat` under masks,\n",
    "we can calculate the cross-entropy loss of the masked language model task in BERT pretraining."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-masked-language-modeling-3"
   },
   "source": [
    "mlm_Y = jnp.array([[7, 8, 9], [10, 20, 30]])\n",
    "mlm_l = optax.softmax_cross_entropy_with_integer_labels(\n",
    "    mlm_Y_hat.reshape((-1, vocab_size)), mlm_Y.reshape(-1))\n",
    "mlm_l.shape"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(6,)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-9 fw=all -->\n",
    "\n",
    "### Next Sentence Prediction\n",
    "\n",
    "Although masked language modeling is able to encode bidirectional context\n",
    "for representing words, it does not explicitly model the logical relationship\n",
    "between text pairs.\n",
    "To help understand the relationship between two text sequences,\n",
    "BERT considers a binary classification task, *next sentence prediction*, in its pretraining.\n",
    "When generating sentence pairs for pretraining,\n",
    "for half of the time they are indeed consecutive sentences with the label \"True\";\n",
    "while for the other half of the time the second sentence is randomly sampled from the corpus with the label \"False\".\n",
    "\n",
    "The following `NextSentencePred` class uses a one-hidden-layer MLP\n",
    "to predict whether the second sentence is the next sentence of the first\n",
    "in the BERT input sequence.\n",
    "Due to self-attention in the Transformer encoder,\n",
    "the BERT representation of the special token “&lt;cls&gt;”\n",
    "encodes both the two sentences from the input.\n",
    "Hence, the output layer (`self.output`) of the MLP classifier takes `X` as input,\n",
    "where `X` is the output of the MLP hidden layer whose input is the encoded “&lt;cls&gt;” token."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-next-sentence-prediction-1"
   },
   "source": [
    "\n",
    "class NextSentencePred(nnx.Module):\n",
    "    \"\"\"The next sentence prediction task of BERT.\"\"\"\n",
    "    def __init__(self, num_hiddens, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.output = nnx.Linear(num_hiddens, 2, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        # `X` shape: (batch size, `num_hiddens`)\n",
    "        return self.output(X)"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-10 fw=all -->\n",
    "\n",
    "We can see that the forward inference of an `NextSentencePred` instance\n",
    "returns binary predictions for each BERT input sequence."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-next-sentence-prediction-2"
   },
   "source": [
    "# Use the `<cls>` token (index 0) as input to NSP\n",
    "# input_shape for NSP: (batch size, `num_hiddens`)\n",
    "cls_X = encoded_X[:, 0, :]\n",
    "nsp = NextSentencePred(num_hiddens)\n",
    "nsp_Y_hat = nsp(cls_X)\n",
    "nsp_Y_hat.shape"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 2)"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-11 fw=all -->\n",
    "\n",
    "The cross-entropy loss of the 2 binary classifications can also be computed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-next-sentence-prediction-3"
   },
   "source": [
    "nsp_y = jnp.array([0, 1])\n",
    "nsp_l = optax.softmax_cross_entropy_with_integer_labels(nsp_Y_hat, nsp_y)\n",
    "nsp_l.shape"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2,)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-12 fw=all -->\n",
    "\n",
    "It is noteworthy that all the labels in both the aforementioned pretraining tasks\n",
    "can be trivially obtained from the pretraining corpus without manual labeling effort.\n",
    "The original BERT has been pretrained on the concatenation of BookCorpus [@Zhu.Kiros.Zemel.ea.2015]\n",
    "and English Wikipedia.\n",
    "These two text corpora are huge:\n",
    "they have 800 million words and 2.5 billion words, respectively.\n",
    "\n",
    "\n",
    "## Putting It All Together\n",
    "\n",
    "When pretraining BERT, the final loss function is a linear combination of\n",
    "both the loss functions for masked language modeling and next sentence prediction.\n",
    "Now we can define the `BERTModel` class by instantiating the three classes\n",
    "`BERTEncoder`, `MaskLM`, and `NextSentencePred`.\n",
    "The forward inference returns the encoded BERT representations `encoded_X`,\n",
    "predictions of masked language modeling `mlm_Y_hat`,\n",
    "and next sentence predictions `nsp_Y_hat`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-putting-it-all-together"
   },
   "source": [
    "\n",
    "class BERTModel(nnx.Module):\n",
    "    \"\"\"The BERT model.\"\"\"\n",
    "    def __init__(self, vocab_size, num_hiddens, ffn_num_hiddens,\n",
    "                 num_heads, num_blks, dropout, max_len=1000, rngs=None):\n",
    "        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs\n",
    "        self.encoder = BERTEncoder(\n",
    "            vocab_size, num_hiddens, ffn_num_hiddens, num_heads, num_blks,\n",
    "            dropout, max_len=max_len, rngs=rngs)\n",
    "        self.hidden = nnx.Linear(num_hiddens, num_hiddens, rngs=rngs)\n",
    "        self.mlm = MaskLM(vocab_size, num_hiddens, rngs=rngs)\n",
    "        self.nsp = NextSentencePred(num_hiddens, rngs=rngs)\n",
    "\n",
    "    def __call__(self, tokens, segments, valid_lens=None, pred_positions=None,\n",
    "                 training=None):\n",
    "        encoded_X = self.encoder(tokens, segments, valid_lens)\n",
    "        if pred_positions is not None:\n",
    "            mlm_Y_hat = self.mlm(encoded_X, pred_positions)\n",
    "        else:\n",
    "            mlm_Y_hat = None\n",
    "        # The hidden layer of the MLP classifier for next sentence prediction.\n",
    "        # 0 is the index of the '<cls>' token\n",
    "        nsp_Y_hat = self.nsp(\n",
    "            jnp.tanh(self.hidden(encoded_X[:, 0, :])))\n",
    "        return encoded_X, mlm_Y_hat, nsp_Y_hat"
   ],
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-13 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Word embedding models such as word2vec and GloVe are context-independent. They assign the same pretrained vector to the same word regardless of the context of the word (if any). It is hard for them to handle well polysemy or complex semantics in natural languages.\n",
    "* For context-sensitive word representations such as ELMo and GPT, representations of words depend on their contexts.\n",
    "* ELMo encodes context bidirectionally but uses task-specific architectures (however, it is practically non-trivial to craft a specific architecture for every natural language processing task); while GPT is task-agnostic but encodes context left-to-right.\n",
    "* BERT combines the best of both worlds: it encodes context bidirectionally and requires minimal architecture changes for a wide range of natural language processing tasks.\n",
    "* The embeddings of the BERT input sequence are the sum of the token embeddings, segment embeddings, and positional embeddings.\n",
    "* Pretraining BERT is composed of two tasks: masked language modeling and next sentence prediction. The former is able to encode bidirectional context for representing words, while the latter explicitly models the logical relationship between text pairs.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. All other things being equal, will a masked language model require more or fewer pretraining steps to converge than a left-to-right language model? Why?\n",
    "1. In the original implementation of BERT, the positionwise feed-forward network in `BERTEncoder` (via `d2l.TransformerEncoderBlock`) and the fully connected layer in `MaskLM` both use the Gaussian error linear unit (GELU) [@Hendrycks.Gimpel.2016] as the activation function. Research into the difference between GELU and ReLU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-md-14 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1490)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
