{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2022 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Notebook autogenerated from walkthrough.py on 08-Sep-2022*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a Hugging Face model on the IPU using a local dataset\n",
    "\n",
    "This PyTorch tutorial will show you how to reuse a Hugging Face model and train it on the IPU using a local dataset.\n",
    "Specifically, we will be fine-tuning a Vision Transformer (ViT) model to detect multiple diseases from chest X-rays. As an X-ray image can have multiple diseases we will be training a multi-label classification model.\n",
    "We will use [Graphcore's Optimum interface](https://github.com/huggingface/optimum-graphcore) to the [Hugging Face Transformers library](https://huggingface.co/docs/transformers/index) to run an existing model on the IPU.\n",
    "We will be using the [google/vit-base-patch16-224-in21k checkpoint](https://huggingface.co/google/vit-base-patch16-224-in21k) pretrained on ImageNet, fine-tuning it on the NIH Chest X-ray Dataset.\n",
    "\n",
    "In this tutorial, you will learn how to:\n",
    "- Repurpose one of the examples from the [Graphcore organisation page on Hugging Face](https://huggingface.co/Graphcore) to fit your use case.\n",
    "- Preprocess a dataset to fit it to an existing model.\n",
    "- Use the Graphcore model cards found on the [Graphcore organisation page on Hugging Face](https://huggingface.co/Graphcore) and reuse checkpoints and config files released by Graphcore.\n",
    "- Maximise IPU utilisation for your specific machine by overriding runtime parameters in the `IPUconfig` object made available in the model cards.\n",
    "\n",
    "If this is your first time using IPUs, read the [IPU Programmer's Guide](https://docs.graphcore.ai/projects/ipu-programmers-guide/en/3.1.0/) to learn the basic concepts.\n",
    "To run your own PyTorch model on the IPU see the [PyTorch basics tutorial](..learning-PyTorch-on-IPU/basics), or to see all existing Graphcore models available from Hugging Face go to the [Graphcore organisation page](https://huggingface.co/Graphcore)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to run this tutorial\n",
    "\n",
    "### Getting the dataset\n",
    "\n",
    "This tutorial uses the NIH Chest X-ray Dataset downloaded from <http://nihcc.app.box.com/v/ChestXray-NIHCC>.\n",
    "\n",
    "### Environment\n",
    "\n",
    "Requirements:\n",
    "\n",
    "- Python packages installed with `python -m pip install -r requirements.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "sst_ignore_md",
     "sst_ignore_code_only"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "579d8a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples_utils import notebook_logging\n",
    "%load_ext gc_logger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graphcore Hugging Face models\n",
    "\n",
    "Hugging Face provides convenient access to pre-trained transformer models.\n",
    "The partnership between Hugging Face and Graphcore allows us to run these models on the IPU.\n",
    "\n",
    "Hugging Face models ported to the IPU can be found on the [Graphcore organisation page on Hugging Face](https://huggingface.co/Graphcore).\n",
    "This tutorial uses the [Vision Transformer model](https://github.com/huggingface/optimum-graphcore/tree/main/examples/image-classification) fine-tuned using the NIH Chest X-ray Dataset, as an example to show how Hugging Face models can be trained with a local dataset on the IPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility imports\n",
    "\n",
    "We start by importing the utilities that will be used\n",
    "later in the tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import contextlib\n",
    "import io\n",
    "from pathlib import Path\n",
    "from scipy.special import softmax\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "import optimum.graphcore as optimum_graphcore\n",
    "import datasets\n",
    "from torchvision import transforms\n",
    "\n",
    "# The `chest-xray-nihcc` directory is assumed to be in the pwd, but may be overridden by the environment variable `DATASETS_DIR`\n",
    "dataset_rootdir = Path(os.environ.get(\"DATASETS_DIR\", \".\")) / \"chest-xray-nihcc-3\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the NIH Chest X-ray Dataset\n",
    "\n",
    "For this tutorial, we are using the NIH Chest X-ray dataset, which contains radiological examinations for the diagnosis of lung diseases.\n",
    "The examinations consist of square, grayscale, X-ray images of 224 pixels with corresponding metadata: Finding Labels, Follow-up #,Patient ID, Patient Age, Patient Gender, View Position, OriginalImage[Width Height] and OriginalImagePixelSpacing[x y].\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the extracted \"images\" directory\n",
    "images_dir = dataset_rootdir / \"images\"\n",
    "tar_file = dataset_rootdir / \"images_01.tar.gz\"\n",
    "bash_file = dataset_rootdir / \"unpack-images.sh\"\n",
    "\n",
    "# If the dataset is packed then unpack the data, this can take up to 10 minutes\n",
    "if not images_dir.exists() and tar_file.exists():\n",
    "    !source {bash_file}\n",
    "\n",
    "\n",
    "# Path to Data_Entry_2017_v2020.csv\n",
    "label_file = dataset_rootdir / \"Data_Entry_2017_v2020.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the labels\n",
    "\n",
    "We are going to train the Graphcore Optimum ViT model to predict the disease (defined by \"Finding Label\") from the  images.\n",
    "\"Finding Label\" can be any number of 14 diseases or a \"No Finding\" label, which indicates that no disease was detected.\n",
    "To be compatible with the `datasets` Hugging Face library, the text labels need to be transformed to N-hot encoded arrays representing the multiple labels which are needed to classify each image.\n",
    "An N-hot encoded array represents the labels as a list of booleans, true if the label corresponds to the image and false if not.\n",
    "\n",
    "First we identify the unique labels in the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset contains the following labels:\n",
      "['Atelectasis', 'Cardiomegaly', 'Consolidation', 'Edema', 'Effusion', 'Emphysema', 'Fibrosis', 'Hernia', 'Infiltration', 'Mass', 'No Finding', 'Nodule', 'Pleural_Thickening', 'Pneumonia', 'Pneumothorax']\n"
     ]
    }
   ],
   "source": [
    "data = pd.read_csv(label_file)\n",
    "\n",
    "# Converts the format of each label in the dataframe from \"LabelA|LabelB|LabelC\"\n",
    "# into [\"LabelA\", \"LabelB\", \"LabelC\"], concatenates the\n",
    "# lists together and removes duplicate labels\n",
    "unique_labels = np.unique(\n",
    "    data[\"Finding Labels\"].str.split(\"|\").aggregate(np.concatenate)\n",
    ").tolist()\n",
    "\n",
    "print(f\"Dataset contains the following labels:\\n{unique_labels}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we transform the labels into N-hot encoded arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_index = {v: i for i, v in enumerate(unique_labels)}\n",
    "\n",
    "\n",
    "def string_to_N_hot(string: str):\n",
    "    true_index = [label_index[cl] for cl in string.split(\"|\")]\n",
    "    label = np.zeros((len(unique_labels),), dtype=float)\n",
    "    label[true_index] = 1\n",
    "    return label\n",
    "\n",
    "\n",
    "data[\"labels\"] = data[\"Finding Labels\"].apply(string_to_N_hot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When loading data using the `datasets.load_dataset` function, labels can be provided either by having folders for each of the labels (see [\"ImageFolder\" documentation](https://huggingface.co/docs/datasets/v2.3.2/en/image_process#imagefolder)) or by having a `metadata.jsonl` file ((see [\"ImageFolder with metadata\" documentation](https://huggingface.co/docs/datasets/v2.3.2/en/image_process#imagefolder-with-metadata))). As the images in this dataset can have multiple labels, we have chosen to use a `metadata.jsonl` file.\n",
    "We write the image file names and their associated labels to the `metadata.jsonl` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_file = images_dir / \"metadata.jsonl\"\n",
    "if not metadata_file.is_file():\n",
    "    data[[\"Image Index\", \"labels\"]].rename(\n",
    "        columns={\"Image Index\": \"file_name\"}\n",
    "    ).to_json(images_dir / \"metadata.jsonl\", orient=\"records\", lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the dataset\n",
    "\n",
    "We are now ready to create the PyTorch dataset and split it into training and validation sets.\n",
    "This step converts the dataset to the [Arrow file format](https://arrow.apache.org) which allows data to be loaded quickly during training and validation ([about Arrow and Hugging Face](https://huggingface.co/docs/datasets/v2.3.2/en/about_arrow)).\n",
    "Because the entire dataset is being loaded and pre-processed it can take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   },
   "outputs": [],
   "source": [
    "train_val_split = 0.05\n",
    "dataset = datasets.load_dataset(\n",
    "    \"imagefolder\",\n",
    "    data_dir=images_dir,\n",
    ")\n",
    "\n",
    "split = dataset[\"train\"].train_test_split(train_val_split)\n",
    "dataset[\"train\"] = split[\"train\"]\n",
    "dataset[\"validation\"] = split[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to import the Vision Transformer(ViT) model from\n",
    "the checkpoint `google/vit-base-patch16-224-in21k`.\n",
    "The checkpoint is a standard model hosted by Hugging Face and is not managed by Graphcore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path = \"google/vit-base-patch16-224-in21k\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fine-tune a pre-trained model, the new dataset must have the same properties as the original dataset used for pre-training.\n",
    "In Hugging Face, the original dataset information is provided in a config file loaded using the `AutoFeatureExtractor`.\n",
    "For this model, the X-ray images are resized to the correct resolution (224x224), converted from grayscale to RGB, and normalized across the RGB channels with a mean (0.5, 0.5, 0.5) and a standard deviation (0.5, 0.5, 0.5)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_extractor = transformers.AutoFeatureExtractor.from_pretrained(\n",
    "    model_name_or_path\n",
    ")\n",
    "\n",
    "\n",
    "class XRayTransform:\n",
    "    \"\"\"\n",
    "    Transforms for pre-processing XRay data across a batch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.transforms = transforms.Compose(\n",
    "            [\n",
    "                transforms.Lambda(lambda pil_img: pil_img.convert(\"RGB\")),\n",
    "                transforms.Resize(feature_extractor.size),\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize(\n",
    "                    mean=feature_extractor.image_mean, std=feature_extractor.image_std\n",
    "                ),\n",
    "                transforms.ConvertImageDtype(dtype=torch.float16),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def __call__(self, example_batch):\n",
    "        example_batch[\"pixel_values\"] = [\n",
    "            self.transforms(pil_img) for pil_img in example_batch[\"image\"]\n",
    "        ]\n",
    "        return example_batch\n",
    "\n",
    "\n",
    "# Set the training transforms\n",
    "dataset[\"train\"].set_transform(XRayTransform())\n",
    "# Set the validation transforms\n",
    "dataset[\"validation\"].set_transform(XRayTransform())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the model to run efficiently images need to be batched. To do this, we define a `batch_sampler` function which returns batches of images and labels in a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_sampler(examples):\n",
    "    pixel_values = torch.stack([example[\"pixel_values\"] for example in examples])\n",
    "    labels = torch.tensor([example[\"labels\"] for example in examples])\n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualising the dataset\n",
    "To examine the dataset we display the first 10 rows of the metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        Image Index          Finding Labels  Follow-up #  Patient ID  \\\n",
      "0  00000001_000.png            Cardiomegaly            0           1   \n",
      "1  00000001_001.png  Cardiomegaly|Emphysema            1           1   \n",
      "2  00000001_002.png   Cardiomegaly|Effusion            2           1   \n",
      "3  00000002_000.png              No Finding            0           2   \n",
      "4  00000003_001.png                  Hernia            0           3   \n",
      "5  00000003_002.png                  Hernia            1           3   \n",
      "6  00000003_003.png     Hernia|Infiltration            2           3   \n",
      "7  00000003_004.png                  Hernia            3           3   \n",
      "8  00000003_005.png                  Hernia            4           3   \n",
      "9  00000003_006.png                  Hernia            5           3   \n",
      "\n",
      "   Patient Age Patient Gender View Position  OriginalImage[Width  Height]  \\\n",
      "0           57              M            PA                 2682     2749   \n",
      "1           58              M            PA                 2894     2729   \n",
      "2           58              M            PA                 2500     2048   \n",
      "3           80              M            PA                 2500     2048   \n",
      "4           74              F            PA                 2500     2048   \n",
      "5           75              F            PA                 2048     2500   \n",
      "6           76              F            PA                 2698     2991   \n",
      "7           77              F            PA                 2500     2048   \n",
      "8           78              F            PA                 2686     2991   \n",
      "9           79              F            PA                 2992     2991   \n",
      "\n",
      "   OriginalImagePixelSpacing[x     y]  \\\n",
      "0                        0.143  0.143   \n",
      "1                        0.143  0.143   \n",
      "2                        0.168  0.168   \n",
      "3                        0.171  0.171   \n",
      "4                        0.168  0.168   \n",
      "5                        0.168  0.168   \n",
      "6                        0.143  0.143   \n",
      "7                        0.168  0.168   \n",
      "8                        0.143  0.143   \n",
      "9                        0.143  0.143   \n",
      "\n",
      "                                              labels  \n",
      "0  [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...  \n",
      "1  [0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, ...  \n",
      "2  [0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, ...  \n",
      "3  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...  \n",
      "4  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ...  \n",
      "5  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ...  \n",
      "6  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, ...  \n",
      "7  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ...  \n",
      "8  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ...  \n",
      "9  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ...  \n"
     ]
    }
   ],
   "source": [
    "print(data.head(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also plot some images from the validation set with their associated labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 15))\n",
    "\n",
    "unique_labels = np.array(unique_labels)\n",
    "\n",
    "convert_image_to_float = transforms.ConvertImageDtype(dtype=torch.float32)\n",
    "for i, data_dict in enumerate(dataset[\"validation\"]):\n",
    "    if i == 12:\n",
    "        break\n",
    "    image = data_dict[\"pixel_values\"]\n",
    "    # Convert image to format supported by imshow\n",
    "    image = convert_image_to_float(image)\n",
    "    label = data_dict[\"labels\"]\n",
    "    ax = plt.subplot(3, 4, i + 1)\n",
    "    ax.set_title(\", \".join(unique_labels[np.argwhere(label).flatten()]))\n",
    "    plt.imshow(image[0])  # Plot only the first channel as they are all identical\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The images are chest X-rays with labels of lung diseases the patient was diagnosed with. Here, we show the transformed images. Our dataset is now ready to be used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the model\n",
    "\n",
    "To train a model on the IPU we need to import it from Hugging Face and define a trainer using the `IPUTrainer` class.\n",
    "The `IPUTrainer` class takes the same arguments as the Hugging Face model along with an `IPUConfig` object which specifies the behaviour for compilation and execution on the IPU.\n",
    "\n",
    "Now we import the ViT model from Hugging Face."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = transformers.AutoModelForImageClassification.from_pretrained(\n",
    "    model_name_or_path, num_labels=len(unique_labels)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use this model on the IPU we need to load the IPU configuration, `IPUConfig`, which gives control to all the parameters specific to Graphcore IPUs.\n",
    "Existing IPU configs can be found at : https://huggingface.co/Graphcore\n",
    "We are going to use `Graphcore/vit-base-ipu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache\") + \"/vit_model_training\"\n",
    "ipu_config = optimum_graphcore.IPUConfig.from_pretrained(\"Graphcore/vit-base-ipu\", executable_cache_dir=executable_cache_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set our training hyperparameters using `IPUTrainingArguments`.\n",
    "This subclasses the Hugging Face `TrainingArguments` class, adding parameters specific to the IPU and its execution characteristics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = optimum_graphcore.IPUTrainingArguments(\n",
    "    output_dir=\".graphcore/vit-model\",\n",
    "    overwrite_output_dir=True,\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    dataloader_num_workers=8,\n",
    "    dataloader_drop_last=True,\n",
    "    num_train_epochs=3,\n",
    "    seed=1337,\n",
    "    logging_steps=50,\n",
    "    save_steps=1000,\n",
    "    remove_unused_columns=False,\n",
    "    warmup_ratio=0.25,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    learning_rate=2e-4,\n",
    "    ignore_data_skip=True,\n",
    "    report_to=\"none\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Performance of multi-label classification models can be assessed using the area under the  ROC (receiver operating characteristic) curve (AUC_ROC). The AUC_ROC represents the ability of the model to separate the different diseases. A score of 0.5 means that it is\n",
    "50% likely to get the correct disease and a score of 1 means that it is able to perfectly separate the diseases.\n",
    "To add this metric to our model evaluation we define a `compute_metrics` function and load the metric from the `datasets` package.\n",
    "The `compute_metrics` function takes the labels predicted by the ViT model and computes the area under the ROC curve.\n",
    "The `compute metrics` function takes an `EvalPrediction` object (a named tuple with a predictions and `label_ids` field), and has to return a dictionary string to float."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_auc = datasets.load_metric(\"roc_auc\", \"multilabel\")\n",
    "\n",
    "def compute_metrics(p):\n",
    "    preds = np.argmax(p.predictions, axis=1)\n",
    "\n",
    "    pred_scores = softmax(p.predictions.astype(\"float32\"), axis=1)\n",
    "    auc = metric_auc.compute(\n",
    "        prediction_scores=pred_scores, references=p.label_ids, multi_class=\"ovo\"\n",
    "    )[\"roc_auc\"]\n",
    "    return {\"roc_auc\": auc}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train the model, we define a trainer using the `IPUTrainer` class which takes care of compiling the model to run on IPUs, and of performing training and evaluation.\n",
    "The `IPUTrainer` class works just like the Hugging Face `Trainer` class, but takes the additional `ipu_config` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = optimum_graphcore.IPUTrainer(\n",
    "    model=model,\n",
    "    ipu_config=ipu_config,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    eval_dataset=dataset[\"validation\"],\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=feature_extractor,\n",
    "    data_collator=batch_sampler,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To accelerate training we will load the last checkpoint if it exists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_checkpoint = None\n",
    "if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir:\n",
    "    last_checkpoint = transformers.trainer_utils.get_last_checkpoint(\n",
    "        training_args.output_dir\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capture the command line output for plotting loss and learning rate\n",
    "output = io.StringIO()\n",
    "\n",
    "with contextlib.redirect_stdout(output):\n",
    "    trainer.train(resume_from_checkpoint=last_checkpoint)\n",
    "\n",
    "# Visualise a fragment of the raw output\n",
    "print(output.getvalue()[:500])\n",
    "print(\"...\")\n",
    "print(output.getvalue()[-500:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting convergence\n",
    "\n",
    "Now that we have completed training we can format and plot the trainer output to evaluate the training behaviour."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   },
   "outputs": [],
   "source": [
    "# Put the trainer logs in a data frame\n",
    "values = []\n",
    "for line in output.getvalue().split(\"\\n\"):\n",
    "    if len(line) > 3 and line[:3] == \"{'l\":\n",
    "        values.append(json.loads(line.replace(\"'\", '\"')))\n",
    "training_records = pd.DataFrame(values)\n",
    "training_records.tail(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We plot the training loss and the learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 1)\n",
    "training_records.plot(x=\"epoch\", y=\"loss\", ax=axs[0])\n",
    "training_records.plot(x=\"epoch\", y=\"learning_rate\", ax=axs[1])\n",
    "fig.set_size_inches(8, 8)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss curve shows a rapid reduction in the loss at the start of training before stabilising around 0.1, showing that the model is learning.\n",
    "The learning rate increases through the warm-up of 25% of the training period, before following a cosine decay."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the evaluation\n",
    "Now that we have trained the model we can evaluate its ability to predict the labels of unseen data using the validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = trainer.evaluate()\n",
    "trainer.log_metrics(\"eval\", metrics)\n",
    "trainer.save_metrics(\"eval\", metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The metrics show the validation AUC_ROC score the tutorial achieves after 3 epochs. The AUC_ROC score is a measure of how well a model performs classification. A score 1 of means the model is able to correctly distinguish between classes for all datapoints. A score of 0.5 means the model is unable to distinguish between classes, either it predicts randomly or the same class for all datapoints. The closer the AUC_ROC is to 1 the better it performs.\n",
    "\n",
    "There are several directions to explore to improve the accuracy of the model including longer training. The validation performance might also be improved through changing optimisers, learning rate, learning rate schedule, loss scaling or using auto-loss scaling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial we have shown how to train a Hugging Face model on the IPU using a local dataset.\n",
    "To do this, we:\n",
    "- loaded the local dataset using the Hugging Face `datasets` package,\n",
    "- defined transforms to ensure the dataset was compatible with the model,\n",
    "- defined the `IPUTrainer` class to allow the model to run on the IPU,\n",
    "- ran the model.\n",
    "\n",
    "The dataset required preprocessing before being compatible with the `datasets` package.\n",
    "We needed to encode the labels as N-hot arrays to make them compatible with the `load_dataset` function.\n",
    "To improve the training of the model, device parameters were loaded and customised in the `IPUConfig` object and hyperparameters were tuned in the `IPUTrainer` class.\n",
    "\n",
    "You are now ready to use Graphcore and Hugging Face models for your own application.\n",
    "To see all the Hugging Face models available for the IPU, see the [Graphcore organisation page](https://huggingface.co/Graphcore) and for more information about\n",
    "how models are run please consult the [Hugging Face documentation](https://huggingface.co/docs) and the [Graphcore Optimum Library](https://github.com/huggingface/optimum-graphcore).\n",
    "For information on how to optimise models for the IPU see the [Memory and Performance Optimisation Guide](https://docs.graphcore.ai/projects/memory-performance-optimisation/).\n",
    "Graphcore also provides reference implementations of model architectures for many applications in the [Graphcore examples repository](https://github.com/graphcore/examples) and in the [Graphcore Model Garden](https://www.graphcore.ai/resources/model-garden)."
   ]
  }
 ],
 "metadata": {
  "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.8.10"
  },
  "traceability": {
   "sdk_version": "3.1.0-EA.1+1118",
   "source_file": "walkthrough.py",
   "sst_version": "0.0.8",
   "timestamp": "2022-09-08T15:27"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
