{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "rEJBSTyZIrIb"
   },
   "source": [
    "# Faster single-label text classification on IPU with Packed BERT - Fine-tuning and Inference\n",
    "\n",
    "This notebook builds on the process in the notebook on \"fine-tuning BERT on a text classification task\" `natural-language-processing/other-use-cases/text_classification.ipynb` and uses [packing](https://www.graphcore.ai/posts/introducing-packed-bert-for-2x-faster-training-in-natural-language-processing), an optimisation method originally used for 2x faster BERT pre-training, which can now also provide massive throughput increases for fine-tuning and batched inference! \n",
    "\n",
    "**So, what *is* packing?** The basic idea of \"packing\" a dataset is to utilise the requirement for constant-shaped inputs into a model. Instead of padding it with empty, unused space, we can recycle this unused space and fill it with more inputs! The architecture of transformer models like BERT supports this, and lets us optimally use this space to process multiple sequences within one input.\n",
    "\n",
    "**And here is why you might want to use it:** Having a single input that contains multiple sequences leads to multiple sequences being processed in parallel in a single pass within a single iteration inside a batch, increasing the \"effective\" batch size of the model by a considerable factor in many cases, and most importantly, increasing model throughput for training and batched inference significantly.\n",
    "\n",
    "This notebook outlines how to easily enable packing for BERT when performing fine-tuning and inference on a text-classification task in 🤗 Optimum, resulting in an impressive 5-6x faster training and inference run-time on the `GLUE/sst2` dataset. \n",
    "\n",
    "You can read more about packing in the original paper on [efficient sequence packing without cross-contamination](https://arxiv.org/abs/2107.02027)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "kTCFado4IrIc"
   },
   "source": [
    "In this notebook, we will see how to fine-tune BERT, a [🤗 Transformers](https://github.com/huggingface/transformers) model to a text classification task of the [GLUE Benchmark](https://gluebenchmark.com/).\n",
    "\n",
    "The GLUE Benchmark is a group of nine classification tasks on sentences or pairs of sentences. The tasks are:\n",
    "\n",
    "- [CoLA](https://nyu-mll.github.io/CoLA/) (Corpus of Linguistic Acceptability) Determine if a sentence is grammatically correct or not.\n",
    "- [MNLI](https://arxiv.org/abs/1704.05426) (Multi-Genre Natural Language Inference) Determine if a sentence entails, contradicts or is unrelated to a given hypothesis. (This dataset has two versions, one with the validation and test sets being from the same distribution and another called \"mismatched\" where the validation and test sets use out-of-domain data.)\n",
    "- [MRPC](https://www.microsoft.com/en-us/download/details.aspx?id=52398) (Microsoft Research Paraphrase Corpus) Determine if two sentences are paraphrases of one another or not.\n",
    "- [QNLI](https://rajpurkar.github.io/SQuAD-explorer/) (Question-answering Natural Language Inference) Determine if the answer to a question is in the second sentence or not. (This dataset is built from the SQuAD dataset.)\n",
    "- [QQP](https://data.quora.com/First-Quora-Dataset-Release-Question-Pairs) (Quora Question Pairs2) Determine if two questions are semantically equivalent or not.\n",
    "- [RTE](https://aclweb.org/aclwiki/Recognizing_Textual_Entailment) (Recognizing Textual Entailment) Determine if a sentence entails a given hypothesis or not.\n",
    "- [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) Determine if the sentence has a positive or negative sentiment.\n",
    "- [STS-B](http://ixa2.si.ehu.es/stswiki/index.php/STSbenchmark) (Semantic Textual Similarity Benchmark) Determine the similarity of two sentences with a score from 1 to 5.\n",
    "- [WNLI](https://cs.nyu.edu/faculty/davise/papers/WinogradSchemas/WS.html) (Winograd Natural Language Inference) Determine if a sentence with an anonymous pronoun and a sentence with this pronoun replaced are entailed or not. (This dataset is built from the Winograd Schema Challenge dataset.)\n",
    "\n",
    "We will see how to easily load the dataset for these tasks and use BERT with packing to fine-tune a model on SST-2. Each task is named using an abbreviation, with `mnli-mm` representing the mismatched version of MNLI (`mnli-mm` has the same training set as `mnli` but has different validation and test sets):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YZbiBDuGIrId",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "GLUE_TASKS = [\"cola\", \"mnli\", \"mnli-mm\", \"mrpc\", \"qnli\", \"qqp\", \"rte\", \"sst2\", \"stsb\", \"wnli\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n",
    "\n",
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies and configuration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the dependencies for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "%pip install optimum-graphcore==0.7\n",
    "%pip install scikit-learn \"datasets>=2.7.0\" evaluate tokenizers matplotlib scipy huggingface_hub\n",
    "\n",
    "%pip install graphcore-cloud-tools[logger]@git+https://github.com/graphcore/graphcore-cloud-tools\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "4RRkXuteIrIh"
   },
   "source": [
    "We use both data parallelism and pipeline parallelism (see this [tutorial on efficient data loading](https://github.com/graphcore/examples/tree/master/tutorials/tutorials/pytorch/efficient_data_loading/walkthrough.ipynb) for more information). Therefore the global batch size, which is the actual number of samples used for the weight update, is determined using four factors:\n",
    "\n",
    "    global batch size = micro_batch_size * gradient accumulation steps * device iterations * replication factor\n",
    "\n",
    "Replication factor is determined by dividing `n_ipu`, the number of IPUs to use, by `ipus_per_replica`.\n",
    "\n",
    "Depending on your model and the IPU Pod you are using, you might need to adjust these batch-size-related arguments.\n",
    "\n",
    "By default this notebook is configured to run on 4 IPUs.\n",
    "\n",
    "We also define `max_seq_length` which is the maximum length a sequence can be, and all sequences will be padded to this length. Therefore, `max_seq_length` should not be larger than the maximum length of the model. Given the small size of the sequences in `sst2`, we can reduce the model maximum input size to `max_seq_length = 256`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this Packed BERT demo, we cover (single-label) sequence classification on the `sst2` dataset. The `task` can be changed to run the other `GLUE` tasks defined above. However, training hyperparameters may need further tuning for these other tasks."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set these parameters and the rest of the notebook should run smoothly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zVvslsfMIrIh",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "task = \"sst2\"\n",
    "model_checkpoint = \"bert-base-uncased\"\n",
    "ipu_config_name = \"Graphcore/bert-base-uncased\"\n",
    "micro_batch_size = 2\n",
    "gradient_accumulation_steps = 32\n",
    "device_iterations = 32\n",
    "max_seq_length = 256"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gradients are not calculated during validation, so gradient accumulation is not applicable, and the global batch size for validation can be defined separately as:\n",
    "\n",
    "```\n",
    "global_validation_batch_size=device_iterations*replication_factor*max_seq_per_pack\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Values for machine size and cache directories can be configured through environment variables or directly in the notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "n_ipu = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\") + \"/packed_bert_slseqcls/\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing your model with the community"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can share your model with the 🤗 community. You do this by completing the following steps:\n",
    "\n",
    "1. Store your authentication token from the 🤗 website. [Sign up to 🤗](https://huggingface.co/join) if you haven't already.\n",
    "2. Execute the following cell and input your username and authentication token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Git-LFS must also be installed to manage large files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!apt install git-lfs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "whPRbBNbIrIl"
   },
   "source": [
    "## Loading the dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "W7QYTpxXIrIl"
   },
   "source": [
    "We will use the [🤗 Datasets library](https://huggingface.co/docs/datasets/index) and the [🤗 Evaluate library](https://huggingface.co/docs/evaluate/index) to download the data and get the metric we will use for evaluation (to compare our model to the benchmark). This can be easily done with the functions `load_dataset` and `evaluate.load()`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IreSlFmlIrIm",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import evaluate"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "CKx2zKs5IrIq"
   },
   "source": [
    "Apart from `mnli-mm` being a special task code, we can directly pass our task name to those functions. `load_dataset` will cache the dataset to avoid downloading it again the next time you run this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "actual_task = \"mnli\" if task == \"mnli-mm\" else task\n",
    "dataset = load_dataset(\"glue\", actual_task)\n",
    "metric = evaluate.load('glue', actual_task)"
   ]
  },
  {
   "attachments": {},
   "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 sets (with more keys for the mismatched validation and test set in the special case of `mnli`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GWiVUF0jIrIv",
    "outputId": "35e3ea43-f397-4a54-c90c-f2cf8d36873e",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "u3EtYfeHIrIz"
   },
   "source": [
    "To access an actual element, you need to select a split (\"train\" in the example below) and specify an index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X6HrpprwIrIz",
    "outputId": "d7670bc0-42e4-4c09-8a6a-5c018ded7d95",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dataset[\"train\"][:10]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "WHUmphG3IrI3"
   },
   "source": [
    "To get a sense of what the data looks like, the following function will show some samples picked randomly from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i3j8APAoIrI3",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "def show_random_elements(dataset, num_examples=10):\n",
    "    assert num_examples <= len(dataset), \"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, datasets.ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SZy5tRB_IrI7",
    "outputId": "ba8f2124-e485-488f-8c0c-254f34f24f13",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "show_random_elements(dataset[\"train\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "lnjDIuQ3IrI-"
   },
   "source": [
    "The metric is an instance of [`datasets.Metric`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Metric):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5o4rUteaIrI_",
    "outputId": "18038ef5-554c-45c5-e00a-133b02ec10f1",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "metric"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "jAWdqcUBIrJC"
   },
   "source": [
    "You can call the `compute` method of metric` with your predictions and labels directly and it will return a dictionary with the metric values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6XN1Rq0aIrJC",
    "outputId": "a4405435-a8a9-41ff-9f79-a13077b587c7",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "fake_preds = np.random.randint(0, 2, size=(64,))\n",
    "fake_labels = np.random.randint(0, 2, size=(64,))\n",
    "metric.compute(predictions=fake_preds, references=fake_labels)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "YOCrQwPoIrJG"
   },
   "source": [
    "Note that `load_metric` has loaded the proper metric associated to your task, which is:\n",
    "\n",
    "- for CoLA: [Matthews Correlation Coefficient](https://en.wikipedia.org/wiki/Matthews_correlation_coefficient)\n",
    "- for MNLI (matched or mismatched): Accuracy\n",
    "- for MRPC: Accuracy and [F1 score](https://en.wikipedia.org/wiki/F1_score)\n",
    "- for QNLI: Accuracy\n",
    "- for QQP: Accuracy and [F1 score](https://en.wikipedia.org/wiki/F1_score)\n",
    "- for RTE: Accuracy\n",
    "- for SST-2: Accuracy\n",
    "- for STS-B: [Pearson Correlation Coefficient](https://en.wikipedia.org/wiki/Pearson_correlation_coefficient) and [Spearman's_Rank_Correlation_Coefficient](https://en.wikipedia.org/wiki/Spearman%27s_rank_correlation_coefficient)\n",
    "- for WNLI: Accuracy\n",
    "\n",
    "so the metric object only computes the relevant metrics needed for your task."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "n9qywopnIrJH"
   },
   "source": [
    "## Preprocessing the data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "YVx71GdAIrJH"
   },
   "source": [
    "Before we can feed the text samples to our model, we need to preprocess them. This is done by a 🤗 Transformers `Tokenizer` which will tokenize the inputs (including converting the tokens to their corresponding IDs in the pre-trained vocabulary) and put it into a format the model expects, as well as generate the other inputs that the 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 pre-training 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": null,
   "metadata": {
    "id": "eXNLu_-nIrJI",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "    \n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "Vl6IidfdIrJK"
   },
   "source": [
    "We pass `use_fast=True` to the call above to use one of the fast tokenizers (backed by Rust) from the 🤗 Tokenizers library. Fast tokenizers are available for almost all models, but if you get an error with the previous call then simply set `use_fast` to False."
   ]
  },
  {
   "attachments": {},
   "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": null,
   "metadata": {
    "id": "a5hBlsrHIrJL",
    "outputId": "acdaa98a-a8cd-4a20-89b8-cc26437bbe90",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "tokenizer(\"Hello, this is one sentence!\", \"And this sentence goes with it.\")"
   ]
  },
  {
   "attachments": {},
   "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 for what we're doing here, but they are required by the model we will instantiate later. You can learn more about keys in this [tutorial on preprocessing](https://huggingface.co/transformers/preprocessing.html).\n",
    "\n",
    "To preprocess our dataset, we will thus need the names of the columns containing the sentence(s). The following dictionary keeps track of the correspondence task to column names:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fyGdtK9oIrJM",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "task_to_keys = {\n",
    "    \"cola\": (\"sentence\", None),\n",
    "    \"mnli\": (\"premise\", \"hypothesis\"),\n",
    "    \"mnli-mm\": (\"premise\", \"hypothesis\"),\n",
    "    \"mrpc\": (\"sentence1\", \"sentence2\"),\n",
    "    \"qnli\": (\"question\", \"sentence\"),\n",
    "    \"qqp\": (\"question1\", \"question2\"),\n",
    "    \"rte\": (\"sentence1\", \"sentence2\"),\n",
    "    \"sst2\": (\"sentence\", None),\n",
    "    \"stsb\": (\"sentence1\", \"sentence2\"),\n",
    "    \"wnli\": (\"sentence1\", \"sentence2\"),\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "xbqtC4MrIrJO"
   },
   "source": [
    "We can double check that this dictionary does work on our current dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "19GG646uIrJO",
    "outputId": "0cb4a520-817e-4f92-8de8-bb45df367657",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "sentence1_key, sentence2_key = task_to_keys[task]\n",
    "\n",
    "if sentence2_key is None:\n",
    "    print(f\"Sentence: {dataset['train'][0][sentence1_key]}\")\n",
    "else:\n",
    "    print(f\"Sentence 1: {dataset['train'][0][sentence1_key]}\")\n",
    "    print(f\"Sentence 2: {dataset['train'][0][sentence2_key]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "2C0hcmp9IrJQ"
   },
   "source": [
    "We can then write the function that will preprocess our samples. We just feed them to the `tokenizer` with the three arguments.`truncation=True` will ensure that an input longer than the maximum length will be truncated to the maximum length. `max_length=max_seq_length` sets the maximum length of a sequence.\n",
    "\n",
    "**Important: since we will use packing later, we don't want to perform any padding in the tokenizer.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vc0BSBLIIrJQ",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# no padding for packing\n",
    "def preprocess_function(examples):\n",
    "    if sentence2_key is None:\n",
    "        return tokenizer(examples[sentence1_key], truncation=True, max_length=max_seq_length)\n",
    "    \n",
    "    return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True, max_length=max_seq_length)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "0lm8ozrJIrJR"
   },
   "source": [
    "This function works with one or several samples. In the case of several samples, the tokenizer will return a list of lists for each key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-b70jh26IrJS",
    "outputId": "acd3a42d-985b-44ee-9daa-af5d944ce1d9",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "preprocess_function(dataset['train'][:5])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "zS-6iXTkIrJT"
   },
   "source": [
    "To apply this function to all the sentences (or pairs of sentences) in our dataset, we use the `map` method of our `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 with one command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DDtsaJeVIrJT",
    "outputId": "aa4734bf-4ef5-4437-9948-2c16363da719",
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "encoded_dataset = dataset.map(preprocess_function, batched=True)\n",
    "len(encoded_dataset['train'])"
   ]
  },
  {
   "attachments": {},
   "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 able to detect when the function you pass to `map` has changed (and thus to not use the cached data). For instance, it will 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 text samples together into batches. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the text samples in a batch concurrently."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Packing the dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To implement packing, we need to pack our dataset first. Each new element will be a \"pack\" containing at most `max_seq_per_pack` sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "max_seq_per_pack = 6"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also define the number of labels in our dataset. `sst2` is a single_label task and it will therefore contain one true class for each example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "num_labels = 3 if task.startswith(\"mnli\") else 1 if task==\"stsb\" else 2\n",
    "problem_type = 'single_label_classification'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Packing algorithm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to pack efficiently, we will use a histogram-based algorithm. The shortest-pack-first histogram packing (SPFHP) was presented in the Graphcore blog post [introducing Packed BERT for a training speedup in natural language processing](https://www.graphcore.ai/posts/introducing-packed-bert-for-2x-faster-training-in-natural-language-processing). We have adapted the [code](https://github.com/graphcore/tutorials/tree/master/blogs_code/packedBERT) from the blog post for this notebook. The full process of packing the dataset consists of four steps:\n",
    "\n",
    "1. Create a histogram of the sequence lengths of the dataset.\n",
    "2. Generate the \"strategy\" for the dataset using one of the state-of-the-art packing algorithms. The strategy maps out the order and indices of the sequences that need to be packed together.\n",
    "3. Use this strategy to create the actual dataset, concatenating the tokenized features together for each column in the dataset, including the labels.\n",
    "4. Finally, pass these new columns into a custom PyTorch dataset, ready to be passed to the PopTorch dataloader!\n",
    "\n",
    "These steps have been simplified through the easy-to-use `utils.packing` package available in Graphcore Optimum. You can simply generate the packed dataset after the usual tokenization and preprocessing by passing all necessary packing configuration to the `PackedDatasetCreator` class, and generate the ready-to-use PyTorch dataset with `.create()`.\n",
    "\n",
    "Within the function, there are some column names used by default. The expected default columns for text classification include:\n",
    "* `input_ids`\n",
    "* `attention_mask`\n",
    "* `token_type_ids`\n",
    "* `labels`\n",
    "\n",
    "These should all be generated automatically when tokenizing any classification dataset for BERT. However, the labels key, as it is not encoded, may have a different name. For this dataset, the column key for the labels for this dataset is `label`, since the dataset creator expects `labels`, we can pass this to the argument `custom_label_key`, so the class can find our labels. \n",
    "\n",
    "The `PackedDatasetCreator` requires different instantiations for different datasets, so it must be called separately for each of our dataset splits. We can set either `training`, `validation` or `inference` to `True` as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from utils.packing.dataset_creator import PackedDatasetCreator\n",
    "\n",
    "train_data_packer = PackedDatasetCreator(\n",
    "    tokenized_dataset = encoded_dataset['train'],\n",
    "    max_sequence_length = max_seq_length,\n",
    "    max_sequences_per_pack = max_seq_per_pack,\n",
    "    training = True,\n",
    "    num_labels = num_labels,\n",
    "    problem_type = problem_type,\n",
    "    algorithm = 'SPFHP',\n",
    "    custom_label_key = 'label'\n",
    ")\n",
    "\n",
    "val_data_packer = PackedDatasetCreator(\n",
    "    tokenized_dataset = encoded_dataset['validation'],\n",
    "    max_sequence_length = max_seq_length,\n",
    "    max_sequences_per_pack = max_seq_per_pack,\n",
    "    validation = True,\n",
    "    num_labels = num_labels,\n",
    "    problem_type = problem_type,\n",
    "    algorithm = 'SPFHP',\n",
    "    custom_label_key = 'label'\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will create the strategy and initialise the necessary parameters for packing the dataset. We can see that the ideal speed-up we have achieved is approximately 5.15x the original dataset, which corresponds directly to the average packing factor: the average number of sequences within one pack.\n",
    "\n",
    "The `PackedDatasetCreator` class also has some other features we do not use here for training, such as `pad_to_global_batch_size`, a feature useful for performing batched inference on a large sample set when we do not want to lose any of the samples. When creating data iterators using the `poptorch.Dataloader`, it applies 'vertical' padding to the dataset, adding filler rows to bring the dataset up to a value divisible by the global batch size, and allows for the largest possible batch sizes to be used without any loss of data."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also view the histogram generated in the first step of the packing process, to observe whether the distribution of sequence lengths in the dataset will benefit from packing. As a general rule, as long as the average length of the sequences in the dataset is 50% or less of the maximum sequence length, packing will offer at least a 2x throughput benefit, in other words: `throughput_increase ≈ max_seq_len/mean_seq_len`\n",
    "\n",
    "Many datasets have distributions with much smaller average lengths, and will benefit much more. We can easily observe this distribution by retrieving and plotting the histogram from the data class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "train_histogram = train_data_packer.histogram\n",
    "\n",
    "plt.hist(train_histogram, bins = [k for k in range(0,max_seq_length,10)]) \n",
    "plt.title(\"Sequence length histogram\") \n",
    "plt.xlabel('Sequence lengths')\n",
    "plt.ylabel('Frequency')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to create the actual packed dataset (step 3 of the packing process outlined above).\n",
    "\n",
    "In this stage, we take the strategy for mapping the sequences by size into packs that were generated by the packing algorithm, and use this to extract the sequences from the tokenized dataset, inserting them into packs for each column in the dataset. Any remaining space in a pack after the sequences have been concatenated is padded to bring all sequences up to the maximum sequence length.\n",
    "\n",
    "Some key features unique to packed datasets are worth mentioning here:\n",
    "\n",
    "- The specific attention mask (`attention_mask`) that is generated contains a unique index for each sequence of the pack and `0` for the remaining padding tokens. This, essentially, tells the model where to look from the perspective of a single token, ignoring any encoded information (such as a different sequence) that is not relevant to that token.\n",
    "    - Example of 3 sequences: `attention_mask = [1,1,1,1,1,1,2,2,2,2,2,3,3,3,3,3,0,...,0,1,2,3]`\n",
    "\n",
    "\n",
    "- The [CLS] tokens of each sequence must be moved to the end of the pack.\n",
    "    - For instance: `[CLS,a,b,c] + [CLS, d,e,f] + [CLS, g,h,i] -> [a,b,c,d,e,f,g,h,i,...,CLS,CLS,CLS]`\n",
    "    \n",
    "\n",
    "- `position_ids` for a pack contains the concatenated `position_ids` of each sequences \n",
    "    - For instance given 3 sequences: `[0,1,2,3,4] + [0,1,2,3] + [0,1,2] -> [1,2,3,4,1,2,3,1,2,...,0,0,0]` (note: the CLS tokens' position ID '0' are also moved to the end of the pack)\n",
    "    \n",
    "- `labels` and `token_type_ids` are also packed to correspond to the `input_ids` pack.\n",
    "\n",
    "\n",
    "To create a dataloader-ready packed dataset, all you need to do is call the `create()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "packed_train_dataset = train_data_packer.create()\n",
    "packed_val_dataset = val_data_packer.create()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize one sample of the new `packed_train_dataset`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "packed_train_dataset[133]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "545PP3o8IrJV"
   },
   "source": [
    "## Fine-tuning the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "FBiW8UpKIrJW"
   },
   "source": [
    "Now that our data is ready, we can download the pre-trained model and fine-tune it."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implement Packed BERT\n",
    "\n",
    "A few model modifications are required to make packing work with BERT.\n",
    "We extend the existing class `BertForSequenceClassification` to `PipelinedPackedBertForSequenceClassification` which incorporates the required changes to the pooler and the model output. The crux of these changes is to modify the generic sequence classification model to handle unpacking multiple sequences in the output stage, treating them as a larger batch size for classification, as well as masking any padding created by packing.\n",
    "\n",
    "First let's load a default BERT configuration using `AutoConfig`. The config includes a new parameter we must set, `max_sequences_per_pack`. This informs the model of the maximum number of sequences it will need to unpack in the model output. It also allows us to clearly define the `num_labels` and `problem_type` for this model.\n",
    "\n",
    "It is essential that the problem type is defines here, as switching between the methods used by different types of classification requires this definition within the custom model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoConfig\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_checkpoint)\n",
    "config.max_sequences_per_pack = max_seq_per_pack\n",
    "config.num_labels = num_labels\n",
    "config.problem_type = problem_type"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can instantiate the model class with the config, loading the weights from the model checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "torch.manual_seed(43)\n",
    "np.random.seed(43)\n",
    "\n",
    "from models.modeling_bert_packed import PipelinedPackedBertForSequenceClassification\n",
    "\n",
    "model = PipelinedPackedBertForSequenceClassification.from_pretrained(\n",
    "    model_checkpoint, config=config).train()\n",
    "\n",
    "print(config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "CczA5lJlIrJX"
   },
   "source": [
    "The warning tells us we are throwing away some weights and randomly initializing others. This is normal in this case, because we are removing the head used to pre-train the model on a masked language modelling objective and replacing it with a new head for which we don't have pre-trained 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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can first test the model on a CPU and observe that the output logits now have the size [batch_size x max_seq_per_pack, 2] = [12, 2] with this notebook default values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from transformers.data.data_collator import default_data_collator\n",
    "import torch\n",
    "\n",
    "model.float()\n",
    "loader = torch.utils.data.DataLoader(packed_train_dataset,\n",
    "                             batch_size=micro_batch_size,\n",
    "                             shuffle=True,\n",
    "                             drop_last=True,\n",
    "                             collate_fn=default_data_collator)\n",
    "data = next(iter(loader))\n",
    "outputs = model(**data)\n",
    "print(outputs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's prepare the model for running on IPUs."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we set the model to half-precision:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "model.half()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For validation, we need to define a function to compute the metrics from the predictions, which will just use `metric` which we loaded earlier. The only preprocessing we have to do is to take `argmax` of our predicted logits (we just squeeze the last axis in the case of STS-B). To ignore the `-100` labels from incomplete packs, we use a boolean mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "metric_name = \"pearson\" if task == \"stsb\" else \"matthews_correlation\" if task == \"cola\" else \"accuracy\"\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    \n",
    "#   Remove the padding labels\n",
    "    mask = (labels != -100)\n",
    "    labels = labels[mask]\n",
    "    \n",
    "    if task != \"stsb\":\n",
    "        predictions = np.argmax(predictions, axis=-1)\n",
    "    else:\n",
    "        predictions = predictions[:, 0]\n",
    "    \n",
    "    predictions = predictions[mask]\n",
    "    \n",
    "    return metric.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "_N8urzhyIrJY"
   },
   "source": [
    "Next, we need to define `IPUConfig`, which is a class that specifies attributes and configuration parameters to compile and put the model on the device. We initialize it with a config name or path, which we set earlier. Then we use it to set the mode attribute `model.ipu_config` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from optimum.graphcore import IPUConfig, IPUTrainer, IPUTrainingArguments\n",
    "\n",
    "ipu_config = IPUConfig.from_pretrained(\n",
    "    ipu_config_name,\n",
    "    executable_cache_dir = executable_cache_dir,\n",
    "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "    replication_factor=1,\n",
    "    device_iterations = device_iterations,\n",
    "    inference_device_iterations= 16,\n",
    "    inference_replication_factor=1\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `IPUTrainingArguments` class defines any custom parameter modification we want to do, such as the initial learning rate for the model. It also allows other options, such as dataloader parameters, micro batch sizes and an automatic push to the Hugging Face Hub (if credentials were set up earlier) to happen at given intervals.\n",
    "\n",
    "These arguments are passed to the `IPUTrainer` class, which wraps the model training and evaluation process into a simple single-line process, doing all of the heavy lifting for us, for example training and evaluation loops, device assignment, optimiser definition and data-loading.\n",
    "\n",
    "Note that only some arbitrary hyperparameter tuning was performed for this task. Other tasks and datasets may require further tuning to get the most optimal results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import default_data_collator\n",
    "\n",
    "args = IPUTrainingArguments(\n",
    "    \"./\"+f\"{model_name}-{task}\",\n",
    "    num_train_epochs=2,\n",
    "    per_device_train_batch_size=micro_batch_size,\n",
    "    per_device_eval_batch_size=2,\n",
    "    learning_rate=9e-5,\n",
    "    warmup_ratio=0.1,\n",
    "    weight_decay=0,\n",
    "    lr_scheduler_type = \"cosine\",\n",
    "    metric_for_best_model=metric_name,\n",
    "    dataloader_drop_last=True,\n",
    "    # dataloader_mode=\"async_rebatched\",\n",
    "    logging_steps=1,\n",
    "    n_ipu=n_ipu,\n",
    "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "    push_to_hub=False\n",
    ")\n",
    "\n",
    "\n",
    "trainer = IPUTrainer(\n",
    "    model,\n",
    "    ipu_config,\n",
    "    args,\n",
    "    train_dataset=packed_train_dataset,\n",
    "    eval_dataset=packed_val_dataset,\n",
    "    data_collator=default_data_collator,\n",
    "    compute_metrics=compute_metrics\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, to train the model we simply call the `train()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***About the performance:*** `IPUTrainer` doesn't take into account that we have packed data samples when computing the speed metrics. It treats a 'sample' as a single input to the model, which means one **pack**.\n",
    "\n",
    "So the actual throughput estimation can be obtained by multiplying `samples_per_second` by the average packing factor (the average number of samples per pack) of the dataset. These were obtained in the `packing_algorithm` section: `5.15` for `sst2` training set and `5.77` for validation set."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can evaluate the model by simply calling the `evaluate()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now push the model to Hugging Face. Simply uncomment and execute the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "#trainer.push_to_hub()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also save the model locally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "trainer.save_model(\"./\"+f\"{model_name}-{task}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have now successfully fine-tuned and evaluated your speed-optimised model for text classification using packing!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fast batched inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Packing can also be used for inference, particularly for performing inference for workloads. This section demonstrates how to perform faster, batched inference with a large number of samples using a super-easy custom pipeline which batches and packs your input data, performs inference and returns post-processed predictions. \n",
    "\n",
    "We need to import the pipeline and initialise a few essential parameters.\n",
    "\n",
    "`model` is the model checkpoint and we are going to use the locally saved checkpoint generated from training `sst2`. `executable_cache_dir`, `problem_type`, `max_seq_length` must be specified. To return predictions organised by class names, the class names for your output must be passed to `label_categories`. \n",
    "\n",
    "The pipeline will automatically determine your model's IPU config, given that the checkpoint was trained using Optimum Graphcore, which will be the case for the model fine-tuned in this notebook.\n",
    "\n",
    "In this example, we pre-load `IPUConfig` and modify some of the default parameters to get the best performance out of inference and leverage the benefits of IPU parallelism. The micro-batch size can also be specified, for which the default is 1.\n",
    "\n",
    "When training, the packing factor affects the convergence in the same way as a large increase in batch size would do. However, for inference, we are free to use a bigger packing factor to speed it up. Let's try it with `max_seq_per_pack = 12`.\n",
    "\n",
    "**Note:** Packing brings huge benefits for performing inference on large amounts of data. For small scale inference tasks, such as those which more suit sequential inference on a single un-batched input, the generic Optimum Graphcore `TextClassificationPipeline` class may be preferred. This won't affect fine-tuning, and the weights generated from fine-tuning using packing will work just the same!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's initialise the `PackedBertTextClassificationPipeline` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "from pipeline.packed_bert import PackedBertTextClassificationPipeline\n",
    "from optimum.graphcore import IPUConfig\n",
    "\n",
    "model = \"./\"+f\"{model_name}-{task}\"\n",
    "# model = 'your_username/{model_name}-{task}' # to load from Hugging Face Hub\n",
    "\n",
    "inference_boosted_ipu_config = IPUConfig.from_pretrained(model, \n",
    "        inference_device_iterations=32,\n",
    "        inference_replication_factor=4,\n",
    "        ipus_per_replica=1,\n",
    "        layers_per_ipu=[12]\n",
    "    )\n",
    "\n",
    "pipeline = PackedBertTextClassificationPipeline(\n",
    "    model = model,\n",
    "    executable_cache_dir = executable_cache_dir,\n",
    "    problem_type='single_label_classification',\n",
    "    max_seq_per_pack=12,\n",
    "    max_seq_length=max_seq_length,\n",
    "    ipu_config=inference_boosted_ipu_config,\n",
    "    micro_batch_size=8,\n",
    "    label_categories=['positive','negative']\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pipeline expects a **list of strings** directly passed to it. There is no need to tokenize, preprocess, pack or post-process the data to use the inference pipeline.\n",
    "\n",
    "As a test, we can load the entire `sst2` dataset and perform packed inference using `.predict()` on the text column to generate predictions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "dataset = datasets.load_dataset(\"glue\", actual_task)\n",
    "preds = pipeline.predict(dataset['train']['sentence'])\n",
    "\n",
    "print(preds.keys())\n",
    "print(f\"Number of predictions: {len(preds['predictions'])}\")\n",
    "print(f\"Preprocessing time: {preds['preprocessing_time']}s\")\n",
    "print(f\"Postprocessing time: {preds['postprocessing_time']}s\")\n",
    "print(f\"Throughput: {preds['throughput']}samples/s\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is minimal overhead from tokenizing and packing the dataset, but the speed benefits are evident. Running the above pipeline, we achieve a throughput of approximately 35 000 samples per second, demonstrating the huge time benefit you can achieve by using packing!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "You may want to try out the following notebooks using BERT with packing for:\n",
    "\n",
    "* Multi-label text classification - `packedBERT_multi_label_text_classification.ipynb`\n",
    "* Question answering - `packedBERT_question_answering.ipynb`\n",
    "\n",
    "Also, check out the full list of [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to get more of a feel for how IPUs perform on other tasks.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Text Classification on GLUE",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}