{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-1 fw=all -->\n",
    "\n",
    "# Natural Language Inference: Fine-Tuning BERT\n",
    "\n",
    "In earlier sections of this chapter,\n",
    "we have designed an attention-based architecture\n",
    "(in that section)\n",
    "for the natural language inference task\n",
    "on the SNLI dataset (as described in that section).\n",
    "Now we revisit this task by fine-tuning BERT.\n",
    "As discussed in that section,\n",
    "natural language inference is a sequence-level text pair classification problem,\n",
    "and fine-tuning BERT only requires an additional MLP-based architecture,\n",
    "as illustrated in the figure.\n",
    "\n",
    "![This section feeds pretrained BERT to an MLP-based architecture for natural language inference.](../img/nlp-map-nli-bert.svg)\n",
    "\n",
    "In this section,\n",
    "we will download a pretrained small version of BERT,\n",
    "then fine-tune it\n",
    "for natural language inference on the SNLI dataset."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-natural-language-inference-fine-tuning-bert"
   },
   "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\n",
    "import json\n",
    "import os"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-2 fw=all -->\n",
    "\n",
    "## Loading Pretrained BERT\n",
    "\n",
    "We have explained how to pretrain BERT on the WikiText-2 dataset in\n",
    "that section and that section\n",
    "(note that the original BERT model is pretrained on much bigger corpora).\n",
    "As discussed in that section,\n",
    "the original BERT model has hundreds of millions of parameters.\n",
    "In the following,\n",
    "we provide two versions of pretrained BERT:\n",
    "\"bert.base\" is about as big as the original BERT base model that requires a lot of computational resources to fine-tune,\n",
    "while \"bert.small\" is a small version to facilitate demonstration."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-loading-pretrained-bert-1"
   },
   "source": [
    "d2l.DATA_HUB['bert.base'] = (d2l.DATA_URL + 'bert.base.torch.zip',\n",
    "                             '225d66f04cae318b841a13d32af3acc165f253ac')\n",
    "d2l.DATA_HUB['bert.small'] = (d2l.DATA_URL + 'bert.small.torch.zip',\n",
    "                              'c72329e68a732bef0452e4b96a1c341c8910f81f')\n",
    "\n",
    "def _load_torch_state_dict(path):\n",
    "    \"\"\"Load a PyTorch state-dict file as {str: numpy.ndarray} without torch.\n",
    "\n",
    "    Works for files saved with ``torch.save(model.state_dict(), path)``\n",
    "    using pickle protocol 2 (the default for PyTorch <= 2.x CPU tensors).\n",
    "    \"\"\"\n",
    "    import pickle, struct, io\n",
    "    from collections import OrderedDict\n",
    "\n",
    "    dtype_info = {\n",
    "        'FloatStorage': (np.float32, 4), 'HalfStorage': (np.float16, 2),\n",
    "        'LongStorage': (np.int64, 8),    'IntStorage': (np.int32, 4),\n",
    "        'DoubleStorage': (np.float64, 8),\n",
    "    }\n",
    "\n",
    "    class _StorageRef:\n",
    "        __slots__ = ('dtype_name', 'key', 'location', 'size')\n",
    "        def __init__(self, dtype_name, key, location, size):\n",
    "            self.dtype_name = dtype_name; self.key = key\n",
    "            self.location = location; self.size = size\n",
    "\n",
    "    class _TensorRef:\n",
    "        __slots__ = ('storage', 'offset', 'size', 'stride')\n",
    "        def __init__(self, storage, offset, size, stride):\n",
    "            self.storage = storage; self.offset = offset\n",
    "            self.size = tuple(size); self.stride = tuple(stride)\n",
    "\n",
    "    class _StorageType:\n",
    "        __slots__ = ('name',)\n",
    "        def __init__(self, name): self.name = name\n",
    "        def __call__(self, _): return self\n",
    "\n",
    "    class _Unpickler(pickle.Unpickler):\n",
    "        def find_class(self, module, name):\n",
    "            if module == 'collections' and name == 'OrderedDict':\n",
    "                return OrderedDict\n",
    "            if module == 'torch._utils' and name == '_rebuild_tensor_v2':\n",
    "                return (lambda s, o, sz, st, *a, **k:\n",
    "                        _TensorRef(s, o, sz, st))\n",
    "            if module == 'torch' and name in dtype_info:\n",
    "                return _StorageType(name)\n",
    "            return super().find_class(module, name)\n",
    "        def persistent_load(self, saved_id):\n",
    "            if isinstance(saved_id, tuple) and saved_id[0] == 'storage':\n",
    "                st = saved_id[1]\n",
    "                return _StorageRef(\n",
    "                    st.name if isinstance(st, _StorageType) else 'FloatStorage',\n",
    "                    saved_id[2], saved_id[3], saved_id[4])\n",
    "            raise RuntimeError(f'Unknown persistent id: {saved_id}')\n",
    "\n",
    "    with open(path, 'rb') as f:\n",
    "        content = f.read()\n",
    "\n",
    "    buf = io.BytesIO(content)\n",
    "    for _ in range(3):                     # skip magic, proto, sysinfo\n",
    "        _Unpickler(buf).load()\n",
    "    data = _Unpickler(buf).load()          # OrderedDict of _TensorRef\n",
    "    storage_keys = _Unpickler(buf).load()  # list of storage keys\n",
    "    pos = buf.tell()\n",
    "\n",
    "    key_dtype = {}\n",
    "    for tref in data.values():\n",
    "        if isinstance(tref, _TensorRef):\n",
    "            key_dtype[tref.storage.key] = tref.storage.dtype_name\n",
    "\n",
    "    storages = {}\n",
    "    for key in storage_keys:\n",
    "        dname = key_dtype.get(key, 'FloatStorage')\n",
    "        np_dt, esz = dtype_info[dname]\n",
    "        n = struct.unpack('<Q', content[pos:pos+8])[0]; pos += 8\n",
    "        storages[key] = np.frombuffer(content, np_dt, count=n, offset=pos)\n",
    "        pos += n * esz\n",
    "\n",
    "    result = OrderedDict()\n",
    "    for name, tref in data.items():\n",
    "        if isinstance(tref, _TensorRef):\n",
    "            s = storages[tref.storage.key]\n",
    "            ne = 1\n",
    "            for d in tref.size:\n",
    "                ne *= d\n",
    "            result[name] = s[tref.offset:tref.offset+ne].reshape(\n",
    "                tref.size).copy()\n",
    "    return result"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-3 fw=all -->\n",
    "\n",
    "Either pretrained BERT model contains a \"vocab.json\" file that defines the vocabulary set\n",
    "and a \"pretrained.params\" file of the pretrained parameters.\n",
    "We implement the following `load_pretrained_model` function to load pretrained BERT parameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-loading-pretrained-bert-2"
   },
   "source": [
    "def load_pretrained_model(pretrained_model, num_hiddens, ffn_num_hiddens,\n",
    "                          num_heads, num_blks, dropout, max_len, devices):\n",
    "    data_dir = d2l.download_extract(pretrained_model)\n",
    "    # Define an empty vocabulary to load the predefined vocabulary\n",
    "    vocab = d2l.Vocab()\n",
    "    vocab.idx_to_token = json.load(open(os.path.join(data_dir, 'vocab.json')))\n",
    "    vocab.token_to_idx = {token: idx for idx, token in enumerate(\n",
    "        vocab.idx_to_token)}\n",
    "    bert = d2l.BERTModel(\n",
    "        len(vocab), num_hiddens, ffn_num_hiddens=ffn_num_hiddens,\n",
    "        num_heads=num_heads, num_blks=num_blks, dropout=dropout,\n",
    "        max_len=max_len)\n",
    "    # Load pretrained PyTorch BERT parameters (as NumPy) into the NNX graph.\n",
    "    pt_state_dict = _load_torch_state_dict(\n",
    "        os.path.join(data_dir, 'pretrained.params'))\n",
    "    _load_torch_into_nnx_bert(bert, pt_state_dict)\n",
    "    return bert, vocab\n",
    "\n",
    "def _load_torch_into_nnx_bert(bert, pt_state_dict):\n",
    "    \"\"\"Load a PyTorch BERT state dict into an NNX model.\"\"\"\n",
    "    p = pt_state_dict\n",
    "    # Encoder: token, segment, position embeddings\n",
    "    enc = bert.encoder\n",
    "    enc.token_embedding.embedding[...] = jnp.array(\n",
    "        p['encoder.token_embedding.weight'])\n",
    "    enc.segment_embedding.embedding[...] = jnp.array(\n",
    "        p['encoder.segment_embedding.weight'])\n",
    "    enc.pos_embedding[...] = jnp.array(p['encoder.pos_embedding'])\n",
    "    # Transformer encoder blocks\n",
    "    for i, blk in enumerate(enc.blks):\n",
    "        prefix = f'encoder.blks.{i}'\n",
    "        # Multi-head attention\n",
    "        attn = blk.attention\n",
    "        for name in ['W_q', 'W_k', 'W_v', 'W_o']:\n",
    "            linear = getattr(attn, name)\n",
    "            linear.kernel[...] = jnp.array(\n",
    "                p[f'{prefix}.attention.{name}.weight'].T)\n",
    "            linear.bias[...] = jnp.array(\n",
    "                p[f'{prefix}.attention.{name}.bias'])\n",
    "        # Addnorm layers (LayerNorm)\n",
    "        for ln_name in ['addnorm1', 'addnorm2']:\n",
    "            ln = getattr(blk, ln_name).ln\n",
    "            ln.scale[...] = jnp.array(p[f'{prefix}.{ln_name}.ln.weight'])\n",
    "            ln.bias[...] = jnp.array(p[f'{prefix}.{ln_name}.ln.bias'])\n",
    "        # FFN\n",
    "        ffn = blk.ffn\n",
    "        ffn.dense1.kernel[...] = jnp.array(\n",
    "            p[f'{prefix}.ffn.dense1.weight'].T)\n",
    "        ffn.dense1.bias[...] = jnp.array(p[f'{prefix}.ffn.dense1.bias'])\n",
    "        ffn.dense2.kernel[...] = jnp.array(\n",
    "            p[f'{prefix}.ffn.dense2.weight'].T)\n",
    "        ffn.dense2.bias[...] = jnp.array(p[f'{prefix}.ffn.dense2.bias'])\n",
    "    # Hidden (tanh) layer\n",
    "    bert.hidden.kernel[...] = jnp.array(p['hidden.0.weight'].T)\n",
    "    bert.hidden.bias[...] = jnp.array(p['hidden.0.bias'])"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-4 fw=all -->\n",
    "\n",
    "To facilitate demonstration on most machines,\n",
    "we will load and fine-tune the small version (\"bert.small\") of the pretrained BERT in this section.\n",
    "In the exercise, we will show how to fine-tune the much larger \"bert.base\" to significantly improve the testing accuracy."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-loading-pretrained-bert-3"
   },
   "source": [
    "devices = d2l.try_all_gpus()\n",
    "bert, vocab = load_pretrained_model(\n",
    "    'bert.small', num_hiddens=256, ffn_num_hiddens=512, num_heads=4,\n",
    "    num_blks=2, dropout=0.1, max_len=512, devices=devices)"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-5 fw=all -->\n",
    "\n",
    "## The Dataset for Fine-Tuning BERT\n",
    "\n",
    "For the downstream task natural language inference on the SNLI dataset,\n",
    "we define a customized dataset class `SNLIBERTDataset`.\n",
    "In each example,\n",
    "the premise and hypothesis form a pair of text sequence\n",
    "and is packed into one BERT input sequence as depicted in the figure.\n",
    "Recall that section that segment IDs\n",
    "are used to distinguish the premise and the hypothesis in a BERT input sequence.\n",
    "With the predefined maximum length of a BERT input sequence (`max_len`),\n",
    "the last token of the longer of the input text pair keeps getting removed until\n",
    "`max_len` is met.\n",
    "To accelerate generation of the SNLI dataset\n",
    "for fine-tuning BERT,\n",
    "we use 4 worker processes to generate training or testing examples in parallel."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-the-dataset-for-fine-tuning-bert-1"
   },
   "source": [
    "class SNLIBERTDataset:\n",
    "    def __init__(self, dataset, max_len, vocab=None):\n",
    "        all_premise_hypothesis_tokens = [[\n",
    "            p_tokens, h_tokens] for p_tokens, h_tokens in zip(\n",
    "            *[d2l.tokenize([s.lower() for s in sentences])\n",
    "              for sentences in dataset[:2]])]\n",
    "        \n",
    "        self.labels = np.asarray(dataset[2], dtype=np.int32)\n",
    "        self.vocab = vocab\n",
    "        self.max_len = max_len\n",
    "        (self.all_token_ids, self.all_segments,\n",
    "         self.valid_lens) = self._preprocess(all_premise_hypothesis_tokens)\n",
    "        print('read ' + str(len(self.all_token_ids)) + ' examples')\n",
    "\n",
    "    def _preprocess(self, all_premise_hypothesis_tokens):\n",
    "        # This Python token/list processing is inexpensive enough here that a\n",
    "        # list comprehension avoids multiprocessing setup and serialization.\n",
    "        out = [self._preprocess_pair(tokens)\n",
    "               for tokens in all_premise_hypothesis_tokens]\n",
    "        all_token_ids = [\n",
    "            token_ids for token_ids, segments, valid_len in out]\n",
    "        all_segments = [segments for token_ids, segments, valid_len in out]\n",
    "        valid_lens = [valid_len for token_ids, segments, valid_len in out]\n",
    "        return (np.asarray(all_token_ids, dtype=np.int32),\n",
    "                np.asarray(all_segments, dtype=np.int32),\n",
    "                np.asarray(valid_lens, dtype=np.float32))\n",
    "\n",
    "    def _preprocess_pair(self, premise_hypothesis_tokens):\n",
    "        p_tokens, h_tokens = premise_hypothesis_tokens\n",
    "        self._truncate_pair_of_tokens(p_tokens, h_tokens)\n",
    "        tokens, segments = d2l.get_tokens_and_segments(p_tokens, h_tokens)\n",
    "        token_ids = self.vocab[tokens] + [self.vocab['<pad>']] \\\n",
    "                             * (self.max_len - len(tokens))\n",
    "        segments = segments + [0] * (self.max_len - len(segments))\n",
    "        valid_len = len(tokens)\n",
    "        return token_ids, segments, valid_len\n",
    "\n",
    "    def _truncate_pair_of_tokens(self, p_tokens, h_tokens):\n",
    "        # Reserve slots for '<cls>', '<sep>', and '<sep>' tokens for the BERT\n",
    "        # input\n",
    "        while len(p_tokens) + len(h_tokens) > self.max_len - 3:\n",
    "            if len(p_tokens) > len(h_tokens):\n",
    "                p_tokens.pop()\n",
    "            else:\n",
    "                h_tokens.pop()\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (self.all_token_ids[idx], self.all_segments[idx],\n",
    "                self.valid_lens[idx]), self.labels[idx]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.all_token_ids)"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-6 fw=all -->\n",
    "\n",
    "After downloading the SNLI dataset,\n",
    "we generate training and testing examples\n",
    "by instantiating the `SNLIBERTDataset` class.\n",
    "Such examples will be read in minibatches during training and testing\n",
    "of natural language inference."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-the-dataset-for-fine-tuning-bert-2"
   },
   "source": [
    "# Reduce `batch_size` if there is an out of memory error. In the original BERT\n",
    "# model, `max_len` = 512\n",
    "batch_size, max_len = 512, 128\n",
    "data_dir = d2l.download_extract('SNLI')\n",
    "train_set = SNLIBERTDataset(d2l.read_snli(data_dir, True), max_len, vocab)\n",
    "test_set = SNLIBERTDataset(d2l.read_snli(data_dir, False), max_len, vocab)\n",
    "train_iter = d2l.load_array(\n",
    "    (train_set.all_token_ids, train_set.all_segments,\n",
    "     train_set.valid_lens, train_set.labels), batch_size, is_train=True)\n",
    "test_iter = d2l.load_array(\n",
    "    (test_set.all_token_ids, test_set.all_segments,\n",
    "     test_set.valid_lens, test_set.labels), batch_size, is_train=False)"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "read 549367 examples\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "read 9824 examples\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-7 fw=all -->\n",
    "\n",
    "## Fine-Tuning BERT\n",
    "\n",
    "As the figure indicates,\n",
    "fine-tuning BERT for natural language inference\n",
    "requires only an extra MLP consisting of two fully connected layers\n",
    "(see `self.hidden` and `self.output`---named `self.output_layer` in the\n",
    "TensorFlow tab to avoid clashing with Keras's reserved `output` property---in\n",
    "the following `BERTClassifier` class).\n",
    "This MLP transforms the\n",
    "BERT representation of the special “&lt;cls&gt;” token,\n",
    "which encodes the information of both the premise and the hypothesis,\n",
    "into three outputs of natural language inference:\n",
    "entailment, contradiction, and neutral."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-fine-tuning-bert-1"
   },
   "source": [
    "class BERTClassifier(nnx.Module):\n",
    "    def __init__(self, bert, rngs=None):\n",
    "        self.bert = bert\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.output = nnx.Linear(bert.hidden.out_features, 3, rngs=rngs)\n",
    "\n",
    "    def __call__(self, tokens_X, segments_X, valid_lens_x):\n",
    "        encoded_X = self.bert.encoder(tokens_X, segments_X, valid_lens_x)\n",
    "        return self.output(jnp.tanh(self.bert.hidden(encoded_X[:, 0, :])))"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-8 fw=all -->\n",
    "\n",
    "In the following,\n",
    "the pretrained BERT model `bert` is fed into the `BERTClassifier` instance `net` for\n",
    "the downstream application.\n",
    "In common implementations of BERT fine-tuning,\n",
    "only the parameters of the output layer of the additional MLP (`net.output`, or `net.output_layer` in the TensorFlow tab) will be learned from scratch.\n",
    "All the parameters of the pretrained BERT encoder (`net.encoder`) and the hidden layer of the additional MLP (`net.hidden`) will be fine-tuned."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-fine-tuning-bert-2"
   },
   "source": [
    "net = BERTClassifier(bert)"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-9 fw=all -->\n",
    "\n",
    "Recall that\n",
    "in that section\n",
    "both the `MaskLM` class and the `NextSentencePred` class\n",
    "have parameters in their employed MLPs.\n",
    "These parameters are part of those in the pretrained BERT model\n",
    "`bert`, and thus part of parameters in `net`.\n",
    "However, such parameters are only for computing\n",
    "the masked language modeling loss\n",
    "and the next sentence prediction loss\n",
    "during pretraining.\n",
    "These two loss functions are irrelevant to fine-tuning downstream applications,\n",
    "thus the parameters of the employed MLPs in \n",
    "`MaskLM` and `NextSentencePred` are not updated (and so their gradients become stale) when BERT is fine-tuned.\n",
    "\n",
    "To allow parameters with stale gradients,\n",
    "the flag `ignore_stale_grad=True` is set in the `step` function of `d2l.train_batch_ch13`.\n",
    "We use this function to train and evaluate the model `net` using the training set\n",
    "(`train_iter`) and the testing set (`test_iter`) of SNLI.\n",
    "Due to the limited computational resources, the training and testing accuracy\n",
    "can be further improved: we leave its discussions in the exercises."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "natural-language-inference-bert-fine-tuning-bert-3"
   },
   "source": [
    "lr, num_epochs = 1e-4, 5\n",
    "optimizer = nnx.Optimizer(net, optax.adam(lr), wrt=nnx.Param)\n",
    "\n",
    "@nnx.jit\n",
    "def train_step(net, optimizer, tokens_X, segments_X, valid_lens_x, labels):\n",
    "    def loss_fn(model):\n",
    "        logits = model(tokens_X, segments_X, valid_lens_x)\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits, labels).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(net)\n",
    "    optimizer.update(net, grads)\n",
    "    return loss\n",
    "\n",
    "@nnx.jit\n",
    "def eval_step(net, tokens_X, segments_X, valid_lens_x, labels):\n",
    "    logits = nnx.view(net, deterministic=True)(\n",
    "        tokens_X, segments_X, valid_lens_x)\n",
    "    return (logits.argmax(axis=-1) == labels).sum()\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss, n_train = jnp.array(0.0), 0\n",
    "    for batch in train_iter:\n",
    "        tokens_X, segments_X, valid_lens_x, labels = (\n",
    "            batch[0], batch[1], batch[2], batch[3])\n",
    "        loss = train_step(net, optimizer, tokens_X, segments_X,\n",
    "                          valid_lens_x, labels)\n",
    "        train_loss += loss * len(labels)\n",
    "        n_train += len(labels)\n",
    "    # Evaluate on test set\n",
    "    n_correct, n_test = jnp.array(0), 0\n",
    "    for batch in test_iter:\n",
    "        tokens_X, segments_X, valid_lens_x, labels = (\n",
    "            batch[0], batch[1], batch[2], batch[3])\n",
    "        n_correct += eval_step(\n",
    "            net, tokens_X, segments_X, valid_lens_x, labels)\n",
    "        n_test += len(labels)\n",
    "    train_loss, n_correct = float(train_loss), int(n_correct)\n",
    "    print(f'epoch {epoch + 1}, loss {train_loss / n_train:.4f}, '\n",
    "          f'test acc {n_correct / n_test:.4f}')"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 1, loss 0.7752, test acc 0.7356\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 2, loss 0.6329, test acc 0.7565\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 3, loss 0.5652, test acc 0.7741\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 4, loss 0.5165, test acc 0.7827\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "epoch 5, loss 0.4798, test acc 0.7856\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-10 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* We can fine-tune the pretrained BERT model for downstream applications, such as natural language inference on the SNLI dataset.\n",
    "* During fine-tuning, the BERT model becomes part of the model for the downstream application. Parameters that are only related to pretraining loss will not be updated during fine-tuning. \n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Fine-tune a much larger pretrained BERT model that is about as big as the original BERT base model if your computational resource allows. Set arguments in the `load_pretrained_model` function as: replacing 'bert.small' with 'bert.base', increasing values of `num_hiddens=256`, `ffn_num_hiddens=512`, `num_heads=4`, and `num_blks=2` to 768, 3072, 12, and 12, respectively. By increasing fine-tuning epochs (and possibly tuning other hyperparameters), can you get a testing accuracy higher than 0.86?\n",
    "1. How to truncate a pair of sequences according to their ratio of length? Compare this pair truncation method and the one used in the `SNLIBERTDataset` class. What are their pros and cons?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=natural-language-inference-bert-md-11 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1526)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
