{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text generation with GPT-J 6B\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 has visibility of previous tokens only.\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 sized models. GPT-J at 6B 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 mentioned above has been fine-tuned on IPU. No finetuning 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.\n",
    "\n",
    "> **Hardware requirements:** This notebook requires at least 16 IPUs, on Paperspace this is available using either an IPU-POD16 or a Bow-POD16 machine."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "In order to run this notebook you will need to be in an environment with the Poplar SDK and PopART installed and enabled - on Paperspace this is handled by default.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
        "%pip install -r requirements.txt"
   ]
  },
  {
   "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\", 16))\n",
    "if number_of_ipus < 16:\n",
    "    raise ValueError(\"This notebook is designed to run with at least 16 IPUs\")\n",
    "\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\")+\"/gptj\"\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 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 chosen a pre-trained model we create a `GPTJPipeline`. 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 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 model gets the correct answer, it 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",
    "Lets try to give a different structure to 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: it lets us 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 16 which means this pipeline can generate answers to 16 questions at a time.\n",
    "\n",
    "Let's use the model to generate countries' 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 off 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 16 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 if two statements agree, disagree or a neutral relative to each other.\n",
    "A common benchmark dataset is the [MNLI GLUE dataset](https://huggingface.co/datasets/glue).\n",
    "\n",
    "It 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 an example 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 examples 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 4 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 3 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 examples can lead to very different results.\n",
    "\n",
    "To evaluate the effectiveness of the engineered prompts we test against a number of unseen examples in the MNLI dataset.\n",
    "In the cell below we test `entailment_instructions_v3`, 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[16:64]\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[:16]"
   ]
  },
  {
   "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 oscillate around 33% accuracy which corresponds to 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",
    "Lets 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.\n",
    "\n",
    "The checkpoint we will be using was fine-tuned on the Graphcore IPU and is hosted on the 🤗 Hugging Face 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",
    "    print_live=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While this model was trained on IPU 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, it 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 pipeline using 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 🤗 Hugging Face 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": {
  "kernelspec": {
   "display_name": "3.1.0+1205_poptorch",
   "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 (default, Jun 22 2022, 20:18:18) \n[GCC 9.4.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b98909bafd6dd6f1e7abe058b85e35058ed7070911f248b1c4db6b998dcd08ee"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}