{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "rEJBSTyZIrIb"
   },
   "source": [
    "Copyright (c) 2023 Graphcore Ltd. All rights reserved.\n",
    "\n",
    "For all available notebooks, check [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to see how IPUs perform on other tasks."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Machine Translation on IPUs using MT5 - Fine-tuning\n",
    "\n",
    "[MT5](https://huggingface.co/docs/transformers/model_doc/MT5) is the multilingual variant of [T5](https://huggingface.co/docs/transformers/model_doc/t5), an encoder-decoder transformer model that leverages a unified text to text approach to represent natural language tasks. A limitation of T5 is that pre-trained checkpoints primarily support only English as the source text for any natural language task so the majority of problems that T5 can solve are English to English tasks. MT5 uses the same architecture as T5 but the encoder is pre-trained on a dataset covering 101 languages. However, unlike T5, MT5's pre-training does not include any supervised training so the MT5 pre-trained checkpoints have limited immediate use. To be performant on any downstream task like machine translation or sentiment analysis, MT5 needs to be fine-tuned.\n",
    "\n",
    "Fine-tuning MT5 is similar to fine-tuning T5 (see Summarization on IPU using T5 Small - Fine-Tuning `summarization.ipynb`). However, since MT5 uses a much larger vocabulary than T5, this needs to be accounted for when placing MT5 on the IPU. This notebook shows how to:\n",
    "- Fine-tune MT5 using the Graphcore IPU with the `IPUSeq2SeqTrainer` for the task of machine translation with a source language that is different from English.\n",
    "- (optional) Place large layers such as an embedding table or projection layer across IPUs when they do not fit on a single IPU.\n",
    "\n",
    "We will use the [WMT dataset](http://www.statmt.org/wmt16/), a machine translation dataset composed of a collection of various sources, including news commentaries and parliament proceedings."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------|--------------|\n",
    "| Natural language processing | Translation | MT5-small | WMT dataset | Fine-tuning | 4 | X6hr |\n",
    "\n",
    "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n",
    "\n",
    "[![Run on Gradient](https://assets.paperspace.io/img/gradient-badge.svg)]()\n",
    "\n",
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies and configuration\n",
    "\n",
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the dependencies for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install \"optimum-graphcore==0.7\" sacrebleu\n",
    "%pip install graphcore-cloud-tools[logger]@git+https://github.com/graphcore/graphcore-cloud-tools\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "n_ipu = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\") + \"/mt5_translation\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For fine-tuning, we will use the `mt5-small` pretrained checkpoint which requires 4 IPUs. Larger MT5 configurations offer a greater capacity of natural language understanding and can better support multiple tasks with a single model at the expense of requiring more memory and compute resource. However, smaller variants such as `mt5-small` still offer good language understanding. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"google/mt5-small\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing your model with the community"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "You can share your model with the 🤗 community. You do this by completing the following steps:\n",
    "1. Store your authentication token from the 🤗 website. [Sign up to 🤗](https://huggingface.co/join) if you haven't already.\n",
    "2. Execute the following cell and input your username and password."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you need to install Git-LFS to manage large files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!apt install git-lfs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "whPRbBNbIrIl"
   },
   "source": [
    "## Loading the dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "W7QYTpxXIrIl"
   },
   "source": [
    "We will use the [🤗 Datasets](https://github.com/huggingface/datasets) library to download the data and define the metric we will use for evaluation. We use the English/Romanian part of the WMT dataset here and the [BLEU](https://huggingface.co/spaces/evaluate-metric/bleu) evaluation metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IreSlFmlIrIm",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset, load_metric\n",
    "\n",
    "raw_datasets = load_dataset(\"wmt16\", \"ro-en\")\n",
    "metric = load_metric(\"sacrebleu\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "WHUmphG3IrI3"
   },
   "source": [
    "To get a sense of what the data looks like, the following function will show some samples picked randomly from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i3j8APAoIrI3",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "def show_random_elements(dataset, num_examples=5):\n",
    "    assert num_examples <= len(dataset), \"Can't pick more elements than there are in the dataset.\"\n",
    "    picks = []\n",
    "    for _ in range(num_examples):\n",
    "        pick = random.randint(0, len(dataset)-1)\n",
    "        while pick in picks:\n",
    "            pick = random.randint(0, len(dataset)-1)\n",
    "        picks.append(pick)\n",
    "    \n",
    "    df = pd.DataFrame(dataset[picks])\n",
    "    for column, typ in dataset.features.items():\n",
    "        if isinstance(typ, datasets.ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SZy5tRB_IrI7",
    "outputId": "ba8f2124-e485-488f-8c0c-254f34f24f13",
    "tags": []
   },
   "outputs": [],
   "source": [
    "show_random_elements(raw_datasets[\"train\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "n9qywopnIrJH"
   },
   "source": [
    "## Preprocessing the data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "YVx71GdAIrJH"
   },
   "source": [
    "To tokenize our input samples we use a 🤗 Transformers `Tokenizer` instantiated using the `AutoTokenizer.from_pretrained` method to ensure that we obtain the same tokenizer and vocabulary used to pre-train MT5. Additionally we define some parameters to configure the behaviour of the tokenizer: \n",
    "- We use the same max sequence length of 128 for fine-tuning as in the [T5 paper](https://arxiv.org/abs/1910.10683). \n",
    "- Sequences with more than 128 tokens will be truncated and those with less than 128 tokens will be padded until the max sequence length is reached.\n",
    "\n",
    "Note that it is necessary to limit all input samples to a fixed length since Graphcore's current Pytorch implementation only runs in static mode. That is, for every forward pass of the model, all inputs need to have the same dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eXNLu_-nIrJI",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from dataclasses import dataclass\n",
    "    \n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\n",
    "\n",
    "@dataclass\n",
    "class TokenizerConfig:\n",
    "    max_length = 128\n",
    "    truncation=True\n",
    "    padding=\"max_length\"\n",
    "\n",
    "tokenizer_config = TokenizerConfig()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "2C0hcmp9IrJQ"
   },
   "source": [
    "Since we are fine-tuning MT5 for Romanian to English translation, we set the prefix to be added to every input sample below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "source_lang = \"ro\"\n",
    "target_lang = \"en\"\n",
    "prefix = f\"translate {source_lang} to {target_lang}: \""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a function to preprocess our input samples and apply it to all elements of all splits in the `dataset`. The first run of the cell below will take more time than subsequent runs since the results are cached by the 🤗 Datasets library. The 🤗 Datasets library is able to detect when the function you pass to `map` has changed (and thus to not use the cached data). For instance, it will detect if you change the task in the first cell and rerun the notebook. 🤗 Datasets warns you when it uses cached files. You can pass `load_from_cache_file=False` in the call to `map` to not use the cached files and force the preprocessing to be applied again.\n",
    "\n",
    "Note that we passed `batched=True` to encode the text samples together into samples. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the text samples in a batch concurrently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vc0BSBLIIrJQ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    inputs = [prefix + ex[source_lang] for ex in examples[\"translation\"]]\n",
    "    targets = [ex[target_lang] for ex in examples[\"translation\"]]\n",
    "    model_inputs = tokenizer(inputs, max_length=tokenizer_config.max_length, padding=tokenizer_config.padding, truncation=tokenizer_config.truncation)\n",
    "\n",
    "    # Setup the tokenizer for targets\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenizer(targets, max_length=tokenizer_config.max_length, padding=tokenizer_config.padding, truncation=tokenizer_config.truncation)\n",
    "        \n",
    "    # Since we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore\n",
    "    # padding in the loss.\n",
    "    labels[\"input_ids\"] = [\n",
    "        [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels[\"input_ids\"]\n",
    "    ]\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\n",
    "tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "545PP3o8IrJV"
   },
   "source": [
    "## Fine-tuning the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "FBiW8UpKIrJW"
   },
   "source": [
    "Now that our data is ready, we can download the pre-trained model and fine-tune it. Since our task is of the sequence-to-sequence kind, we use the `AutoModelForSeq2SeqLM` class. Like with the tokenizer, the `from_pretrained` method will download and cache the model for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TlqNaB8jIrJW",
    "outputId": "84916cf3-6e6c-47f3-d081-032ec30a4132",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq\n",
    "from optimum.graphcore import IPUConfig, IPUSeq2SeqTrainer, IPUSeq2SeqTrainingArguments\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To instantiate a `IPUSeq2SeqTrainer`, we will need to define: \n",
    "* `IPUConfig`, which is a class that specifies attributes and configuration parameters to compile and put the model on the device.\n",
    "* `IPUSeq2SeqTrainingArguments`, which is a class that contains all the attributes to customize the training.\n",
    "* Data collator.\n",
    "* How to compute the metrics from the predictions.\n",
    "\n",
    "We begin by firstly instantiating an `IPUConfig`. For the models supported in Optimum Graphcore, IPU configurations are provided on the [Graphcore Huggingface hub](https://huggingface.co/Graphcore) that can be used for fine-tuning or inference. For an introduction to the options available in the `IPUConfig` please view the `natural-language-processing/introduction_to_optimum_graphcore.ipynb` notebook. We can initialise an `IPUConfig` by loading a prepared configuration from the Huggingface hub as shown below. For detail on how the `IPUConfig` for `mt5-small` is constructed, read the section below. It is optional and you can proceed to fine-tuning MT5 however it provides some background on how you can construct your own `IPUConfig` for your own models as well as how to use options such as [`embedding_serialization_factor`](https://huggingface.co/docs/optimum/main/graphcore/ipu_config#optimum.graphcore.IPUConfig.embedding_serialization_factor) to fit large embedding tables on one or more IPUs."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "    <summary><b>Constructing an IPUConfig for mt5-small</b></summary>\n",
    "    \n",
    "In order to place `mt5-small` for fine-tuning on the IPU, an approximation of its memory\n",
    "footprint is required. To do so, firstly note that `mt5-small` has 300M parameters. If we choose to use mixed precision training, we can obtain a partial approximation for the memory footprint by considering the memory required to store the model weights and gradients (see\n",
    "[Anatomy of Model's Memory](https://huggingface.co/docs/transformers/v4.20.1/en/perf_train_gpu_one#anatomy-of-models-memory) for the calculation details below):\n",
    "    \n",
    "- 6 bytes * 300M = 1717MB for the model weights\n",
    "- 4 bytes * 300M = 1145MB for gradients\n",
    "\n",
    "Excluding the memory requirements for other aspects such as optimizer states, activations, temporary variables and model code to name a few, atleast 2862MB is required to fine-tune `mt5-small`. Since a single MK2 IPU has 900MB of SRAM, we need atleast 4 IPUs to hold the model in memory. \n",
    "\n",
    "To decide on how to place the constituent layers of `mt5-small` on the IPU, observe that the official configuration of `mt5-small` can be summarized as consisting of:\n",
    "\n",
    "- An input embedding table of size (250112 x 512)\n",
    "- 8 Encoder layers\n",
    "- 8 Decoder layers\n",
    "- A language modelling head or linear projection layer with weights of size (250112 x 512). Note that by default input and output embeddings are _not_ tied for MT5 so the linear projection layer has a separate set of weights. \n",
    "\n",
    "Since the input embedding table and output projection layer each use approximately 128M parameters when weights are _not_ tied, we can try to reserve IPU-0 and IPU-3 for these layers and split the 16 encoder and decoder layers equally across IPU-1 and IPU-2. This configuration can be specified in the `IPUConfig` via `layers_per_ipu:=[0, 8, 8, 0]` where the value 0 is used to reserve specific IPUs for non-transformer layers. For MT5 in Optimum Graphcore this translates to placing the embedding table and language modelling head on the first and last IPUs respectively. However if we choose to tie input and output embedding weights, we no longer need to reserve the last IPU for the language modelling head as it can be placed on IPU-0: `layers_per_ipu:=[0, 8, 4, 4]`.\n",
    "\n",
    "If we construct the `IPUConfig` by simply setting `ipu_config=IPUConfig(layers_per_ipu=[0, 8, 8, 0])` and proceeding with defining the `IPUSeq2SeqTrainer` as done in the [Fine-tuning the model](#fine_tuning_the_model) section above, unfortunately the model will not compile due to an Out of Memory error. To understand why the model does not fit, the [PopVision Graph Analyser](https://docs.graphcore.ai/projects/graphcore-popvision-user-guide/en/2.0.0/popvision.html) (see getting started [video](https://share.vidyard.com/watch/jbEaiQqkEinNnFBBFwEHNV)) can be used to analyse the memory footprint. A summary of the memory use per IPU is shown below: \n",
    "\n",
    "![image-2.png](images/mt5_oom.png)\n",
    "\n",
    "It can be seen that IPU-0 and IPU-3 exceed the max memory permitted. For cases like the above, it is typically due to the embedding weights being large and the corresponding weight update process being expensive. For example if we use the AdamW optimiser, an additional 8 bytes * 250112 * 512 = 977MB of memory is required during the update step of the parameters on IPU-0 and IPU-3 since the AdamW optimiser uses additional state parameters. If the choice of optimiser cannot be changed, the optimiser states can be chosen to be placed off chip in the [IPU Gateway DRAM](https://docs.graphcore.ai/projects/ipu-programmers-guide/en/latest/about_ipu.html) by setting the option `optimizer_state_offchip=True` in the `IPUConfig`, reducing the memory used on the IPU during the forward pass. However when updating the large embedding weights during the backwards pass, the optimiser states for the embeddings will still need to be moved to the IPU and in this case they are of size 977MB, exceeding the IPU max memory. The amount of memory moved to the IPU when updating the embedding weights can be reduced if we serialize the computation of the embeddings so that a large embedding computation is split into N sequential sub-embedding computations. By serializing the computation of the embeddings into N parts, the update step for the embeddings in the backwards pass will also be serialised into N parts, meaning only 1/N of the optimiser states for the embeddings need to be moved to the IPU at a given step.\n",
    "\n",
    "The `IPUConfig` provides options to serialize the input embedding table and the computation of the language modelling head via:\n",
    "- `embedding_serialization_factor: int`, for example `embedding_serialization_factor=4` will serialize the input embedding computation into 4 partial calculations.\n",
    "- `projection_serialization_factor: int`\n",
    "\n",
    "In addition to serializing the computation of the input embeddings and the projection layer, it is also possible to place the serialized computation across a selection of contiguous IPUs via options below. Note that these options are mutually exclusive with the above options which restrict serializing computations to a single IPU.\n",
    "- `serialized_embedding_splits_per_ipu: List[int]` for instance `serialized_embedding_splits_per_ipu=[3, 1, 0, 0]`, will serialize the embedding computation into 4 parts, with 3 parts occuring on IPU-0 and the last computation on IPU-1. Note that the `list` specifying the number of sub computations to place on each IPU needs to be of the same length as the `layers_per_ipu` parameter introduced above.\n",
    "- `serialized_projection_splits_per_ipu: List[int]` \n",
    "\n",
    "By setting `optimizer_state_offchip=True` and serializing the input and projection layer embeddings, `mt5-small` with the [Graphcore/mt5-small-ipu](https://huggingface.co/Graphcore/mt5-small-ipu) `IPUConfig` can be fine-tuned on the IPU. Other memory management options are available in the `IPUConfig` and can be viewed in the [documentation](https://huggingface.co/docs/optimum/main/graphcore/ipu_config#optimum.graphcore.IPUConfig). With the above information hopefully it is clearer why [Graphcore/mt5-small-ipu](https://huggingface.co/Graphcore/mt5-small-ipu) has particular options enabled. The above information details the memory requirements for fine-tuning `mt5-small`, can you follow a similar process for defining an `IPUConfig` to run `mt5-small` in inference mode (note that inference variants for options like `layers_per_ipu` can be set via `inference_layers_per_ipu`)?\n",
    "- How many IPUs are required to run `mt5-small` in inference mode?\n",
    "- Is it necessary to serialize input embeddings and the projection layer for inference?\n",
    "- If the input and output weight embeddings are tied, is it possible to serialize embedding / projection layer computation?\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ipu_config = IPUConfig.from_pretrained(\"Graphcore/mt5-small-ipu\", executable_cache_dir=executable_cache_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "_N8urzhyIrJY"
   },
   "source": [
    "Next we define the training arguments that will be used with the `IPUSeq2SeqTrainingArguments` class. A folder name is required as that will be used to save the checkpoints of the model. All other arguments are optional but we tweak them for our fine-tuning.\n",
    "\n",
    "We set the evaluation to be done at the end of each epoch and tweak the learning rate. The effective batch size is also configured by providing arguments to the  `micro_batch_size`, `gradient_accumulation_steps` and `n_ipu` parameters. The `n_ipu` parameter sets the number of IPUs to be used, which the trainer will use to define the number of replicas for data parallel training. To make sure the effective batch size remains the same we adjust gradient accumulation steps to take into account additional training replicas. Since `IPUSeq2SeqTrainer` will save the model regularly and our dataset is quite large, we tell it to make a maximum of three.\n",
    "\n",
    "The `push_to_hub` parameter in `IPUSeq2SeqTrainer` is necessary if we want to push the model to the [🤗 Models Hub](https://huggingface.co/models) regularly during training. You can enable it if you did follow the installation steps at the beginning of this notebook. If you want to save your model locally to a name that is different to the name of the repository it will be pushed to, or if you want to push your model under an organization and not your name space, use the `hub_model_id` argument to set the repo name (it needs to be the full name, including your namespace: for instance `\"sgugger/marian-finetuned-ro-to-en\"` or `\"huggingface/marian-finetuned-ro-to-en\"`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Bliy8zgjIrJY",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "# Gradient accumulation steps = 128 for 1 replica.\n",
    "# Increase replicas if we have additional ipus available and\n",
    "# adjust gradient accumulation steps. The trainer will automatically \n",
    "# use data parallel training by using the `n_ipu` argument\n",
    "replication_factor = n_ipu // ipu_config.ipus_per_replica\n",
    "gradient_accumulation_steps = 128 // replication_factor\n",
    "\n",
    "args = IPUSeq2SeqTrainingArguments(\n",
    "    f\"{model_name}-finetuned-{source_lang}-to-{target_lang}\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=1e-3,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "    n_ipu=n_ipu,\n",
    "    save_total_limit=3,\n",
    "    num_train_epochs=2,\n",
    "    predict_with_generate=True,\n",
    "    generation_max_length=TokenizerConfig.max_length,\n",
    "    dataloader_drop_last=True,\n",
    "    logging_steps=10,\n",
    "    push_to_hub=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "km3pGVdTIrJc"
   },
   "source": [
    "Since we are using a Seq2Seq model, a special kind of data collator is required that will be used by the trainer when collating samples with a Pytorch DataLoader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "7sZOdRlRIrJd"
   },
   "source": [
    "The last thing to define for our `IPUSeq2SeqTrainer` is how to compute the metrics from the predictions. We need to define a function for this, which will just use the `metric` defined earlier. We have to do a bit of pre-processing to decode the predictions into text samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UmvbnJ9JIrJd",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def postprocess_text(preds, labels):\n",
    "    preds = [pred.strip() for pred in preds]\n",
    "    labels = [[label.strip()] for label in labels]\n",
    "    return preds, labels\n",
    "\n",
    "def compute_metrics(eval_preds):\n",
    "    preds, labels = eval_preds\n",
    "    if isinstance(preds, tuple):\n",
    "        preds = preds[0]\n",
    "    # Replace -100 in the labels as we can't decode them.\n",
    "    preds = np.where(preds != -100, preds, tokenizer.pad_token_id)\n",
    "    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # Some simple post-processing\n",
    "    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)\n",
    "\n",
    "    result = metric.compute(predictions=decoded_preds, references=decoded_labels)\n",
    "    result = {\"bleu\": result[\"score\"]}\n",
    "\n",
    "    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]\n",
    "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
    "    result = {k: round(v, 4) for k, v in result.items()}\n",
    "    return result"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "rXuFTAzDIrJe"
   },
   "source": [
    "Then we just need to pass all of this together with our datasets to the `IPUSeq2SeqTrainer` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "imY1oC3SIrJf",
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = IPUSeq2SeqTrainer(\n",
    "    model,\n",
    "    ipu_config,\n",
    "    args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "CdzABDVcIrJg"
   },
   "source": [
    "We now fine-tune our model by calling the `train` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uNx5pyRlIrJh",
    "outputId": "077e661e-d36c-469b-89b8-7ff7f73541ec",
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer.model.config.use_cache=False\n",
    "trainer.train()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can upload the result of the training to the 🤗 Hub:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer.push_to_hub()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also share this model and other users can load it with the identifier \"your-username/the-name-you-picked\" so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import AutoModelForSeq2SeqLM\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\"sgugger/my-awesome-model\")\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion and next steps\n",
    "\n",
    "In this notebook we have demonstrated how to fine-tune MT5 for the task of machine translation using a source language that is different from English, a limitation of the T5 model. To reduce the time to train MT5 for translation we can use data parallel training. To use more than 1 replica more IPUs are required. On Paperspace, this is available using either an IPU-POD16 or a BoW-IPU-POD16. Please contact Graphcore if you need assistance running on larger platforms.\n",
    "\n",
    "Want to see how to use MT5 for zero-shot text classification? Try out Zero-Shot Text Classification on IPUs using MT5 - Inference `mt5_xnli.ipynb` notebook. For all available notebooks, check [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to see how IPUs perform on other tasks.\n",
    "\n",
    "Have a question? Please contact us on our [Graphcore community channel](https://www.graphcore.ai/join-community).\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Translation",
   "provenance": []
  },
  "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
