notebooks/Infer-masked-token-with-BERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Infer masked token with BERT\n",
    "\n",
    "The task is simple: given a sentence with a masked token, the model should predict which token can replace the masked position. For example, given `The car is [MASK]`, a possible solution could be `red`. This task is part of the BERT pretaining described by its authors (see: https://arxiv.org/pdf/1810.04805.pdf), thus we do not need to fine-tune the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare and import modules\n",
    "\n",
    "With your environment configured, you can now prepare and import the BERT modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torchvision in c:\\users\\dimas\\anaconda3\\lib\\site-packages (0.5.0)\n",
      "Requirement already satisfied: pillow>=4.1.1 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from torchvision) (7.0.0)\n",
      "Requirement already satisfied: six in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from torchvision) (1.14.0)\n",
      "Requirement already satisfied: numpy in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from torchvision) (1.18.1)\n",
      "Requirement already satisfied: torch==1.4.0 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from torchvision) (1.4.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: transformers in c:\\users\\dimas\\anaconda3\\lib\\site-packages (3.0.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (4.42.1)\n",
      "Requirement already satisfied: sentencepiece!=0.1.92 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (0.1.91)\n",
      "Requirement already satisfied: requests in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (2.22.0)\n",
      "Requirement already satisfied: numpy in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (1.18.1)\n",
      "Requirement already satisfied: packaging in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (20.1)\n",
      "Requirement already satisfied: tokenizers==0.8.1.rc1 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (0.8.1rc1)\n",
      "Requirement already satisfied: filelock in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (3.0.12)\n",
      "Requirement already satisfied: regex!=2019.12.17 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (2020.7.14)\n",
      "Requirement already satisfied: sacremoses in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from transformers) (0.0.43)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from requests->transformers) (2.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from requests->transformers) (2019.11.28)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from requests->transformers) (1.24.3)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from requests->transformers) (3.0.4)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from packaging->transformers) (2.4.6)\n",
      "Requirement already satisfied: six in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from packaging->transformers) (1.14.0)\n",
      "Requirement already satisfied: click in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from sacremoses->transformers) (7.0)\n",
      "Requirement already satisfied: joblib in c:\\users\\dimas\\anaconda3\\lib\\site-packages (from sacremoses->transformers) (0.14.1)\n"
     ]
    }
   ],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import BertTokenizer, BertForMaskedLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get model\n",
    "\n",
    "Visit https://huggingface.co/transformers/pretrained_models.html to see the full list of pretrained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BERT_MODEL = 'bert-base-multilingual-cased'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForMaskedLM were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['cls.predictions.decoder.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(BERT_MODEL)\n",
    "model = BertForMaskedLM.from_pretrained(BERT_MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize & encode\n",
    "\n",
    "BERT expects 3 different inputs:\n",
    "- Token IDs: the tokens transformed into numbers.\n",
    "- Mask: sequence of `0` (if there are PAD tokens in that position) and `1` (otherwise).\n",
    "- Segments or Type IDs: sequence of `0` and `1` to distinguish between the first and the second sentence in NSP tasks. In this notebook, we do not need this input, so it will be always `0`.\n",
    "\n",
    "The only constraint is that the **maximum number of tokens is 512**. Please, note that there are extra tokens which we are going to add (see the next section).\n",
    "\n",
    "For example:\n",
    "```\n",
    "Text:       Is this jacksonville?\n",
    "---------------------------------------------------------------------------------\n",
    "Tokens:     [CLS] Is    this  ja    ##cks ##on  ##ville ?   [SEP] [PAD] [PAD] ...\n",
    "Token IDs:  101   12034 10531 10201 18676 10263 12043   136 102   100   100   ...\n",
    "Mask:       1     1     1     1     1     1     1       1   1     0     0     ...\n",
    "Type IDs:   0     0     0     0     0     0     0       0   0     0     0     ...\n",
    "```\n",
    "\n",
    "Note: Token IDs may be different depending on the tokenizer.\n",
    "\n",
    "For further details, see BERT implementation `convert_single_example()` at https://github.com/google-research/bert/blob/eedf5716ce1268e56f0a50264a88cafad334ac61/run_classifier.py#L377 (permalink)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manual encoding\n",
    "\n",
    "This section is an example of how to get the encoded input of BERT.\n",
    "\n",
    "We use the BERT Tokenizer to split the text into tokens. Then, according to the paper, it is necessary to add two tokens at the beginning and the ending: `[CLS]` and `[SEP]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_masks(tokens, max_len=128):\n",
    "    \"\"\"Mask for padding\"\"\"\n",
    "    if len(tokens) > max_len:\n",
    "        raise IndexError(\"Token length more than max length!\")\n",
    "    return [1] * len(tokens) + [0] * (max_len - len(tokens))\n",
    "\n",
    "\n",
    "def get_segments(tokens, max_len=128):\n",
    "    \"\"\"Segments: 0 for the first sequence, 1 for the second\"\"\"\n",
    "    if len(tokens) > max_len:\n",
    "        raise IndexError(\"Token length more than max 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_len - len(tokens))\n",
    "\n",
    "\n",
    "def get_ids(tokens, tokenizer, max_len=128):\n",
    "    \"\"\"Token ids from Tokenizer vocab\"\"\"\n",
    "    token_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "    input_ids = token_ids + [tokenizer.pad_token_id] * (max_len - len(token_ids))\n",
    "    return input_ids\n",
    "\n",
    "\n",
    "def encode_input(sentence, max_len=128):\n",
    "    stokens = tokenize(sentence)\n",
    "    \n",
    "    input_ids = get_ids(stokens, tokenizer, max_len)\n",
    "    input_masks = get_masks(stokens, max_len)\n",
    "    input_segments = get_segments(stokens, max_len)\n",
    "    \n",
    "    model_input = {\n",
    "        'input_word_ids': np.array([input_ids]),\n",
    "        'input_masks': np.array([input_masks]),\n",
    "        'input_segments': np.array([input_segments])\n",
    "    }\n",
    "    \n",
    "    mask_pos = np.array(np.array(stokens) == '[MASK]', dtype='int')\n",
    "    mask_pos = np.concatenate((mask_pos, np.zeros(max_len - len(mask_pos))))\n",
    "    mask_pos = mask_pos.astype('int')\n",
    "    \n",
    "    return model_input, mask_pos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(sentence):\n",
    "    stokens = tokenizer.tokenize(sentence)\n",
    "    i = 0\n",
    "    while i < len(stokens) - 2:\n",
    "        if stokens[i] == '[' and stokens[i+1] == 'mask' and stokens[i+2] == ']':\n",
    "            stokens[i] = '[MASK]'\n",
    "            stokens.pop(i+2)\n",
    "            stokens.pop(i+1)\n",
    "        i = i + 1\n",
    "    \n",
    "    stokens = ['[CLS]'] + stokens + ['[SEP]']\n",
    "    return stokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS]', 'Hello', ',', 'I', \"'\", 'm', 'a', '[MASK]', 'model', '.', '[SEP]']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"Hello, I'm a [MASK] model.\"\n",
    "tokenize(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'input_word_ids': array([[  101, 31178,   117,   146,   112,   181,   169,   103, 13192,\n",
       "            119,   102,     0,     0,     0,     0]]),\n",
       "  'input_masks': array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]),\n",
       "  'input_segments': array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])},\n",
       " array([0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of BERT inputs\n",
    "text = \"Hello, I'm a [MASK] model.\"\n",
    "encode_input(text, max_len=15) # max_len=15 for display purpose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Torch encoder\n",
    "\n",
    "Torch has already implemented similar code as above, thus you should use these functions if you do not need any custom behaviour.\n",
    "\n",
    "Here you can see the result of `tokenize`. Note that it does not include `[CLS]` and `[SEP]`, but the result of `encoded_input` does include them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Hello', ',', 'I', \"'\", 'm', 'a', '[MASK]', 'model', '.']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"Hello, I'm a [MASK] model.\"\n",
    "tokenizer.tokenize(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[  101, 31178,   117,   146,   112,   181,   169,   103, 13192,   119,\n",
       "           102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"Hello, I'm a [MASK] model.\"\n",
    "\n",
    "encoded_input = tokenizer(text, return_tensors='pt')\n",
    "encoded_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infer masked word\n",
    "\n",
    "Now, we can use the described encoded data as input of our BERT model and, then, we apply softmax to the output. The result are the probabilities that a token fit in a certain position, so we should focus on the position of the masked token and get the top probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_topk_predictions(model, tokenizer, text, topk=5):\n",
    "    encoded_input = tokenizer(text, return_tensors='pt')\n",
    "    logits = model(encoded_input['input_ids'],\n",
    "                   encoded_input['token_type_ids'],\n",
    "                   encoded_input['attention_mask'],\n",
    "                   masked_lm_labels=None)[0]\n",
    "\n",
    "    logits = logits.squeeze(0)\n",
    "    probs = torch.softmax(logits, dim=-1)\n",
    "\n",
    "    mask_cnt = 0\n",
    "    token_ids = encoded_input['input_ids'][0]\n",
    "    \n",
    "    top_preds = []\n",
    "\n",
    "    for idx, _ in enumerate(token_ids):\n",
    "        if token_ids[idx] == tokenizer.mask_token_id:\n",
    "            mask_cnt += 1\n",
    "            \n",
    "            topk_prob, topk_indices = torch.topk(probs[idx, :], topk)\n",
    "            topk_indices = topk_indices.cpu().numpy()\n",
    "            topk_tokens = tokenizer.convert_ids_to_tokens(topk_indices)\n",
    "            for prob, tok_str, tok_id in zip(topk_prob, topk_tokens, topk_indices):\n",
    "                top_preds.append({'token_str': tok_str,\n",
    "                                  'token_id': tok_id,\n",
    "                                  'probability': float(prob)})\n",
    "    \n",
    "    return top_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_topk_predictions(model, tokenizer, text, pretty_prob=False):\n",
    "    top_preds = get_topk_predictions(model, tokenizer, text)\n",
    "    \n",
    "    print(text)\n",
    "    print('=' * 40)\n",
    "    for item in top_preds:\n",
    "        if not pretty_prob:\n",
    "            print('%s %.4f' % (item['token_str'], item['probability']))\n",
    "        else:\n",
    "            probability = item['probability'] * 100\n",
    "            print('%s %.2f%%' % (item['token_str'], probability))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, I'm a [MASK] model.\n",
      "========================================\n",
      "model 6.66%\n",
      "real 4.59%\n",
      "business 3.30%\n",
      "mathematical 3.16%\n",
      "new 2.78%\n"
     ]
    }
   ],
   "source": [
    "text = \"Hello, I'm a [MASK] model.\"\n",
    "display_topk_predictions(model, tokenizer, text, pretty_prob=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The doctor ran to the emergency room to see a [MASK].\n",
      "========================================\n",
      "doctor 7.59%\n",
      "woman 3.52%\n",
      "fire 2.40%\n",
      "patient 2.38%\n",
      "problem 1.87%\n"
     ]
    }
   ],
   "source": [
    "text = 'The doctor ran to the emergency room to see a [MASK].'\n",
    "display_topk_predictions(model, tokenizer, text, pretty_prob=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The predictions in English seems to work quite good.\n",
    "\n",
    "Since I am using a multilingual model, it should have nice predictions in other languages too, e.g. Spanish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Este coche es [MASK].\n",
      "========================================\n",
      "perdido 2.83%\n",
      "coupé 2.72%\n",
      "retirado 2.69%\n",
      "motor 1.57%\n",
      "abandonado 1.56%\n"
     ]
    }
   ],
   "source": [
    "text = 'Este coche es [MASK].'\n",
    "display_topk_predictions(model, tokenizer, text, pretty_prob=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Despite I introduced a very simple sentence (`This car is [MASK]`), the outputs do not make sense in a sentence (although they are related to cars). I think that the model is confusing the two Spanish verbs `ser` and `estar` (both mean `to be` in English) or it does not have enough context to output a good result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I added the adverb \"*very*\" to the previous sentence (`This car is very [MASK].`), so I expect that it outputs better tokens since the context is better too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Este coche es muy [MASK].\n",
      "========================================\n",
      "popular 59.34%\n",
      "sencillo 3.47%\n",
      "raro 2.98%\n",
      "vendido 1.42%\n",
      "pequeño 1.41%\n"
     ]
    }
   ],
   "source": [
    "text = 'Este coche es muy [MASK].'\n",
    "display_topk_predictions(model, tokenizer, text, pretty_prob=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, the model success to return tokens which make sense.\n",
    "\n",
    "Here is one more example but, in this case, the sentence is written in Russian: `I think that Nastya is a very [MASK] person`. Spoiler: the output words make sense :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Я считаю, что Настя очень [MASK] человек.\n",
      "========================================\n",
      "молодой 54.84%\n",
      "большой 20.15%\n",
      "великий 5.03%\n",
      "лучший 4.91%\n",
      "новый 1.64%\n"
     ]
    }
   ],
   "source": [
    "text = 'Я считаю, что Настя очень [MASK] человек.'\n",
    "display_topk_predictions(model, tokenizer, text, pretty_prob=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}