Experiments/huggingface/transformers/question_answering-tf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X4cRE8IbIrIV"
   },
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets. Uncomment the following cell and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "MOsHUjgdIrIW",
    "outputId": "f84a093e-147f-470e-aad9-80fb51193c8e"
   },
   "outputs": [],
   "source": [
    "#! pip install transformers datasets huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this notebook locally, make sure your environment has an install from the last version of those libraries.\n",
    "\n",
    "To be able to share your model with the community and generate results like the one shown in the picture below via the inference API, there are a few more steps to follow.\n",
    "\n",
    "First you have to store your authentication token from the Hugging Face website (sign up [here](https://huggingface.co/join) if you haven't already!) then uncomment the following cell and input your token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you need to install Git-LFS and setup Git if you haven't already. Uncomment the following instructions and adapt with your name and email:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !apt install git-lfs\n",
    "# !git config --global user.email \"you@example.com\"\n",
    "# !git config --global user.name \"Your Name\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure your version of Transformers is at least 4.16.0 since the functionality was introduced in that version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.21.0.dev0\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "print(transformers.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HFASsisvIrIb"
   },
   "source": [
    "You can find a script version of this notebook to fine-tune your model in a distributed fashion using multiple GPUs or TPUs [here](https://github.com/huggingface/transformers/tree/master/examples/question-answering)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rEJBSTyZIrIb"
   },
   "source": [
    "# Fine-tuning a model on a question-answering task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will see how to fine-tune one of the [🤗 Transformers](https://github.com/huggingface/transformers) model to a question answering task, which is the task of extracting the answer to a question from a given context. We will see how to easily load a dataset for these kinds of tasks and use Keras to fine-tune a model on it. Note that this model **does not generate new text!** Instead, it selects a span of the input passage as the answer.\n",
    "\n",
    "![Widget inference representing the QA task](images/question_answering.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4RRkXuteIrIh"
   },
   "source": [
    "This notebook is built to run on any question answering task with the same format as SQUAD (version 1 or 2), with any model checkpoint from the [Model Hub](https://huggingface.co/models) as long as that model has a version with a token classification head and a fast tokenizer (check on [this table](https://huggingface.co/transformers/index.html#bigtable) if this is the case). It might, however, need some small adjustments if you decide to use a different dataset than the one used here. Depending on your model and the GPU you are using, you might need to adjust the batch size to avoid out-of-memory errors. Set those three parameters, then the rest of the notebook should run smoothly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "zVvslsfMIrIh"
   },
   "outputs": [],
   "source": [
    "# This flag is the difference between SQUAD v1 or 2 (if you're using another dataset, it indicates if impossible\n",
    "# answers are allowed or not).\n",
    "squad_v2 = False\n",
    "model_checkpoint = \"distilbert-base-uncased\"\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "whPRbBNbIrIl"
   },
   "source": [
    "## Loading the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W7QYTpxXIrIl"
   },
   "source": [
    "We will use the [🤗 Datasets](https://github.com/huggingface/datasets) library to download the data and get the metric we need to use for evaluation (to compare our model to the benchmark). This can be easily done with the functions `load_dataset` and `load_metric`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "IreSlFmlIrIm"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset, load_metric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CKx2zKs5IrIq"
   },
   "source": [
    "For our example here, we'll use the [SQUAD dataset](https://rajpurkar.github.io/SQuAD-explorer/). The notebook should work with any question answering dataset in the 🤗 Datasets library. If you're using your own dataset in a JSON or CSV file (see the [Datasets documentation](https://huggingface.co/docs/datasets/loading_datasets.html#from-local-files) on how to load them), it might need some adjustments to the column names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 270,
     "referenced_widgets": [
      "69caab03d6264fef9fc5649bffff5e20",
      "3f74532faa86412293d90d3952f38c4a",
      "50615aa59c7247c4804ca5cbc7945bd7",
      "fe962391292a413ca55dc932c4279fa7",
      "299f4b4c07654e53a25f8192bd1d7bbd",
      "ad04ed1038154081bbb0c1444784dcc2",
      "7c667ad22b5740d5a6319f1b1e3a8097",
      "46c2b043c0f84806978784a45a4e203b",
      "80e2943be35f46eeb24c8ab13faa6578",
      "de5956b5008d4fdba807bae57509c393",
      "931db1f7a42f4b46b7ff8c2e1262b994",
      "6c1db72efff5476e842c1386fadbbdba",
      "ccd2f37647c547abb4c719b75a26f2de",
      "d30a66df5c0145e79693e09789d96b81",
      "5fa26fc336274073abbd1d550542ee33",
      "2b34de08115d49d285def9269a53f484",
      "d426be871b424affb455aeb7db5e822e",
      "160bf88485f44f5cb6eaeecba5e0901f",
      "745c0d47d672477b9bb0dae77b926364",
      "d22ab78269cd4ccfbcf70c707057c31b",
      "d298eb19eeff453cba51c2804629d3f4",
      "a7204ade36314c86907c562e0a2158b8",
      "e35d42b2d352498ca3fc8530393786b2",
      "75103f83538d44abada79b51a1cec09e",
      "f6253931d90543e9b5fd0bb2d615f73a",
      "051aa783ff9e47e28d1f9584043815f5",
      "0984b2a14115454bbb009df71c1cf36f",
      "8ab9dfce29854049912178941ef1b289",
      "c9de740e007141958545e269372780a4",
      "cbea68b25d6d4ba09b2ce0f27b1726d5",
      "5781fc45cf8d486cb06ed68853b2c644",
      "d2a92143a08a4951b55bab9bc0a6d0d3",
      "a14c3e40e5254d61ba146f6ec88eae25",
      "c4ffe6f624ce4e978a0d9b864544941a",
      "1aca01c1d8c940dfadd3e7144bb35718",
      "9fbbaae50e6743f2aa19342152398186",
      "fea27ca6c9504fc896181bc1ff5730e5",
      "940d00556cb849b3a689d56e274041c2",
      "5cdf9ed939fb42d4bf77301c80b8afca",
      "94b39ccfef0b4b08bf2fb61bb0a657c1",
      "9a55087c85b74ea08b3e952ac1d73cbe",
      "2361ab124daf47cc885ff61f2899b2af",
      "1a65887eb37747ddb75dc4a40f7285f2",
      "3c946e2260704e6c98593136bd32d921",
      "50d325cdb9844f62a9ecc98e768cb5af",
      "aa781f0cfe454e9da5b53b93e9baabd8",
      "6bb68d3887ef43809eb23feb467f9723",
      "7e29a8b952cf4f4ea42833c8bf55342f",
      "dd5997d01d8947e4b1c211433969b89b",
      "2ace4dc78e2f4f1492a181bcd63304e7",
      "bbee008c2791443d8610371d1f16b62b",
      "31b1c8a2e3334b72b45b083688c1a20c",
      "7fb7c36adc624f7dbbcb4a831c1e4f63",
      "0b7c8f1939074794b3d9221244b1344d",
      "a71908883b064e1fbdddb547a8c41743",
      "2f5223f26c8541fc87e91d2205c39995"
     ]
    },
    "id": "s_AY1ATSIrIq",
    "outputId": "fd0578d1-8895-443d-b56f-5908de9f1b6b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset squad (/home/matt/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "354ff956c46a4157b2b39ec80be14272",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "datasets = load_dataset(\"squad_v2\" if squad_v2 else \"squad\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzfPtOMoIrIu"
   },
   "source": [
    "The `datasets` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "GWiVUF0jIrIv",
    "outputId": "35e3ea43-f397-4a54-c90c-f2cf8d36873e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 87599\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 10570\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the training, validation and test sets all have a column for the context, the question and the answers to those questions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u3EtYfeHIrIz"
   },
   "source": [
    "To access an actual element, you need to select a split first, then give an index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "X6HrpprwIrIz",
    "outputId": "d7670bc0-42e4-4c09-8a6a-5c018ded7d95"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '5733be284776f41900661182',\n",
       " 'title': 'University_of_Notre_Dame',\n",
       " 'context': 'Architecturally, the school has a Catholic character. Atop the Main Building\\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \"Venite Ad Me Omnes\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.',\n",
       " 'question': 'To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?',\n",
       " 'answers': {'text': ['Saint Bernadette Soubirous'], 'answer_start': [515]}}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the answers are indicated by their start position in the text (here at character 515) and their full text, which is a substring of the context as we mentioned above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WHUmphG3IrI3"
   },
   "source": [
    "To get a sense of what the data looks like, the following function will show some examples picked randomly from the dataset and decoded back to strings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "i3j8APAoIrI3"
   },
   "outputs": [],
   "source": [
    "from datasets import ClassLabel, Sequence\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "\n",
    "def show_random_elements(dataset, num_examples=10):\n",
    "    assert num_examples <= len(\n",
    "        dataset\n",
    "    ), \"Can't pick more elements than there are in the dataset.\"\n",
    "    picks = []\n",
    "    for _ in range(num_examples):\n",
    "        pick = random.randint(0, len(dataset) - 1)\n",
    "        while pick in picks:\n",
    "            pick = random.randint(0, len(dataset) - 1)\n",
    "        picks.append(pick)\n",
    "\n",
    "    df = pd.DataFrame(dataset[picks])\n",
    "    for column, typ in dataset.features.items():\n",
    "        if isinstance(typ, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):\n",
    "            df[column] = df[column].transform(\n",
    "                lambda x: [typ.feature.names[i] for i in x]\n",
    "            )\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "SZy5tRB_IrI7",
    "outputId": "ba8f2124-e485-488f-8c0c-254f34f24f13",
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>context</th>\n",
       "      <th>question</th>\n",
       "      <th>answers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>56f71ecf711bf01900a449a5</td>\n",
       "      <td>Treaty</td>\n",
       "      <td>When a state limits its treaty obligations through reservations, other states party to that treaty have the option to accept those reservations, object to them, or object and oppose them. If the state accepts them (or fails to act at all), both the reserving state and the accepting state are relieved of the reserved legal obligation as concerns their legal obligations to each other (accepting the reservation does not change the accepting state's legal obligations as concerns other parties to the treaty). If the state opposes, the parts of the treaty affected by the reservation drop out completely and no longer create any legal obligations on the reserving and accepting state, again only as concerns each other. Finally, if the state objects and opposes, there are no legal obligations under that treaty between those two state parties whatsoever. The objecting and opposing state essentially refuses to acknowledge the reserving state is a party to the treaty at all.</td>\n",
       "      <td>Who remains unaffected when a party's reservation is accepted by a second party?</td>\n",
       "      <td>{'text': ['other parties to the treaty'], 'answer_start': [480]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5730ec8305b4da19006bcc47</td>\n",
       "      <td>United_States_Air_Force</td>\n",
       "      <td>Due to the Budget sequestration in 2013, the USAF was forced to ground many of its squadrons. The Commander of Air Combat Command, General Mike Hostage indicated that the USAF must reduce its F-15 and F-16 fleets and eliminate platforms like the A-10 in order to focus on a fifth-generation jet fighter future. In response to squadron groundings and flight time reductions, many Air Force pilots have opted to resign from active duty and enter the Air Force Reserve and Air National Guard while pursuing careers in the commercial airlines where they can find flight hours on more modern aircraft.</td>\n",
       "      <td>Who was the Commander of Air Combat Command in 2013?</td>\n",
       "      <td>{'text': ['General Mike Hostage'], 'answer_start': [131]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>56fb7d7a8ddada1400cd6476</td>\n",
       "      <td>Middle_Ages</td>\n",
       "      <td>Under the Capetian dynasty France slowly began to expand its authority over the nobility, growing out of the Île-de-France to exert control over more of the country in the 11th and 12th centuries. They faced a powerful rival in the Dukes of Normandy, who in 1066 under William the Conqueror (duke 1035–1087), conquered England (r. 1066–87) and created a cross-channel empire that lasted, in various forms, throughout the rest of the Middle Ages. Normans also settled in Sicily and southern Italy, when Robert Guiscard (d. 1085) landed there in 1059 and established a duchy that later became the Kingdom of Sicily. Under the Angevin dynasty of Henry II (r. 1154–89) and his son Richard I (r. 1189–99), the kings of England ruled over England and large areas of France,[W] brought to the family by Henry II's marriage to Eleanor of Aquitaine (d. 1204), heiress to much of southern France.[X] Richard's younger brother John (r. 1199–1216) lost Normandy and the rest of the northern French possessions in 1204 to the French King Philip II Augustus (r. 1180–1223). This led to dissension among the English nobility, while John's financial exactions to pay for his unsuccessful attempts to regain Normandy led in 1215 to Magna Carta, a charter that confirmed the rights and privileges of free men in England. Under Henry III (r. 1216–72), John's son, further concessions were made to the nobility, and royal power was diminished. The French monarchy continued to make gains against the nobility during the late 12th and 13th centuries, bringing more territories within the kingdom under their personal rule and centralising the royal administration. Under Louis IX (r. 1226–70), royal prestige rose to new heights as Louis served as a mediator for most of Europe.[Y]</td>\n",
       "      <td>During what period did William reign over England?</td>\n",
       "      <td>{'text': ['1066–87'], 'answer_start': [331]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>572e9a3fc246551400ce43c6</td>\n",
       "      <td>Steven_Spielberg</td>\n",
       "      <td>Studio producers Richard D. Zanuck and David Brown offered Spielberg the director's chair for Jaws, a thriller-horror film based on the Peter Benchley novel about an enormous killer shark. Spielberg has often referred to the gruelling shoot as his professional crucible. Despite the film's ultimate, enormous success, it was nearly shut down due to delays and budget over-runs. But Spielberg persevered and finished the film. It was an enormous hit, winning three Academy Awards (for editing, original score and sound) and grossing more than $470 million worldwide at the box office. It also set the domestic record for box office gross, leading to what the press described as \"Jawsmania.\":248 Jaws made Spielberg a household name and one of America's youngest multi-millionaires, allowing him a great deal of autonomy for his future projects.:250 It was nominated for Best Picture and featured Spielberg's first of three collaborations with actor Richard Dreyfuss.</td>\n",
       "      <td>How many Academy Awards did the film \"Jaws\" win?</td>\n",
       "      <td>{'text': ['three'], 'answer_start': [458]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>56d4c4e72ccc5a1400d8321b</td>\n",
       "      <td>Beyoncé</td>\n",
       "      <td>On January 7, 2012, Beyoncé gave birth to her first child, a daughter, Blue Ivy Carter, at Lenox Hill Hospital in New York. Five months later, she performed for four nights at Revel Atlantic City's Ovation Hall to celebrate the resort's opening, her first performances since giving birth to Blue Ivy.</td>\n",
       "      <td>Where was Beyoncé's first public performance after giving birth?</td>\n",
       "      <td>{'text': ['Revel Atlantic City's Ovation Hall'], 'answer_start': [176]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>573212a80fdd8d15006c6758</td>\n",
       "      <td>Party_leaders_of_the_United_States_House_of_Representatives</td>\n",
       "      <td>During this early period, it was more usual that neither major party grouping (Federalists and Democratic-Republicans) had an official leader. In 1813, for instance, a scholar recounts that the Federalist minority of 36 Members needed a committee of 13 \"to represent a party comprising a distinct minority\" and \"to coordinate the actions of men who were already partisans in the same cause.\" In 1828, a foreign observer of the House offered this perspective on the absence of formal party leadership on Capitol Hill:</td>\n",
       "      <td>In early 19th century, what were 2 common parties?</td>\n",
       "      <td>{'text': ['(Federalists and Democratic-Republicans)'], 'answer_start': [78]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>572fd29ea23a5019007fca49</td>\n",
       "      <td>Bacteria</td>\n",
       "      <td>In ordinary circumstances, transduction, conjugation, and transformation involve transfer of DNA between individual bacteria of the same species, but occasionally transfer may occur between individuals of different bacterial species and this may have significant consequences, such as the transfer of antibiotic resistance. In such cases, gene acquisition from other bacteria or the environment is called horizontal gene transfer and may be common under natural conditions. Gene transfer is particularly important in antibiotic resistance as it allows the rapid transfer of resistance genes between different pathogens.</td>\n",
       "      <td>What is horizontal gene transfer?</td>\n",
       "      <td>{'text': ['gene acquisition from other bacteria or the environment'], 'answer_start': [339]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>5727e3c24b864d1900163f55</td>\n",
       "      <td>Oklahoma</td>\n",
       "      <td>More than 12,000 miles (19,000 km) of roads make up the state's major highway skeleton, including state-operated highways, ten turnpikes or major toll roads, and the longest drivable stretch of Route 66 in the nation. In 2008, Interstate 44 in Oklahoma City was Oklahoma's busiest highway, with a daily traffic volume of 123,300 cars. In 2010, the state had the nation's third highest number of bridges classified as structurally deficient, with nearly 5,212 bridges in disrepair, including 235 National Highway System Bridges.</td>\n",
       "      <td>Oklahoma has the longest drivable stretch of what famous highway?</td>\n",
       "      <td>{'text': ['Route 66'], 'answer_start': [194]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>5733ecdb4776f41900661524</td>\n",
       "      <td>Portugal</td>\n",
       "      <td>The President, who is elected to a five-year term, has an executive role: the current President is Aníbal Cavaco Silva. The Assembly of the Republic is a single chamber parliament composed of 230 deputies elected for a four-year term. The Government is headed by the Prime Minister (currently António Costa) and includes Ministers and Secretaries of State. The Courts are organized into several levels, among the judicial, administrative and fiscal branches. The Supreme Courts are institutions of last resort/appeal. A thirteen-member Constitutional Court oversees the constitutionality of the laws.</td>\n",
       "      <td>How many members sit on the Constitutional Court?</td>\n",
       "      <td>{'text': ['thirteen'], 'answer_start': [520]}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>5733796c4776f41900660b65</td>\n",
       "      <td>Saint_Barth%C3%A9lemy</td>\n",
       "      <td>Marine mammals are many, such as the dolphins, porpoises and whales, which are seen here during the migration period from December till May. Turtles are a common sight along the coastline of the island. They are a protected species and in the endangered list. It is stated that it will take 15–50 years for this species to attain reproductive age. Though they live in the sea, the females come to the shore to lay eggs and are protected by private societies. Three species of turtles are particularly notable. These are: The leatherback sea turtles which have leather skin instead of a shell and are the largest of the type found here, some times measuring a much as 3 m (average is about 1.5 m) and weighing about 450 kg (jellyfish is their favourite diet); the hawksbill turtles, which have hawk-like beaks and found near reefs, generally about 90 cm in diameter and weigh about 60 kg and their diet consists of crabs and snails; and the green turtles, herbivores which have rounded heads, generally about 90 cm in diameter and live amidst tall sea grasses.</td>\n",
       "      <td>Where do green turtles live?</td>\n",
       "      <td>{'text': ['amidst tall sea grasses'], 'answer_start': [1035]}</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_random_elements(datasets[\"train\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n9qywopnIrJH"
   },
   "source": [
    "## Preprocessing the training data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YVx71GdAIrJH"
   },
   "source": [
    "Before we can feed those texts to our model, we need to preprocess them. This is done by a 🤗 Transformers `Tokenizer` which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pretrained vocabulary) and put it in a format the model expects, as well as generate the other inputs that model requires.\n",
    "\n",
    "To do all of this, we instantiate our tokenizer with the `AutoTokenizer.from_pretrained` method, which will ensure:\n",
    "\n",
    "- we get a tokenizer that corresponds to the model architecture we want to use,\n",
    "- we download the vocabulary used when pretraining this specific checkpoint.\n",
    "\n",
    "That vocabulary will be cached, so it's not downloaded again the next time we run the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "eXNLu_-nIrJI"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vl6IidfdIrJK"
   },
   "source": [
    "The following assertion ensures that our tokenizer is a fast tokenizer (backed by Rust) from the 🤗 Tokenizers library. Those fast tokenizers are available for almost all models, and we will need some of the special features they have for our preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "\n",
    "assert isinstance(tokenizer, transformers.PreTrainedTokenizerFast)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check which type of models have a fast tokenizer available and which don't in the [big table of models](https://huggingface.co/transformers/index.html#bigtable)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rowT4iCLIrJK"
   },
   "source": [
    "You can directly call this tokenizer on two sentences (one for the answer, one for the context):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "a5hBlsrHIrJL",
    "outputId": "acdaa98a-a8cd-4a20-89b8-cc26437bbe90"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101, 2054, 2003, 2115, 2171, 1029, 102, 2026, 2171, 2003, 25353, 22144, 2378, 1012, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer(\"What is your name?\", \"My name is Sylvain.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on the model you selected, you will see different keys in the dictionary returned by the cell above. They don't matter much for what we're doing here (just know they are required by the model we will instantiate later), you can learn more about them in [this tutorial](https://huggingface.co/transformers/preprocessing.html) if you're interested.\n",
    "\n",
    "Now one specific thing for the preprocessing in question answering is how to deal with very long documents. We usually truncate them in other tasks, when they are longer than the model maximum sentence length, but here, removing part of the the context might result in losing the answer we are looking for. To deal with this, we will allow one (long) example in our dataset to give several input features, each of length shorter than the maximum length of the model (or the one we set as a hyper-parameter). Also, just in case the answer lies at the point we split a long context, we allow some overlap between the features we generate controlled by the hyper-parameter `doc_stride`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_length = 384  # The maximum length of a feature (question and context)\n",
    "doc_stride = 128  # The allowed overlap between two part of the context when splitting is performed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's find one long example in our dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, example in enumerate(datasets[\"train\"]):\n",
    "    if len(tokenizer(example[\"question\"], example[\"context\"])[\"input_ids\"]) > 384:\n",
    "        break\n",
    "example = datasets[\"train\"][i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without any truncation, we get the following length for the input IDs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "396"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer(example[\"question\"], example[\"context\"])[\"input_ids\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, if we just truncate, we will lose information (and possibly the answer to our question):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "384"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(\n",
    "    tokenizer(\n",
    "        example[\"question\"],\n",
    "        example[\"context\"],\n",
    "        max_length=max_length,\n",
    "        truncation=\"only_second\",\n",
    "    )[\"input_ids\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we never want to truncate the question, only the context, and so we use the `only_second` truncation method. Our tokenizer can automatically return a list of features capped by a certain maximum length, with the overlap we talked about above, we just have to tell it to do so with `return_overflowing_tokens=True` and by passing the stride:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_example = tokenizer(\n",
    "    example[\"question\"],\n",
    "    example[\"context\"],\n",
    "    max_length=max_length,\n",
    "    truncation=\"only_second\",\n",
    "    return_overflowing_tokens=True,\n",
    "    stride=doc_stride,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we don't have one list of `input_ids`, but several: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[384, 157]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[len(x) for x in tokenized_example[\"input_ids\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And if we decode them, we can see the overlap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS] how many wins does the notre dame men's basketball team have? [SEP] the men's basketball team has over 1, 600 wins, one of only 12 schools who have reached that mark, and have appeared in 28 ncaa tournaments. former player austin carr holds the record for most points scored in a single game of the tournament with 61. although the team has never won the ncaa tournament, they were named by the helms athletic foundation as national champions twice. the team has orchestrated a number of upsets of number one ranked teams, the most notable of which was ending ucla's record 88 - game winning streak in 1974. the team has beaten an additional eight number - one teams, and those nine wins rank second, to ucla's 10, all - time in wins against the top team. the team plays in newly renovated purcell pavilion ( within the edmund p. joyce center ), which reopened for the beginning of the 2009 – 2010 season. the team is coached by mike brey, who, as of the 2014 – 15 season, his fifteenth at notre dame, has achieved a 332 - 165 record. in 2009 they were invited to the nit, where they advanced to the semifinals but were beaten by penn state who went on and beat baylor in the championship. the 2010 – 11 team concluded its regular season ranked number seven in the country, with a record of 25 – 5, brey's fifth straight 20 - win season, and a second - place finish in the big east. during the 2014 - 15 season, the team went 32 - 6 and won the acc conference tournament, later advancing to the elite 8, where the fighting irish lost on a missed buzzer - beater against then undefeated kentucky. led by nba draft picks jerian grant and pat connaughton, the fighting irish beat the eventual national champion duke blue devils twice during the season. the 32 wins were [SEP]\n",
      "[CLS] how many wins does the notre dame men's basketball team have? [SEP] championship. the 2010 – 11 team concluded its regular season ranked number seven in the country, with a record of 25 – 5, brey's fifth straight 20 - win season, and a second - place finish in the big east. during the 2014 - 15 season, the team went 32 - 6 and won the acc conference tournament, later advancing to the elite 8, where the fighting irish lost on a missed buzzer - beater against then undefeated kentucky. led by nba draft picks jerian grant and pat connaughton, the fighting irish beat the eventual national champion duke blue devils twice during the season. the 32 wins were the most by the fighting irish team since 1908 - 09. [SEP]\n"
     ]
    }
   ],
   "source": [
    "for x in tokenized_example[\"input_ids\"][:2]:\n",
    "    print(tokenizer.decode(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's going to take some work to properly label the answers here: we need to find in which of those features the answer actually is, and where exactly in that feature. The models we will use require the start and end positions of these answers in the tokens, so we will also need to to map parts of the original context to some tokens. Thankfully, the tokenizer we're using can help us with that by returning an `offset_mapping`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0, 0), (0, 3), (4, 8), (9, 13), (14, 18), (19, 22), (23, 28), (29, 33), (34, 37), (37, 38), (38, 39), (40, 50), (51, 55), (56, 60), (60, 61), (0, 0), (0, 3), (4, 7), (7, 8), (8, 9), (10, 20), (21, 25), (26, 29), (30, 34), (35, 36), (36, 37), (37, 40), (41, 45), (45, 46), (47, 50), (51, 53), (54, 58), (59, 61), (62, 69), (70, 73), (74, 78), (79, 86), (87, 91), (92, 96), (96, 97), (98, 101), (102, 106), (107, 115), (116, 118), (119, 121), (122, 126), (127, 138), (138, 139), (140, 146), (147, 153), (154, 160), (161, 165), (166, 171), (172, 175), (176, 182), (183, 186), (187, 191), (192, 198), (199, 205), (206, 208), (209, 210), (211, 217), (218, 222), (223, 225), (226, 229), (230, 240), (241, 245), (246, 248), (248, 249), (250, 258), (259, 262), (263, 267), (268, 271), (272, 277), (278, 281), (282, 285), (286, 290), (291, 301), (301, 302), (303, 307), (308, 312), (313, 318), (319, 321), (322, 325), (326, 330), (330, 331), (332, 340), (341, 351), (352, 354), (355, 363), (364, 373), (374, 379), (379, 380), (381, 384), (385, 389), (390, 393), (394, 406), (407, 408), (409, 415), (416, 418)]\n"
     ]
    }
   ],
   "source": [
    "tokenized_example = tokenizer(\n",
    "    example[\"question\"],\n",
    "    example[\"context\"],\n",
    "    max_length=max_length,\n",
    "    truncation=\"only_second\",\n",
    "    return_overflowing_tokens=True,\n",
    "    return_offsets_mapping=True,\n",
    "    stride=doc_stride,\n",
    ")\n",
    "print(tokenized_example[\"offset_mapping\"][0][:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives the corresponding start and end character in the original text for each token in our input IDs. The very first token (`[CLS]`) has (0, 0) because it doesn't correspond to any part of the question/answer, then the second token is the same as the characters 0 to 3 of the question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "how How\n"
     ]
    }
   ],
   "source": [
    "first_token_id = tokenized_example[\"input_ids\"][0][1]\n",
    "offsets = tokenized_example[\"offset_mapping\"][0][1]\n",
    "print(\n",
    "    tokenizer.convert_ids_to_tokens([first_token_id])[0],\n",
    "    example[\"question\"][offsets[0] : offsets[1]],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we can use this mapping to find the position of the start and end tokens of our answer in a given feature. We just have to distinguish which parts of the offsets correspond to the question and which part correspond to the context, this is where the `sequence_ids` method of our `tokenized_example` can be useful:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, None]\n"
     ]
    }
   ],
   "source": [
    "sequence_ids = tokenized_example.sequence_ids()\n",
    "print(sequence_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It returns `None` for the special tokens, then 0 or 1 depending on whether the corresponding token comes from the first sentence past (the question) or the second (the context). Now with all of this, we can find the first and last token of the answer in one of our input feature (or if the answer is not in this feature):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23 26\n"
     ]
    }
   ],
   "source": [
    "answers = example[\"answers\"]\n",
    "start_char = answers[\"answer_start\"][0]\n",
    "end_char = start_char + len(answers[\"text\"][0])\n",
    "\n",
    "# Start token index of the current span in the text.\n",
    "token_start_index = 0\n",
    "while sequence_ids[token_start_index] != 1:\n",
    "    token_start_index += 1\n",
    "\n",
    "# End token index of the current span in the text.\n",
    "token_end_index = len(tokenized_example[\"input_ids\"][0]) - 1\n",
    "while sequence_ids[token_end_index] != 1:\n",
    "    token_end_index -= 1\n",
    "\n",
    "# Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).\n",
    "offsets = tokenized_example[\"offset_mapping\"][0]\n",
    "if (\n",
    "    offsets[token_start_index][0] <= start_char\n",
    "    and offsets[token_end_index][1] >= end_char\n",
    "):\n",
    "    # Move the token_start_index and token_end_index to the two ends of the answer.\n",
    "    # Note: we could go after the last offset if the answer is the last word (edge case).\n",
    "    while (\n",
    "        token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char\n",
    "    ):\n",
    "        token_start_index += 1\n",
    "    start_position = token_start_index - 1\n",
    "    while offsets[token_end_index][1] >= end_char:\n",
    "        token_end_index -= 1\n",
    "    end_position = token_end_index + 1\n",
    "    print(start_position, end_position)\n",
    "else:\n",
    "    print(\"The answer is not in this feature.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can double check that it is indeed the correct answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "over 1, 600\n",
      "over 1,600\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    tokenizer.decode(\n",
    "        tokenized_example[\"input_ids\"][0][start_position : end_position + 1]\n",
    "    )\n",
    ")\n",
    "print(answers[\"text\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this notebook to work with any kind of model, we need to account for the special case where the model expects padding on the left (in which case we switch the order of the question and the context):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_on_right = tokenizer.padding_side == \"right\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's put everything together in one function we will apply to our training set. In the case of impossible answers (the answer is in another feature given by an example with a long context), we set the cls index for both the start and end position. We could also simply discard those examples from the training set if the flag `allow_impossible_answers` is `False`. Since the preprocessing is already complex enough as it is, we've kept is simple for this part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_train_features(examples):\n",
    "    # Tokenize our examples with truncation and padding, but keep the overflows using a stride. This results\n",
    "    # in one example possible giving several features when a context is long, each of those features having a\n",
    "    # context that overlaps a bit the context of the previous feature.\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[\"question\" if pad_on_right else \"context\"],\n",
    "        examples[\"context\" if pad_on_right else \"question\"],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_length,\n",
    "        stride=doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "\n",
    "    # Since one example might give us several features if it has a long context, we need a map from a feature to\n",
    "    # its corresponding example. This key gives us just that.\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "    # The offset mappings will give us a map from token to character position in the original context. This will\n",
    "    # help us compute the start_positions and end_positions.\n",
    "    offset_mapping = tokenized_examples.pop(\"offset_mapping\")\n",
    "\n",
    "    # Let's label those examples!\n",
    "    tokenized_examples[\"start_positions\"] = []\n",
    "    tokenized_examples[\"end_positions\"] = []\n",
    "\n",
    "    for i, offsets in enumerate(offset_mapping):\n",
    "        # We will label impossible answers with the index of the CLS token.\n",
    "        input_ids = tokenized_examples[\"input_ids\"][i]\n",
    "        cls_index = input_ids.index(tokenizer.cls_token_id)\n",
    "\n",
    "        # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "\n",
    "        # One example can give several spans, this is the index of the example containing this span of text.\n",
    "        sample_index = sample_mapping[i]\n",
    "        answers = examples[\"answers\"][sample_index]\n",
    "        # If no answers are given, set the cls_index as answer.\n",
    "        if len(answers[\"answer_start\"]) == 0:\n",
    "            tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "            tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "        else:\n",
    "            # Start/end character index of the answer in the text.\n",
    "            start_char = answers[\"answer_start\"][0]\n",
    "            end_char = start_char + len(answers[\"text\"][0])\n",
    "\n",
    "            # Start token index of the current span in the text.\n",
    "            token_start_index = 0\n",
    "            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):\n",
    "                token_start_index += 1\n",
    "\n",
    "            # End token index of the current span in the text.\n",
    "            token_end_index = len(input_ids) - 1\n",
    "            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):\n",
    "                token_end_index -= 1\n",
    "\n",
    "            # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).\n",
    "            if not (\n",
    "                offsets[token_start_index][0] <= start_char\n",
    "                and offsets[token_end_index][1] >= end_char\n",
    "            ):\n",
    "                tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "                tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "            else:\n",
    "                # Otherwise move the token_start_index and token_end_index to the two ends of the answer.\n",
    "                # Note: we could go after the last offset if the answer is the last word (edge case).\n",
    "                while (\n",
    "                    token_start_index < len(offsets)\n",
    "                    and offsets[token_start_index][0] <= start_char\n",
    "                ):\n",
    "                    token_start_index += 1\n",
    "                tokenized_examples[\"start_positions\"].append(token_start_index - 1)\n",
    "                while offsets[token_end_index][1] >= end_char:\n",
    "                    token_end_index -= 1\n",
    "                tokenized_examples[\"end_positions\"].append(token_end_index + 1)\n",
    "\n",
    "    return tokenized_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0lm8ozrJIrJR"
   },
   "source": [
    "This function works with one or several examples. In the case of several examples, the tokenizer will return a list of lists for each key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "-b70jh26IrJS",
    "outputId": "acd3a42d-985b-44ee-9daa-af5d944ce1d9"
   },
   "outputs": [],
   "source": [
    "features = prepare_train_features(datasets[\"train\"][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zS-6iXTkIrJT"
   },
   "source": [
    "To apply this function on all the sentences (or pairs of sentences) in our dataset, we just use the `map` method of the `dataset` object we created earlier. This will apply the function on all the elements of all the splits in `dataset`, so our training, validation and testing data will be preprocessed in one single command. Since our preprocessing changes the number of samples, we need to remove the old columns when applying it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "DDtsaJeVIrJT",
    "outputId": "aa4734bf-4ef5-4437-9948-2c16363da719"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-ad89cfc588b4b5ad.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-123d7bb970edffa2.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenized_datasets = datasets.map(\n",
    "    prepare_train_features, batched=True, remove_columns=datasets[\"train\"].column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "voWiw8C7IrJV"
   },
   "source": [
    "Even better, the results are automatically cached by the 🤗 Datasets library to avoid spending time on this step the next time you run your notebook. The 🤗 Datasets library is normally smart enough to detect when the function you pass to map has changed (and thus requires to not use the cache data). For instance, it will properly detect if you change the task in the first cell and rerun the notebook. 🤗 Datasets warns you when it uses cached files, you can pass `load_from_cache_file=False` in the call to `map` to not use the cached files and force the preprocessing to be applied again.\n",
    "\n",
    "Note that we passed `batched=True` to encode the texts by batches together. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the texts in a batch concurrently."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "545PP3o8IrJV"
   },
   "source": [
    "## Fine-tuning the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FBiW8UpKIrJW"
   },
   "source": [
    "Now that our data is ready for training, we can download the pretrained model and fine-tune it. Since our task is question answering, we use the `TFAutoModelForQuestionAnswering` class. Like with the tokenizer, the `from_pretrained` method will download and cache the model for us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "TlqNaB8jIrJW",
    "outputId": "84916cf3-6e6c-47f3-d081-032ec30a4132"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-21 15:10:11.409257: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.415291: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.415996: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.417100: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-07-21 15:10:11.419725: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.420421: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.421095: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.747224: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.747919: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.748580: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-07-21 15:10:11.749220: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 21699 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:21:00.0, compute capability: 8.6\n",
      "2022-07-21 15:10:12.356985: I tensorflow/stream_executor/cuda/cuda_blas.cc:1786] TensorFloat-32 will be used for the matrix multiplication. This will only be logged once.\n",
      "Some layers from the model checkpoint at distilbert-base-uncased were not used when initializing TFDistilBertForQuestionAnswering: ['vocab_layer_norm', 'vocab_projector', 'activation_13', 'vocab_transform']\n",
      "- This IS expected if you are initializing TFDistilBertForQuestionAnswering 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 TFDistilBertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some layers of TFDistilBertForQuestionAnswering were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['dropout_19', 'qa_outputs']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import TFAutoModelForQuestionAnswering\n",
    "\n",
    "model = TFAutoModelForQuestionAnswering.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CczA5lJlIrJX"
   },
   "source": [
    "The warning is telling us we are throwing away some weights (the `vocab_transform` and `vocab_layer_norm` layers) and randomly initializing some other (the `pre_classifier` and `classifier` layers). This is absolutely normal in this case, because we are removing the head used to pretrain the model on a masked language modeling objective and replacing it with a new head for which we don't have pretrained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_N8urzhyIrJY"
   },
   "source": [
    "To train our model, we will need to define a few more things. The first two arguments are to setup everything so we can push the model to the [Hub](https://huggingface.co/models) at the end of training. Remove the two of them if you didn't follow the installation steps at the top of the notebook, otherwise you can change the value of `push_to_hub_model_id` to something you would prefer.\n",
    "\n",
    "We also tweak the learning rate, use the `batch_size` defined at the top of the notebook and customize the number of epochs for training, as well as the weight decay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "Bliy8zgjIrJY"
   },
   "outputs": [],
   "source": [
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "push_to_hub_model_id = f\"{model_name}-finetuned-squad\"\n",
    "learning_rate = 2e-5\n",
    "num_train_epochs = 2\n",
    "weight_decay = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we convert our datasets to `tf.data.Dataset`, which Keras understands natively. There are two ways to do this - we can use the slightly more low-level [`Dataset.to_tf_dataset()`](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset.to_tf_dataset) method, or we can use [`Model.prepare_tf_dataset()`](https://huggingface.co/docs/transformers/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset). The main difference between these two is that the `Model` method can inspect the model to determine which column names it can use as input, which means you don't need to specify them yourself. It also supplies a default data collator that will work fine for us, as our samples are already padded to the same length and ready to go."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = model.prepare_tf_dataset(\n",
    "    tokenized_datasets[\"train\"],\n",
    "    shuffle=True,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "validation_set = model.prepare_tf_dataset(\n",
    "    tokenized_datasets[\"validation\"],\n",
    "    shuffle=False,\n",
    "    batch_size=batch_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can create an optimizer and specify a loss function. The `create_optimizer` function gives us a very solid `AdamW` optimizer with weight decay and a learning rate schedule, but it needs us to compute the number of training steps to build that schedule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import create_optimizer\n",
    "\n",
    "total_train_steps = len(train_set) * num_train_epochs\n",
    "\n",
    "optimizer, schedule = create_optimizer(\n",
    "    init_lr=learning_rate, num_warmup_steps=0, num_train_steps=total_train_steps\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that most Transformers models compute loss internally, so we actually don't have to specify anything there! You can of course set your own loss function if you want, but by default our models will choose the 'obvious' loss that matches their task, such as cross-entropy in the case of language modelling. The built-in loss will also correctly handle things like masking the loss on padding tokens, or unlabelled tokens in the case of masked language modelling, so we recommend using it unless you're an advanced user!\n",
    "\n",
    "In addition, because the outputs and loss for this model class are quite straightforward, we can use built-in Keras metrics - these are liable to misbehave in other contexts (for example, they don't know about the masking in masked language modelling) but work well here.\n",
    "\n",
    "We can also use `jit_compile` to compile the model with [XLA](https://www.tensorflow.org/xla). In other cases, we should be careful about that - if our inputs might have variable sequence lengths, we may end up having to do a new XLA compilation for each possible length, because XLA compilation expects a static input shape! In this notebook, however, we have padded all examples to exactly the same length. This makes it perfect for XLA, which will give us a nice performance boost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No loss specified in compile() - the model's internal loss computation will be used as the loss. Don't panic - this is a common way to train TensorFlow models in Transformers! To disable this behaviour please pass a loss argument, or explicitly pass `loss=None` if you do not want your model to compute a loss.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "model.compile(optimizer=optimizer, jit_compile=True, metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rXuFTAzDIrJe"
   },
   "source": [
    "We will evaluate our model and compute metrics in the next section (this is a very long operation, so we will only compute the evaluation loss during training). For now, let's just train our model. We can also add a callback to sync up our model with the Hub - this allows us to resume training from other machines and even test the model's inference quality midway through training! If you don't want to do this, simply remove the callbacks argument in the call to `fit()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "imY1oC3SIrJf"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matt/PycharmProjects/notebooks/examples/qa_model_save is already a clone of https://huggingface.co/Rocketknight1/distilbert-base-uncased-finetuned-squad. Make sure you pull the latest changes with `repo.git_pull()`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-21 15:10:19.739992: I tensorflow/compiler/xla/service/service.cc:170] XLA service 0x7f12c8010180 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:\n",
      "2022-07-21 15:10:19.740026: I tensorflow/compiler/xla/service/service.cc:178]   StreamExecutor device (0): NVIDIA GeForce RTX 3090, Compute Capability 8.6\n",
      "2022-07-21 15:10:19.883966: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:263] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
      "2022-07-21 15:10:19.914890: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:57] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. tf_distil_bert_for_question_answering/distilbert/embeddings/dropout/dropout/random_uniform/RandomUniform\n",
      "2022-07-21 15:10:32.499915: I tensorflow/compiler/jit/xla_compilation_cache.cc:478] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   6/5532 [..............................] - ETA: 11:03 - loss: 5.8639 - end_logits_accuracy: 0.0104 - start_logits_accuracy: 0.0000e+00    WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0370s vs `on_train_batch_end` time: 0.1004s). Check your callbacks.\n",
      "5532/5532 [==============================] - 718s 127ms/step - loss: 1.5124 - end_logits_accuracy: 0.6041 - start_logits_accuracy: 0.5680 - val_loss: 1.1534 - val_end_logits_accuracy: 0.6849 - val_start_logits_accuracy: 0.6443\n",
      "Epoch 2/2\n",
      "5532/5532 [==============================] - 697s 126ms/step - loss: 0.9726 - end_logits_accuracy: 0.7301 - start_logits_accuracy: 0.6915 - val_loss: 1.1130 - val_end_logits_accuracy: 0.7014 - val_start_logits_accuracy: 0.6644\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f15b01ce350>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers.keras_callbacks import PushToHubCallback\n",
    "from tensorflow.keras.callbacks import TensorBoard\n",
    "\n",
    "push_to_hub_callback = PushToHubCallback(\n",
    "    output_dir=\"./qa_model_save\",\n",
    "    tokenizer=tokenizer,\n",
    "    hub_model_id=push_to_hub_model_id,\n",
    ")\n",
    "\n",
    "tensorboard_callback = TensorBoard(log_dir=\"./qa_model_save/logs\")\n",
    "\n",
    "callbacks = [tensorboard_callback, push_to_hub_callback]\n",
    "\n",
    "model.fit(\n",
    "    train_set,\n",
    "    validation_data=validation_set,\n",
    "    epochs=num_train_epochs,\n",
    "    callbacks=callbacks,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluating our model will require a bit more work, as we will need to map the predictions of our model back to parts of the context. The model itself predicts logits for the start and end position of our answers: if we take a batch from our validation dataset, here is the output our model gives us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['start_logits', 'end_logits'])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch = next(iter(validation_set))\n",
    "output = model.predict_on_batch(batch)\n",
    "output.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the model is a dict-like object that contains the loss (since we provided labels), the start and end logits. We won't need the loss for our predictions, let's have a look a the logits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((16, 384), (16, 384))"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.start_logits.shape, output.end_logits.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have one logit for each feature and each token. The most obvious thing to predict an answer for each feature is to take the index for the maximum of the start logits as a start position and the index of the maximum of the end logits as an end position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 46,  57,  78,  43, 118, 108,  72,  35, 108,  34,  73,  41,  80,\n",
       "         91, 156,  35]),\n",
       " array([ 47,  58,  81,  44, 118, 109,  75,  37, 109,  36,  76,  42,  83,\n",
       "         94, 158,  35]))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.argmax(output.start_logits, -1), np.argmax(output.end_logits, -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will work great in a lot of cases, but what if this prediction gives us something impossible: the start position could be greater than the end position, or point to a span of text in the question instead of the answer. In that case, we might want to look at the second best prediction to see if it gives a possible answer and select that instead.\n",
    "\n",
    "However, picking the second best answer is not as easy as picking the best one: is it the second best index in the start logits with the best index in the end logits? Or the best index in the start logits with the second best index in the end logits? And if that second best answer is not possible either, it gets even trickier for the third best answer.\n",
    "\n",
    "\n",
    "To classify our answers, we will use the score obtained by adding the start and end logits. We won't try to order all the possible answers and limit ourselves to with a hyper-parameter we call `n_best_size`. We'll pick the best indices in the start and end logits and gather all the answers this predicts. After checking if each one is valid, we will sort them by their score and keep the best one. Here is how we would do this on the first feature in the batch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_best_size = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "start_logits = output.start_logits[0]\n",
    "end_logits = output.end_logits[0]\n",
    "# Gather the indices the best start/end logits:\n",
    "start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "valid_answers = []\n",
    "for start_index in start_indexes:\n",
    "    for end_index in end_indexes:\n",
    "        if (\n",
    "            start_index <= end_index\n",
    "        ):  # We need to refine that test to check the answer is inside the context\n",
    "            valid_answers.append(\n",
    "                {\n",
    "                    \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                    \"text\": \"\",  # We need to find a way to get back the original substring corresponding to the answer in the context\n",
    "                }\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then we can sort the `valid_answers` according to their `score` and only keep the best one. The only point left is how to check a given span is inside the context (and not the question) and how to get back the text inside. To do this, we need to add two things to our validation features:\n",
    "- the ID of the example that generated the feature (since each example can generate several features, as seen before);\n",
    "- the offset mapping that will give us a map from token indices to character positions in the context.\n",
    "\n",
    "That's why we will re-process the validation set with the following function, slightly different from `prepare_train_features`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_validation_features(examples):\n",
    "    # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results\n",
    "    # in one example possible giving several features when a context is long, each of those features having a\n",
    "    # context that overlaps a bit the context of the previous feature.\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[\"question\" if pad_on_right else \"context\"],\n",
    "        examples[\"context\" if pad_on_right else \"question\"],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_length,\n",
    "        stride=doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "\n",
    "    # Since one example might give us several features if it has a long context, we need a map from a feature to\n",
    "    # its corresponding example. This key gives us just that.\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "\n",
    "    # We keep the example_id that gave us this feature and we will store the offset mappings.\n",
    "    tokenized_examples[\"example_id\"] = []\n",
    "\n",
    "    for i in range(len(tokenized_examples[\"input_ids\"])):\n",
    "        # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "        context_index = 1 if pad_on_right else 0\n",
    "\n",
    "        # One example can give several spans, this is the index of the example containing this span of text.\n",
    "        sample_index = sample_mapping[i]\n",
    "        tokenized_examples[\"example_id\"].append(examples[\"id\"][sample_index])\n",
    "\n",
    "        # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token\n",
    "        # position is part of the context or not.\n",
    "        tokenized_examples[\"offset_mapping\"][i] = [\n",
    "            (o if sequence_ids[k] == context_index else None)\n",
    "            for k, o in enumerate(tokenized_examples[\"offset_mapping\"][i])\n",
    "        ]\n",
    "\n",
    "    return tokenized_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And like before, we can apply that function to our validation set easily:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-fb6eddd5466a5d8b.arrow\n"
     ]
    }
   ],
   "source": [
    "validation_features = datasets[\"validation\"].map(\n",
    "    prepare_validation_features,\n",
    "    batched=True,\n",
    "    remove_columns=datasets[\"validation\"].column_names,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And turn the dataset into a `tf.data.Dataset` as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_dataset = model.prepare_tf_dataset(\n",
    "     validation_features,\n",
    "     shuffle=False,\n",
    "     batch_size=batch_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can grab the predictions for all features by using the `model.predict` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "674/674 [==============================] - 730s 1s/step\n"
     ]
    }
   ],
   "source": [
    "raw_predictions = model.predict(validation_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TFQuestionAnsweringModelOutput(loss=None, start_logits=array([[-7.6091537, -8.416484 , -8.680255 , ..., -9.14007  , -9.111827 ,\n",
       "        -9.095983 ],\n",
       "       [-7.444468 , -8.380038 , -8.667503 , ..., -9.137428 , -9.108936 ,\n",
       "        -9.093311 ],\n",
       "       [-6.9005785, -6.8720226, -7.9134784, ..., -9.034261 , -9.037982 ,\n",
       "        -9.096408 ],\n",
       "       ...,\n",
       "       [-5.2020297, -7.8957915, -8.570762 , ..., -9.107713 , -9.104418 ,\n",
       "        -9.130447 ],\n",
       "       [-3.5689824, -6.5537987, -7.506344 , ..., -9.123406 , -9.082561 ,\n",
       "        -9.090938 ],\n",
       "       [-3.1193774, -7.857499 , -8.254857 , ..., -9.13785  , -9.124171 ,\n",
       "        -9.154689 ]], dtype=float32), end_logits=array([[-6.6888733, -8.407975 , -8.232986 , ..., -8.920444 , -8.954519 ,\n",
       "        -8.96373  ],\n",
       "       [-6.5028334, -8.344763 , -8.204293 , ..., -8.921926 , -8.957324 ,\n",
       "        -8.965635 ],\n",
       "       [-6.7003646, -7.6689005, -8.584744 , ..., -8.975139 , -8.959321 ,\n",
       "        -8.931076 ],\n",
       "       ...,\n",
       "       [-4.556566 , -8.519492 , -8.63644  , ..., -8.948354 , -8.937384 ,\n",
       "        -8.92393  ],\n",
       "       [-3.1397572, -7.477724 , -7.658732 , ..., -8.916698 , -8.921466 ,\n",
       "        -8.956884 ],\n",
       "       [-2.569292 , -8.102478 , -8.430557 , ..., -8.907092 , -8.90624  ,\n",
       "        -8.899828 ]], dtype=float32), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "raw_predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now refine the test we had before: since we set `None` in the offset mappings when it corresponds to a part of the question, it's easy to check if an answer is fully inside the context. We also eliminate very long answers from our considerations (with an hyper-parameter we can tune)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_answer_length = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 14.749695, 'text': 'Denver Broncos'},\n",
       " {'score': 12.7321825,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers'},\n",
       " {'score': 10.229258, 'text': 'Carolina Panthers'},\n",
       " {'score': 9.956345, 'text': 'Broncos'},\n",
       " {'score': 9.882124,\n",
       "  'text': 'American Football Conference (AFC) champion Denver Broncos'},\n",
       " {'score': 9.054972,\n",
       "  'text': 'The American Football Conference (AFC) champion Denver Broncos'},\n",
       " {'score': 8.676702, 'text': 'Denver'},\n",
       " {'score': 8.392237,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC)'},\n",
       " {'score': 7.9388328,\n",
       "  'text': 'Broncos defeated the National Football Conference (NFC) champion Carolina Panthers'},\n",
       " {'score': 7.8646116,\n",
       "  'text': 'American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers'},\n",
       " {'score': 7.692413,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference'},\n",
       " {'score': 7.0697618,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24–10'},\n",
       " {'score': 7.037459,\n",
       "  'text': 'The American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers'},\n",
       " {'score': 6.493333,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC'},\n",
       " {'score': 6.1674027,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC) champion'},\n",
       " {'score': 6.023146,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC) champion Carolina'},\n",
       " {'score': 6.008625, 'text': 'champion Denver Broncos'},\n",
       " {'score': 5.645863, 'text': 'Panthers'},\n",
       " {'score': 5.2083697,\n",
       "  'text': 'National Football Conference (NFC) champion Carolina Panthers'},\n",
       " {'score': 4.8981276,\n",
       "  'text': 'Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24–10 to earn their third Super Bowl title.'}]"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "start_logits = output.start_logits[0]\n",
    "end_logits = output.end_logits[0]\n",
    "offset_mapping = validation_features[0][\"offset_mapping\"]\n",
    "# The first feature comes from the first example. For the more general case, we will need to be match the example_id to\n",
    "# an example index\n",
    "context = datasets[\"validation\"][0][\"context\"]\n",
    "\n",
    "# Gather the indices the best start/end logits:\n",
    "start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "valid_answers = []\n",
    "for start_index in start_indexes:\n",
    "    for end_index in end_indexes:\n",
    "        # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond\n",
    "        # to part of the input_ids that are not in the context.\n",
    "        if (\n",
    "            start_index >= len(offset_mapping)\n",
    "            or end_index >= len(offset_mapping)\n",
    "            or offset_mapping[start_index] is None\n",
    "            or offset_mapping[end_index] is None\n",
    "        ):\n",
    "            continue\n",
    "        # Don't consider answers with a length that is either < 0 or > max_answer_length.\n",
    "        if end_index < start_index or end_index - start_index + 1 > max_answer_length:\n",
    "            continue\n",
    "        if (\n",
    "            start_index <= end_index\n",
    "        ):  # We need to refine that test to check the answer is inside the context\n",
    "            start_char = offset_mapping[start_index][0]\n",
    "            end_char = offset_mapping[end_index][1]\n",
    "            valid_answers.append(\n",
    "                {\n",
    "                    \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                    \"text\": context[start_char:end_char],\n",
    "                }\n",
    "            )\n",
    "\n",
    "valid_answers = sorted(valid_answers, key=lambda x: x[\"score\"], reverse=True)[\n",
    "    :n_best_size\n",
    "]\n",
    "valid_answers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare to the actual ground-truth answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': ['Denver Broncos', 'Denver Broncos', 'Denver Broncos'],\n",
       " 'answer_start': [177, 177, 177]}"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"validation\"][0][\"answers\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our model's most likely answer is correct!\n",
    "\n",
    "As we mentioned in the code above, this was easy on the first feature because we knew it comes from the first example. For the other features, we will need a map between examples and their corresponding features. Also, since one example can give several features, we will need to gather together all the answers in all the features generated by a given example, then pick the best one. The following code builds a map from example index to its corresponding features indices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "examples = datasets[\"validation\"]\n",
    "features = validation_features\n",
    "\n",
    "example_id_to_index = {k: i for i, k in enumerate(examples[\"id\"])}\n",
    "features_per_example = collections.defaultdict(list)\n",
    "for i, feature in enumerate(features):\n",
    "    features_per_example[example_id_to_index[feature[\"example_id\"]]].append(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're almost ready for our post-processing function. The last bit to deal with is the impossible answer (when `squad_v2 = True`). The code above only keeps answers that are inside the context, we need to also grab the score for the impossible answer (which has start and end indices corresponding to the index of the CLS token). When one example gives several features, we have to predict the impossible answer when all the features give a high score to the impossible answer (since one feature could predict the impossible answer just because the answer isn't in the part of the context it has access too), which is why the score of the impossible answer for one example is the *minimum* of the scores for the impossible answer in each feature generated by the example.\n",
    "\n",
    "We then predict the impossible answer when that score is greater than the score of the best non-impossible answer. All combined together, this gives us this post-processing function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "def postprocess_qa_predictions(\n",
    "    examples,\n",
    "    features,\n",
    "    all_start_logits,\n",
    "    all_end_logits,\n",
    "    n_best_size=20,\n",
    "    max_answer_length=30,\n",
    "):\n",
    "    # Build a map example to its corresponding features.\n",
    "    example_id_to_index = {k: i for i, k in enumerate(examples[\"id\"])}\n",
    "    features_per_example = collections.defaultdict(list)\n",
    "    for i, feature in enumerate(features):\n",
    "        features_per_example[example_id_to_index[feature[\"example_id\"]]].append(i)\n",
    "\n",
    "    # The dictionaries we have to fill.\n",
    "    predictions = collections.OrderedDict()\n",
    "\n",
    "    # Logging.\n",
    "    print(\n",
    "        f\"Post-processing {len(examples)} example predictions split into {len(features)} features.\"\n",
    "    )\n",
    "\n",
    "    # Let's loop over all the examples!\n",
    "    for example_index, example in enumerate(tqdm(examples)):\n",
    "        # Those are the indices of the features associated to the current example.\n",
    "        feature_indices = features_per_example[example_index]\n",
    "\n",
    "        min_null_score = None  # Only used if squad_v2 is True.\n",
    "        valid_answers = []\n",
    "\n",
    "        context = example[\"context\"]\n",
    "        # Looping through all the features associated to the current example.\n",
    "        for feature_index in feature_indices:\n",
    "            # We grab the predictions of the model for this feature.\n",
    "            start_logits = all_start_logits[feature_index]\n",
    "            end_logits = all_end_logits[feature_index]\n",
    "            # This is what will allow us to map some the positions in our logits to span of texts in the original\n",
    "            # context.\n",
    "            offset_mapping = features[feature_index][\"offset_mapping\"]\n",
    "\n",
    "            # Update minimum null prediction.\n",
    "            cls_index = features[feature_index][\"input_ids\"].index(\n",
    "                tokenizer.cls_token_id\n",
    "            )\n",
    "            feature_null_score = start_logits[cls_index] + end_logits[cls_index]\n",
    "            if min_null_score is None or min_null_score < feature_null_score:\n",
    "                min_null_score = feature_null_score\n",
    "\n",
    "            # Go through all possibilities for the `n_best_size` greater start and end logits.\n",
    "            start_indexes = np.argsort(start_logits)[\n",
    "                -1 : -n_best_size - 1 : -1\n",
    "            ].tolist()\n",
    "            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "            for start_index in start_indexes:\n",
    "                for end_index in end_indexes:\n",
    "                    # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond\n",
    "                    # to part of the input_ids that are not in the context.\n",
    "                    if (\n",
    "                        start_index >= len(offset_mapping)\n",
    "                        or end_index >= len(offset_mapping)\n",
    "                        or not offset_mapping[start_index]\n",
    "                        or not offset_mapping[end_index]\n",
    "                    ):\n",
    "                        continue\n",
    "                    # Don't consider answers with a length that is either < 0 or > max_answer_length.\n",
    "                    if (\n",
    "                        end_index < start_index\n",
    "                        or end_index - start_index + 1 > max_answer_length\n",
    "                    ):\n",
    "                        continue\n",
    "                    start_char = offset_mapping[start_index][0]\n",
    "                    end_char = offset_mapping[end_index][1]\n",
    "                    valid_answers.append(\n",
    "                        {\n",
    "                            \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                            \"text\": context[start_char:end_char],\n",
    "                        }\n",
    "                    )\n",
    "\n",
    "        if len(valid_answers) > 0:\n",
    "            best_answer = sorted(valid_answers, key=lambda x: x[\"score\"], reverse=True)[\n",
    "                0\n",
    "            ]\n",
    "        else:\n",
    "            # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid\n",
    "            # failure.\n",
    "            best_answer = {\"text\": \"\", \"score\": 0.0}\n",
    "\n",
    "        # Let's pick our final answer: the best one or the null answer (only for squad_v2)\n",
    "        if not squad_v2:\n",
    "            predictions[example[\"id\"]] = best_answer[\"text\"]\n",
    "        else:\n",
    "            answer = (\n",
    "                best_answer[\"text\"] if best_answer[\"score\"] > min_null_score else \"\"\n",
    "            )\n",
    "            predictions[example[\"id\"]] = answer\n",
    "\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can apply our post-processing function to our raw predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Post-processing 10570 example predictions split into 10784 features.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9bf775b07074ce894069add7ed27b14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10570 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "final_predictions = postprocess_qa_predictions(\n",
    "    datasets[\"validation\"],\n",
    "    validation_features,\n",
    "    raw_predictions[\"start_logits\"],\n",
    "    raw_predictions[\"end_logits\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can load the metric from the datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = load_metric(\"squad_v2\" if squad_v2 else \"squad\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can call compute on it. We just need to format predictions and labels a bit as it expects a list of dictionaries and not one big dictionary. In the case of squad_v2, we also have to set a `no_answer_probability` argument (which we set to 0.0 here as we have already set the answer to empty if we picked it)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'exact_match': 76.46168401135289, 'f1': 84.61642863605374}"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if squad_v2:\n",
    "    formatted_predictions = [\n",
    "        {\"id\": k, \"prediction_text\": v, \"no_answer_probability\": 0.0}\n",
    "        for k, v in final_predictions.items()\n",
    "    ]\n",
    "else:\n",
    "    formatted_predictions = [\n",
    "        {\"id\": k, \"prediction_text\": v} for k, v in final_predictions.items()\n",
    "    ]\n",
    "references = [\n",
    "    {\"id\": ex[\"id\"], \"answers\": ex[\"answers\"]} for ex in datasets[\"validation\"]\n",
    "]\n",
    "metric.compute(predictions=formatted_predictions, references=references)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you ran the callback above, you can now share this model with all your friends, family or favorite pets: they can all load it with the identifier `\"your-username/the-name-you-picked\"` so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import TFAutoModelForQuestionAnswering\n",
    "\n",
    "model = TFAutoModelForQuestionAnswering.from_pretrained(\"your-username/my-awesome-model\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we've trained our model, let's see how we could load it and use it to answer questions in future! First, let's load it from the hub. This means we can resume the code from here without needing to rerun everything above every time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59da7ff386e24bf280e6942cef82966a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading config.json:   0%|          | 0.00/538 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "983af5c90b634617ab108dc1dcbe2ee3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tf_model.h5:   0%|          | 0.00/253M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some layers from the model checkpoint at Rocketknight1/distilbert-base-uncased-finetuned-squad were not used when initializing TFDistilBertForQuestionAnswering: ['dropout_19']\n",
      "- This IS expected if you are initializing TFDistilBertForQuestionAnswering 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 TFDistilBertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some layers of TFDistilBertForQuestionAnswering were not initialized from the model checkpoint at Rocketknight1/distilbert-base-uncased-finetuned-squad and are newly initialized: ['dropout_76']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5273be5a734045f6bd5f790889dd66c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tokenizer_config.json:   0%|          | 0.00/360 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa2ab1172b6e4795b1f3474e79e69487",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading vocab.txt:   0%|          | 0.00/226k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8906b2f4561e48129370990b32bafdca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tokenizer.json:   0%|          | 0.00/695k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26e8396607874e6d8f724d0bc0f9f419",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, TFAutoModelForQuestionAnswering\n",
    "\n",
    "# You can, of course, use your own username and model name here \n",
    "# once you've pushed your model using the code above!\n",
    "checkpoint = \"Rocketknight1/distilbert-base-uncased-finetuned-squad\"\n",
    "model = TFAutoModelForQuestionAnswering.from_pretrained(checkpoint)\n",
    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's get some sample text and ask a question. Feel free to substitute your own text and question!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\"\"The dominant sequence transduction models are based on complex recurrent or convolutional \n",
    "neural networks in an encoder-decoder configuration. The best performing models also connect the encoder\n",
    "and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, \n",
    "based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on \n",
    "two machine translation tasks show these models to be superior in quality while being more parallelizable \n",
    "and requiring significantly less time to train.\"\"\"\n",
    "question = \"What kind of mechanisms is Transformer based on?\"\n",
    "\n",
    "inputs = tokenizer([context], [question], return_tensors=\"np\")\n",
    "\n",
    "outputs = model(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The outputs are logits, so let's use argmax to find the largest logit, which represents the model's best guess for the right answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n",
      "65\n",
      "[ 3086 10595]\n"
     ]
    }
   ],
   "source": [
    "start_position = np.argmax(outputs.start_logits[0])\n",
    "end_position = np.argmax(outputs.end_logits[0])\n",
    "print(start_position)\n",
    "print(end_position)\n",
    "\n",
    "# Extract this substring from the inputs\n",
    "answer = inputs[\"input_ids\"][0, start_position: end_position + 1]\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well, these are definitely tokens. Let's decode them back to text:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'attention mechanisms'"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An alternative way to quickly perform inference with any model on the hub is to use the [Pipeline API](https://huggingface.co/docs/transformers/main_classes/pipelines), which abstracts away all the steps we did manually above. It will perform the preprocessing, forward pass and postprocessing all in a single object.\n",
    "\n",
    "Let's showcase this for our trained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some layers from the model checkpoint at Rocketknight1/distilbert-base-uncased-finetuned-squad were not used when initializing TFDistilBertForQuestionAnswering: ['dropout_19']\n",
      "- This IS expected if you are initializing TFDistilBertForQuestionAnswering 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 TFDistilBertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some layers of TFDistilBertForQuestionAnswering were not initialized from the model checkpoint at Rocketknight1/distilbert-base-uncased-finetuned-squad and are newly initialized: ['dropout_116']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "question_answerer = pipeline(\"question-answering\",  \"Rocketknight1/distilbert-base-uncased-finetuned-squad\", framework=\"tf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'score': 0.534810483455658,\n",
       " 'start': 320,\n",
       " 'end': 340,\n",
       " 'answer': 'attention mechanisms'}"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question_answerer(context=context, question=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Question Answering on SQUAD",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}