examples/multiple_choice-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, 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": 1,
   "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 some of the functionality we use 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": {},
   "source": [
    "We also quickly upload some telemetry - this tells us which examples and software versions are getting used so we know where to prioritize our maintenance efforts. We don't collect (or care about) any personally identifiable information, but if you'd prefer not to be counted, feel free to skip this step or delete this cell entirely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.utils import send_example_telemetry\n",
    "\n",
    "send_example_telemetry(\"multiple_choice_notebook\", framework=\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rEJBSTyZIrIb"
   },
   "source": [
    "# Fine-tuning a model on a multiple choice task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTCFado4IrIc"
   },
   "source": [
    "In this notebook, we will see how to fine-tune one of the [🤗 Transformers](https://github.com/huggingface/transformers) model on a multiple-choice task. In a multiple-choice task, multiple answers or continuations are provided for each input, and the model must guess which is most plausible.  The dataset used here is [SWAG](https://www.aclweb.org/anthology/D18-1009/) but you can adapt the pre-processing to any other multiple choice dataset you like, or your own data. SWAG is a dataset about commonsense reasoning, where each example describes a situation and proposes four continuations that could follow it. \n",
    "\n",
    "This notebook is built to run  with any model checkpoint from the [Model Hub](https://huggingface.co/models) as long as that model has a version with a mutiple choice head. 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 two parameters, then the rest of the notebook should run smoothly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "zVvslsfMIrIh"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"bert-base-cased\"\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. This can be easily done with the `load_dataset` function.  "
   ]
  },
  {
   "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": [
    "`load_dataset` will cache the dataset to avoid downloading it again the next time you run this cell."
   ]
  },
  {
   "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 swag (/home/matt/.cache/huggingface/datasets/swag/regular/0.0.0/9640de08cdba6a1469ed3834fcab4b8ad8e38caf5d1ba5e7436d8b1fd067ad4c)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f73b371cb74c4e399faea25052830c5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "datasets = load_dataset(\"swag\", \"regular\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzfPtOMoIrIu"
   },
   "source": [
    "The `dataset` 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: ['video-id', 'fold-ind', 'startphrase', 'sent1', 'sent2', 'gold-source', 'ending0', 'ending1', 'ending2', 'ending3', 'label'],\n",
       "        num_rows: 73546\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['video-id', 'fold-ind', 'startphrase', 'sent1', 'sent2', 'gold-source', 'ending0', 'ending1', 'ending2', 'ending3', 'label'],\n",
       "        num_rows: 20006\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['video-id', 'fold-ind', 'startphrase', 'sent1', 'sent2', 'gold-source', 'ending0', 'ending1', 'ending2', 'ending3', 'label'],\n",
       "        num_rows: 20005\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets"
   ]
  },
  {
   "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": [
       "{'video-id': 'anetv_jkn6uvmqwh4',\n",
       " 'fold-ind': '3416',\n",
       " 'startphrase': 'Members of the procession walk down the street holding small horn brass instruments. A drum line',\n",
       " 'sent1': 'Members of the procession walk down the street holding small horn brass instruments.',\n",
       " 'sent2': 'A drum line',\n",
       " 'gold-source': 'gold',\n",
       " 'ending0': 'passes by walking down the street playing their instruments.',\n",
       " 'ending1': 'has heard approaching them.',\n",
       " 'ending2': \"arrives and they're outside dancing and asleep.\",\n",
       " 'ending3': 'turns the lead singer watches the performance.',\n",
       " 'label': 0}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"train\"][0]"
   ]
  },
  {
   "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 in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "i3j8APAoIrI3"
   },
   "outputs": [],
   "source": [
    "from datasets import ClassLabel\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",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "SZy5tRB_IrI7",
    "outputId": "ba8f2124-e485-488f-8c0c-254f34f24f13"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>video-id</th>\n",
       "      <th>fold-ind</th>\n",
       "      <th>startphrase</th>\n",
       "      <th>sent1</th>\n",
       "      <th>sent2</th>\n",
       "      <th>gold-source</th>\n",
       "      <th>ending0</th>\n",
       "      <th>ending1</th>\n",
       "      <th>ending2</th>\n",
       "      <th>ending3</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>lsmdc0005_Chinatown-48562</td>\n",
       "      <td>957</td>\n",
       "      <td>In a moment Cross can be seen, looking toward camera. He</td>\n",
       "      <td>In a moment Cross can be seen, looking toward camera.</td>\n",
       "      <td>He</td>\n",
       "      <td>gen</td>\n",
       "      <td>closes the pantry door a little bit more.</td>\n",
       "      <td>crosses to a door to a thick dark case that closes clear.</td>\n",
       "      <td>looks up at an empty window in the club, then hangs up, embarrassed on the cold.</td>\n",
       "      <td>is carrying a trail, which wraps around his ankle.</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>anetv_W6y6Vmk5edg</td>\n",
       "      <td>11587</td>\n",
       "      <td>A gymnast is seen leaning across a long beam and begins performing a gymnastics routine. The girl</td>\n",
       "      <td>A gymnast is seen leaning across a long beam and begins performing a gymnastics routine.</td>\n",
       "      <td>The girl</td>\n",
       "      <td>gen</td>\n",
       "      <td>performs several kickboxing while moving her arms in and down.</td>\n",
       "      <td>does a gymnastic routine on the balance beam and herself on to the balance beam.</td>\n",
       "      <td>move through several flips and flips on the beam.</td>\n",
       "      <td>begins jumping up and down on the mat.</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>lsmdc3042_KARATE_KID-19795</td>\n",
       "      <td>12868</td>\n",
       "      <td>He stares off, vacantly. Someone</td>\n",
       "      <td>He stares off, vacantly.</td>\n",
       "      <td>Someone</td>\n",
       "      <td>gold</td>\n",
       "      <td>puts a three - slim guard into his mouth and kisses him.</td>\n",
       "      <td>is in slow motion, his knees still scarred by the battle.</td>\n",
       "      <td>flips across the dark to the car.</td>\n",
       "      <td>paces toward him and glances at the wreck.</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>lsmdc0026_The_Big_Fish-62739</td>\n",
       "      <td>18560</td>\n",
       "      <td>Someone walks into the river, up to his knees. He</td>\n",
       "      <td>Someone walks into the river, up to his knees.</td>\n",
       "      <td>He</td>\n",
       "      <td>gold</td>\n",
       "      <td>is pulled away, and his shoulders upset.</td>\n",
       "      <td>finds his hand seething, standing around the cab, raising his head.</td>\n",
       "      <td>turns back so his father can face the crowd.</td>\n",
       "      <td>moves vigorously he then turns and does a little push the bubbles into his mouth, leaving his mouth open, breathing ringed.</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>anetv_ZPVrC5185NM</td>\n",
       "      <td>11852</td>\n",
       "      <td>She pushes the baby back and forth on a swing. The baby</td>\n",
       "      <td>She pushes the baby back and forth on a swing.</td>\n",
       "      <td>The baby</td>\n",
       "      <td>gold</td>\n",
       "      <td>laughs and smiles as she swings.</td>\n",
       "      <td>comes back and takes a few flips.</td>\n",
       "      <td>hops backwards into the swing.</td>\n",
       "      <td>walks back inside and throws it out.</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>lsmdc3024_EASY_A-11617</td>\n",
       "      <td>11559</td>\n",
       "      <td>She unbuckles her seat belt and they share a tight embrace. Someone</td>\n",
       "      <td>She unbuckles her seat belt and they share a tight embrace.</td>\n",
       "      <td>Someone</td>\n",
       "      <td>gen</td>\n",
       "      <td>pulls a knapsack from a shelf by someone's coat pocket as she opens one of the boxes.</td>\n",
       "      <td>sits in the limo.</td>\n",
       "      <td>slips his arm around her waist.</td>\n",
       "      <td>blinks, then leans forward and presses his lips to her chin.</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>lsmdc0050_Indiana_Jones_and_the_last_crusade-70715</td>\n",
       "      <td>3232</td>\n",
       "      <td>Now he sees the pendulum has been guarding a small corridor which turns a corner to the left fifty yards ahead. Wooden wheels</td>\n",
       "      <td>Now he sees the pendulum has been guarding a small corridor which turns a corner to the left fifty yards ahead.</td>\n",
       "      <td>Wooden wheels</td>\n",
       "      <td>gold</td>\n",
       "      <td>turn - - the mechanism controlling the spinning blades.</td>\n",
       "      <td>follow through the rising foam.</td>\n",
       "      <td>cranks back in an electronic system.</td>\n",
       "      <td>press against a gold and metal gate.</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>anetv_fZQS02Ypca4</td>\n",
       "      <td>847</td>\n",
       "      <td>A young girl is seen sitting and speaking to the camera while using a brush to powder her face. A color pallet is held up next to her and she</td>\n",
       "      <td>A young girl is seen sitting and speaking to the camera while using a brush to powder her face.</td>\n",
       "      <td>A color pallet is held up next to her and she</td>\n",
       "      <td>gold</td>\n",
       "      <td>begins rubbing the powder all over her eyes.</td>\n",
       "      <td>uses the brush on brush drys on her hair while continuing to speak to the camera.</td>\n",
       "      <td>applies the sides and pans far away.</td>\n",
       "      <td>begins brushing her teeth and showing an image with the brush in her long hair.</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>lsmdc0051_Men_in_black-70855</td>\n",
       "      <td>14113</td>\n",
       "      <td>Someone barks a few orders to them. He</td>\n",
       "      <td>Someone barks a few orders to them.</td>\n",
       "      <td>He</td>\n",
       "      <td>gen</td>\n",
       "      <td>has a hair mustache.</td>\n",
       "      <td>pretends to sit back.</td>\n",
       "      <td>puts his slacks back on.</td>\n",
       "      <td>sends aiming a heft torch.</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>anetv_76RoR_LbIzQ</td>\n",
       "      <td>16378</td>\n",
       "      <td>Pictures of an office is shown. A woman</td>\n",
       "      <td>Pictures of an office is shown.</td>\n",
       "      <td>A woman</td>\n",
       "      <td>gold</td>\n",
       "      <td>is explaining how she is wiping the car drive.</td>\n",
       "      <td>is sitting on a long swing holding a stick.</td>\n",
       "      <td>is doing another woman's hair.</td>\n",
       "      <td>is seen playing a song on a stage.</td>\n",
       "      <td>2</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": {},
   "source": [
    "Each example in the dataset has a context composed of a first sentence (`sent1`) and an introduction to the second sentence (`sent2`). Then four possible endings are given (`ending0`, `ending1`, `ending2` and `ending3`) and the model must pick the right one (`label`). The following function lets us visualize a given example a bit better:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_one(example):\n",
    "    print(f\"Context: {example['sent1']}\")\n",
    "    print(f\"  A - {example['sent2']} {example['ending0']}\")\n",
    "    print(f\"  B - {example['sent2']} {example['ending1']}\")\n",
    "    print(f\"  C - {example['sent2']} {example['ending2']}\")\n",
    "    print(f\"  D - {example['sent2']} {example['ending3']}\")\n",
    "    print(f\"\\nGround truth: option {['A', 'B', 'C', 'D'][example['label']]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: Members of the procession walk down the street holding small horn brass instruments.\n",
      "  A - A drum line passes by walking down the street playing their instruments.\n",
      "  B - A drum line has heard approaching them.\n",
      "  C - A drum line arrives and they're outside dancing and asleep.\n",
      "  D - A drum line turns the lead singer watches the performance.\n",
      "\n",
      "Ground truth: option A\n"
     ]
    }
   ],
   "source": [
    "show_one(datasets[\"train\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: Now it's someone's turn to rain blades on his opponent.\n",
      "  A - Someone pats his shoulder and spins wildly.\n",
      "  B - Someone lunges forward through the window.\n",
      "  C - Someone falls to the ground.\n",
      "  D - Someone rolls up his fast run from the water and tosses in the sky.\n",
      "\n",
      "Ground truth: option C\n"
     ]
    }
   ],
   "source": [
    "show_one(datasets[\"train\"][15])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n9qywopnIrJH"
   },
   "source": [
    "## Preprocessing the 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, convert 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": 12,
   "metadata": {
    "id": "eXNLu_-nIrJI"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rowT4iCLIrJK"
   },
   "source": [
    "You can directly call this tokenizer on one sentence or a pair of sentences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "a5hBlsrHIrJL",
    "outputId": "acdaa98a-a8cd-4a20-89b8-cc26437bbe90"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101, 8667, 117, 1142, 1110, 170, 5650, 106, 102, 1262, 1142, 5650, 2947, 1114, 1122, 119, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer(\"Hello, this is a sentence!\", \"And this sentence goes with it.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qo_0B1M2IrJM"
   },
   "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2C0hcmp9IrJQ"
   },
   "source": [
    "We can now write the function that will preprocess our samples. The tricky part is to put all the possible pairs of sentences into two big lists before passing them to the tokenizer, then un-flatten the result so that each example has four input ids, attentions masks, etc.\n",
    "\n",
    "When calling the `tokenizer`, we use the argument `truncation=True`. This will ensure that an input longer that what the model selected can handle will be truncated to the maximum length accepted by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "vc0BSBLIIrJQ"
   },
   "outputs": [],
   "source": [
    "ending_names = [\"ending0\", \"ending1\", \"ending2\", \"ending3\"]\n",
    "\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    # Repeat each first sentence four times to go with the four possibilities of second sentences.\n",
    "    first_sentences = [[context] * 4 for context in examples[\"sent1\"]]\n",
    "    # Grab all second sentences possible for each context.\n",
    "    question_headers = examples[\"sent2\"]\n",
    "    second_sentences = [\n",
    "        [f\"{header} {examples[end][i]}\" for end in ending_names]\n",
    "        for i, header in enumerate(question_headers)\n",
    "    ]\n",
    "\n",
    "    # Flatten everything\n",
    "    first_sentences = sum(first_sentences, [])\n",
    "    second_sentences = sum(second_sentences, [])\n",
    "\n",
    "    # Tokenize\n",
    "    tokenized_examples = tokenizer(first_sentences, second_sentences, truncation=True)\n",
    "    # Un-flatten\n",
    "    return {\n",
    "        k: [v[i : i + 4] for i in range(0, len(v), 4)]\n",
    "        for k, v in tokenized_examples.items()\n",
    "    }"
   ]
  },
  {
   "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 of lists for each key: a list of all examples (here 5), then a list of all choices (4) and a list of input IDs (length varying here since we did not apply any padding):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5 4 [30, 25, 30, 28]\n"
     ]
    }
   ],
   "source": [
    "examples = datasets[\"train\"][:5]\n",
    "features = preprocess_function(examples)\n",
    "print(\n",
    "    len(features[\"input_ids\"]),\n",
    "    len(features[\"input_ids\"][0]),\n",
    "    [len(x) for x in features[\"input_ids\"][0]],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check we didn't do anything wrong when grouping all possibilites and unflattening them, let's have a look at the decoded inputs for a given example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS] A drum line passes by walking down the street playing their instruments. [SEP] Members of the procession are playing ping pong and celebrating one left each in quick. [SEP]',\n",
       " '[CLS] A drum line passes by walking down the street playing their instruments. [SEP] Members of the procession wait slowly towards the cadets. [SEP]',\n",
       " '[CLS] A drum line passes by walking down the street playing their instruments. [SEP] Members of the procession makes a square call and ends by jumping down into snowy streets where fans begin to take their positions. [SEP]',\n",
       " '[CLS] A drum line passes by walking down the street playing their instruments. [SEP] Members of the procession play and go back and forth hitting the drums while the audience claps for them. [SEP]']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx = 3\n",
    "[tokenizer.decode(features[\"input_ids\"][idx][i]) for i in range(4)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare it to the ground truth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: A drum line passes by walking down the street playing their instruments.\n",
      "  A - Members of the procession are playing ping pong and celebrating one left each in quick.\n",
      "  B - Members of the procession wait slowly towards the cadets.\n",
      "  C - Members of the procession makes a square call and ends by jumping down into snowy streets where fans begin to take their positions.\n",
      "  D - Members of the procession play and go back and forth hitting the drums while the audience claps for them.\n",
      "\n",
      "Ground truth: option D\n"
     ]
    }
   ],
   "source": [
    "show_one(datasets[\"train\"][3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zS-6iXTkIrJT"
   },
   "source": [
    "This seems alright, so we can apply this function on all the examples in our dataset. All we need to do is to 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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/swag/regular/0.0.0/9640de08cdba6a1469ed3834fcab4b8ad8e38caf5d1ba5e7436d8b1fd067ad4c/cache-1931e2c0368a1bc4.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/swag/regular/0.0.0/9640de08cdba6a1469ed3834fcab4b8ad8e38caf5d1ba5e7436d8b1fd067ad4c/cache-df4adf2eee309953.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/swag/regular/0.0.0/9640de08cdba6a1469ed3834fcab4b8ad8e38caf5d1ba5e7436d8b1fd067ad4c/cache-41eb40ca99e099f0.arrow\n"
     ]
    }
   ],
   "source": [
    "encoded_datasets = datasets.map(preprocess_function, batched=True)"
   ]
  },
  {
   "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, so 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 handle 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, we can download the pretrained model and fine-tune it. Since all our task is about multiple choice, we use the `AutoModelForMultipleChoice` class. Like with the tokenizer, the `from_pretrained` method will download and cache the model for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "TlqNaB8jIrJW",
    "outputId": "84916cf3-6e6c-47f3-d081-032ec30a4132"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-21 13:43:42.021257: 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 13:43:42.060063: 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 13:43:42.061084: 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 13:43:42.062943: 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 13:43:42.066328: 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 13:43:42.067335: 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 13:43:42.068316: 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 13:43:42.798900: 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 13:43:42.799597: 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 13:43:42.800270: 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 13:43:42.800895: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 21750 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:21:00.0, compute capability: 8.6\n",
      "2022-07-21 13:43:43.912195: 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",
      "All model checkpoint layers were used when initializing TFBertForMultipleChoice.\n",
      "\n",
      "Some layers of TFBertForMultipleChoice were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier']\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 TFAutoModelForMultipleChoice\n",
    "\n",
    "model = TFAutoModelForMultipleChoice.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 others (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, and 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": [
    "Next, we set some names and hyperparameters for the model. The first two variables are used 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "Bliy8zgjIrJY"
   },
   "outputs": [],
   "source": [
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "push_to_hub_model_id = f\"{model_name}-finetuned-swag\"\n",
    "\n",
    "learning_rate = 5e-5\n",
    "batch_size = batch_size\n",
    "num_train_epochs = 2\n",
    "weight_decay = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "km3pGVdTIrJc"
   },
   "source": [
    "Next we need to tell our `Dataset` how to form batches from the pre-processed inputs.  We haven't done any padding yet because we will pad each batch to the maximum length inside the batch (instead of doing so with the maximum length of the whole dataset). This will be the job of the *data collator*. A data collator takes a list of examples and converts them to a batch (by, in our case, applying padding). Since there is no data collator in the library that works on our specific problem, we will write one, adapted from the `DataCollatorWithPadding`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from transformers.tokenization_utils_base import (\n",
    "    PreTrainedTokenizerBase,\n",
    "    PaddingStrategy,\n",
    ")\n",
    "from typing import Optional, Union\n",
    "import tensorflow as tf\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataCollatorForMultipleChoice:\n",
    "    \"\"\"\n",
    "    Data collator that will dynamically pad the inputs for multiple choice received.\n",
    "    \"\"\"\n",
    "\n",
    "    tokenizer: PreTrainedTokenizerBase\n",
    "    padding: Union[bool, str, PaddingStrategy] = True\n",
    "    max_length: Optional[int] = None\n",
    "    pad_to_multiple_of: Optional[int] = None\n",
    "\n",
    "    def __call__(self, features):\n",
    "        label_name = \"label\" if \"label\" in features[0].keys() else \"labels\"\n",
    "        labels = [feature.pop(label_name) for feature in features]\n",
    "        batch_size = len(features)\n",
    "        num_choices = len(features[0][\"input_ids\"])\n",
    "        flattened_features = [\n",
    "            [{k: v[i] for k, v in feature.items()} for i in range(num_choices)]\n",
    "            for feature in features\n",
    "        ]\n",
    "        flattened_features = sum(flattened_features, [])\n",
    "\n",
    "        batch = self.tokenizer.pad(\n",
    "            flattened_features,\n",
    "            padding=self.padding,\n",
    "            max_length=self.max_length,\n",
    "            pad_to_multiple_of=self.pad_to_multiple_of,\n",
    "            return_tensors=\"np\",\n",
    "        )\n",
    "\n",
    "        # Un-flatten\n",
    "        batch = {\n",
    "            k: tf.reshape(v, (batch_size, num_choices, -1)) for k, v in batch.items()\n",
    "        }\n",
    "        # Add back labels\n",
    "        batch[\"labels\"] = tf.convert_to_tensor(labels, dtype=tf.int64)\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When called on a list of examples, it will flatten all the inputs/attentions masks etc. in big lists that it will pass to the `tokenizer.pad` method. This will return a dictionary with big tensors (of shape `(batch_size * 4) x seq_length`) that we then unflatten.\n",
    "\n",
    "We can check this data collator works on a list of features, we just have to make sure to remove all features that are not inputs accepted by our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "accepted_keys = [\"input_ids\", \"attention_mask\", \"label\"]\n",
    "features = [\n",
    "    {k: v for k, v in encoded_datasets[\"train\"][i].items() if k in accepted_keys}\n",
    "    for i in range(10)\n",
    "]\n",
    "batch = DataCollatorForMultipleChoice(tokenizer)(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Value(dtype='int8', id=None)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_datasets[\"train\"].features[\"attention_mask\"].feature.feature"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, all those flatten/un-flattens are sources of potential errors so let's make another sanity check on our inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS] Someone walks over to the radio. [SEP] Someone hands her another phone. [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]',\n",
       " '[CLS] Someone walks over to the radio. [SEP] Someone takes the drink, then holds it. [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]',\n",
       " '[CLS] Someone walks over to the radio. [SEP] Someone looks off then looks at someone. [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]',\n",
       " '[CLS] Someone walks over to the radio. [SEP] Someone stares blearily down at the floor. [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[tokenizer.decode(batch[\"input_ids\"][8][i].numpy().tolist()) for i in range(4)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: Someone walks over to the radio.\n",
      "  A - Someone hands her another phone.\n",
      "  B - Someone takes the drink, then holds it.\n",
      "  C - Someone looks off then looks at someone.\n",
      "  D - Someone stares blearily down at the floor.\n",
      "\n",
      "Ground truth: option D\n"
     ]
    }
   ],
   "source": [
    "show_one(datasets[\"train\"][8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7sZOdRlRIrJd"
   },
   "source": [
    "All good! Now we can use this collator as a collation function for our dataset. \n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PrefetchDataset element_spec=({'input_ids': TensorSpec(shape=(16, 4, None), dtype=tf.int64, name=None), 'token_type_ids': TensorSpec(shape=(16, 4, None), dtype=tf.int64, name=None), 'attention_mask': TensorSpec(shape=(16, 4, None), dtype=tf.int64, name=None)}, TensorSpec(shape=(16,), dtype=tf.int64, name=None))>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_collator = DataCollatorForMultipleChoice(tokenizer)\n",
    "\n",
    "train_set = model.prepare_tf_dataset(\n",
    "    encoded_datasets[\"train\"],\n",
    "    shuffle=True,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=data_collator,\n",
    ")\n",
    "\n",
    "validation_set = model.prepare_tf_dataset(\n",
    "    encoded_datasets[\"validation\"],\n",
    "    shuffle=False,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=data_collator,\n",
    ")\n",
    "\n",
    "train_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, our dataset will output a 2-tuple where the first element is a dict containing `input_ids`, `token_type_ids` and `attention_mask`, and the second element is the label. This is exactly what we want for our model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compile our model. First, we specify an optimizer. Using the `create_optimizer` function we can get a nice `AdamW` optimizer with weight decay and a learning rate decay schedule set up for free - but to compute that schedule, it needs to know how long training will take."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import create_optimizer\n",
    "\n",
    "total_train_steps = (len(encoded_datasets[\"train\"]) // batch_size) * 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",
    "In some of our other examples, we use `jit_compile` to compile the model with [XLA](https://www.tensorflow.org/xla). In this case, we should be careful about that - because our inputs 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! For small datasets, this will probably result in spending more time on XLA compilation than actually training, which isn't very helpful.\n",
    "\n",
    "If you really want to use XLA without these problems (for example, if you're training on TPU), you can create a tokenizer with `padding=\"max_length\"`. This will pad all of your samples to the same length, ensuring that a single XLA compilation will suffice for your entire dataset. Note that depending on the nature of your dataset, this may result in a lot of wasted computation on padding tokens!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "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(\n",
    "    optimizer=optimizer,\n",
    "    metrics=[\"accuracy\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can 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! Make sure to change the `username` if you do. If you don't want to do this, simply remove the callbacks argument in the call to `fit()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matt/PycharmProjects/notebooks/examples/mc_model_save is already a clone of https://huggingface.co/Rocketknight1/bert-base-cased-finetuned-swag. Make sure you pull the latest changes with `repo.git_pull()`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "4596/4596 [==============================] - ETA: 0s - loss: 0.8709 - accuracy: 0.6465"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Several commits (2) will be pushed upstream.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4596/4596 [==============================] - 827s 178ms/step - loss: 0.8709 - accuracy: 0.6465 - val_loss: 0.6167 - val_accuracy: 0.7590\n",
      "Epoch 2/2\n",
      "4596/4596 [==============================] - 820s 178ms/step - loss: 0.3868 - accuracy: 0.8568 - val_loss: 0.6014 - val_accuracy: 0.7795\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fa5b8519780>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers.keras_callbacks import PushToHubCallback\n",
    "from tensorflow.keras.callbacks import TensorBoard\n",
    "\n",
    "tensorboard_callback = TensorBoard(log_dir=\"./mc_model_save/logs\")\n",
    "\n",
    "push_to_hub_callback = PushToHubCallback(\n",
    "    output_dir=\"./mc_model_save\",\n",
    "    tokenizer=tokenizer,\n",
    "    hub_model_id=push_to_hub_model_id,\n",
    ")\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": [
    "If you used 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 TFAutoModelForMultipleChoice\n",
    "\n",
    "model = TFAutoModelForMultipleChoice.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": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7d543f74e1647af83acce4fa2a8acdc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading config.json:   0%|          | 0.00/645 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa7bc00e99654f1daf48f4b4603374a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tf_model.h5:   0%|          | 0.00/413M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some layers from the model checkpoint at Rocketknight1/bert-base-cased-finetuned-swag were not used when initializing TFBertForMultipleChoice: ['dropout_37']\n",
      "- This IS expected if you are initializing TFBertForMultipleChoice 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 TFBertForMultipleChoice from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the layers of TFBertForMultipleChoice were initialized from the model checkpoint at Rocketknight1/bert-base-cased-finetuned-swag.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFBertForMultipleChoice for predictions without further training.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4065cdaf16ba4b128263fea2f67335cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tokenizer_config.json:   0%|          | 0.00/347 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d61f34bd35094f6f9ecef5cfc6245309",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading vocab.txt:   0%|          | 0.00/208k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87d5c9eb3c604f1fbf99953e03a5b620",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tokenizer.json:   0%|          | 0.00/653k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69545af63f5b4df885d2eeddecd56aa3",
       "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, TFAutoModelForMultipleChoice\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/bert-base-cased-finetuned-swag\"\n",
    "model = TFAutoModelForMultipleChoice.from_pretrained(checkpoint)\n",
    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see how to use this model for inference. The SWAG task we trained on is a commonsense inference benchmark, where we ask the model to indicate which of four completions of a sentence is realistic and makes sense in context. Let's use a sample input from SWAG and see how we can get predictions for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_start = \"Members of the procession walk down the street holding small horn brass instruments. A drum line\"\n",
    "endings = [\n",
    "    'passes by walking down the street playing their instruments.',\n",
    "    'has heard approaching them.',\n",
    "    \"arrives and they're outside dancing and asleep.\",\n",
    "    'turns the lead singer watches the performance.',\n",
    "]\n",
    "full_sentences = [f\"{input_start} {ending}\" for ending in endings]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we tokenize this input. Note that our inputs need to be reshaped a little - multiple choice models expect inputs to have the shape `(num_samples, num_choices, num_tokens)` - this means we will need to add a sample/batch dimension of length 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "tokenized = tokenizer(full_sentences, padding=\"longest\", return_tensors=\"np\")\n",
    "tokenized = {key: np.expand_dims(array, 0) for key, array in tokenized.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we run these inputs through our model and see what it guesses!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The answer is choice 0: passes by walking down the street playing their instruments.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "outputs = model(tokenized).logits\n",
    "answer = np.argmax(outputs)\n",
    "\n",
    "print(f\"The answer is choice {answer}: {endings[answer]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Multiple choice on SWAG",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}