lecture06.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d715d6a9-c3d7-492f-aa7d-6465b264d01c",
   "metadata": {},
   "source": [
    "## Exercise 3 - Generating hep-ph titles with Transformers "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda90d94-084b-46a7-98e0-19c4c3e63c34",
   "metadata": {},
   "source": [
    "> Too lazy to think of a catchy title for your next paper? Then this exercise is for you!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33cd450c-b40e-4116-8619-6647371d0661",
   "metadata": {},
   "source": [
    "## Learning objectives\n",
    "\n",
    "* Gain hands-on experience using the Hugging Face Transformers library\n",
    "* Understand how to prepare and tokenizer a text dataset for Transformer models\n",
    "* Learn how to fine-tune a Transformer for text summarization and how to evaluate it's performance\n",
    "\n",
    "Understanding the main steps involved in fine-tuning Transformers for NLP, will put you in good stead to understand advanced particle physics applications such as the [Particle Transformer](https://arxiv.org/abs/2202.03772)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df3227f-1c01-4552-804f-d8324dc539db",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Chapter 6 of [_Natural Language Processing with Transformers_](https://transformersbook.com/) by L. Tunstall, L. von Werra, and T. Wolf\n",
    "* [_On the Use of ArXiv as a Dataset_](https://arxiv.org/abs/1905.00075) by C. Clement et al."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "517b33cb-9796-464c-96a2-14da6eb96aef",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a60f6f46-f4a7-49c8-acbb-df68c1318751",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and run this cell if using Colab, Kaggle etc\n",
    "# %pip install transformers datasets evaluate rouge_score nltk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f2d15b3-0f44-4135-b6b4-83776de269cf",
   "metadata": {},
   "source": [
    "To be able to share your model with the community there are a few more steps to follow.\n",
    "\n",
    "First you have to store your authentication token from the Hugging Face website (sign up here if you haven't already!) then execute the following cell and input your username and password:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a1a131-d4ce-46bd-a00f-cc4307052124",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23eeedfe-2f01-47c0-b4e0-06cdc990832c",
   "metadata": {},
   "source": [
    "Then you need to install Git-LFS. Uncomment and execute the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e41fb47a-6667-4d49-a4a6-d4a27cab00eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !apt install git-lfs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c23f0cd-7efe-4dd4-b40e-a73d0c0a03b8",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bdcfefa-9b67-480c-b15b-320e7d2d0c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from evaluate import load\n",
    "from IPython.display import YouTubeVideo\n",
    "from transformers import (\n",
    "    AutoModelForSeq2SeqLM,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForSeq2Seq,\n",
    "    Seq2SeqTrainer,\n",
    "    Seq2SeqTrainingArguments,\n",
    "    pipeline,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c210f7-1dc4-4ace-9ba8-88f48e5bad4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "# Suppress logs to keep things tidy\n",
    "datasets.logging.set_verbosity_error()\n",
    "# Download special package for computing metrics\n",
    "nltk.download(\"punkt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e42b0bd-7d57-4c5a-9792-8cdbb1fb310e",
   "metadata": {},
   "source": [
    "## Loading the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b6b06d8-478f-4038-b475-ccf2c7dd8d53",
   "metadata": {},
   "source": [
    "Until now, we've been using the Top Tagging Dataset to explore various neural network architectures, from MLPs to CNNs. Today we'll do something completely different and explore an application of Transformers to a natural language processing (NLP) task called _text summarization_! As the name suggests, text summarization involves condensing long documents into crips summaries. To give the task a physics flavour, we'll use a dump of arXiv hep-ph papers and train a model to summarise the abstract into a title. This way, the next time you're lacking inspiration on your next big paper, you can just get the model to cook it up for you!\n",
    "\n",
    "To get started, we'll need a dataset of arXiv papers. Fortunately, someone from the community has uploaded a dump of papers to the Hugging Face Hub, so we can download it using the familiar `load_dataset()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d7d1fe-dfc4-43e3-a297-54944038d056",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_dataset = load_dataset(\"gfissore/arxiv-abstracts-2021\", split=\"train\")\n",
    "raw_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "213c8389-b640-481f-9628-b5885f80ab19",
   "metadata": {},
   "source": [
    "Okay, this is quite a lot of papers! Let's take a look at one example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b986bbc9-0db1-49b7-9567-f7f82f86f5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "621f6862-7790-4cae-b92f-baf5ed0d7304",
   "metadata": {},
   "source": [
    "We can see that each example consists of an `abstract` and `title`, along with various metadata about the submission. To keep things focused, let's filter the dataset for just those papers which have a `hep-ph` category. To do so, we can use the `filter()` method in the `datasets` library. This method applies a boolean function to every row in the dataset, and removes rows where the function evaluates to false. In our case, we're interested in the `categories` column, so let's check how many hep-ph papers we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9418e7a2-1614-4c2b-bfda-7507f8944488",
   "metadata": {},
   "outputs": [],
   "source": [
    "category = \"hep-ph\"\n",
    "hep_dataset = raw_dataset.filter(lambda x: category in x[\"categories\"])\n",
    "hep_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3df0b032-8d2a-4202-b8cc-dd95c9de7b97",
   "metadata": {},
   "source": [
    "Great, this is a much more manageable dataset to work with! As a sanity check, you can pick your name or that of a colleague to see if any relevant papers are found in the corpus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "204a4fd2-81ec-4d54-96e0-5c0f1f66e314",
   "metadata": {},
   "outputs": [],
   "source": [
    "author = \"tunstall\"\n",
    "sample_dataset = hep_dataset.filter(lambda x: \"tunstall\" in x[\"authors\"].lower())\n",
    "for row in sample_dataset:\n",
    "    print(f\"Title: {row['title']} \\n\")\n",
    "    print(f\"Authors: {row['authors']} \\n\")\n",
    "    print(f\"Abstract: {row['abstract']}\")\n",
    "    print(\"=\" * 50, \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94255774-aad8-47af-980d-b96c73ff8fd8",
   "metadata": {},
   "source": [
    "These look like hep-ph papers, so now let's process the raw text in the abstracts and titles into a format that's suitable for neural networks!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69a2fb39-5252-47b8-91a6-bc7a3d88872d",
   "metadata": {},
   "source": [
    "## From text to tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f41239d1-073b-4502-a575-6473d1f8b432",
   "metadata": {},
   "source": [
    "Like other machine learning models, Transformers expect their inputs in the form of numbers (not strings) and so some form of preprocessing is required. For NLP, this preprocessing step is called tokenization. Tokenization converts strings into atomic chunks called tokens, and these tokens are subsequently encoded as numerical vectors.\n",
    "\n",
    "For more information about tokenizers, check out the following video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee2d1cf7-ff90-4114-9db2-c8e07f6ba06f",
   "metadata": {},
   "outputs": [],
   "source": [
    "YouTubeVideo(\"VFp38yj8h3A\", width=600, height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa588fb-88e6-440f-abe4-64c2fdab199d",
   "metadata": {},
   "source": [
    "Each pretrained model comes with its own tokenizer, so to get started let's download the tokenizer of a popular model called T5 from the Hub:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17fd8c82-dfb2-4075-8d2f-218f64179982",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint = \"t5-small\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a040c8bb-cf24-4bdd-ab70-98edf94e3db7",
   "metadata": {},
   "source": [
    "The tokenizer has a few interesting attributes such as the vocabulary size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78c38779-661a-4467-b487-6face3cf9b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.vocab_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f980d97b-1e18-4c95-9160-8827b976a51d",
   "metadata": {},
   "source": [
    "This tells us that T5 has 32,100 tokens that is can use to represent text with. Some of the tokens are called special tokens to indicate whether a token is the start or end of a sentence, or corresponds to the mask that is associated with language modeling. \n",
    "\n",
    "When you feed strings to the tokenizer, you'll get at least two fields (some models have more, depending on how they're trained):\n",
    "\n",
    "* `input_ids`: These correspond to the numerical encodings that map each token to an integer\n",
    "* `attention_mask`: This indicates to the model which tokens should be ignored when computing self-attention\n",
    "\n",
    "Let's see how this works with a simple example. First we encode the string:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c55db134-396f-4f97-b21e-6c8d5a4547c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoded_str = tokenizer(\"Albert Einstein lived in Bern\")\n",
    "encoded_str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb7347d-2c94-418c-a04c-6b3b443da3f4",
   "metadata": {},
   "source": [
    "and then decode the input IDs to see the mapping explicitly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8416c56-0b56-4536-a0f7-2ed3afc02ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for token in encoded_str[\"input_ids\"]:\n",
    "    print(token, tokenizer.decode([token]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8acf8a5d-ef31-465f-a9b4-e2b69c5bb69c",
   "metadata": {},
   "source": [
    "So to prepare our inputs, we simply need to apply the tokenizer to each example in our corpus. The only subtlety is that our targets are the paper titles, and these are also strings! So, we'll also need to tokenize them as well. The following function takes care of both these preprocessing steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7970df54-6de2-4825-a28f-62d440e2120c",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_input_length = 1024  # Truncate abstracts longer than this\n",
    "max_target_length = 128  # Truncate titles longer than this\n",
    "prefix = \"summarize: \"  # A special feature of T5 to indicate which task to condition the model on\n",
    "\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    inputs = [prefix + doc for doc in examples[\"abstract\"]]\n",
    "    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)\n",
    "\n",
    "    # Setup the tokenizer for targets\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenizer(\n",
    "            examples[\"title\"], max_length=max_target_length, truncation=True\n",
    "        )\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9707101e-671c-45d0-88a5-b0dae7ed03b5",
   "metadata": {},
   "source": [
    "With this function we can tokenize the whole dataset with a `map()` operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99856309-ca1e-4263-9a30-f6e4f98cfa89",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset = hep_dataset.map(\n",
    "    preprocess_function, batched=True, remove_columns=hep_dataset.column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88bc290-5c1f-47af-8c04-33c512f91a9c",
   "metadata": {},
   "source": [
    "Now that we've tokenized our corpus, it's time to load a pretrained model!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "808a7cf1-1351-4508-8e73-6b22db1357ad",
   "metadata": {},
   "source": [
    "## Loading a pretrained model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfb9538c-d802-44b7-8d06-750752004bbf",
   "metadata": {},
   "source": [
    "To load a pretrained model from the Hub is quite simple: just select the appropriate `AutoModelForXxx` class and use the `from_pretrained()` function with the model checkpoint. In our case, we're dealing with a sequence-to-sequence task (mapping abstracts to titles), so the corresponding autoclass is as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2c91cd-8266-4f27-89d5-97608e76b31a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21687824-6161-4a70-b819-2acb8c96a618",
   "metadata": {},
   "source": [
    "These warnings are perfectly normal - they are telling us that the weights in the head of the network are randomly initialised and so we should fine-tune the model on a downstream task.\n",
    "\n",
    "Now that we have a model, the next step is to initialise a `Trainer` that will take care of the training loop for us (similar to the `Learner` in fastai). Let's do that next."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9fdc81b-e994-400b-9242-27ff12cbec1b",
   "metadata": {},
   "source": [
    "## Creating a Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a8afd86-2fcb-4002-99d2-7b38878c6d41",
   "metadata": {},
   "source": [
    "To create a `Trainer`, we usually need a few basic ingredients:\n",
    "\n",
    "* A `TrainingArguments` class to define all the hyperparameters\n",
    "* A `compute_metrics()` function to compute metrics during evaluation\n",
    "* Datasets to train and evaluate on\n",
    "\n",
    "For more information about the `Trainer` check out the following video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7319d911-6f4a-4ac3-8072-d282ca93e51d",
   "metadata": {},
   "outputs": [],
   "source": [
    "YouTubeVideo(\"nvBXf7s7vTI\", width=600, height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1ea458e-eeaa-4a3c-854a-9f91f7017ce1",
   "metadata": {},
   "source": [
    "Let's start with the `TrainingArguments`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "081d927f-95a8-4ff3-96e4-fd1e66a6aa5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "args = Seq2SeqTrainingArguments(\n",
    "    f\"{model_name}-finetuned-arxiv\",\n",
    "    overwrite_output_dir=True,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    weight_decay=0.01,\n",
    "    save_steps=2000,\n",
    "    num_train_epochs=1,\n",
    "    predict_with_generate=True,\n",
    "    fp16=True,\n",
    "    push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6332391d-44b8-4609-9b60-81006fed17e5",
   "metadata": {},
   "source": [
    "Here we've defined output_dir to save our checkpoints and tweaked some of the default hyperparameters like the learning rate and weight decay. The push_to_hub argument will push each checkpoint to the Hub automatically for us, so we can reuse the model at any point in the future!\n",
    "\n",
    "Now that we've defined the hyperparameters, the next step is to define the metrics. Measuring the performance of text generation tasks like summarization or translation is not as straightforward as classification/regression tasks. For example, given a review like “I loved reading the Hunger Games”, there are multiple valid summaries, like “I loved the Hunger Games” or “Hunger Games is a great read”. Clearly, applying some sort of exact match between the generated summary and the label is not a good solution — even humans would fare poorly under such a metric, because we all have our own writing style.\n",
    "\n",
    "For summarization, one of the most commonly used metrics is the ROUGE score (short for Recall-Oriented Understudy for Gisting Evaluation). The basic idea behind this metric is to compare a generated summary against a set of reference summaries that are typically created by humans. If you want to learn more about this metric, check out the video below - for now it is enough to know that higher ROUGE scores are associated with \"better\" summaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beef7f70-88df-431b-b499-635f842b4192",
   "metadata": {},
   "outputs": [],
   "source": [
    "YouTubeVideo(\"TMshhnrEXlg\", width=600, height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "019ddd58-505c-4528-be69-508ae46518a9",
   "metadata": {},
   "source": [
    "To load a metric, we'll use the `evaluate` library which hosts a wide variety of metrics for machine learning. To load a metric is quite simple:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9afa4ad3-39b5-4381-b41a-46437afacf80",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = load(\"rouge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "931e2522-1b39-4bc1-a5dc-36f99d5aba64",
   "metadata": {},
   "source": [
    "And once we have a metric, we can now compute the ROUGE scores using the `compute()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f03e0c7-db45-4e79-97c6-900aab1ee003",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_summary = \"I absolutely loved reading the Hunger Games\"\n",
    "reference_summary = \"I loved reading the Hunger Games\"\n",
    "scores = metric.compute(predictions=[generated_summary], references=[reference_summary])\n",
    "scores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2754542a-6289-4b28-96fd-89a993ee54dc",
   "metadata": {},
   "source": [
    "Whoa, there’s a lot of information in that output — what does it all mean? First, `evaluate` actually computes confidence intervals for precision, recall, and F1-score; these are the low, mid, and high attributes you can see here. Moreover, `evaluate` computes a variety of ROUGE scores which are based on different types of text granularity when comparing the generated and reference summaries. The rouge1 variant is the overlap of unigrams — this is just a fancy way of saying the overlap of words and is exactly the metric we’ve discussed above. To verify this, let’s pull out the mid value of our scores:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc039540-0d5d-42ff-8234-fcd5babbf208",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores[\"rouge1\"].mid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c6d9575-9e06-475b-8af5-8409ac607c7c",
   "metadata": {},
   "source": [
    "Now what about those other ROUGE scores? rouge2 measures the overlap between bigrams (think the overlap of pairs of words), while rougeL and rougeLsum measure the longest matching sequences of words by looking for the longest common substrings in the generated and reference summaries. The “sum” in rougeLsum refers to the fact that this metric is computed over a whole summary, while rougeL is computed as the average over individual sentences.\n",
    "\n",
    "Next, let's define the `compute_metrics()` function that we'll use to evaluate our model. For summarization this is a bit more involved than simply calling `metric.compute()` on the model’s predictions, since we need to decode the outputs and labels into text before we can compute the ROUGE scores. The following function does exactly that, and also makes use of the `sent_tokenize()` function from nltk to separate the summary sentences with newlines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6c0390-3ef5-4c0d-a20a-42febafec8bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "    # Replace -100 in the labels as we can't decode them.\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # Rouge expects a newline after each sentence\n",
    "    decoded_preds = [\n",
    "        \"\\n\".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds\n",
    "    ]\n",
    "    decoded_labels = [\n",
    "        \"\\n\".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels\n",
    "    ]\n",
    "\n",
    "    result = metric.compute(\n",
    "        predictions=decoded_preds, references=decoded_labels, use_stemmer=True\n",
    "    )\n",
    "    # Extract a few results\n",
    "    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
    "\n",
    "    # Add mean generated length\n",
    "    prediction_lens = [\n",
    "        np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions\n",
    "    ]\n",
    "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
    "\n",
    "    return {k: round(v, 4) for k, v in result.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6244f4b4-c362-4be1-8984-b661e9518801",
   "metadata": {},
   "source": [
    "Next, we need to define a data collator for our sequence-to-sequence task. Since T5 is an encoder-decoder Transformer model, one subtlety with preparing our batches is that during decoding we need to shift the labels to the right by one. This is required to ensure that the decoder only sees the previous ground truth labels and not the current or future ones, which would be easy for the model to memorize. This is similar to how masked self-attention is applied to the inputs in a task like causal language modeling.\n",
    "\n",
    "Luckily, 🤗 Transformers provides a `DataCollatorForSeq2Seq` collator that will dynamically pad the inputs and the labels for us. To instantiate this collator, we simply need to provide the tokenizer and model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61639973-4bba-460c-a713-30f65133ee0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5059e8c2-909a-4a5a-9c84-3ab05017e081",
   "metadata": {},
   "source": [
    "We finally have all the ingredients we need to train with! We now simply need to create a train and test split and instantiate the trainer with the standard arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a18a367-a5f7-4ccd-b191-4303f5fe2d0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env TOKENIZERS_PARALLELISM=false"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ae1b4c-0060-4a76-812f-11b06c524b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = tokenized_dataset.train_test_split(seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "220eaec7-cfe0-41e9-994c-19dfcbea801e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Seq2SeqTrainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"test\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080c81e0-b763-4136-b053-ea0625a12dbd",
   "metadata": {},
   "source": [
    "The only thing left to do is launch our training run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12765a8-cd75-40c1-8f38-a31aa4364da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e98c531-84f7-4fc0-90f7-62d8f3bd7bdf",
   "metadata": {},
   "source": [
    "The final thing to do is push the model weights to the Hub, as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ef6765-3b97-42d3-9127-b92a21edba39",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298e205c-6e4e-4ad8-988a-06b3eba9fa47",
   "metadata": {},
   "source": [
    "## Using your fine-tuned model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "869a3821-9046-463e-8738-71feb796b632",
   "metadata": {},
   "source": [
    "Now that we've pushed our model to the Hub, let's see how it fares on some abstracts that it's never seen before! First, we'll load the model using the `pipeline()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9701fc68-ad6a-4392-b5f6-d5cbf678e127",
   "metadata": {},
   "outputs": [],
   "source": [
    "summarizer = pipeline(\"summarization\", model=\"lewtun/t5-small-finetuned-arxiv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26563de-1646-4def-b662-d4bacf646a7d",
   "metadata": {},
   "source": [
    "And then we can feed it an abstract to summarise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aed0b3c8-559c-48cd-8e34-ef446b23593d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from https://arxiv.org/abs/2206.00678\n",
    "abstract = \"\"\"We demonstrate that the observed cosmological excess of matter over antimatter\n",
    "may originate from a heavy QCD axion that solves the strong CP problem but has a\n",
    "mass much larger than that given by the Standard Model QCD strong dynamics. We\n",
    "investigate a rotation of the heavy QCD axion in field space, which is\n",
    "transferred into a baryon asymmetry through weak and strong sphaleron processes.\n",
    "This provides a strong cosmological motivation for heavy QCD axions, which are\n",
    "of high experimental interest. The viable parameter space has an axion mass ma\n",
    "between 1~MeV and 10 GeV and a decay constant fa<105 GeV, which can be probed by\n",
    "accelerator-based direct axion searches and observations of the cosmic microwave\n",
    "background.\"\"\"\n",
    "\n",
    "summarizer(abstract)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3fd16b-37b1-407b-b070-a8f96cee7727",
   "metadata": {},
   "source": [
    "The original title is _Axiogenesis with a Heavy QCD Axion_, so our model has done a pretty good job!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdbb5628-c819-4469-810e-a55fb4e2c621",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "* Try fine-tuning the model on a different arXiv category like hep-lat. Do the generated titles still make sense?"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}