notebooks/qa_step_by_step.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "controversial-writer",
   "metadata": {},
   "source": [
    "# Open Domain Question Answering based on Explain Anything Like I'm Five Blog Post\n",
    "\n",
    "This notebook is a step by step implementation of the open domain question answering system\n",
    "described in the blog post https://yjernite.github.io/lfqa.html. You should read that first as there is a lot \n",
    "of good detail there on the how and whys (particuarly the model training) which I won't be replicating here. \n",
    "\n",
    "The purpose of this notebook, is to work through all the details of setting up the QA system assuming we already have pre-trained models. There are a lot of details and I wanted to go through the exercise and understand each\n",
    "step. \n",
    "\n",
    "No training of models is done. I used the pre-trained models available on the Huggingface Hub.  The code in this noteook is largely a refactoring and commenting of the original code that you can find here: https://github.com/huggingface/notebooks/blob/master/longform-qa/lfqa_utils.py."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "massive-longitude",
   "metadata": {},
   "source": [
    "# Preliminaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "combined-karma",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do some imports. There are not many.\n",
    "\n",
    "import time\n",
    "import numpy as np \n",
    "import pandas as pd\n",
    "\n",
    "from datasets import load_dataset # The blog used the nlp library, this has been superceeded by datasets\n",
    "\n",
    "import torch as torch\n",
    "from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM\n",
    "\n",
    "import faiss # For doing fast similarity search: https://github.com/facebookresearch/faiss\n",
    "\n",
    "# add the src path\n",
    "import sys, os\n",
    "sys.path.append(\"../src\") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "significant-treasurer",
   "metadata": {},
   "source": [
    "### Import the dataset we will need\n",
    "\n",
    "To set up the QA system we actually only need the wiki_snippets dataset. This is what we use for making the dense index.  It's basically wikipedia cut up into 100 word chunks.  \n",
    "\n",
    "The eli5 dataset is actually only used for model training. It is however pretty interesting and worth examining. I won't do it here though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "pleasant-portsmouth",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset wiki_snippets (/home/rob/.cache/huggingface/datasets/wiki_snippets/wiki40b_en_100_0/1.0.0/d152a0e6a420c02b9b26e7f75f45fb54c818cae1d83e8f164f0b1a13ac7998ae)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfb7c4daedec445c8c5ad4fb6cb24522",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# the wiki_snippets dataset only has a training data partition \n",
    "# It was made by dicing up the wiki40b dataset\n",
    "\n",
    "# load the dataset and grab the training partition\n",
    "wiki40b_snippets = load_dataset(\"wiki_snippets\", name = \"wiki40b_en_100_0\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dimensional-render",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of rows: 17553713\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'_id': '{\"datasets_id\": 22758, \"wiki_id\": \"Q1199856\", \"sp\": 6, \"sc\": 1727, \"ep\": 6, \"ec\": 2363}',\n",
       " 'article_title': 'The Tears',\n",
       " 'datasets_id': 22758,\n",
       " 'end_character': 2363,\n",
       " 'end_paragraph': 6,\n",
       " 'passage_text': 'reinvigorated by each other\\'s company. Anderson talks excitedly of Tears songs like the ballad Asylum, inspired by his father\\'s struggle with depression, as having moved away from ‘Suede cliches or Brett Anderson cliches ... it\\'s not, you know, opiated fop territory.’\"\\nFrom the start, Anderson was insistent that the band would not be playing any songs by Suede. Things would change over time, however, as the band ended up playing the B-side, \"The Living Dead\", to an enthusiastic reception, during an encore for their show at the Sheffield Leadmill in April. In April 2005, the band\\'s first single, \"Refugees\", was released.',\n",
       " 'section_title': 'History',\n",
       " 'start_character': 1727,\n",
       " 'start_paragraph': 6,\n",
       " 'wiki_id': 'Q1199856'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check the number of rows/elements in this dataset\n",
    "print(\"number of rows:\", wiki40b_snippets.num_rows)\n",
    "print()\n",
    "# And take a look at an element\n",
    "wiki40b_snippets[123456]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "royal-warrant",
   "metadata": {},
   "source": [
    "There's a lot of info in here, but we will only be using the 'passage_text.'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "jewish-initial",
   "metadata": {},
   "source": [
    "# Get the Retribert Model and Tokenizer to make the passage text embeddings\n",
    "\n",
    "This QA system works by taking a question, embedding it using a BERT type model, and then comparing that embedding with a large set of previously calclulated embeddings over all passages in the wiki_snippets dataset. The results of this similarity search are used as a context for a BART type sequence to sequence model that generates the answer. Note that the blog also discusses sparse retrival using elastic search ... I'm not going to do that. I'm only going to make the dense index. \n",
    "\n",
    "The dense embedings are made using a so called \"RetriBERT\" model which is a small BERT model that has two different embedding heads. One head is for embedding questions, the other is for embedding passages that contain context for answers. Otherwise the BERT weights are identical. \n",
    "\n",
    "This model is trained on the ELI5 dataset. As explained in the blog, the assumption is made that the ELI5 answers are similar to the wikisnippets passages. Thus, by training the retribert model so that ELI5 questions and answers are close in the embedding space, we can assume that questions will also be close to contextually appropriate passages of wikisnippets.  \n",
    "\n",
    "Also note that the embedding space is only 128 dimensions, which is relatively small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "powered-aruba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU available. default_device set to cuda:0\n"
     ]
    }
   ],
   "source": [
    "# first check that the GPU is working. I always like to do this in case some update has messed up my PyTorch\n",
    "if torch.cuda.is_available():\n",
    "    default_device = \"cuda:0\"\n",
    "    print(\"GPU available. default_device set to cuda:0\")\n",
    "else:\n",
    "    default_device = 'cpu'\n",
    "    print(\"GPU not available. default_device set to cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "specific-start",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RetriBertModel were not initialized from the model checkpoint at yjernite/retribert-base-uncased and are newly initialized: ['bert_query.embeddings.position_ids']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# now download the RetriBERT tokenizer from the huggingface hub\n",
    "qar_tokenizer = AutoTokenizer.from_pretrained('yjernite/retribert-base-uncased')\n",
    "\n",
    "# and the RetriBERT model\n",
    "qar_model = AutoModel.from_pretrained('yjernite/retribert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bored-binary",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RetriBertConfig {\n",
      "  \"_name_or_path\": \"yjernite/retribert-base-uncased\",\n",
      "  \"architectures\": [\n",
      "    \"RetriBertModel\"\n",
      "  ],\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"model_type\": \"retribert\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 8,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"projection_dim\": 128,\n",
      "  \"share_encoders\": true,\n",
      "  \"transformers_version\": \"4.11.3\",\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# it's interesting to look at the model config. This is a pretty small embedding model ... \n",
    "# which makes it really neat that the embedding works so well!\n",
    "\n",
    "print(qar_model.config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "indonesian-phase",
   "metadata": {},
   "source": [
    "# Play with the embedding model a bit to understand how to use it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "signed-villa",
   "metadata": {},
   "source": [
    "I played a bit with embedding passages and questions and comparing the resulting embedding vectors directly. \n",
    "To figure out how to do this, I looked in the model code (specifically the RetrievalQAEmbedder class in the lfqa_utils.py file). In addition to the forward method you'll also see two methods: embed_questions and embed_answers. These allow you to do just that for questions and answers independently. The relevant bits \n",
    "of these methods only differ in the embedding head used. (Note that the forward method embeds both questions and answers and only returns the loss for training. So you don't want to use that for embedding the knowledge base.)  \n",
    "\n",
    "Another thing to note is that the Huggingface tokenizer class has changed a bit in how it can be called. So some of the code in the blog (explicitly changing to longTensors, using batch_encode_plus etc is no longer needed. You can just call the tokenizer directly.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "delayed-rabbit",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "#T_1f826858_2864_11ec_9fad_0cdd2483efb4row0_col0,#T_1f826858_2864_11ec_9fad_0cdd2483efb4row1_col0,#T_1f826858_2864_11ec_9fad_0cdd2483efb4row2_col0,#T_1f826858_2864_11ec_9fad_0cdd2483efb4row3_col0,#T_1f826858_2864_11ec_9fad_0cdd2483efb4row4_col0{\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >0</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4row0_col0\" class=\"data row0 col0\" >Ági Szalóki Life She started singing as a toddler, considering Márta Sebestyén a role model. Her musical background is traditional folk music; she first won recognition for singing with Ökrös in a traditional folk style, and Besh o droM, a Balkan gypsy brass band. With these ensembles she toured around the world from the Montreal Jazz Festival, through Glastonbury Festival to the Théatre de la Ville in Paris, from New York to Beijing.\n",
       "Since 2005, she began to pursue her solo career and explore various genres, such as jazz, thirties ballads, or children's songs.\n",
       "Until now, three of her six released albums</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4row1_col0\" class=\"data row1 col0\" >C English In English orthography, ⟨c⟩ generally represents the \"soft\" value of /s/ before the letters ⟨e⟩ (including the Latin-derived digraphs ⟨ae⟩ and ⟨oe⟩, or the corresponding ligatures ⟨æ⟩ and ⟨œ⟩), ⟨i⟩, and ⟨y⟩, and a \"hard\" value of /k/ before any other letters or at the end of a word. However, there are a number of exceptions in English: \"soccer\" and \"Celt\" are words that have /k/ where /s/ would be expected.\n",
       "The \"soft\" ⟨c⟩ may represent the /ʃ/ sound in the digraph ⟨ci⟩ when this precedes a vowel, as in the words 'delicious' and 'appreciate', and also in the</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4row2_col0\" class=\"data row2 col0\" >17th century writers. Later 19th century mumming Old Father Christmas continued to make his annual appearance in Christmas folk plays throughout the 19th century, his appearance varying considerably according to local custom. Sometimes, as in Hervey's book of 1836, he was portrayed (below left) as a hunchback.\n",
       "One unusual portrayal (below centre) was described several times by William Sandys between 1830 and 1852, all in essentially the same terms: \"Father Christmas is represented as a grotesque old man, with a large mask and comic wig, and a huge club in his hand.\" This representation is considered by the folklore scholar Peter</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4row3_col0\" class=\"data row3 col0\" >evidently been employed; I will only tell you, that we passed by winding paths, over acres and acres, with a constant varying surface, where on all sides were growing every variety of shrubs and flowers, with more than natural grace, all set in borders of greenest, closest turf, and all kept with consummate neatness.\n",
       "Birkenhead Park was used as a template for the creation of Sefton Park, which opened in Liverpool in 1872. Points of interest The Grand Entrance is at the northeast entrance to Birkenhead Park. It consists of three arches flanked by lodges and is in Ionic style. </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_1f826858_2864_11ec_9fad_0cdd2483efb4row4_col0\" class=\"data row4 col0\" >one way or another. The teams Each team in the NCBL recruits its own players, and raises its own funds however they wish to. Some teams have sponsors while other teams don't. 2006 Tier I Championship Game 3 Bottom of the 7th, the Nepean Brewers complete their sweep of Marc Sports.\n",
       "2006 NCBL Tier I Championship Game 3 on YouTube</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f30b44a5eb8>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# so let's grab some random passages and make up some questions based on them to see if the embeddings work\n",
    "idxs = [0,5000,10000,100000,1000000]\n",
    "test_passages = [p for p in wiki40b_snippets[idxs][\"passage_text\"]]\n",
    "df_passages = pd.DataFrame(test_passages)\n",
    "df_passages.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "historical-closure",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "#T_1f826859_2864_11ec_9fad_0cdd2483efb4row0_col0,#T_1f826859_2864_11ec_9fad_0cdd2483efb4row1_col0,#T_1f826859_2864_11ec_9fad_0cdd2483efb4row2_col0,#T_1f826859_2864_11ec_9fad_0cdd2483efb4row3_col0,#T_1f826859_2864_11ec_9fad_0cdd2483efb4row4_col0{\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >0</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4row0_col0\" class=\"data row0 col0\" >Who  did Ági Szalóki tour with?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4row1_col0\" class=\"data row1 col0\" >Does a digraph precede a vowel?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4row2_col0\" class=\"data row2 col0\" >where did Father Christmas make his appearance?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4row3_col0\" class=\"data row3 col0\" >What was used for the creation of Sefton park?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_1f826859_2864_11ec_9fad_0cdd2483efb4row4_col0\" class=\"data row4 col0\" >How are players recruited?</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f30b44a59b0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now make up some questions based on each of these\n",
    "test_questions = [\"Who  did Ági Szalóki tour with?\", \n",
    "             \"Does a digraph precede a vowel?\", \n",
    "             \"where did Father Christmas make his appearance?\",\n",
    "            \"What was used for the creation of Sefton park?\",\n",
    "            \"How are players recruited?\"]\n",
    "\n",
    "df_questions = pd.DataFrame(test_questions)\n",
    "df_questions.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "civic-bermuda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next tokenize the questions and the passages. Note that Huggingface tokenizers assume you are\n",
    "# passing in a list of strings to be tokenized. So we can pass in all the passages and questions\n",
    "# at the same time.\n",
    "#\n",
    "# Also note that some of the syntax regarding padding has changed since the blog was written. \n",
    "\n",
    "tokenized_test_questions = qar_tokenizer(test_questions,\n",
    "                                   max_length=128, \n",
    "                                   padding=\"max_length\", \n",
    "                                   truncation = True, \n",
    "                                   return_tensors='pt')\n",
    "\n",
    "tokenized_test_passages = qar_tokenizer(test_passages,\n",
    "                                   max_length=128, \n",
    "                                   padding=\"max_length\", \n",
    "                                   truncation = True, \n",
    "                                   return_tensors='pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "defensive-germany",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of embedded questions (5, 128)\n",
      "Size of embedded passages (5, 128)\n"
     ]
    }
   ],
   "source": [
    "# Now embed the passage and questions \n",
    "\n",
    "# First put the embedding model on the GPU and set it to eval mode\n",
    "# Note the GPU isn't really necessary for a couple passages ... it is necessary later on when we are making\n",
    "# the entire index\n",
    "qar_model.to(default_device).eval()\n",
    "\n",
    "# Next embed the questions. Use torch.no_grad() so we don't calculate gradients and have to detach them\n",
    "with torch.no_grad():\n",
    "    embedded_test_questions = qar_model.embed_questions(tokenized_test_questions[\"input_ids\"].to(default_device), \n",
    "                                                      tokenized_test_questions[\"attention_mask\"].to(default_device))\n",
    "\n",
    "# put back on the cpu and convert to numpy. Note we don't have to detach gradients because we used .no_grad()\n",
    "embedded_test_questions = embedded_test_questions.cpu().numpy()\n",
    "\n",
    "# and the passages ... note the different method used\n",
    "with torch.no_grad():\n",
    "    embedded_test_passages = qar_model.embed_answers(tokenized_test_passages[\"input_ids\"].to(default_device), \n",
    "                                                      tokenized_test_passages[\"attention_mask\"].to(default_device))\n",
    "\n",
    "embedded_test_passages = embedded_test_passages.cpu().numpy()\n",
    "\n",
    "# and print the shape of the resultant numpy arrays\n",
    "print(\"Size of embedded questions\", np.shape(embedded_test_questions))\n",
    "print(\"Size of embedded passages\", np.shape(embedded_test_passages))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "freelance-suite",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[21.93355    7.2254443 10.973549  10.100238  13.696762 ]\n",
      " [ 2.0702991 17.253456   4.4251814  2.614605  -3.751867 ]\n",
      " [10.714686   0.5131312 24.815298   9.250219  -2.1042044]\n",
      " [ 6.839991   3.0273747  5.2295    15.840267   9.030561 ]\n",
      " [ 5.577931   8.142668  -1.9331653  1.7066447 18.850311 ]]\n"
     ]
    }
   ],
   "source": [
    "# And finally, let's find the dot product of these two matricies. This *should* be largest on the diagonal\n",
    "\n",
    "test_scores = np.dot(embedded_test_questions,embedded_test_passages.T)\n",
    "print(test_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compressed-group",
   "metadata": {},
   "source": [
    "OK that seems to work. The diagonal has the largest elements in each row (or column, it's symmetrical). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "confused-veteran",
   "metadata": {},
   "source": [
    "# Embed the entire Wikipedia data set\n",
    "\n",
    "\n",
    "To make the embeddings for the entire knowledge base we essentially repeat the above steps but looping over the entire dataset and save the results. Code for doing this is in the src/qa_utils.py file, specifically the embed_passages and embed_passage_batch functions. \n",
    "\n",
    "You'll probably want to run this overnight. **It took me about 14 hours on a Nvidia 2070 Super GPU.** You can do a shorter test run by setting the n_batches argument of embed_passages to a small integer and then figure out how long the full run will take with some math. \n",
    "\n",
    "Also note that the entire embedded dataset takes up a little over 8GB of memory. This is actually relatively small because retriebert only uses an 128 dimension embedding. For reference, DPR/RAG uses BERT base which has a 768 dimensional embedding. The knowledge base takes around 70GB to store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "continent-california",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "starting to embed\n",
      "embedding 512 passages in 2 separate batches\n",
      "embedding batch 0 of 2 batches 256 passages in batch at time 0.0028543472290039062\n"
     ]
    }
   ],
   "source": [
    "# NOTE: MAKE SURE YOU WANT TO ACTUALLY RUN THIS BECAUSE IF YOU AREN'T CAREFUL YOU MAY OVERWRITE YOUR PREVIOUSLY\n",
    "# CALCULATED RESULTS!!!!!\n",
    "\n",
    "from qa_utils import embed_passages # the code is fully commented\n",
    "\n",
    "test_filename = \"../wiki_embeddings/test_embed.dat\" # change this to the full file path where you want to store the index\n",
    "\n",
    "#Note, if the file exists, running embed passages will overwrite it. Be careful and don't do that!\n",
    "if not os.path.isfile(test_filename): # here as a safety measure\n",
    "    print(\"starting to embed\")\n",
    "    embed_passages(wiki40b_snippets, \n",
    "                   qar_tokenizer, \n",
    "                   qar_model, \n",
    "                   n_batches = 2, # remove this line if you want to embed the entire dataset\n",
    "                   index_file_name = test_filename)\n",
    "else:\n",
    "    print(\"file already exists are you sure you want to overwrite and re-embed?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "subtle-wilderness",
   "metadata": {},
   "source": [
    "# Create a Faiss Index\n",
    "\n",
    "Faiss is a library out of Facebook that allows for super efficient vector similarity search on the scale of searching through billions of vectors. It does both exact search as well as various approximate searches that are quicker and applicable to extremely large sets of vectors. The tutorials  https://github.com/facebookresearch/faiss/wiki/Getting-started are a pretty quick and easy read and there are a bunch of linked papers.\n",
    "\n",
    "Search algorithms use the concept of an \"index\" which as far as I can tell is a fancy lookup table. They are easy to make with faiss. Here we just use the simplest index which does exact search. It's still pretty fast, even on CPU!  That's good because I only have an 8GB GPU so I can't fit the index there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "blind-fifth",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CHANGE THIS TO THE FILE THAT YOU STORED ALL THE EMBEDDINGS IN!\n",
    "embedding_filename =  \"../wiki_embeddings/wiki_index_embed.dat\"\n",
    "\n",
    "# make a numpy memmap of the embeddings in *read* mode. memmap is just a way of storing large numpy\n",
    "# arrays on disk and accessing them as if they were in RAM\n",
    "# you need to pass the filename, datatype, mode and the shape of the stored embeddings\n",
    "wiki40b_passage_embedding = np.memmap(embedding_filename,\n",
    "                                      dtype='float32', \n",
    "                                      mode='r',\n",
    "                                      shape=(wiki40b_snippets.num_rows, 128)\n",
    "                                     )\n",
    "\n",
    "# Create the index on the CPU\n",
    "wiki40b_index_flat = faiss.IndexFlatIP(128)\n",
    "wiki40b_index_flat.add(wiki40b_passage_embedding)\n",
    "\n",
    "# OR create the index on the GPU (If you have a enough VRAM 10GB or above probably)\n",
    "# faiss_res = faiss.StandardGpuResources()\n",
    "# wiki40b_index_flat = faiss.IndexFlatIP(128)\n",
    "# wiki40b_gpu_index = faiss.index_cpu_to_gpu(faiss_res, 0, wiki40b_index_flat)\n",
    "# wiki40b_gpu_index.add(wiki40b_passage_embedding[0:10000,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "listed-organic",
   "metadata": {},
   "source": [
    "Now we want to query the index, here's a function to do that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "taken-spanking",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_index(question, embedding_model, tokenizer, wiki_dataset, kb_index, \n",
    "                n_results=10, max_length=128, min_passage_length=20, device=\"cpu\"):\n",
    "    \n",
    "    \"\"\" This is a refactoring of blog function 'query_qa_dense_index'\"\"\"\n",
    "    \n",
    "    embedding_model.to(device)\n",
    "    \n",
    "    # embed the question\n",
    "    tokenized_question = tokenizer([question],\n",
    "                                   max_length=max_length, \n",
    "                                   padding=\"max_length\", \n",
    "                                   truncation = True, \n",
    "                                   return_tensors='pt')\n",
    "    \n",
    "    \n",
    "    with torch.no_grad():\n",
    "        embedded_question = embedding_model.embed_questions(tokenized_question[\"input_ids\"].to(device), \n",
    "                                                      tokenized_question[\"attention_mask\"].to(device))\n",
    "    \n",
    "    # now put on the cpu as numpy so we can do faiss. default, it should be on the cpu already    \n",
    "    embedded_question = embedded_question.cpu().numpy()\n",
    "    \n",
    "    # now query the index, using faiss. getting more than we need to make sure the text is long enough\n",
    "    D, I = kb_index.search(embedded_question, 2* n_results)\n",
    "   \n",
    "    # get the results of the query\n",
    "    all_wikidata = [wiki_dataset[int(k)] for k in I[0]]\n",
    "    all_passages = \"<P> \" + \" <P> \".join([p[\"passage_text\"] for p in all_wikidata])\n",
    "    \n",
    "    # this is just to make a dictionary we can look at in pandas\n",
    "    res_list = [dict([(k, p[k]) for k in wiki_dataset.column_names]) for p in all_wikidata]\n",
    "    res_list = [res for res in res_list if len(res[\"passage_text\"].split()) > min_passage_length][:n_results]\n",
    "    # add the faiss score\n",
    "    for r, sc in zip(res_list, D[0]):\n",
    "        r[\"score\"] = float(sc)\n",
    "    \n",
    "    \n",
    "    return all_passages, res_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "concerned-notice",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "#T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col2,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col0,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col1,#T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col2{\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >Article</th>        <th class=\"col_heading level0 col1\" >Sections</th>        <th class=\"col_heading level0 col2\" >Text</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col0\" class=\"data row0 col0\" >---</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col1\" class=\"data row0 col1\" >---</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row0_col2\" class=\"data row0 col2\" >--- How did the second world war end?</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col0\" class=\"data row1 col0\" >Europe</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col1\" class=\"data row1 col1\" >20th century to the present</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row1_col2\" class=\"data row1 col2\" >was to encircle Germany and cut the Germans off from Scandinavian resources. Around the same time, Germany moved troops into Denmark. The Phoney War continued.\n",
       "In May 1940, Germany attacked France through the Low Countries. France capitulated in June 1940. By August Germany began a bombing offensive on Britain, but failed to convince the Britons to give up. In 1941, Germany invaded the Soviet Union in Operation Barbarossa. On 7 December 1941 Japan's attack on Pearl Harbor drew the United States into the conflict as allies of the British Empire and other allied forces.\n",
       "After the staggering Battle of Stalingrad in 1943,</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col0\" class=\"data row2 col0\" >Military strategy</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col1\" class=\"data row2 col1\" >European Allies</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row2_col2\" class=\"data row2 col2\" >invasion of French North-Africa), Sicily and southern Italy were invaded, leading to the defeat of Fascist Italy. Churchill especially favoured a Southern strategy, aiming to attack the \"soft underbelly\" of Axis Europe through Italy, Greece and the Balkans in a strategy similar to the First World War idea of \"knocking out the supports\".  Roosevelt favoured a more direct approach through northern Europe, and with the Invasion of Normandy in June 1944, the weight of Allied effort shifted to the direct conquest of Germany.\n",
       "From 1944, as German defeat became more and more inevitable, the shape of post-war Europe assumed greater</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col0\" class=\"data row3 col0\" >Battle of Arras (1917)</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col1\" class=\"data row3 col1\" >Home fronts</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row3_col2\" class=\"data row3 col2\" >the war. Hundreds of thousands of casualties had been suffered at the battles of Gallipoli, the Somme and Verdun, with little prospect of victory in sight. The British Prime Minister, H. H. Asquith, resigned in early December 1916 and was succeeded by David Lloyd George. In France, Prime Minister Aristide Briand, along with Minister of Defence Hubert Lyautey were politically diminished and resigned in March 1917, following disagreements over the prospective Nivelle Offensive. The United States was close to declaring war on Germany; American public opinion was growing increasingly incensed by U-boat attacks upon civilian shipping, which had begun with</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col0\" class=\"data row4 col0\" >Germany–United Kingdom relations</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col1\" class=\"data row4 col1\" >World War II & Occupation</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row4_col2\" class=\"data row4 col2\" >Germany, but the United States greatly funded and supplied the British. In December 1941, United States entered the war against Germany and Japan after the attack on Pearl Harbor by Japan, which also later overwhelmed British outposts in the Pacific from Hong Kong to Singapore.\n",
       "The Allied invasion of France on D-Day in June 1944 as well as strategic bombing and land forces all contributed to the final defeat of Germany. Occupation As part of the Yalta and Potsdam agreements, Britain took control of its own sector in occupied Germany. It soon merged its sector with the American and French sectors,</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row5\" class=\"row_heading level0 row5\" >5</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col0\" class=\"data row5 col0\" >Italian invasion of France</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col1\" class=\"data row5 col1\" >Italian imperial ambitions & Battle of France</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row5_col2\" class=\"data row5 col2\" >impression of weakness\". Germany supplied Italy with about one million tons of coal a month beginning in the spring of 1940, an amount that even exceeded Mussolini's demand of August 1939 that Italy receive six million tons of coal for its first twelve months of war. Battle of France On 1 September 1939, Germany invaded Poland. Following a month of war, Poland was defeated. A period of inaction, called the Phoney War, then followed between the Allies and Germany. On 10 May 1940, this inactivity ended as Germany began Fall Gelb (Case Yellow) against France and the neutral nations of</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row6\" class=\"row_heading level0 row6\" >6</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col0\" class=\"data row6 col0\" >Military history of Germany</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col1\" class=\"data row6 col1\" >First World War (1914–18) & Weimar Republic and the Third Reich (1918–39)</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row6_col2\" class=\"data row6 col2\" >the British tank attack at the Battle of Cambrai.\n",
       "In March 1918 the German army Spring Offensive began an impressive advance creating a salient in the allied line. The offensive stalled as the British and French fell back and then counterattacked. The Germans did not have the airpower or tanks to secure their battlefield gains. The Allies, invigorated by American manpower, money, and food, counterattacked in late summer and rolled over the depleted German lines, as the German navy rebelled and support for the war on the homefront evaporated. Weimar Republic and the Third Reich (1918–39) The treaty of Versailles imposed</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row7\" class=\"row_heading level0 row7\" >7</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col0\" class=\"data row7 col0\" >Battle of France</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col1\" class=\"data row7 col1\" >Occupation</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row7_col2\" class=\"data row7 col2\" >docked at ports in Vichy France and North Africa and use them in an invasion of Britain (Operation Sea Lion). Within a month, the Royal Navy attacked the French naval forces stationed in North Africa in the Attack on Mers-el-Kébir. The British Chiefs of Staff Committee had concluded in May 1940 that if France collapsed, \"we do not think we could continue the war with any chance of success\" without \"full economic and financial support\" from the United States. Churchill's desire for American aid led in September to the Destroyers for Bases agreement that began the wartime Anglo-American partnership.\n",
       "The occupation</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row8\" class=\"row_heading level0 row8\" >8</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col0\" class=\"data row8 col0\" >Spanish Civil War</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col1\" class=\"data row8 col1\" >Francoist repression after the war and Republican exile & International relations</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row8_col2\" class=\"data row8 col2\" >returned. The third wave occurred after the War, at the end of March 1939, when thousands of Republicans tried to board ships to exile, although few succeeded. International relations The political and emotional repercussions of the War transcended the National scale, becoming a precursor to World War II. The war has frequently been described as the \"prelude to\" or the \"opening round of\" the Second World War, as part of an international battle against fascism. However Stanley Payne suggests this isn't accurate, arguing that the international alliance that was created in December 1941, once the United States entered WW2, was</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row9\" class=\"row_heading level0 row9\" >9</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col0\" class=\"data row9 col0\" >World War II</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col1\" class=\"data row9 col1\" >Allies gain momentum (1943–44)</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row9_col2\" class=\"data row9 col2\" >named the Italian Social Republic, causing an Italian civil war. The Western Allies fought through several lines until reaching the main German defensive line in mid-November.\n",
       "German operations in the Atlantic also suffered. By May 1943, as Allied counter-measures became increasingly effective, the resulting sizeable German submarine losses forced a temporary halt of the German Atlantic naval campaign. In November 1943, Franklin D. Roosevelt and Winston Churchill met with Chiang Kai-shek in Cairo and then with Joseph Stalin in Tehran. The former conference determined the post-war return of Japanese territory and the military planning for the Burma Campaign, while the latter</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4level0_row10\" class=\"row_heading level0 row10\" >10</th>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col0\" class=\"data row10 col0\" >France–United Kingdom relations</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col1\" class=\"data row10 col1\" >Second World War & 1945-1956</td>\n",
       "                        <td id=\"T_33417c94_2864_11ec_9fad_0cdd2483efb4row10_col2\" class=\"data row10 col2\" >Vichy (until October 1942) and avoided recognition of de Gaulle. Churchill, caught between the U.S. and de Gaulle, tried to find a compromise. 1945-1956 Following D-Day, relations between the two peoples were at a high, as the British were greeted as liberators. Following the surrender of Germany in May 1945, the UK and France became close as both feared the Americans would withdraw from Europe leaving them vulnerable to the Soviet Union's expanding communist bloc. The UK was successful in strongly advocating that France be given a zone of occupied Germany. Both states were amongst the five Permanent Members of</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f30b5a49a90>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# and we can ask a question and return some contextual passages\n",
    "\n",
    "test_question = \"How did the second world war end?\"\n",
    "\n",
    "all_passages, res_list = query_index(test_question, qar_model, qar_tokenizer, \n",
    "                   wiki40b_snippets, wiki40b_index_flat, device='cpu')\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Article': ['---'] + [res['article_title'] for res in res_list],\n",
    "    'Sections': ['---'] + [res['section_title'] if res['section_title'].strip() != '' else res['article_title']\n",
    "                 for res in res_list],\n",
    "    'Text': ['--- ' + test_question] + [res['passage_text'] for res in res_list],\n",
    "})\n",
    "df.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "certified-extraction",
   "metadata": {},
   "source": [
    "# Now that we have the index, we can ask a question and get a generated response.\n",
    "\n",
    "The contextual passages are inputs into a BART-large sequence-to-sequence model. Essentially the question and contextual passages are tokenized, concatenated together and then  then truncated to 1024 tokens. This is then passed through the BART model to create the response. The blog post goes through the fine-tuning proceedure.\n",
    "\n",
    "Here are the steps that you can take to answer a list of questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "large-indian",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First load up the BART model \n",
    "qa_s2s_tokenizer = AutoTokenizer.from_pretrained('yjernite/bart_eli5')\n",
    "qa_s2s_model = AutoModelForSeq2SeqLM.from_pretrained('yjernite/bart_eli5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "wooden-death",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# put the BERT embedding model and the BART seq-2-seq model where you want them\n",
    "# I put the embedding model on my cpu (because the index is large) and the bart model on my GPU\n",
    "# also put them in eval mode\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "qar_device = \"cpu\"\n",
    "qa_s2s_device = \"cuda\"\n",
    "\n",
    "qar_model.to(qar_device).eval()\n",
    "qa_s2s_model.to(qa_s2s_device).eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "neural-arizona",
   "metadata": {},
   "outputs": [],
   "source": [
    "# come up with a list of questions:\n",
    "questions = [\"How do birds fly?\",\n",
    "             \"How do fish breath water?\",\n",
    "             \"How did the polynesians travel across the ocean?\",\n",
    "            \"When did we learn to cook food?\",\n",
    "            \"Will be ever be able to live on Mars? Why or why not?\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "expressed-storage",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style  type=\"text/css\" >\n",
       "#T_464af522_2864_11ec_9fad_0cdd2483efb4row0_col0,#T_464af522_2864_11ec_9fad_0cdd2483efb4row0_col1,#T_464af522_2864_11ec_9fad_0cdd2483efb4row1_col0,#T_464af522_2864_11ec_9fad_0cdd2483efb4row1_col1,#T_464af522_2864_11ec_9fad_0cdd2483efb4row2_col0,#T_464af522_2864_11ec_9fad_0cdd2483efb4row2_col1,#T_464af522_2864_11ec_9fad_0cdd2483efb4row3_col0,#T_464af522_2864_11ec_9fad_0cdd2483efb4row3_col1,#T_464af522_2864_11ec_9fad_0cdd2483efb4row4_col0,#T_464af522_2864_11ec_9fad_0cdd2483efb4row4_col1{\n",
       "            text-align:  left;\n",
       "        }</style><table id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4\" ><thead>    <tr>        <th class=\"blank level0\" ></th>        <th class=\"col_heading level0 col0\" >Question</th>        <th class=\"col_heading level0 col1\" >Answer</th>    </tr></thead><tbody>\n",
       "                <tr>\n",
       "                        <th id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row0_col0\" class=\"data row0 col0\" >How do birds fly?</td>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row0_col1\" class=\"data row0 col1\" >Birds use their wings to generate lift. When they flap their wings, the up and down motion of the wing pushes the air in front of them, creating lift.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4level0_row1\" class=\"row_heading level0 row1\" >1</th>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row1_col0\" class=\"data row1 col0\" >How do fish breath water?</td>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row1_col1\" class=\"data row1 col1\" >Fish don't \"breath\" water. Their gills absorb the oxygen in the water and exhale it as gaseous oxygen. It's the same way we exhale CO2.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4level0_row2\" class=\"row_heading level0 row2\" >2</th>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row2_col0\" class=\"data row2 col0\" >How did the polynesians travel across the ocean?</td>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row2_col1\" class=\"data row2 col1\" >The Polynesians didn't travel across the ocean. They traveled across the Polynesian Triangle, which is a huge stretch of ocean between the islands of Hawaii and New Caledonia.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4level0_row3\" class=\"row_heading level0 row3\" >3</th>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row3_col0\" class=\"data row3 col0\" >When did we learn to cook food?</td>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row3_col1\" class=\"data row3 col1\" >Cooking has been around for as long as humans have been around. It's just a matter of getting it right the first time you do it. If you don't know how to cook, you won't be able to do it the next time.</td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                        <th id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4level0_row4\" class=\"row_heading level0 row4\" >4</th>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row4_col0\" class=\"data row4 col0\" >Will be ever be able to live on Mars? Why or why not?</td>\n",
       "                        <td id=\"T_464af522_2864_11ec_9fad_0cdd2483efb4row4_col1\" class=\"data row4 col1\" >Yes, we will be able to live on Mars in the future. The problem is that we don't have the technology to do so right now. We don't even know how to get there.</td>\n",
       "            </tr>\n",
       "    </tbody></table>"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f30b405b278>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# now loop through these and answer them\n",
    "\n",
    "answers = []\n",
    "\n",
    "for i in range(len(questions)):\n",
    "    # create support document with the dense index\n",
    "    question = questions[i]\n",
    "    \n",
    "    doc, res_list = query_index(question, qar_model, qar_tokenizer, \n",
    "                   wiki40b_snippets, wiki40b_index_flat, device=qar_device)\n",
    "    \n",
    "    # concatenate question and support document into BART input\n",
    "    question_doc = \"question: {} context: {}\".format(question, doc)\n",
    "    \n",
    "    # concatenate question and support document into BART input\n",
    "    question_doc = \"question: {} context: {}\".format(question, doc)\n",
    "\n",
    "    # tokenize this with the BART tokenizer. Note we are truncating the total inputs at max_input_length\n",
    "    s2s_inputs = qa_s2s_tokenizer([question_doc],\n",
    "                                       max_length=1024, \n",
    "                                       padding=\"max_length\", \n",
    "                                       truncation = True, \n",
    "                                       return_tensors='pt').to(qa_s2s_device)\n",
    "\n",
    "    # now feed this into the BART sequence to sequence model and generate answers\n",
    "    generated_ids = qa_s2s_model.generate(input_ids=s2s_inputs[\"input_ids\"],\n",
    "                                              attention_mask=s2s_inputs[\"attention_mask\"],\n",
    "                                              min_length=32,\n",
    "                                              max_length=128,\n",
    "                                              do_sample=False,\n",
    "                                              early_stopping=True,\n",
    "                                              num_beams=8,\n",
    "                                              temperature=1.0,\n",
    "                                              top_k=None,\n",
    "                                              top_p=None,\n",
    "                                              eos_token_id=qa_s2s_tokenizer.eos_token_id,\n",
    "                                              no_repeat_ngram_size=3,\n",
    "                                              num_return_sequences=1,\n",
    "                                              decoder_start_token_id=qa_s2s_tokenizer.bos_token_id)[0]\n",
    "\n",
    "    # and make the list of answers by decoding the generated ids\n",
    "    answer = qa_s2s_tokenizer.decode(generated_ids, skip_special_tokens=True).strip()\n",
    "\n",
    "    answers += [answer]\n",
    "\n",
    "# And make a pretty dataframe to show the results\n",
    "df = pd.DataFrame({\n",
    "    'Question': questions,\n",
    "    'Answer': answers,\n",
    "})\n",
    "df.style.set_properties(**{'text-align': 'left'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "allied-meeting",
   "metadata": {},
   "source": [
    "And that's basically it. I found it impressive that one could do so well with such a simple implementation using relatively small models. One can think of all sorts of tricks that could make this better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "legitimate-leadership",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}