examples/language_modeling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X4cRE8IbIrIV"
   },
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets. Uncomment the following cell and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "MOsHUjgdIrIW",
    "outputId": "f84a093e-147f-470e-aad9-80fb51193c8e"
   },
   "outputs": [],
   "source": [
    "#! pip install datasets transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this notebook locally, make sure your environment has an install from the last version of those libraries.\n",
    "\n",
    "To be able to share your model with the community and generate results like the one shown in the picture below via the inference API, 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](https://huggingface.co/join) if you haven't already!) then execute the following cell and input your username and password:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you need to install Git-LFS. Uncomment the following instructions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !apt install git-lfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure your version of Transformers is at least 4.11.0 since the functionality was introduced in that version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "\n",
    "print(transformers.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HFASsisvIrIb"
   },
   "source": [
    "You can find a script version of this notebook to fine-tune your model in a distributed fashion using multiple GPUs or TPUs [here](https://github.com/huggingface/transformers/tree/master/examples/language-modeling)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also quickly upload some telemetry - this tells us which examples and software versions are getting used so we know where to prioritize our maintenance efforts. We don't collect (or care about) any personally identifiable information, but if you'd prefer not to be counted, feel free to skip this step or delete this cell entirely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.utils import send_example_telemetry\n",
    "\n",
    "send_example_telemetry(\"language_modeling_notebook\", framework=\"pytorch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a3KD3WXU3l-O"
   },
   "source": [
    "# Fine-tuning a language model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JAscNNUD3l-P"
   },
   "source": [
    "In this notebook, we'll see how to fine-tune one of the [🤗 Transformers](https://github.com/huggingface/transformers) model on a language modeling tasks. We will cover two types of language modeling tasks which are:\n",
    "\n",
    "- Causal language modeling: the model has to predict the next token in the sentence (so the labels are the same as the inputs shifted to the right). To make sure the model does not cheat, it gets an attention mask that will prevent it to access the tokens after token i when trying to predict the token i+1 in the sentence.\n",
    "\n",
    "![Widget inference representing the causal language modeling task](images/causal_language_modeling.png)\n",
    "\n",
    "- Masked language modeling: the model has to predict some tokens that are masked in the input. It still has access to the whole sentence, so it can use the tokens before and after the tokens masked to predict their value.\n",
    "\n",
    "![Widget inference representing the masked language modeling task](images/masked_language_modeling.png)\n",
    "\n",
    "We will see how to easily load and preprocess the dataset for each one of those tasks, and how to use the `Trainer` API to fine-tune a model on it.\n",
    "\n",
    "A script version of this notebook you can directly run on a distributed environment or on TPU is available in our [examples folder](https://github.com/huggingface/transformers/tree/master/examples)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1r_n9OWV3l-Q"
   },
   "source": [
    "## Preparing the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kswRMhPc3l-Q"
   },
   "source": [
    "For each of those tasks, we will use the [Wikitext 2]() dataset as an example. You can load it very easily with the 🤗 Datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n2ZRs1cL3l-R",
    "outputId": "11151c56-be90-4d11-e7df-db85e745ca5c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset wikitext (/home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f1-9jepM3l-W"
   },
   "source": [
    "You can replace the dataset above with any dataset hosted on [the hub](https://huggingface.co/datasets) or use your own files. Just uncomment the following cell and replace the paths with values that will lead to your files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uxSaGa_l3l-W"
   },
   "outputs": [],
   "source": [
    "# datasets = load_dataset(\"text\", data_files={\"train\": path_to_train.txt, \"validation\": path_to_validation.txt}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jY1SwIrY3l-a"
   },
   "source": [
    "You can also load datasets from a csv or a JSON file, see the [full documentation](https://huggingface.co/docs/datasets/loading_datasets.html#from-local-files) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u3EtYfeHIrIz"
   },
   "source": [
    "To access an actual element, you need to select a split first, then give an index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X6HrpprwIrIz",
    "outputId": "d7670bc0-42e4-4c09-8a6a-5c018ded7d95"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': ' The game \\'s battle system , the BliTZ system , is carried over directly from Valkyira Chronicles . During missions , players select each unit using a top @-@ down perspective of the battlefield map : once a character is selected , the player moves the character around the battlefield in third @-@ person . A character can only act once per @-@ turn , but characters can be granted multiple turns at the expense of other characters \\' turns . Each character has a field and distance of movement limited by their Action Gauge . Up to nine characters can be assigned to a single mission . During gameplay , characters will call out if something happens to them , such as their health points ( HP ) getting low or being knocked out by enemy attacks . Each character has specific \" Potentials \" , skills unique to each character . They are divided into \" Personal Potential \" , which are innate skills that remain unaltered unless otherwise dictated by the story and can either help or impede a character , and \" Battle Potentials \" , which are grown throughout the game and always grant boons to a character . To learn Battle Potentials , each character has a unique \" Masters Table \" , a grid @-@ based skill table that can be used to acquire and link different skills . Characters also have Special Abilities that grant them temporary boosts on the battlefield : Kurt can activate \" Direct Command \" and move around the battlefield without depleting his Action Point gauge , the character Reila can shift into her \" Valkyria Form \" and become invincible , while Imca can target multiple enemy units with her heavy weapon . \\n'}"
      ]
     },
     "execution_count": null,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"train\"][10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WHUmphG3IrI3"
   },
   "source": [
    "To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ur5sNUcZ3l-g"
   },
   "outputs": [],
   "source": [
    "from datasets import ClassLabel\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "def show_random_elements(dataset, num_examples=10):\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, 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": "1Uk8NROQ3l-k",
    "outputId": "a822dcec-51e3-4dba-c73c-dba9e0301726"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DuVall joined Alice in Chains as lead singer during the band 's reunion concerts . Duff McKagan again joined the band for the reunion tour , playing rhythm guitar on selected songs . Before the tour , Kinney mentioned in an interview that he would be interested in writing new material , but not as Alice in Chains . However , AliceinChains.com reported that the band had begun writing new material , with DuVall on lead vocals . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>In chess tournaments and matches , the frequency with which each player receives white and black is an important consideration . In matches , the players ' colors in the first game are determined by drawing lots , and alternated thereafter . In round robin tournaments with an odd number of players , each player receives an equal number of whites and blacks ; with an even number of players , each receives one extra white or black . Where one or more players withdraws from the tournament , the tournament director may change the assigned colors in some games so that no player receives two more blacks than whites , or vice versa . The double @-@ round robin tournament is considered to give the most reliable final standings , since each player receives the same number of whites and blacks , and plays both White and Black against each opponent . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>= = = Rich Rodriguez era = = = \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>= = = = First siege of Kaifeng = = = = \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Furthermore , it appears that Iguanodon became more quadrupedal as it got older and heavier ; juvenile I. bernissartensis have shorter arms than adults ( 60 % of hindlimb length versus 70 % for adults ) . When walking as a quadruped , the animal 's hands would have been held so that the palms faced each other , as shown by iguanodontian trackways and the anatomy of this genus 's arms and hands . The three toed pes ( foot ) of Iguanodon was relatively long , and when walking , both the hand and the foot would have been used in a digitigrade fashion ( walking on the fingers and toes ) . The maximum speed of Iguanodon has been estimated at 24 km / h ( 15 mph ) , which would have been as a biped ; it would not have been able to gallop as a quadruped . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>With little guidance on water allocation from the Supreme Court , proponents of the dam feared endless litigation . A Colorado attorney proposed that the seven states which fell within the river 's basin ( California , Nevada , Arizona , Utah , New Mexico , Colorado and Wyoming ) form an interstate compact , with the approval of Congress . Such compacts were authorized by Article I of the United States Constitution but had never been concluded among more than two states . In 1922 , representatives of seven states met with then @-@ Secretary of Commerce Herbert Hoover . Initial talks produced no result , but when the Supreme Court handed down the Wyoming v. Colorado decision undermining the claims of the upstream states , they became anxious to reach an agreement . The resulting Colorado River Compact was signed on November 24 , 1922 . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>He goes on to note that , once freed from the possession , the victim will never again be able to eat tofu , azukimeshi , or other foods favored by foxes : \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>The only legal marriage was between Kody and his first wife Meri , until their legal divorce in September 2014 . ( He legally married fourth wife Robyn in December 2014 in order to legally adopt her three children ) . The other marriages are considered spiritual unions . As of 2015 Kody has been married to Meri for 25 years , Janelle for 22 years , Christine for 21 years , and Robyn for 5 years . Kody and Meri have a daughter named Mariah , their only child . Kody and Janelle have six children : daughters Madison and Savanah and sons Logan , Hunter , Garrison , and Gabriel . Kody and Christine have six children : daughters Aspyn , Mykelti , Gwendlyn , Ysabel , and Truely and son Paedon . Robyn had three children from her first marriage , which was monogamous : Dayton , Aurora , and Breanna . Kody legally adopted them in June 2015 . Kody and Robyn have two children : son Solomon and daughter Ariella . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>= = Ranks and insignia = = \\n</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_random_elements(datasets[\"train\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CKerdF353l-o"
   },
   "source": [
    "As we can see, some of the texts are a full paragraph of a Wikipedia article while others are just titles or empty lines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JEA1ju653l-p"
   },
   "source": [
    "## Causal Language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v5GTGKZS3l-q"
   },
   "source": [
    "For causal language modeling (CLM) we are going to take all the texts in our dataset and concatenate them after they are tokenized. Then we will split them in examples of a certain sequence length. This way the model will receive chunks of contiguous text that may look like:\n",
    "```\n",
    "part of text 1\n",
    "```\n",
    "or \n",
    "```\n",
    "end of text 1 [BOS_TOKEN] beginning of text 2\n",
    "```\n",
    "depending on whether they span over several of the original texts in the dataset or not. The labels will be the same as the inputs, shifted to the left.\n",
    "\n",
    "We will use the [`distilgpt2`](https://huggingface.co/distilgpt2) model for this example. You can pick any of the checkpoints listed [here](https://huggingface.co/models?filter=causal-lm) instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-WGBCO343l-q"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"distilgpt2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5io6fY_d3l-u"
   },
   "source": [
    "To tokenize all our texts with the same vocabulary that was used when training the model, we have to download a pretrained tokenizer. This is all done by the `AutoTokenizer` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iAYlS40Z3l-v"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "    \n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rpOiBrJ13l-y"
   },
   "source": [
    "We can now call the tokenizer on all our texts. This is very simple, using the [`map`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map) method from the Datasets library. First we define a function that call the tokenizer on our texts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lS2m25YM3l-z"
   },
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M9xVAa3s3l-2"
   },
   "source": [
    "Then we apply it to all the splits in our `datasets` object, using `batched=True` and 4 processes to speed up the preprocessing. We won't need the `text` column afterward, so we discard it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NVAO0H8u3l-3",
    "outputId": "30d88b8a-e353-4e13-f709-8e5e06ef747b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-0a686d6f64cb210f.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-659bcb80cad0097c.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-7f22912475d34c88.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-b3566e2fe9c5c036.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-a7e5c438d154ad4d.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-012753d728f48615.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-1d058a410941ea13.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-d3f134f9a6ca08b5.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-1fac954e1a542924.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-f3cc2c4468f35c7e.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-997632d4bc92501b.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-657bf68aa1caa201.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8qik3J_C3l-7"
   },
   "source": [
    "If we now look at an element of our datasets, we will see the text have been replaced by the `input_ids` the model will need:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nYv_mcKk3l-7",
    "outputId": "8334734c-0f86-4e18-ec17-4216a2d5dd18"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       " 'input_ids': [796, 569, 18354, 7496, 17740, 6711, 796, 220, 198]}"
      ]
     },
     "execution_count": null,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets[\"train\"][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "obvgcXda3l--"
   },
   "source": [
    "Now for the harder part: we need to concatenate all our texts together then split the result in small chunks of a certain `block_size`. To do this, we will use the `map` method again, with the option `batched=True`. This option actually lets us change the number of examples in the datasets by returning a different number of examples than we got. This way, we can create our new samples from a batch of examples.\n",
    "\n",
    "First, we grab the maximum length our model was pretrained with. This might be a big too big to fit in your GPU RAM, so here we take a bit less at just 128."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DVHs5aCA3l-_"
   },
   "outputs": [],
   "source": [
    "# block_size = tokenizer.model_max_length\n",
    "block_size = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RpNfGiMw3l_A"
   },
   "source": [
    "Then we write the preprocessing function that will group our texts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iaAJy5Hu3l_B"
   },
   "outputs": [],
   "source": [
    "def group_texts(examples):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n",
    "        # customize this part to your needs.\n",
    "    total_length = (total_length // block_size) * block_size\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LGJWXtNv3l_C"
   },
   "source": [
    "First note that we duplicate the inputs for our labels. This is because the model of the 🤗 Transformers library apply the shifting to the right, so we don't need to do it manually.\n",
    "\n",
    "Also note that by default, the `map` method will send a batch of 1,000 examples to be treated by the preprocessing function. So here, we will drop the remainder to make the concatenated tokenized texts a multiple of `block_size` every 1,000 examples. You can adjust this behavior by passing a higher batch size (which will also be processed slower). You can also speed-up the preprocessing by using multiprocessing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gXUSfBrq3l_C",
    "outputId": "34e55885-3d8f-4f05-cbdb-706ce56a25f8"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-da77bf362d4c6fa4.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-7d08a6d62516c9ff.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-a985b575c96ddae3.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-47fffef35acafddb.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-9281db6733359e2e.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-6c06eb34e3ec189c.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-21ab2ae5196fbb0c.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-b014c0146a937326.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-3dbc3c09d8c329a2.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-a92a7c440f530426.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-9f8b3fe14e22eebc.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-91c30819fb1d1769.arrow\n"
     ]
    }
   ],
   "source": [
    "lm_datasets = tokenized_datasets.map(\n",
    "    group_texts,\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6n84V8Gc3l_G"
   },
   "source": [
    "And we can check our datasets have changed: now the samples contain chunks of `block_size` contiguous tokens, potentially spanning over several of our original texts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hTeGCLl_3l_G",
    "outputId": "ab381a07-f92e-4b14-f7b6-e4af5513d7c4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' game and follows the \" Nameless \", a penal military unit serving the nation of Gallia during the Second Europan War who perform secret black operations and are pitted against the Imperial unit \" Calamaty Raven \". \\n The game began development in 2010, carrying over a large portion of the work done on Valkyria Chronicles II. While it retained the standard features of the series, it also underwent multiple adjustments, such as making the game more forgiving for series newcomers. Character designer Raita Honjou and composer Hitoshi Sakimoto both returned from previous entries, along with Valkyria Chronicles II director Takeshi Oz'"
      ]
     },
     "execution_count": null,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(lm_datasets[\"train\"][1][\"input_ids\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iEmeQ7Xm3l_H"
   },
   "source": [
    "Now that the data has been cleaned, we're ready to instantiate our `Trainer`. We will a model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sPqQA3TT3l_I"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "model = AutoModelForCausalLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VyPQTOF_3l_J"
   },
   "source": [
    "And some `TrainingArguments`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jElf8LJ33l_K"
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YbSwEhQ63l_L"
   },
   "outputs": [],
   "source": [
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "training_args = TrainingArguments(\n",
    "    f\"{model_name}-finetuned-wikitext2\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last argument to setup everything so we can push the model to the [Hub](https://huggingface.co/models) regularly during training. Remove it if you didn't follow the installation steps at the top of the notebook. If you want to save your model locally in a name that is different than the name of the repository it will be pushed, 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/gpt-finetuned-wikitext2\"` or `\"huggingface/gpt-finetuned-wikitext2\"`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sZRbT9ui3l_N"
   },
   "source": [
    "We pass along all of those to the `Trainer` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OEuqwIra3l_N"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=lm_datasets[\"train\"],\n",
    "    eval_dataset=lm_datasets[\"validation\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Vvz34Td3l_O"
   },
   "source": [
    "And we can train our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NyZvu_MF3l_P",
    "outputId": "b69d0931-7f1f-4f2d-fdb8-09d37c7418bb"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "        </style>\n",
       "      \n",
       "      <progress value='7002' max='7002' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7002/7002 09:05, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>3.767790</td>\n",
       "      <td>3.664047</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>3.648033</td>\n",
       "      <td>3.645540</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>3.594949</td>\n",
       "      <td>3.641928</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=7002, training_loss=3.69408126695944)"
      ]
     },
     "execution_count": null,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3APq-vUc3l_R"
   },
   "source": [
    "Once the training is completed, we can evaluate our model and get its perplexity on the validation set like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "diKZnB1I3l_R",
    "outputId": "9b3ac725-0117-4830-f380-a555ee57c8cf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Perplexity: 38.17\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "eval_results = trainer.evaluate()\n",
    "print(f\"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wY82caEX3l_i"
   },
   "source": [
    "You can now upload the result of the training to the Hub, just execute this instruction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now share this model with all your friends, family, favorite pets: they can all load it with the identifier `\"your-username/the-name-you-picked\"` so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"sgugger/my-awesome-model\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q-EIELH43l_T"
   },
   "source": [
    "## Masked language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LWk97-Ny3l_T"
   },
   "source": [
    "For masked language modeling (MLM) we are going to use the same preprocessing as before for our dataset with one additional step: we will randomly mask some tokens (by replacing them by `[MASK]`) and the labels will be adjusted to only include the masked tokens (we don't have to predict the non-masked tokens).\n",
    "\n",
    "We will use the [`distilroberta-base`](https://huggingface.co/distilroberta-base) model for this example. You can pick any of the checkpoints listed [here](https://huggingface.co/models?filter=masked-lm) instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QRTpmyCc3l_T"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"distilroberta-base\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "12F1ulgT3l_V"
   },
   "source": [
    "We can apply the same tokenization function as before, we just need to update our tokenizer to use the checkpoint we just picked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h8RCYcvr3l_V",
    "outputId": "a5ffeb0a-71da-4b27-e57a-c62f1927562e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-333e4baa6f280a66.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-23acd0930cc16da7.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-56ae8ad41a9fdf19.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-599a47a0e666ad65.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-5ba855475cae1614.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-3aa432d7bf07952c.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-75711d299cf5397c.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-46595b48d78f0089.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-9bd1f40499aceff5.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-8dcfc6147d5fa31f.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-2995af94c706ccfc.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-6e6c35e31dcb428c.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)\n",
    "tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MTuy8UUs3l_X"
   },
   "source": [
    "And like before, we group texts together and chunk them in samples of length `block_size`. You can skip that step if your dataset is composed of individual sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LVYPMwEs3l_X",
    "outputId": "e71ed7f1-b182-4643-a8fb-3d731c70e40b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-661796332aa2b576.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-e019d91824c225fd.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-b5875c725d0e5cb7.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-a8e3eeaa703ca023.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-5b01f983be82f7d4.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-f21374c1e117e44a.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-9423e584c93db0b1.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-519591986bd5de0e.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-1567968e11ef94e5.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-505a19c89d31abc3.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-1ab0ccc328411d63.arrow\n",
      "Loading cached processed dataset at /home/sgugger/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/47c57a6745aa5ce8e16a5355aaa4039e3aa90d1adad87cef1ad4e0f29e74ac91/cache-1c817d18f58aea63.arrow\n"
     ]
    }
   ],
   "source": [
    "lm_datasets = tokenized_datasets.map(\n",
    "    group_texts,\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nFJ49iHJ3l_Z"
   },
   "source": [
    "The rest is very similar to what we had, with two exceptions. First we use a model suitable for masked LM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PM10A9Za3l_Z",
    "outputId": "fff2d5bb-397d-4d5d-9aa9-933090cb6680"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaForMaskedLM were not initialized from the model checkpoint at distilroberta-base and are newly initialized: ['lm_head.decoder.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForMaskedLM\n",
    "model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We redefine our `TrainingArguments`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YbSwEhQ63l_L"
   },
   "outputs": [],
   "source": [
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "training_args = TrainingArguments(\n",
    "    f\"{model_name}-finetuned-wikitext2\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like before, the last argument to setup everything so we can push the model to the [Hub](https://huggingface.co/models) regularly during training. Remove it if you didn't follow the installation steps at the top of the notebook. If you want to save your model locally in a name that is different than the name of the repository it will be pushed, 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/bert-finetuned-wikitext2\"` or `\"huggingface/bert-finetuned-wikitext2\"`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z6uuUnvz3l_b"
   },
   "source": [
    "Finally, we use a special `data_collator`. The `data_collator` is a function that is responsible of taking the samples and batching them in tensors. In the previous example, we had nothing special to do, so we just used the default for this argument. Here we want to do the random-masking. We could do it as a pre-processing step (like the tokenization) but then the tokens would always be masked the same way at each epoch. By doing this step inside the `data_collator`, we ensure this random masking is done in a new way each time we go over the data.\n",
    "\n",
    "To do this masking for us, the library provides a `DataCollatorForLanguageModeling`. We can adjust the probability of the masking:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nRZ-5v_P3l_b"
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bqHnWcYC3l_d"
   },
   "source": [
    "Then we just have to pass everything to `Trainer` and begin training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V-Y3gNqV3l_d"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=lm_datasets[\"train\"],\n",
    "    eval_dataset=lm_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y9TFqDG_3l_e",
    "outputId": "2e0c8bca-0e04-4b4f-ad06-8dd320af6c37"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "        </style>\n",
       "      \n",
       "      <progress value='7218' max='7218' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7218/7218 09:11, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.109144</td>\n",
       "      <td>1.919501</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.986566</td>\n",
       "      <td>1.884295</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.952879</td>\n",
       "      <td>1.861102</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=7218, training_loss=2.0377309222603213)"
      ]
     },
     "execution_count": null,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KDBi0reX3l_g"
   },
   "source": [
    "Like before, we can evaluate our model on the validation set. The perplexity is much lower than for the CLM objective because for the MLM objective, we only have to make predictions for the masked tokens (which represent 15% of the total here) while having access to the rest of the tokens. It's thus an easier task for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4hSaANqj3l_g",
    "outputId": "eeeb8727-2e27-4aeb-ac71-c98123214661"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "        </style>\n",
       "      \n",
       "      <progress value='249' max='249' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [249/249 00:04]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Perplexity: 6.37\n"
     ]
    }
   ],
   "source": [
    "eval_results = trainer.evaluate()\n",
    "print(f\"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wY82caEX3l_i"
   },
   "source": [
    "You can now upload the result of the training to the Hub, just execute this instruction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now share this model with all your friends, family, favorite pets: they can all load it with the identifier `\"your-username/the-name-you-picked\"` so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import AutoModelForMaskedLM\n",
    "\n",
    "model = AutoModelForMaskedLM.from_pretrained(\"sgugger/my-awesome-model\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Fine-tune a language model",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}