{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-1 fw=all -->\n",
    "\n",
    "# The Dataset for Pretraining BERT\n",
    "\n",
    "To pretrain the BERT model as implemented in that section,\n",
    "we need to generate the dataset in the ideal format to facilitate\n",
    "the two pretraining tasks:\n",
    "masked language modeling and next sentence prediction.\n",
    "On the one hand,\n",
    "the original BERT model is pretrained on the concatenation of\n",
    "two huge corpora BookCorpus and English Wikipedia (see that section),\n",
    "making it hard to run for most readers of this book.\n",
    "On the other hand,\n",
    "the off-the-shelf pretrained BERT model\n",
    "may not fit for applications from specific domains like medicine.\n",
    "Thus, it is getting popular to pretrain BERT on a customized dataset.\n",
    "To facilitate the demonstration of BERT pretraining,\n",
    "we use a smaller corpus WikiText-2 [@Merity.Xiong.Bradbury.ea.2016].\n",
    "\n",
    "Comparing with the PTB dataset used for pretraining word2vec in that section,\n",
    "WikiText-2 (i) retains the original punctuation, making it suitable for next sentence prediction; (ii) retains the original case and numbers; (iii) is over twice larger."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-the-dataset-for-pretraining-bert-1"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np\n",
    "import os\n",
    "import random"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-2 fw=all -->\n",
    "\n",
    "In the WikiText-2 dataset,\n",
    "each line represents a paragraph where\n",
    "space is inserted between any punctuation and its preceding token.\n",
    "Paragraphs with at least two sentences are retained.\n",
    "To split sentences, we only use the period as the delimiter for simplicity.\n",
    "We leave discussions of more complex sentence splitting techniques in the exercises\n",
    "at the end of this section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-the-dataset-for-pretraining-bert-2"
   },
   "source": [
    "\n",
    "WIKITEXT_2_URL = ('https://huggingface.co/datasets/Salesforce/wikitext/'\n",
    "                  'resolve/main/wikitext-2-v1/train-00000-of-00001.parquet')\n",
    "WIKITEXT_2_SHA1 = '98ee727e59fcc34fddaadae93e15b1f8ed5561a4'\n",
    "\n",
    "\n",
    "def _read_wiki(data_dir=None):\n",
    "    import contextlib\n",
    "    import io\n",
    "    import pandas as pd\n",
    "    with contextlib.redirect_stdout(io.StringIO()):\n",
    "        fname = d2l.download(WIKITEXT_2_URL, folder='../data',\n",
    "                              sha1_hash=WIKITEXT_2_SHA1)\n",
    "    lines = pd.read_parquet(fname)['text'].tolist()\n",
    "    # Uppercase letters are converted to lowercase ones\n",
    "    paragraphs = [line.strip().lower().split(' . ')\n",
    "                  for line in lines if len(line.split(' . ')) >= 2]\n",
    "    random.shuffle(paragraphs)\n",
    "    return paragraphs"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-3 fw=all -->\n",
    "\n",
    "## Defining Helper Functions for Pretraining Tasks\n",
    "\n",
    "In the following,\n",
    "we begin by implementing helper functions for the two BERT pretraining tasks:\n",
    "next sentence prediction and masked language modeling.\n",
    "These helper functions will be invoked later\n",
    "when transforming the raw text corpus\n",
    "into the dataset of the ideal format to pretrain BERT.\n",
    "\n",
    "### Generating the Next Sentence Prediction Task\n",
    "\n",
    "According to descriptions of that section,\n",
    "the `_get_next_sentence` function generates a training example\n",
    "for the binary classification task."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-generating-the-next-sentence-prediction-task-1"
   },
   "source": [
    "\n",
    "def _get_next_sentence(sentence, next_sentence, paragraphs):\n",
    "    if random.random() < 0.5:\n",
    "        is_next = True\n",
    "    else:\n",
    "        # `paragraphs` is a list of lists of lists\n",
    "        next_sentence = random.choice(random.choice(paragraphs))\n",
    "        is_next = False\n",
    "    return sentence, next_sentence, is_next"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-4 fw=all -->\n",
    "\n",
    "The following function generates training examples for next sentence prediction\n",
    "from the input `paragraph` by invoking the `_get_next_sentence` function.\n",
    "Here `paragraph` is a list of sentences, where each sentence is a list of tokens.\n",
    "The argument `max_len` specifies the maximum length of a BERT input sequence during pretraining."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-generating-the-next-sentence-prediction-task-2"
   },
   "source": [
    "\n",
    "def _get_nsp_data_from_paragraph(paragraph, paragraphs, vocab, max_len):\n",
    "    nsp_data_from_paragraph = []\n",
    "    for i in range(len(paragraph) - 1):\n",
    "        tokens_a, tokens_b, is_next = _get_next_sentence(\n",
    "            paragraph[i], paragraph[i + 1], paragraphs)\n",
    "        # Consider 1 '<cls>' token and 2 '<sep>' tokens\n",
    "        if len(tokens_a) + len(tokens_b) + 3 > max_len:\n",
    "            continue\n",
    "        tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b)\n",
    "        nsp_data_from_paragraph.append((tokens, segments, is_next))\n",
    "    return nsp_data_from_paragraph"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-5 fw=all -->\n",
    "\n",
    "### Generating the Masked Language Modeling Task\n",
    "\n",
    "In order to generate training examples\n",
    "for the masked language modeling task\n",
    "from a BERT input sequence,\n",
    "we define the following `_replace_mlm_tokens` function.\n",
    "In its inputs, `tokens` is a list of tokens representing a BERT input sequence,\n",
    "`candidate_pred_positions` is a list of token indices of the BERT input sequence\n",
    "excluding those of special tokens (special tokens are not predicted in the masked language modeling task),\n",
    "and `num_mlm_preds` indicates the number of predictions (recall 15% random tokens to predict).\n",
    "Following the definition of the masked language modeling task in that section,\n",
    "at each prediction position, the input may be replaced by\n",
    "a special “&lt;mask&gt;” token or a random token, or remain unchanged.\n",
    "In the end, the function returns the input tokens after possible replacement,\n",
    "the token indices where predictions take place and labels for these predictions."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-generating-the-masked-language-modeling-task-1"
   },
   "source": [
    "\n",
    "def _replace_mlm_tokens(tokens, candidate_pred_positions, num_mlm_preds,\n",
    "                        vocab):\n",
    "    # For the input of a masked language model, make a new copy of tokens and\n",
    "    # replace some of them by '<mask>' or random tokens\n",
    "    mlm_input_tokens = [token for token in tokens]\n",
    "    pred_positions_and_labels = []\n",
    "    # Shuffle for getting 15% random tokens for prediction in the masked\n",
    "    # language modeling task\n",
    "    random.shuffle(candidate_pred_positions)\n",
    "    for mlm_pred_position in candidate_pred_positions:\n",
    "        if len(pred_positions_and_labels) >= num_mlm_preds:\n",
    "            break\n",
    "        masked_token = None\n",
    "        # 80% of the time: replace the word with the '<mask>' token\n",
    "        if random.random() < 0.8:\n",
    "            masked_token = '<mask>'\n",
    "        else:\n",
    "            # 10% of the time: keep the word unchanged\n",
    "            if random.random() < 0.5:\n",
    "                masked_token = tokens[mlm_pred_position]\n",
    "            # 10% of the time: replace the word with a random word\n",
    "            else:\n",
    "                masked_token = random.choice(vocab.idx_to_token)\n",
    "        mlm_input_tokens[mlm_pred_position] = masked_token\n",
    "        pred_positions_and_labels.append(\n",
    "            (mlm_pred_position, tokens[mlm_pred_position]))\n",
    "    return mlm_input_tokens, pred_positions_and_labels"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-6 fw=all -->\n",
    "\n",
    "By invoking the aforementioned `_replace_mlm_tokens` function,\n",
    "the following function takes a BERT input sequence (`tokens`)\n",
    "as an input and returns indices of the input tokens\n",
    "(after possible token replacement as described in that section),\n",
    "the token indices where predictions take place,\n",
    "and label indices for these predictions."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-generating-the-masked-language-modeling-task-2"
   },
   "source": [
    "\n",
    "def _get_mlm_data_from_tokens(tokens, vocab):\n",
    "    candidate_pred_positions = []\n",
    "    # `tokens` is a list of strings\n",
    "    for i, token in enumerate(tokens):\n",
    "        # Special tokens are not predicted in the masked language modeling\n",
    "        # task\n",
    "        if token in ['<cls>', '<sep>']:\n",
    "            continue\n",
    "        candidate_pred_positions.append(i)\n",
    "    # 15% of random tokens are predicted in the masked language modeling task\n",
    "    num_mlm_preds = max(1, round(len(tokens) * 0.15))\n",
    "    mlm_input_tokens, pred_positions_and_labels = _replace_mlm_tokens(\n",
    "        tokens, candidate_pred_positions, num_mlm_preds, vocab)\n",
    "    pred_positions_and_labels = sorted(pred_positions_and_labels,\n",
    "                                       key=lambda x: x[0])\n",
    "    pred_positions = [v[0] for v in pred_positions_and_labels]\n",
    "    mlm_pred_labels = [v[1] for v in pred_positions_and_labels]\n",
    "    return vocab[mlm_input_tokens], pred_positions, vocab[mlm_pred_labels]"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-7 fw=all -->\n",
    "\n",
    "## Transforming Text into the Pretraining Dataset\n",
    "\n",
    "Now we are almost ready to customize a `Dataset` class for pretraining BERT.\n",
    "Before that, \n",
    "we still need to define a helper function `_pad_bert_inputs`\n",
    "to append the special “&lt;pad&gt;” tokens to the inputs.\n",
    "Its argument `examples` contain the outputs from the helper functions `_get_nsp_data_from_paragraph` and `_get_mlm_data_from_tokens` for the two pretraining tasks."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-transforming-text-into-the-pretraining-dataset-1"
   },
   "source": [
    "\n",
    "def _pad_bert_inputs(examples, max_len, vocab):\n",
    "    max_num_mlm_preds = round(max_len * 0.15)\n",
    "    all_token_ids, all_segments, valid_lens,  = [], [], []\n",
    "    all_pred_positions, all_mlm_weights, all_mlm_labels = [], [], []\n",
    "    nsp_labels = []\n",
    "    for (token_ids, pred_positions, mlm_pred_label_ids, segments,\n",
    "         is_next) in examples:\n",
    "        all_token_ids.append(jnp.array(token_ids + [vocab['<pad>']] * (\n",
    "            max_len - len(token_ids)), dtype=jnp.int32))\n",
    "        all_segments.append(jnp.array(segments + [0] * (\n",
    "            max_len - len(segments)), dtype=jnp.int32))\n",
    "        # `valid_lens` excludes count of '<pad>' tokens\n",
    "        valid_lens.append(jnp.array(len(token_ids), dtype=jnp.float32))\n",
    "        all_pred_positions.append(jnp.array(pred_positions + [0] * (\n",
    "            max_num_mlm_preds - len(pred_positions)), dtype=jnp.int32))\n",
    "        # Predictions of padded tokens will be filtered out in the loss via\n",
    "        # multiplication of 0 weights\n",
    "        all_mlm_weights.append(\n",
    "            jnp.array([1.0] * len(mlm_pred_label_ids) + [0.0] * (\n",
    "                max_num_mlm_preds - len(pred_positions)),\n",
    "                dtype=jnp.float32))\n",
    "        all_mlm_labels.append(jnp.array(mlm_pred_label_ids + [0] * (\n",
    "            max_num_mlm_preds - len(mlm_pred_label_ids)), dtype=jnp.int32))\n",
    "        nsp_labels.append(jnp.array(is_next, dtype=jnp.int32))\n",
    "    return (all_token_ids, all_segments, valid_lens, all_pred_positions,\n",
    "            all_mlm_weights, all_mlm_labels, nsp_labels)"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-8 fw=all -->\n",
    "\n",
    "Putting the helper functions for\n",
    "generating training examples of the two pretraining tasks,\n",
    "and the helper function for padding inputs together,\n",
    "we customize the following `_WikiTextDataset` class as the WikiText-2 dataset for pretraining BERT.\n",
    "By implementing the `__getitem__ `function,\n",
    "we can arbitrarily access the pretraining (masked language modeling and next sentence prediction) examples \n",
    "generated from a pair of sentences from the WikiText-2 corpus.\n",
    "\n",
    "The original BERT model uses WordPiece embeddings whose vocabulary size is 30000 [@Wu.Schuster.Chen.ea.2016].\n",
    "The tokenization method of WordPiece is a slight modification of\n",
    "the original byte pair encoding algorithm in that section.\n",
    "For simplicity, we use the `d2l.tokenize` function for tokenization.\n",
    "Infrequent tokens that appear less than five times are filtered out."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-transforming-text-into-the-pretraining-dataset-2"
   },
   "source": [
    "\n",
    "class _WikiTextDataset:\n",
    "    def __init__(self, paragraphs, max_len):\n",
    "        # Input `paragraphs[i]` is a list of sentence strings representing a\n",
    "        # paragraph; while output `paragraphs[i]` is a list of sentences\n",
    "        # representing a paragraph, where each sentence is a list of tokens\n",
    "        paragraphs = [d2l.tokenize(\n",
    "            paragraph, token='word') for paragraph in paragraphs]\n",
    "        sentences = [sentence for paragraph in paragraphs\n",
    "                     for sentence in paragraph]\n",
    "        self.vocab = d2l.Vocab(sentences, min_freq=5, reserved_tokens=[\n",
    "            '<pad>', '<mask>', '<cls>', '<sep>'])\n",
    "        # Get data for the next sentence prediction task\n",
    "        examples = []\n",
    "        for paragraph in paragraphs:\n",
    "            examples.extend(_get_nsp_data_from_paragraph(\n",
    "                paragraph, paragraphs, self.vocab, max_len))\n",
    "        # Get data for the masked language model task\n",
    "        examples = [(_get_mlm_data_from_tokens(tokens, self.vocab)\n",
    "                      + (segments, is_next))\n",
    "                     for tokens, segments, is_next in examples]\n",
    "        # Pad inputs\n",
    "        (self.all_token_ids, self.all_segments, self.valid_lens,\n",
    "         self.all_pred_positions, self.all_mlm_weights,\n",
    "         self.all_mlm_labels, self.nsp_labels) = _pad_bert_inputs(\n",
    "            examples, max_len, self.vocab)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (self.all_token_ids[idx], self.all_segments[idx],\n",
    "                self.valid_lens[idx], self.all_pred_positions[idx],\n",
    "                self.all_mlm_weights[idx], self.all_mlm_labels[idx],\n",
    "                self.nsp_labels[idx])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.all_token_ids)"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-9 fw=all -->\n",
    "\n",
    "By using the `_read_wiki` function and the `_WikiTextDataset` class,\n",
    "we define the following `load_data_wiki` to download the WikiText-2 dataset\n",
    "and generate pretraining examples from it."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-transforming-text-into-the-pretraining-dataset-3"
   },
   "source": [
    "\n",
    "def load_data_wiki(batch_size, max_len):\n",
    "    \"\"\"Load the WikiText-2 dataset.\"\"\"\n",
    "    paragraphs = _read_wiki()\n",
    "    train_set = _WikiTextDataset(paragraphs, max_len)\n",
    "    # Create an index array and shuffle it\n",
    "    indices = list(range(len(train_set)))\n",
    "    random.shuffle(indices)\n",
    "\n",
    "    # Return a callable so each call yields a fresh iterator (one-shot\n",
    "    # generators can't be re-entered for multi-epoch training).\n",
    "    def data_iter():\n",
    "        for i in range(0, len(indices), batch_size):\n",
    "            batch_indices = indices[i:i + batch_size]\n",
    "            if len(batch_indices) < batch_size:\n",
    "                continue\n",
    "            batch = [train_set[idx] for idx in batch_indices]\n",
    "            yield (jnp.stack([b[0] for b in batch]),\n",
    "                   jnp.stack([b[1] for b in batch]),\n",
    "                   jnp.stack([b[2] for b in batch]),\n",
    "                   jnp.stack([b[3] for b in batch]),\n",
    "                   jnp.stack([b[4] for b in batch]),\n",
    "                   jnp.stack([b[5] for b in batch]),\n",
    "                   jnp.stack([b[6] for b in batch]))\n",
    "    return data_iter, train_set.vocab"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-10 fw=all -->\n",
    "\n",
    "Setting the batch size to 512 and the maximum length of a BERT input sequence to be 64,\n",
    "we print out the shapes of a minibatch of BERT pretraining examples.\n",
    "Note that in each BERT input sequence,\n",
    "$10$ ($64 \\times 0.15$) positions are predicted for the masked language modeling task."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-transforming-text-into-the-pretraining-dataset-4"
   },
   "source": [
    "batch_size, max_len = 512, 64\n",
    "# train_iter is a callable returning a fresh iterator on each call.\n",
    "train_iter, vocab = load_data_wiki(batch_size, max_len)\n",
    "\n",
    "for (tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X,\n",
    "     mlm_Y, nsp_y) in train_iter():\n",
    "    print(tokens_X.shape, segments_X.shape, valid_lens_x.shape,\n",
    "          pred_positions_X.shape, mlm_weights_X.shape, mlm_Y.shape,\n",
    "          nsp_y.shape)\n",
    "    break"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(512, 64) (512, 64) (512,) (512, 10) (512, 10) (512, 10) (512,)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-11 fw=all -->\n",
    "\n",
    "In the end, let's take a look at the vocabulary size.\n",
    "Even after filtering out infrequent tokens,\n",
    "it is still over twice larger than that of the PTB dataset."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bert-dataset-transforming-text-into-the-pretraining-dataset-5"
   },
   "source": [
    "len(vocab)"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "20257"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-12 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Comparing with the PTB dataset, the WikiText-2 dataset retains the original punctuation, case and numbers, and is over twice larger.\n",
    "* We can arbitrarily access the pretraining (masked language modeling and next sentence prediction) examples generated from a pair of sentences from the WikiText-2 corpus.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. For simplicity, the period is used as the only delimiter for splitting sentences. Try other sentence splitting techniques, such as the spaCy and NLTK. Take NLTK as an example. You need to install NLTK first: `pip install nltk`. In the code, first `import nltk`. Then, download the Punkt sentence tokenizer: `nltk.download('punkt')`. To split sentences such as `sentences = 'This is great ! Why not ?'`, invoking `nltk.tokenize.sent_tokenize(sentences)` will return a list of two sentence strings: `['This is great !', 'Why not ?']`.\n",
    "1. What is the vocabulary size if we do not filter out any infrequent token?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=bert-dataset-md-13 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1496)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
