{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning GPT-J for text entailment\n",
    "\n",
    "Copyright (c) 2023 Graphcore Ltd."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[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 on 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",
    "In this example we apply this idea and fine-tune GPT-J as a Causal Language Model (CLM) for Text Entailment on [GLUE MNLI dataset](https://huggingface.co/datasets/glue#mnli).\n",
    "\n",
    "You can easily adapt this example to do your custom fine-tuning on several downstream tasks, such as Question Answering, Named Entity Recognition, Sentiment Analysis, Text Classification in general: you just need to prepare data in the right way.\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",
    "Our weights are also available as an Hugging Face checkpoint at [Graphcore/gptj-mnli]( https://huggingface.co/Graphcore/gptj-mnli).\n",
    "\n",
    "> **Note:** This notebook requires at least 16 IPUs, on Paperspace this is available using either an IPU-POD16 or a Bow-POD16 machine. Fine tuning the model takes more than 2 hours, however you can assess the fine tuned model by downloading the [Graphcore/gptj-mnli]( https://huggingface.co/Graphcore/gptj-mnli) checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paperspace setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8140e787",
   "metadata": {},
   "source": [
    "This notebook saves checkpoints during training to allow you to resume the fine-tuning from intermediate states.\n",
    "On Paperspace these checkpoints are stored in `/storage`, this is a shared storage space with your team, allowing you to reuse the checkpoint on a different notebook instance."
   ]
  },
  {
   "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(f\"This example need 16 IPUs to work. Detected {number_of_ipus}\")\n",
    "\n",
    "os.environ[\"POPART_CACHE_DIR\"] = os.getenv(\n",
    "    \"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\"\n",
    ")\n",
    "checkpoint_dir = os.getenv(\"PERSISTENT_CHECKPOINT_DIR\", \"checkpoints\") + \"/gpt-j\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial setup\n",
    "### Fine-tuning configuration\n",
    "First of all, we need to load a base configuration, defined in `config/finetuning.yml`.\n",
    "This file has optimised configurations to run the model on IPUs.\n",
    "We need to pick the one suitable for a Pod16.\n",
    "\n",
    "This configuration uses a sequence length of 1024 tokens. GPT-J layers are split across 16 IPUs, using [Tensor Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf). No data parallelism is used (this extra optimization is available when using a Pod64).\n",
    "The `gptj_config_setup` sets up the specified configuration and configures logging and Weight and Biases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.setup import gptj_config_setup\n",
    "from config import CONFIG_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **W&B**: We support logging to Weights & Biases.\n",
    "If you want to use it, you will first need to manually log in (see the quickstart guide [here](https://docs.wandb.ai/quickstart)).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set this to True if you want to use W&B. Be sure to be logged in.\n",
    "wandb_setup = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose a configuration\n",
    "config, *_ = gptj_config_setup(\n",
    "    CONFIG_DIR / \"finetuning.yml\", \"release\", \"gptj_6B_1024_pod16\", wandb_setup\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(config.dumps_yaml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validation configuration\n",
    "Configurations for inference-only are available in `config/inference.yml`:\n",
    "- `gpt-j` is the base configuration, and is guaranteed to fit into memory with 1024 sequence length.\n",
    "- `gpt-j-mnli` is optimised for the MNLI dataset. This selects a bigger batch size but requires the sequence length to be reduced.\n",
    "\n",
    "In this example we will start from the base one, and manually reduce the sequence length and increase the batch size later on.\n",
    "You can do the same on your custom dataset to find the best configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.setup import gptj_config_setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **W&B** We support logging to Weights & Biases.\n",
    "If you want to use it, you will first need to manually log in (see the quickstart guide [here](https://docs.wandb.ai/quickstart))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_setup_on_eval = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_config, args, _ = gptj_config_setup(\n",
    "    CONFIG_DIR / \"inference.yml\",\n",
    "    \"release\",\n",
    "    \"gpt-j\",\n",
    "    hf_model_setup=False,\n",
    "    wandb_setup=wandb_setup,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(eval_config.dumps_yaml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "The MNLI 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",
    "Data splits for the MNLI dataset are the following:\n",
    "\n",
    "|train |validation_matched|validation_mismatched|\n",
    "|-----:|-----------------:|--------------------:|\n",
    "|392702|              9815|                 9832|\n",
    "\n",
    "\n",
    "You can explore it [on Hugging Face](https://huggingface.co/datasets/glue/viewer/mnli/train).\n",
    "![MNLI dataset](imgs/mnli_dataset.png)\n",
    "\n",
    "\n",
    "### Training pre-processing\n",
    "The columns we are interested in are `hypothesis`, `premise` and `label`.\n",
    "\n",
    "The first step consists of forming input prompts with the format\n",
    "```bash\n",
    "mnli hypothesis: {hypothesis} premise: {premise} target: {label} <|endoftext|>\n",
    "```\n",
    "For example:\n",
    "```\n",
    "mnli hypothesis: Product and geography are what make cream skimming work.  premise: Conceptually cream skimming has two basic dimensions - product and geography. target: neutral<|endoftext|>\n",
    "```\n",
    "\n",
    "Then, prompt sentences are tokenized and packed together to form 1024 token sequences, following the [Hugging Face packing algorithm](https://github.com/huggingface/transformers/blob/v4.20.1/examples/pytorch/language-modeling/run_clm.py). No padding is used.\n",
    "\n",
    "Finally, the prompt is split into `input_ids` and `labels`. The input consists of the full sentence but for the last token (`prompt[:-1]`), and the label is the sentence shifted by one (`prompt[1:]`).\n",
    "Given the training format, no extra care is needed to account for different sequences: the model does not need to know which sentence a token belongs to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer\n",
    "import data.hf_data_utils as hf_data_utils\n",
    "import data.mnli_data as mnli_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next two cells are the ones you want to change when doing a custom fine-tuning.\n",
    "\n",
    "We first load the MNLI dataset, and then create a custom pre-processing function to build prompts suitable for a\n",
    "text-to-text setup.\n",
    "For a custom fine-tuning, you will need to choose a format for your prompts and change the `form_training_prompts` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load HF dataset\n",
    "dataset = load_dataset(\"glue\", \"mnli\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Form prompts in the format mnli hypothesis: {hypothesis} premise: {premise} target: {class_label} <|endoftext|>\n",
    "def form_training_prompts(example):\n",
    "    hypothesis = example[\"hypothesis\"]\n",
    "    premise = example[\"premise\"]\n",
    "    class_label = [\"entailment\", \"neutral\", \"contradiction\"][example[\"label\"]]\n",
    "\n",
    "    example[\n",
    "        \"text\"\n",
    "    ] = f\"mnli hypothesis: {hypothesis} premise: {premise} target: {class_label}<|endoftext|>\"\n",
    "    return example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.map(\n",
    "    form_training_prompts,\n",
    "    remove_columns=[\"hypothesis\", \"premise\", \"label\", \"idx\"],\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Generating text prompt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shows first textual prompt\n",
    "print(dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After that, we tokenize the prompts. You won't need to change this step for a custom fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6b\")\n",
    "tokenizer.add_special_tokens({\"pad_token\": \"<|extratoken_1|>\"})  # index 50257"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tokenize prompts\n",
    "dataset = dataset.map(\n",
    "    mnli_data.tokenizes_text(tokenizer),\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=1,\n",
    "    remove_columns=dataset.column_names,\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Tokenizing text\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shows first tokenized prompt\n",
    "print(dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we use the Hugging Face packing algorithm (`group_text`) to create packed sentences of the specified sequence length,\n",
    "and separate inputs and labels.\n",
    "Again, this is a step you are not going to change for a custom fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pack tokenized prompts into sequences and split sequences in input_ids and labels\n",
    "dataset = dataset.map(\n",
    "    hf_data_utils.group_texts(config),\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=1,\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Packing sequences\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show a portion of first sentence. You can see that the label is the input shifted by one.\n",
    "print(\"first 10 tokens of first sentence\")\n",
    "print(\"input_ids\")\n",
    "print(dataset[\"input_ids\"][0][:10])\n",
    "print(\"labels - shifted by one\")\n",
    "print(dataset[\"labels\"][0][:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note** If you want to adapt this code for another dataset, be sure to call the inputs and labels in the same way: `input_ids` and `labels`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validation pre-processing\n",
    "For validation, we use the [mnli validation_mismatched](https://huggingface.co/datasets/glue/viewer/mnli_mismatched/validation) split.\n",
    "\n",
    "\"Mismatched\" means that the validation examples are not derived from the same sources as those in the training set and therefore don't closely resemble any of the examples seen at training time.\n",
    "\n",
    "Similar to what we did for training, the first pre-processing step is creating prompts, this time without including the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def form_validation_prompts(example):\n",
    "    hypothesis = example[\"hypothesis\"]\n",
    "    premise = example[\"premise\"]\n",
    "    class_label = [\"entailment\", \"neutral\", \"contradiction\"][example[\"label\"]]\n",
    "\n",
    "    example[\"text\"] = f\"mnli hypothesis: {hypothesis} premise: {premise} target:\"\n",
    "    return example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset = load_dataset(\"glue\", \"mnli\", split=\"validation_mismatched\")\n",
    "eval_dataset = eval_dataset.map(\n",
    "    form_validation_prompts,\n",
    "    remove_columns=[\"hypothesis\", \"premise\", \"idx\"],\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Generating text prompt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(eval_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, input prompts are tokenized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_validation_features(dataset, tokenizer):\n",
    "    tokenized_examples = []\n",
    "    for example in dataset[\"text\"]:\n",
    "        tokenized_example = tokenizer.encode(example, return_tensors=\"pt\").squeeze()\n",
    "        tokenized_examples.append(tokenized_example)\n",
    "    return {\"input_ids\": tokenized_examples, \"label\": dataset[\"label\"]}\n",
    "\n",
    "\n",
    "eval_dataset = eval_dataset.map(\n",
    "    prepare_validation_features,\n",
    "    batched=True,\n",
    "    remove_columns=eval_dataset.column_names,\n",
    "    load_from_cache_file=False,\n",
    "    fn_kwargs={\"tokenizer\": tokenizer},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(eval_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note** If you want to adapt this code for another dataset, be sure to call the inputs and labels in the same way: `input_ids` and `label`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customise configuration and create a Trainer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right at the beginning of the notebook we loaded the base configurations for training and inference. We are now going to show how to customise some of the parameters for your needs.\n",
    "\n",
    "### Customise training configuration\n",
    "In the cells below we list the parameters you are most likely to play around with when doing a custom fine-tuning.\n",
    "\n",
    "These are the training steps, dropout probability and optimizer/learning rate parameters.\n",
    "\n",
    "Moreover, it is important that you specify **checkpoint** parameters, namely a folder to save the fine-tuned weights and a periodicity for checkpointing. Be aware that saving checkpoints takes time, so you don't want to save them too often.\n",
    "To disable intermediate checkpoints set `config.checkpoint.steps = 0`. The final checkpoint is always saved provided the `config.checkpoint.save` directory is given. Set it to `None` if you don't want to save weights, but it's unlikely you want to disable the last checkpoint.\n",
    "\n",
    "If you are not resuming training and you don't care about resuming the training later on you can reduce the time and memory required to save checkpoints by specifying `optim_state=False`. In this case, only the model weights will be saved, while the optimiser state will be discarded.\n",
    "\n",
    "Checkpoints will be saved in the directory given by the environment variable `PERSISTENT_CHECKPOINT_DIR`, which we saved in `checkpoint_dir` at the beginning. You can have a look at the path by printing it out in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(checkpoint_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customise training arguments\n",
    "config.model.dropout_prob = 0.0\n",
    "config.training.steps = 400"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customise optimiser and learning rate schedule\n",
    "config.training.optimizer.learning_rate.maximum = 5e-06\n",
    "config.training.optimizer.learning_rate.warmup_proportion = 0.005995\n",
    "config.training.optimizer.learning_rate.beta1 = 0.9\n",
    "config.training.optimizer.learning_rate.beta2 = 0.999\n",
    "config.training.optimizer.learning_rate.weight_decay = 0.0\n",
    "config.training.optimizer.learning_rate.gradient_clipping = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customise checkpoints\n",
    "config.checkpoint.save = (\n",
    "    checkpoint_dir  # where the model is saved. None means don't save any checkpoint\n",
    ")\n",
    "config.checkpoint.steps = (\n",
    "    100  # how often you save the model. 0 means only the final checkpoint is saved\n",
    ")\n",
    "config.checkpoint.to_keep = 4  # maximum number of checkpoints kept on disk\n",
    "config.checkpoint.optim_state = (\n",
    "    False  # Whether to include the optimiser state in checkpoints\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resume training\n",
    "config.checkpoint.load = (\n",
    "    None  # you can specify a directory containing a previous checkpoint\n",
    ")\n",
    "# os.path.join(checkpoint_dir, ...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Customise validation configuration\n",
    "For validation there are fewer relevant parameters.\n",
    "\n",
    "You can control the maximum number of tokens generated by the model using the  `output_length` parameter. If you know that the targets are only a few tokens long, it is convenient to set it to a small number. Generation stops if the output token is `<|endoftext|>` or after `output_length` tokens are generated.\n",
    "\n",
    "In our case, we set the `output_length` to 5 to accommodate all class labels and the `<|endoftext|>` token.\n",
    "\n",
    "You can also reduce the model `sequence_len` to account for the maximum length of sentences encountered in the validation dataset, plus the `output_length` specified.\n",
    "In our case of the MNLI example, `sequence_len = 229` and this allows us to increase the batch size to `16`.\n",
    "\n",
    "If you adapt this example to another dataset, you can compare your dataset `max_len` with ours and decide if you can safely use the increased batch size. Otherwise, stick with the default one or create your own.\n",
    "\n",
    "You can specify a checkpoint to be used for validation. If `None` is provided, the latest weights will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify maximum number of tokens that can be generated.\n",
    "eval_config.inference.output_length = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reducing sequence length\n",
    "max_len = reduce(lambda l, e: max(l, len(e[\"input_ids\"])), eval_dataset, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Maximum length in mnli-mismatched dataset: {max_len}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_config.model.sequence_length = max_len + eval_config.inference.output_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Setting sequence length to {eval_config.model.sequence_length}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Increase batch size\n",
    "eval_config.execution.micro_batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can specify a directory containing a previous checkpoint\n",
    "# None means latest weights are used\n",
    "config.checkpoint.load = None  # os.path.join(checkpoint_dir, ...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(config.dumps_yaml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we need to define a validation metric. We will use `accuracy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_metric = evaluate.load(\"accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to define a function to convert our generated labels to integer indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def postprocess_mnli_predictions(generated_sentences):\n",
    "    labels_to_ids = {\"entailment\": 0, \"neutral\": 1, \"contradiction\": 2, \"unknown\": -1}\n",
    "    predictions = []\n",
    "    for s in generated_sentences:\n",
    "        answer = mnli_data.extract_class_label(s)\n",
    "        predictions.append(labels_to_ids[answer])\n",
    "    return predictions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "717507f8",
   "metadata": {},
   "source": [
    "### Create a Trainer\n",
    "Once a config is specified, we are ready to create the training session with the help of the \n",
    "`GPTJTrainer` class.\n",
    "You need to provide the following arguments:\n",
    "\n",
    "- *config*: the training configuration.\n",
    "- *pretrained*: the Hugging Face pre-trained model, used to initialise the weights.\n",
    "- *dataset*: the training dataset.\n",
    "\n",
    "Moreover, you can specify:\n",
    "\n",
    "- *eval_dataset*: the validation dataset.\n",
    "- *eval_config*: the inference configuration, to be used in validation.\n",
    "- *tokenizer*: the tokenizer, needed by validation.\n",
    "- *metric*: the metric for validation. An Hugging Face metric from the `evaluate` module. We use the `accuracy` metric.\n",
    "- *process_answers_func*: a function to convert the generated answers to the format required by the metric and the labels. For example we need to convert textual categories `[entailment, contradiction,neutral]` to indices.\n",
    "\n",
    "These extra arguments can also be provided later on when calling `trainer.evaluate(...)`.\n",
    "\n",
    "If you want to run fine-tuning, the pre-trained model should be `EleutherAI/gpt-j-6b`. If you want to just run validation on the existing Hugging Face checkpoint, you should change it to the fine-tuned model `Graphcore/gptj-mnli`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.trainer import GPTJTrainer\n",
    "from transformers.models.gptj.modeling_gptj import GPTJForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change to Graphcore/gptj-mnli for validation only\n",
    "pretrained = GPTJForCausalLM.from_pretrained(r\"EleutherAI/gpt-j-6b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = GPTJTrainer(\n",
    "    config,\n",
    "    pretrained,\n",
    "    dataset,\n",
    "    eval_dataset,\n",
    "    eval_config,\n",
    "    tokenizer,\n",
    "    accuracy_metric,\n",
    "    postprocess_mnli_predictions,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run fine-tuning\n",
    "We can now run training for the number of steps you set in the config.\n",
    "Checkpoints will be saved in the folder you specified in `config.checkpoint.save`, with the periodicity set in `config.checkpoint.steps`.\n",
    "\n",
    "The first time you run `trainer.train()` it takes around 10 minutes to compile the training model.\n",
    "\n",
    "Training takes around 8 minutes to start because the pre-trained weights need to be downloaded to the IPUs.\n",
    "After that, each step takes around 20-22 seconds.\n",
    "This time does not include checkpointing time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run validation\n",
    "Finally, we validate our model on the [mnli_mismatched](https://huggingface.co/datasets/glue/viewer/mnli_mismatched/test) split of the MNLI dataset.\n",
    "\n",
    "Generative inference is performed token-by-token using a greedy heuristic: the next token is chosen based on the highest logits.\n",
    "\n",
    "We run token-by-token inference in batches to generate labels for multiple examples at once.\n",
    "We then compute the accuracy by comparing the model answers with the true labels.\n",
    "\n",
    "The resulting model matches SOTA performance with 82.5% accuracy.\n",
    "\n",
    "```\n",
    "Total number of examples                 9832\n",
    "Number with badly formed result          0\n",
    "Number with incorrect result             1725\n",
    "Number with correct result               8107 [82.5%]\n",
    " ```\n",
    "\n",
    "The first time you run `trainer.evaluate()` it takes around 4 minutes to compile the inference model.\n",
    "\n",
    "Running validation on the whole dataset takes around 7 minutes.\n",
    "\n",
    "Note that:\n",
    "\n",
    "- If you have compiled the training model and you don't specify an `eval_config.checkpoint.load` folder, the latest weights will be used.\n",
    "\n",
    "- If you instead specify a repository in `eval_config.checkpoint.load`, you will be evaluating that specific set of weights.\n",
    "\n",
    "- If none of the above holds or if you specify `trainer.evaluate(use_pretrained=True)`, weights from `pretrained` will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you want to change the pretrained model to run validation on the HF checkpoint, uncomment and run below\n",
    "# pretrained = GPTJForCausalLM.from_pretrained(\"Graphcore/gptj-mnli\")\n",
    "# trainer.pretrained = pretrained\n",
    "# trainer.evaluate(use_pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Hugging Face checkpoint\n",
    "You can save the trained weights so that they can be uploaded to Hugging Face and used with Hugging Face's PyTorch model.\n",
    "You can specify a checkpoint path if you want to convert a specific checkpoint instead of the latest weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_checkpoint_path = os.path.join(checkpoint_dir, \"hf_checkpoint\")\n",
    "ckpt_path = None  # os.path.join(checkpoint_dir, ...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finetuned = trainer.save_hf_checkpoint(hf_checkpoint_path, ckpt_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the model with Hugging Face pipeline\n",
    "The same model can later be used with the standard Hugging Face pipeline on any hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6b\")\n",
    "tokenizer.add_special_tokens({\"pad_token\": \"<|extratoken_1|>\"})\n",
    "hf_model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"Graphcore/gptj-mnli\", pad_token_id=tokenizer.pad_token_id\n",
    ")\n",
    "generator = pipeline(\"text-generation\", model=hf_model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = (\n",
    "    \"mnli hypothesis: Your contributions were of no help with our students' education.\"\n",
    "    \"premise: Your contribution helped make it possible for us to provide our students with a quality education. target:\"\n",
    ")\n",
    "\n",
    "out = generator(prompt, return_full_text=False, max_new_tokens=5, top_k=1)\n",
    "print(out)\n",
    "# [{'generated_text': ' contradiction'}]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}