GPTJ-generative-inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text Generation with GPT-J 6B on IPUs\n",
    "\n",
    "[GPT-J](https://huggingface.co/EleutherAI/gpt-j-6B) is a causal decoder-only transformer model which can be used for text-generation.\n",
    "Causal means that a causal mask is used in the decoder attention, so that each token only has visibility of previous tokens.\n",
    "\n",
    "Language models are very powerful because a huge variety of tasks can be formulated as text-to-text problems and thus adapted to fit the generative setup, where the model is asked to correctly predict future tokens. This idea has been widely explored in the [T5 paper: Exploring the Limits of Transfer Learning with a Unified\n",
    "Text-to-Text Transformer](https://arxiv.org/pdf/1910.10683.pdf).\n",
    "\n",
    "Note that for these kind of tasks you don't need GPT-3 175B parameter sized models. GPT-J at 6B parameters has very good language understanding and is suitable for most of these scenarios. Larger models give only a small improvement in language understanding. Mainly they add more world knowledge and better performance at free text generation as might be used in an AI Assistant or chatbot.\n",
    "\n",
    "In this notebook you will:\n",
    "\n",
    "- generate text in 5 lines of code with GPT-J on the Graphcore IPU;\n",
    "- use GPT-J to answer questions and build prompts to reliably use text-generation for more specific NLP tasks;\n",
    "- explore the effect on the model of the prompt format and compare the model performance with 0-shot (no examples included in the prompt) and few-shot prompting (where a few typical examples are included in the prompt);\n",
    "- improve text generation throughput using batched inference;\n",
    "- understand the limitations of the base GPT-J checkpoint when it comes to more complex NLP tasks;\n",
    "- use the model to identify whether statements agree or disagree (entailment). For this more complex task, we show the benefit of fine-tuning and load a checkpoint fine-tuned on the MNLI dataset from the Hugging Face Hub which achieves much better performance on this specific task.\n",
    "\n",
    "\n",
    "The MNLI checkpoint on Hugging Face has been fine-tuned on IPUs. No fine-tuning is performed in this notebook but you can learn more on fine-tuning GPT-J in the finetuning notebook `finetuning.ipynb`.\n",
    "By exploring this notebook, you will gain insight into how various tasks can be formulated in a text-to-text format and how this flexibility can be utilized to fine-tune the model using your own dataset."
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To ensure smooth execution of the notebook, we load and check environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "number_of_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "if number_of_ipus < 4:\n",
    "    raise ValueError(\"This notebook is designed to run with at least 4 IPUs\")\n",
    "\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\")\n",
    "os.environ[\"POPART_CACHE_DIR\"] = executable_cache_dir\n",
    "checkpoint_directory = os.getenv(\"CHECKPOINT_DIR\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running GPT-J on the IPU\n",
    "\n",
    "We start by running the original [GPT-J model published by Eleuther AI](https://huggingface.co/EleutherAI/gpt-j-6B) on the Graphcore IPU. \n",
    "In a few lines of code we load the configuration and use it to create a pipeline object which will allow us to interactively run the model and use it for text generation.\n",
    "\n",
    "While this application is written in Graphcore's PopXL framework, no knowledge of the framework is required to use or train this application as all parameters are controlled through configuration options.\n",
    "Base configurations are available in the `config/inference.yml` file and can be loaded as follows:\n",
    " <!-- PopXL is a framework which provides fine grained control of execution, memory and parallelism.  -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Setup ---\n",
    "import run_inference\n",
    "\n",
    "config, *_ = run_inference.gptj_config_setup(\n",
    "    \"config/inference.yml\", \"release\", \"gpt-j-mnli\"\n",
    ")\n",
    "print(config.dumps_yaml())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The configuration can be edited and stored in a new file to suit your needs. It contains all the arguments which define the model and control the execution of the application on the IPU.\n",
    "\n",
    "For inference, the main arguments we are interested in modifying are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The number of tokens generated before stopping\n",
    "# Note the model will stop before this if it generates an <|endoftext|> token\n",
    "config.inference.output_length = 10\n",
    "# The number of prompts which will be processed at once\n",
    "config.execution.micro_batch_size = 12\n",
    "# The maximum tokenized sequence length (input + generated) handled by the model\n",
    "config.model.sequence_length = 512"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we're going to combine this configuration with pre-trained weights.\n",
    "The `pipeline` utility accepts directly the name of a pre-trained checkpoint from the Hugging Face Hub.\n",
    "The reference pre-trained checkpoint is the [6 billion parameter GPT-J checkpoint from EleutherAI](https://huggingface.co/EleutherAI/gpt-j-6B), which has been trained on [the Pile](https://pile.eleuther.ai/) open source dataset and has not undergone fine-tuning on any specific task. Hence, it is suitable for general text generation but, as you will see, is not very good on downstream tasks (such as question answering or entailment).\n",
    "\n",
    "\n",
    "Once we have a config and have chosen a pre-trained model we create a `GPTJPipeline` pipeline. If you are not planning to use long\n",
    "prompts, you can reduce the sequence length either by changing the config or by providing the `sequence_length` argument to the pipeline. Here, we reduce it from the\n",
    "default value of 1024 to 512. Reducing the sequence length allows you to fit more batches into memory. As you will see in the batched inference section, you can maximise performance by making the model process several prompts at the same time.\n",
    "\n",
    "Creating the pipeline takes a few minutes as the checkpoint is downloaded and loaded into the session:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import api\n",
    "\n",
    "general_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"EleutherAI/gpt-j-6b\",\n",
    "    sequence_length=512,\n",
    "    micro_batch_size=12,\n",
    "    output_length=20,\n",
    "    print_live=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can explore the attributes of the `general_model` pipeline:\n",
    "\n",
    "- `pretrained` contains the `GPTJForCausalLM` class from the Hugging Face Transformers library which is used to load the weights;\n",
    "- `tokenizer` contains the tokenizer loaded with the pre-trained checkpoint from the Hugging Face Hub;\n",
    "- `config` has the input config;\n",
    "- `session` is the PopXL session which can be run on the IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "general_model.tokenizer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the pipeline to do standard text generation starting from an input prompt. In this demo we use **greedy generation**, meaning the most likely token is chosen at each step. This corresponds to running an Hugging Face pipeline with the option `sample=False`.\n",
    "The first execution of the pipeline takes a few minutes because the program is compiled and the weights are downloaded from the Hugging Face Hub and copied to the IPU. Subsequent executions will be much faster.\n",
    "In the cell below, we ask the model to answer a simple question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = general_model(\"What is the capital of France?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the answer is correct, the model includes it in a long form answer and continues generating irrelevant text instead of the `<|endoftext|>` token.\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sensitivity to prompt format\n",
    "\n",
    "The model output is very sensitive to the prompt format, even spaces can make a difference. To get the best results, you have to experiment with the format of your input prompts. Structured prompts can make it easy to post-process the model outputs to extract the relevant pieces of information.\n",
    "\n",
    "Let's try to use a different structure for our initial question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = general_model(\n",
    "    \"\"\"Question: What is the capital of Country?\n",
    "Answer: City\n",
    "Question: What is the capital of France?\n",
    "Answer:\"\"\",\n",
    ")\n",
    "out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we get the answer in a short format, which is nice. Irrelevant extra text is still generated though. \n",
    "However, we can easily extract the relevant information. First of all, we can reduce the amount of text generated by the model to 10 tokens, which is enough for our answer.\n",
    "We then extract only the string which answers the question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = general_model(\n",
    "    f\"\"\"Question: What is the capital of Country?\n",
    "Answer: city\n",
    "Question: What is the capital of France?\n",
    "Answer:\"\"\",\n",
    "    output_length=10,\n",
    ")\n",
    "capitals = [answer.splitlines()[0].strip() for answer in out]\n",
    "capitals"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batched text generation\n",
    "\n",
    "This model configuration supports batched generation where we can generate text based on multiple prompts at the same time.\n",
    "This can be useful to increase throughput when text prompts are queued for processing in a production environment.\n",
    "This config has a batch size of 12 which means this pipeline can generate answers to 12 questions at a time.\n",
    "\n",
    "Let's use the model to generate country names we can later use to build batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_for_countries = \"List countries in Europe, America, Asia and Africa: France, \"\n",
    "out = general_model(prompt_for_countries, output_length=120, print_live=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "countries = [\n",
    "    c.strip() for c in (prompt_for_countries + out[0]).split(\":\")[1].split(\",\")\n",
    "]\n",
    "countries"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using these generated countries, we compose a batch of structured prompts and pass it to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = general_model(\n",
    "    [\n",
    "        f\"\"\"Question: What is the capital of China?\n",
    "Answer: Beijing\n",
    "Question: What is the capital of {country}?\n",
    "Answer:\"\"\"\n",
    "        for country in countries\n",
    "    ],\n",
    "    print_live=False,\n",
    "    output_length=10,\n",
    ")\n",
    "capitals = [answer.splitlines()[0].strip() for answer in out]\n",
    "list(zip(countries, capitals))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to batched inference 12 prompts can be processed at the same time, which leads to much higher throughput."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Determining entailment\n",
    "\n",
    "In this section, we evaluate the ability of our generative pipeline to perform a more complex natural language processing task.\n",
    "This task will show the limits of the general model, even when few-shot prompting is used.\n",
    "\n",
    "**Entailment** is the task of determining whether two statements agree, disagree or are neutral relative to each other.\n",
    "A common benchmark dataset is the [MNLI GLUE dataset](https://huggingface.co/datasets/glue).\n",
    "\n",
    "This dataset consists of pairs of sentences, a *premise* and a *hypothesis*.\n",
    "The task is to predict the relation between the premise and the hypothesis, which can be:\n",
    "- `entailment`: hypothesis follows from the premise,\n",
    "- `contradiction`: hypothesis contradicts the premise,\n",
    "- `neutral`: hypothesis and premise are unrelated.\n",
    "\n",
    "As we did for question answering, we can try to use our generative model to tackle this task by creating prompts:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def entailment_prompt(hypothesis, premise, target=\"\"):\n",
    "    sep = \".\\n\" if target else \"\"\n",
    "    return f\"mnli hypothesis: {hypothesis} premise: {premise} target: {target}{sep}\"\n",
    "\n",
    "\n",
    "entailment_prompt(\"The person is leaving.\", \"Hello, welcome to the country.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Valid prompts for this task are not as easy to come up with so we load validation data from the MNLI task from the GLUE dataset using the 🤗 Datasets library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "dataset = datasets.load_dataset(\"glue\", \"mnli\", split=\"validation_mismatched\")\n",
    "dataset[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the `label_to_target` mapping below to turn integer class labels from the dataset into the name of the class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_label_to_target = [\"entailment\", \"neutral\", \"contradiction\", \"unknown\"]\n",
    "mnli_label_to_target[dataset[1][\"label\"]]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try passing a single prompt to the model. For question answering this was enough to get the right answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a check in the pipeline for sequence length - get one of each\n",
    "general_model(\n",
    "    entailment_prompt(dataset[1][\"hypothesis\"], dataset[1][\"premise\"]),\n",
    "    print_live=True,\n",
    "    output_length=10,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We were hoping for the model to predict \"contradiction\", or an answer which expresses that idea.\n",
    "The model does not get it right: determining entailment from a single formatted prompt is not enough to get it to perform the task.\n",
    "\n",
    "We can try to improve the performance by adding instructions and providing examples using the **few-shot prompting** technique. To help with that we create a helper function `mnli_data_to_example` which will turn a single data entry from the MNLI dataset into a sample prompt complete with the target prediction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnli_data_to_example(data: dict):\n",
    "    return entailment_prompt(\n",
    "        data[\"hypothesis\"], data[\"premise\"], mnli_label_to_target[data[\"label\"]]\n",
    "    )\n",
    "\n",
    "\n",
    "mnli_data_to_example(dataset[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The idea is to form input prompts that contain a few samples of the task we want the model to perform. \n",
    "\n",
    "Coming up with a prompt which will reliably generate the intended results can be challenging especially for a more complex task like entailment classification.\n",
    "In the cell below we show four possible ways of generating entailment instructions in the prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try asking for the task to be completed\n",
    "entailment_instructions_v1 = \"How are the statements related?\\n\"\n",
    "\n",
    "# Try explaining the task with some explicit \"Examples\"\n",
    "entailment_instructions_v2 = (\n",
    "    \"Tell me if the next statements entailment, contradiction, neutral.\\n\"\n",
    "    + \"Example - \"\n",
    "    + entailment_prompt(\"Goodbye.\", \"Hey there.\", \"contradiction\")\n",
    "    + \"Example - \"\n",
    "    + entailment_prompt(\"Hello.\", \"Hey there.\", \"entailment\")\n",
    ")\n",
    "\n",
    "# Reinforce the instructions with every example\n",
    "entailment_instructions_v3 = (\n",
    "    \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + entailment_prompt(\"Goodbye.\", \"Hey there.\", \"contradiction\")\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + entailment_prompt(\"Hello.\", \"Hey there.\", \"entailment\")\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + entailment_prompt(\"The person is traveling.\", \"The cat is black.\", \"neutral\")\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    ")\n",
    "\n",
    "#  Use the dataset instead of handcrafted prompts\n",
    "entailment_instructions_v4 = (\n",
    "    \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + mnli_data_to_example(dataset[0])\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + mnli_data_to_example(dataset[2])\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    "    + mnli_data_to_example(dataset[8])\n",
    "    + \"Tell me if the statements entailment, contradiction, neutral.\\n\"\n",
    ")\n",
    "\n",
    "general_model(\n",
    "    entailment_instructions_v3\n",
    "    + entailment_prompt(\"Hello, welcome.\", \"The person is leaving.\"),\n",
    "    print_live=True,\n",
    "    output_length=40,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the model does pick one of the three options, it does not select the right one.\n",
    "The big challenge with this approach is obtaining reliable results. Apparently inconsequential changes in the phrasing of the samples can lead to very different results.\n",
    "\n",
    "To evaluate the effectiveness of the engineered prompts we test against a number of unseen samples in the MNLI dataset.\n",
    "In the cell below we test `entailment_instructions_v3`. You can try out the other prompts and see how they perform, or come up with your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_sample = dataset[12:48]\n",
    "out = general_model(\n",
    "    [\n",
    "        entailment_instructions_v4 + entailment_prompt(hypothesis, premise)\n",
    "        for hypothesis, premise in zip(\n",
    "            dataset_sample[\"hypothesis\"], dataset_sample[\"premise\"]\n",
    "        )\n",
    "    ],\n",
    "    print_live=False,\n",
    "    output_length=10,\n",
    ")\n",
    "# Strip out everything in the output after new lines\n",
    "processed = [\n",
    "    (mnli_label_to_target[label], answer.splitlines()[0].strip())\n",
    "    for label, answer in zip(dataset_sample[\"label\"], out)\n",
    "]\n",
    "processed = [(label in answer, label, answer) for label, answer in processed]\n",
    "print(f\"Got {sum(p for p, *_ in processed)} / {len(processed)} correct\")\n",
    "processed[:12]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The instructions in prompt `entailment_instructions_v3` and `entailment_instructions_v4` are enough to capture some elements of the task and let the model guess one of the three target classes every time.\n",
    "However, they are not enough for the model to correctly classify the entailment of the hypothesis and the premise. They accuracy oscillates around 33% which corresponds to a random choice between the three classes.\n",
    "\n",
    "As we can see few-shot prompting is not sufficient for the GPT-J model to complete this task.\n",
    "Let's release the IPUs to let us try a checkpoint for the model fine-tuned to perform this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "general_model.detach()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using a fine-tuned model\n",
    "\n",
    "In order to complete the entailment task we are going to use a fine-tuned model on the MNLI task of the GLUE dataset. For this model we have a shorter sequence length of 256, and so can fit a larger batch size of 16 into memory.\n",
    "\n",
    "The checkpoint we will be using was fine-tuned on the Graphcore IPU and is hosted on the 🤗 Hub at [Graphcore/gptj-mnli](https://huggingface.co/Graphcore/gptj-mnli). To see how this checkpoint was generated check out the [fine-tuning notebook](finetuning.ipynb).\n",
    "As we did before we can load this checkpoint with a single command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"Graphcore/gptj-mnli\",\n",
    "    sequence_length=256,\n",
    "    micro_batch_size=16,\n",
    "    print_live=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While this model was trained on IPUs, a checkpoint trained on GPUs could just as well be loaded in a multi-accelerator workflow.\n",
    "\n",
    "Just like the previous checkpoint, the model can handle arbitrary text generation questions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_model(\"Hey there\", output_length=30)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, on the entailment task the performance is much better, even without instructions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_model(entailment_prompt(\"The person is leaving.\", \"Hello, welcome.\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It got it right! Those sentences are contradictory. Now let's try our samples from the GLUE dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = mnli_model(\n",
    "    [\n",
    "        entailment_prompt(hypothesis, premise)\n",
    "        for hypothesis, premise in zip(\n",
    "            dataset[:16][\"hypothesis\"], dataset[:16][\"premise\"]\n",
    "        )\n",
    "    ],\n",
    "    print_live=True,\n",
    "    output_length=10,\n",
    ")\n",
    "# Strip out everything in the output after new lines\n",
    "mnli_label_to_target = [\"entailment\", \"neutral\", \"contradiction\", \"unknown\"]\n",
    "[\n",
    "    (\n",
    "        mnli_label_to_target[label],\n",
    "        answer.splitlines()[0].strip().replace(\"<|endoftext|>\", \"\"),\n",
    "    )\n",
    "    for label, answer in zip(dataset[:16][\"label\"], out)\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It gets almost all of them right and clearly has some knowledge of the task we need it to complete.\n",
    "\n",
    "We can create a pipeline specific to this task which handles the prompt pre-processing and the post-processing of the generated text.\n",
    "We can change pipelines with the `from_gptj_pipeline` factory method, which gives us a ready-to-use pipeline for entailment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_pipeline = api.GPTJEntailmentPipeline.from_gptj_pipeline(mnli_model)\n",
    "mnli_pipeline(\"Hey there.\", \"Goodbye.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate the model, we now run 200 samples of the dataset through our pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 200\n",
    "out = mnli_pipeline(\n",
    "    premise=dataset[:sample_size][\"premise\"],\n",
    "    hypothesis=dataset[:sample_size][\"hypothesis\"],\n",
    "    print_live=False,\n",
    "    output_length=5,\n",
    ")\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "results = pd.DataFrame(\n",
    "    [\n",
    "        (mnli_label_to_target[label] == answer, mnli_label_to_target[label], answer)\n",
    "        for label, answer in zip(dataset[:sample_size][\"label\"], out)\n",
    "    ],\n",
    "    columns=[\"correct\", \"label\", \"prediction\"],\n",
    ")\n",
    "results.head(16)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets check the performance on the 200 samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Got {results['correct'].sum()}/{len(results)} correct\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get approximately 80% correct as expected from the model card."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_pipeline.detach()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook has demonstrated how easy it is to run GPT-J on the Graphcore IPU using this implementation of the model and 🤗 Hub checkpoints of the model weights. While not as powerful as larger models for free text generation, medium-size, auto-regressive models like GPT-J can still be successfully fine-tuned to handle a range of NLP tasks such as question answering, sentiment analysis, and named entity recognition.\n",
    "\n",
    "In less than 10 lines of code, we were able to load the model onto the IPU and perform NLP tasks. We showed how the prompt format influences the model output and built structured prompts for question answering. We demonstrated how batched inference can be used to increase throughput by generating multiple answers at a time.\n",
    "\n",
    "We also tackled the more complex task of determining entailment between statements, and found that the standard GPT-J checkpoint was not effective.\n",
    "We tried few-shot prompting, concatenating several examples in the input prompt, to give the model instructions for the task. Even with some improvement, this technique wasn't enough for the general model to successfully classify entailment.\n",
    "Finally, we showed that a model fine-tuned on this specific downstream task (using the MNLI GLUE dataset) performs much better, achieving approximately 80% accuracy on 200 validation samples.\n",
    "\n",
    "Overall, this notebook showcases the potential for GPT-J to be used effectively and efficiently on several downstream tasks after a simple fine-tuning.\n",
    "Next, find out how to fine-tune GPT-J on the IPU in our notebook on \"fine-tuning GPT-J on the MNLI dataset\" `finetuning.ipynb`."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}