10_transformers-from-scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and run this cell if you're on Colab or Kaggle\n",
    "# !git clone https://github.com/nlp-with-transformers/notebooks.git\n",
    "# %cd notebooks\n",
    "# from install import *\n",
    "# install_requirements(is_chapter10=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "from utils import *\n",
    "setup_chapter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Transformers from Scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note:** In this chapter a large dataset and the script to train a large language model on a distributed infrastructure are built. As such not all the steps in this notebook are executable on platforms such as Colab or Kaggle. Either downscale the steps at critical points or use this notebook as an inspiration when building a script for distributed training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Large Datasets and Where to Find Them"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Challenges of Building a Large-Scale Corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide_output\n",
    "from transformers import pipeline, set_seed\n",
    "\n",
    "generation_gpt = pipeline(\"text-generation\", model=\"openai-gpt\")\n",
    "generation_gpt2 = pipeline(\"text-generation\", model=\"gpt2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT  size: 116.5M parameters\n",
      "GPT2 size: 124.4M parameters\n"
     ]
    }
   ],
   "source": [
    "def model_size(model):\n",
    "    return sum(t.numel() for t in model.parameters())\n",
    "\n",
    "print(f\"GPT  size: {model_size(generation_gpt.model)/1000**2:.1f}M parameters\")\n",
    "print(f\"GPT2 size: {model_size(generation_gpt2.model)/1000**2:.1f}M parameters\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "set_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT completions:\n",
      "1.\n",
      "When they came back.\n",
      " \" we need all we can get, \" jason said once they had settled into the back of\n",
      "the truck without anyone stopping them. \" after getting out here, it 'll be up\n",
      "to us what to find. for now\n",
      "2.\n",
      "When they came back.\n",
      " his gaze swept over her body. he 'd dressed her, too, in the borrowed clothes\n",
      "that she 'd worn for the journey.\n",
      " \" i thought it would be easier to just leave you there. \" a woman like\n",
      "3.\n",
      "When they came back to the house and she was sitting there with the little boy.\n",
      " \" don't be afraid, \" he told her. she nodded slowly, her eyes wide. she was so\n",
      "lost in whatever she discovered that tom knew her mistake\n",
      "\n",
      "GPT-2 completions:\n",
      "1.\n",
      "When they came back we had a big dinner and the other guys went to see what\n",
      "their opinion was on her. I did an hour and they were happy with it.\n",
      "2.\n",
      "When they came back to this island there had been another massacre, but he could\n",
      "not help but feel pity for the helpless victim who had been left to die, and\n",
      "that they had failed that day. And so was very, very grateful indeed.\n",
      "3.\n",
      "When they came back to our house after the morning, I asked if she was sure. She\n",
      "said, \"Nope.\" The two kids were gone that morning. I thought they were back to\n",
      "being a good friend.\n",
      "\n",
      "When Dost\n"
     ]
    }
   ],
   "source": [
    "def enum_pipeline_ouputs(pipe, prompt, num_return_sequences):\n",
    "    out = pipe(prompt, num_return_sequences=num_return_sequences,\n",
    "               clean_up_tokenization_spaces=True)\n",
    "    return \"\\n\".join(f\"{i+1}.\" + s[\"generated_text\"] for i, s in enumerate(out))\n",
    "\n",
    "prompt = \"\\nWhen they came back\"\n",
    "print(\"GPT completions:\\n\" + enum_pipeline_ouputs(generation_gpt, prompt, 3))\n",
    "print(\"\")\n",
    "print(\"GPT-2 completions:\\n\" + enum_pipeline_ouputs(generation_gpt2, prompt, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a Custom Code Dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating a dataset with Google BigQuery"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#sidebar To Filter the Noise or Not?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Working with Large Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Memory mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note:** The following code block assumes that you have downloaded the BigQuery dataset to a folder called `codeparrot`. We suggest skipping this step since it will unpack the compressed files and require ~180GB of disk space. This code is just for demonstration purposes and you can just continue below with the streamed dataset which will not consume that much disk space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide_output\n",
    "from datasets import load_dataset, DownloadConfig\n",
    "\n",
    "download_config = DownloadConfig(delete_extracted=True)\n",
    "dataset = load_dataset(\"./codeparrot\", split=\"train\",\n",
    "                       download_config=download_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of python files code in dataset : 18695559\n",
      "Dataset size (cache file) : 183.68 GB\n",
      "RAM memory used: 4924 MB\n"
     ]
    }
   ],
   "source": [
    "import psutil, os\n",
    "\n",
    "print(f\"Number of python files code in dataset : {len(dataset)}\")\n",
    "ds_size = sum(os.stat(f[\"filename\"]).st_size for f in dataset.cache_files)\n",
    "# os.stat.st_size is expressed in bytes, so we convert to GB\n",
    "print(f\"Dataset size (cache file) : {ds_size / 2**30:.2f} GB\")\n",
    "# Process.memory_info is expressed in bytes, so we convert to MB\n",
    "print(f\"RAM used: {psutil.Process(os.getpid()).memory_info().rss >> 20} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-cae7a1d2f0dbde67\n"
     ]
    }
   ],
   "source": [
    "# hide_output\n",
    "streamed_dataset = load_dataset('./codeparrot', split=\"train\", streaming=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "iterator = iter(streamed_dataset)\n",
    "\n",
    "print(dataset[0] == next(iterator))\n",
    "print(dataset[1] == next(iterator))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remote_dataset = load_dataset('transformersbook/codeparrot', split=\"train\",\n",
    "                              streaming=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Datasets to the Hugging Face Hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29ced71e91434126970160a03cc006a5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.17k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f437f06babc4f01b5f02ed2e11a274f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/773k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "086fad17475145a2960cf393d6da4da5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.32M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83b90218ddd54563b5abc524ed820741",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/508 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8fb62d059f0a4ab397767ec6497ab5ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/792k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1fd604f95db44748a1b53a31be9ff5a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.33M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# hide_output\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "def tok_list(tokenizer, string):\n",
    "    input_ids = tokenizer(string, add_special_tokens=False)[\"input_ids\"]\n",
    "    return [tokenizer.decode(tok) for tok in input_ids]\n",
    "\n",
    "tokenizer_T5 = AutoTokenizer.from_pretrained(\"t5-base\")\n",
    "tokenizer_camembert = AutoTokenizer.from_pretrained(\"camembert-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "T5 tokens for \"sex\": ['', 's', 'ex']\n",
      "CamemBERT tokens for \"being\": ['be', 'ing']\n"
     ]
    }
   ],
   "source": [
    "print(f'T5 tokens for \"sex\": {tok_list(tokenizer_T5,\"sex\")}')\n",
    "print(f'CamemBERT tokens for \"being\": {tok_list(tokenizer_camembert,\"being\")}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Tokenizer Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Measuring Tokenizer Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Tokenizer for Python "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['def', 'Ġsay', '_', 'hello', '():', 'Ċ', 'Ġ', 'Ġ', 'Ġ', 'Ġprint', '(\"',\n",
      "'Hello', ',', 'ĠWorld', '!\"', ')', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_',\n",
      "'hello', '()', 'Ċ']\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "python_code = r\"\"\"def say_hello():\n",
    "    print(\"Hello, World!\")\n",
    "# Print it\n",
    "say_hello()\n",
    "\"\"\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "print(tokenizer(python_code).tokens())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.backend_tokenizer.normalizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('def', (0, 3)), ('Ġsay', (3, 7)), ('_', (7, 8)), ('hello', (8, 13)), ('():',\n",
      "(13, 16)), ('ĊĠĠĠ', (16, 20)), ('Ġprint', (20, 26)), ('(\"', (26, 28)), ('Hello',\n",
      "(28, 33)), (',', (33, 34)), ('ĠWorld', (34, 40)), ('!\")', (40, 43)), ('Ġ#', (43,\n",
      "45)), ('ĠPrint', (45, 51)), ('Ġit', (51, 54)), ('Ċ', (54, 55)), ('Ċ', (55, 56)),\n",
      "('say', (56, 59)), ('_', (59, 60)), ('hello', (60, 65)), ('()', (65, 67)), ('Ċ',\n",
      "(67, 68))]\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(python_code))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "`a` is encoded as `b'a'` with a single byte: 97\n",
      "`€` is encoded as `b'\\xe2\\x82\\xac'` with three bytes: [226, 130, 172]\n"
     ]
    }
   ],
   "source": [
    "a, e = u\"a\", u\"€\"\n",
    "byte = ord(a.encode(\"utf-8\"))\n",
    "print(f'`{a}` is encoded as `{a.encode(\"utf-8\")}` with a single byte: {byte}')\n",
    "byte = [ord(chr(i)) for i in e.encode(\"utf-8\")]\n",
    "print(f'`{e}` is encoded as `{e.encode(\"utf-8\")}` with three bytes: {byte}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of our base vocabulary: 256\n",
      "First element: `!`, last element: `Ń`\n"
     ]
    }
   ],
   "source": [
    "from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode\n",
    "\n",
    "byte_to_unicode_map = bytes_to_unicode()\n",
    "unicode_to_byte_map = dict((v, k) for k, v in byte_to_unicode_map.items())\n",
    "base_vocab = list(unicode_to_byte_map.keys())\n",
    "\n",
    "print(f'Size of our base vocabulary: {len(base_vocab)}')\n",
    "print(f'First element: `{base_vocab[0]}`, last element: `{base_vocab[-1]}`')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Description</th>\n",
       "      <th>Character</th>\n",
       "      <th>Bytes</th>\n",
       "      <th>Mapped bytes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>Regular characters</td>\n",
       "      <td>`a` and `?`</td>\n",
       "      <td>97 and 63</td>\n",
       "      <td>`a` and `?`</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Non-printable control character (CARRIAGE RETURN)</td>\n",
       "      <td>`U+000D`</td>\n",
       "      <td>13</td>\n",
       "      <td>`č`</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>A space</td>\n",
       "      <td>` `</td>\n",
       "      <td>32</td>\n",
       "      <td>`Ġ`</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>A non-breakable space</td>\n",
       "      <td>`\\xa0`</td>\n",
       "      <td>160</td>\n",
       "      <td>`ł`</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>A newline character</td>\n",
       "      <td>`\\n`</td>\n",
       "      <td>10</td>\n",
       "      <td>`Ċ`</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# hide_input\n",
    "#id unicode_mapping\n",
    "#caption Examples of character mappings in BPE\n",
    "#hide_input\n",
    "import pandas as pd\n",
    "from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode\n",
    "\n",
    "byte_to_unicode_map = bytes_to_unicode()\n",
    "unicode_to_byte_map = dict((v, k) for k, v in byte_to_unicode_map.items())\n",
    "base_vocab = list(unicode_to_byte_map.keys())\n",
    "\n",
    "examples = [\n",
    "    ['Regular characters', '`a` and `?`', f'{ord(\"a\")} and {ord(\"?\")}' , f'`{byte_to_unicode_map[ord(\"a\")]}` and `{byte_to_unicode_map[ord(\"?\")]}`'],\n",
    "    ['Nonprintable control character (carriage return)', '`U+000D`', f'13', f'`{byte_to_unicode_map[13]}`'],\n",
    "    ['A space', '` `', f'{ord(\" \")}', f'`{byte_to_unicode_map[ord(\" \")]}`'],\n",
    "    ['A nonbreakable space', '`\\\\xa0`', '160', f'`{byte_to_unicode_map[ord(chr(160))]}`'],\n",
    "    ['A newline character', '`\\\\n`', '10', f'`{byte_to_unicode_map[ord(chr(10))]}`'],\n",
    "]\n",
    "\n",
    "pd.DataFrame(examples, columns = ['Description', 'Character', 'Bytes', 'Mapped bytes'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('def', (0, 3)), ('Ġsay', (3, 7)), ('_', (7, 8)), ('hello', (8, 13)), ('():',\n",
      "(13, 16)), ('ĊĠĠĠ', (16, 20)), ('Ġprint', (20, 26)), ('(\"', (26, 28)), ('Hello',\n",
      "(28, 33)), (',', (33, 34)), ('ĠWorld', (34, 40)), ('!\")', (40, 43)), ('Ġ#', (43,\n",
      "45)), ('ĠPrint', (45, 51)), ('Ġit', (51, 54)), ('Ċ', (54, 55)), ('Ċ', (55, 56)),\n",
      "('say', (56, 59)), ('_', (59, 60)), ('hello', (60, 65)), ('()', (65, 67)), ('Ċ',\n",
      "(67, 68))]\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(python_code))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of the vocabulary: 50257\n"
     ]
    }
   ],
   "source": [
    "print(f\"Size of the vocabulary: {len(tokenizer)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['def', 'Ġsay', '_', 'hello', '():', 'Ċ', 'Ġ', 'Ġ', 'Ġ', 'Ġprint', '(\"',\n",
      "'Hello', ',', 'ĠWorld', '!\"', ')', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_',\n",
      "'hello', '()', 'Ċ']\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer(python_code).tokens())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training a Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂ', '\n",
      "=================================================================', '\n",
      "----------------------------------------------------------------',\n",
      "'................................................................',\n",
      "'ÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂ',\n",
      "'----------------------------------------------------------------',\n",
      "'================================================================',\n",
      "'________________________________________________________________']\n"
     ]
    }
   ],
   "source": [
    "tokens = sorted(tokenizer.vocab.items(), key=lambda x: len(x[0]), reverse=True)\n",
    "print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[:8]]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<|endoftext|>', ' gazed', ' informants', ' Collider', ' regress', 'ominated',\n",
      "' amplification', 'Compar', '….\"', ' (/', 'Commission', ' Hitman']\n"
     ]
    }
   ],
   "source": [
    "tokens = sorted(tokenizer.vocab.items(), key=lambda x: x[1], reverse=True)\n",
    "print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[:12]]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "743bca69d71649908db9ca5760af61d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Check remote data files:   0%|          | 0/183 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration codeparrot-train-99775fd6743284b5\n"
     ]
    }
   ],
   "source": [
    "#hide_output\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "length = 100000\n",
    "dataset_name = 'transformersbook/codeparrot-train'\n",
    "dataset = load_dataset(dataset_name, split=\"train\", streaming=True)\n",
    "iter_dataset = iter(dataset)\n",
    "\n",
    "def batch_iterator(batch_size=10):\n",
    "    for _ in tqdm(range(0, length, batch_size)):\n",
    "        yield [next(iter_dataset)['content'] for _ in range(batch_size)]\n",
    "\n",
    "new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(), \n",
    "                                                  vocab_size=12500,\n",
    "                                                  initial_alphabet=base_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['  ', '    ', '   ', '        ', 'se', 'in', '       ', 're', 'on', 'te', '\\n\n",
      "', '\\n        ', 'or', 'st', 'de', '\\n   ', 'th', 'le', ' =', 'lf', 'self',\n",
      "'me', 'al']\n"
     ]
    }
   ],
   "source": [
    "tokens = sorted(new_tokenizer.vocab.items(), key=lambda x: x[1], reverse=False)\n",
    "print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[257:280]]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[' capt', ' embedded', ' regarding', 'Bundle', '355', ' recv', ' dmp', ' vault',\n",
      "' Mongo', ' possibly', 'implementation', 'Matches']\n"
     ]
    }
   ],
   "source": [
    "print([f'{new_tokenizer.convert_tokens_to_string(t)}' for t,_ in tokens[-12:]]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['def', 'Ġs', 'ay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '(\"', 'Hello', ',',\n",
      "'ĠWor', 'ld', '!\")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 's', 'ay', '_', 'hello',\n",
      "'()', 'Ċ']\n"
     ]
    }
   ],
   "source": [
    "print(new_tokenizer(python_code).tokens())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are in total 35 Python keywords.\n",
      "No, keyword `await` is not in the vocabulary\n",
      "No, keyword `finally` is not in the vocabulary\n",
      "No, keyword `nonlocal` is not in the vocabulary\n"
     ]
    }
   ],
   "source": [
    "import keyword\n",
    "\n",
    "print(f'There are in total {len(keyword.kwlist)} Python keywords.')\n",
    "for keyw in keyword.kwlist:\n",
    "    if keyw not in new_tokenizer.vocab:\n",
    "        print(f'No, keyword `{keyw}` is not in the vocabulary')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [05:08<00:00,  1.54s/it]\n"
     ]
    }
   ],
   "source": [
    "# hide_output\n",
    "length = 200000\n",
    "new_tokenizer_larger = tokenizer.train_new_from_iterator(batch_iterator(),\n",
    "    vocab_size=32768, initial_alphabet=base_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['lineEdit', 'spik', ' BC', 'pective', 'OTA', 'theus', 'FLUSH', ' excutils',\n",
      "'00000002', ' DIVISION', 'CursorPosition', ' InfoBar']\n"
     ]
    }
   ],
   "source": [
    "tokens = sorted(new_tokenizer_larger.vocab.items(), key=lambda x: x[1],\n",
    "                reverse=False)\n",
    "print([f'{tokenizer.convert_tokens_to_string(t)}' for t, _ in tokens[-12:]]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['def', 'Ġsay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '(\"', 'Hello', ',',\n",
      "'ĠWorld', '!\")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_', 'hello', '()',\n",
      "'Ċ']\n"
     ]
    }
   ],
   "source": [
    "print(new_tokenizer_larger(python_code).tokens())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No, keyword `nonlocal` is not in the vocabulary\n"
     ]
    }
   ],
   "source": [
    "for keyw in keyword.kwlist:\n",
    "    if keyw not in new_tokenizer_larger.vocab:\n",
    "        print(f'No, keyword `{keyw}` is not in the vocabulary')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving a Custom Tokenizer on the Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/transformersbook/codeparrot into local empty directory.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'https://huggingface.co/transformersbook/codeparrot/commit/1c284adaa3cc9f8635ae7e3377bd3739f48bc09a'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#hide_output\n",
    "model_ckpt = \"codeparrot\"\n",
    "org = \"transformersbook\"\n",
    "new_tokenizer_larger.push_to_hub(model_ckpt, organization=org)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['def', 'Ġsay', '_', 'hello', '():', 'ĊĠĠĠ', 'Ġprint', '(\"', 'Hello', ',',\n",
      "'ĠWorld', '!\")', 'Ġ#', 'ĠPrint', 'Ġit', 'Ċ', 'Ċ', 'say', '_', 'hello', '()',\n",
      "'Ċ']\n"
     ]
    }
   ],
   "source": [
    "reloaded_tokenizer = AutoTokenizer.from_pretrained(org + \"/\" + model_ckpt)\n",
    "print(reloaded_tokenizer(python_code).tokens())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/transformersbook/codeparrot-small-vocabulary into local empty directory.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'https://huggingface.co/transformersbook/codeparrot-small-vocabulary/commit/0b37bed9956d95d0b79ada169f6a281e15c63381'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#hide_output\n",
    "new_tokenizer.push_to_hub(model_ckpt+ \"-small-vocabulary\", organization=org)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a Model from Scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Tale of Pretraining Objectives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Causal language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Masked language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sequence-to-sequence training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **NOTE**: In the following code block, a large GPT-2 checkpoint is loaded into memory. On platforms like Colab and Kaggle, this can cause the instance to crash due to insufficient RAM or GPU memory. You can still run the example if you use the small checkpoint by replacing the configuration with `config = AutoConfig.from_pretrained(\"gpt2\", vocab_size=len(tokenizer))`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be84ca77ca144954af8ae4820ec6685b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/787 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#hide_output\n",
    "from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(org + \"/\" + model_ckpt)\n",
    "config = AutoConfig.from_pretrained(\"gpt2-xl\", vocab_size=len(tokenizer))\n",
    "model = AutoModelForCausalLM.from_config(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT-2 (xl) size: 1529.6M parameters\n"
     ]
    }
   ],
   "source": [
    "print(f'GPT-2 (xl) size: {model_size(model)/1000**2:.1f}M parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide_output\n",
    "model.save_pretrained(\"models/\" + model_ckpt, push_to_hub=True,\n",
    "                      organization=org)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_ckpt)\n",
    "config_small = AutoConfig.from_pretrained(\"gpt2\", vocab_size=len(tokenizer))\n",
    "model_small = AutoModelForCausalLM.from_config(config_small)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT-2 size: 111.0M parameters\n"
     ]
    }
   ],
   "source": [
    "print(f'GPT-2 size: {model_size(model_small)/1000**2:.1f}M parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/transformersbook/codeparrot-small into local empty directory.\n"
     ]
    }
   ],
   "source": [
    "#hide_output\n",
    "model_small.save_pretrained(\"models/\" + model_ckpt + \"-small\", push_to_hub=True,\n",
    "                            organization=org)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing the Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/500 [00:00<01:16,  6.54it/s]Token indices sequence length is longer than the specified maximum sequence length for this model (2605 > 1024). Running this sequence through the model will result in indexing errors\n",
      "100%|██████████| 500/500 [00:04<00:00, 122.59it/s]\n"
     ]
    }
   ],
   "source": [
    "#hide_output\n",
    "examples, total_characters, total_tokens = 500, 0, 0\n",
    "dataset = load_dataset('transformersbook/codeparrot-train', split='train',\n",
    "                       streaming=True)\n",
    "\n",
    "for _, example in tqdm(zip(range(examples), iter(dataset)), total=examples):\n",
    "    total_characters += len(example['content'])\n",
    "    total_tokens += len(tokenizer(example['content']).tokens())\n",
    "\n",
    "characters_per_token = total_characters / total_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6233025034779565\n"
     ]
    }
   ],
   "source": [
    "print(characters_per_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import IterableDataset\n",
    "\n",
    "class ConstantLengthDataset(IterableDataset):\n",
    "    \n",
    "    def __init__(self, tokenizer, dataset, seq_length=1024,\n",
    "                 num_of_sequences=1024, chars_per_token=3.6):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.concat_token_id = tokenizer.eos_token_id\n",
    "        self.dataset = dataset\n",
    "        self.seq_length = seq_length\n",
    "        self.input_characters = seq_length * chars_per_token * num_of_sequences\n",
    "    \n",
    "    def __iter__(self):\n",
    "        iterator = iter(self.dataset)\n",
    "        more_examples = True\n",
    "        while more_examples:\n",
    "            buffer, buffer_len = [], 0\n",
    "            while True:\n",
    "                if buffer_len >= self.input_characters:\n",
    "                    m=f\"Buffer full: {buffer_len}>={self.input_characters:.0f}\"\n",
    "                    print(m)\n",
    "                    break\n",
    "                try:\n",
    "                    m=f\"Fill buffer: {buffer_len}<{self.input_characters:.0f}\"\n",
    "                    print(m)\n",
    "                    buffer.append(next(iterator)[\"content\"])\n",
    "                    buffer_len += len(buffer[-1])\n",
    "                except StopIteration:\n",
    "                    iterator = iter(self.dataset)\n",
    "\n",
    "            all_token_ids = []\n",
    "            tokenized_inputs = self.tokenizer(buffer, truncation=False)\n",
    "            for tokenized_input in tokenized_inputs['input_ids']:\n",
    "                all_token_ids.extend(tokenized_input + [self.concat_token_id])\n",
    "            \n",
    "            for i in range(0, len(all_token_ids), self.seq_length):\n",
    "                input_ids = all_token_ids[i : i + self.seq_length]\n",
    "                if len(input_ids) == self.seq_length:\n",
    "                    yield torch.tensor(input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fill buffer: 0<36864\n",
      "Fill buffer: 3311<36864\n",
      "Fill buffer: 9590<36864\n",
      "Fill buffer: 22177<36864\n",
      "Fill buffer: 25530<36864\n",
      "Fill buffer: 31098<36864\n",
      "Fill buffer: 32232<36864\n",
      "Fill buffer: 33867<36864\n",
      "Buffer full: 41172>=36864\n",
      "Lengths of the sequences: [1024, 1024, 1024, 1024, 1024]\n"
     ]
    }
   ],
   "source": [
    "shuffled_dataset = dataset.shuffle(buffer_size=100)\n",
    "constant_length_dataset = ConstantLengthDataset(tokenizer, shuffled_dataset,\n",
    "                                                num_of_sequences=10)\n",
    "dataset_iterator = iter(constant_length_dataset)\n",
    "\n",
    "lengths = [len(b) for _, b in zip(range(5), dataset_iterator)]\n",
    "print(f\"Lengths of the sequences: {lengths}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import Namespace\n",
    "\n",
    "# Commented parameters correspond to the small model\n",
    "config = {\"train_batch_size\": 2, # 12\n",
    "          \"valid_batch_size\": 2, # 12\n",
    "          \"weight_decay\": 0.1,\n",
    "          \"shuffle_buffer\": 1000,\n",
    "          \"learning_rate\": 2e-4, # 5e-4\n",
    "          \"lr_scheduler_type\": \"cosine\",\n",
    "          \"num_warmup_steps\": 750, # 2000\n",
    "          \"gradient_accumulation_steps\": 16, # 1\n",
    "          \"max_train_steps\": 50000, # 150000\n",
    "          \"max_eval_steps\": -1,\n",
    "          \"seq_length\": 1024,\n",
    "          \"seed\": 1,\n",
    "          \"save_checkpoint_steps\": 50000} # 15000\n",
    "\n",
    "args = Namespace(**config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import logging\n",
    "import wandb\n",
    "\n",
    "def setup_logging(project_name):\n",
    "    logger = logging.getLogger(__name__)\n",
    "    logging.basicConfig(\n",
    "        format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
    "        datefmt=\"%m/%d/%Y %H:%M:%S\", level=logging.INFO, handlers=[\n",
    "        logging.FileHandler(f\"log/debug_{accelerator.process_index}.log\"),\n",
    "        logging.StreamHandler()])\n",
    "    if accelerator.is_main_process: # We only want to set up logging once\n",
    "        wandb.init(project=project_name, config=args)\n",
    "        run_name = wandb.run.name\n",
    "        tb_writer = SummaryWriter()\n",
    "        tb_writer.add_hparams(vars(args), {'0': 0})\n",
    "        logger.setLevel(logging.INFO)\n",
    "        datasets.utils.logging.set_verbosity_debug()\n",
    "        transformers.utils.logging.set_verbosity_info()\n",
    "    else:\n",
    "        tb_writer = None\n",
    "        run_name = ''\n",
    "        logger.setLevel(logging.ERROR)\n",
    "        datasets.utils.logging.set_verbosity_error()\n",
    "        transformers.utils.logging.set_verbosity_error()\n",
    "    return logger, tb_writer, run_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_metrics(step, metrics):\n",
    "    logger.info(f\"Step {step}: {metrics}\")\n",
    "    if accelerator.is_main_process:\n",
    "        wandb.log(metrics)\n",
    "        [tb_writer.add_scalar(k, v, step) for k, v in metrics.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "328dc6d7d05c452e8d8e2cab5b4b9c4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Check remote data files:   0%|          | 0/183 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration codeparrot-train-938ce362e6f661b1\n",
      "Using custom data configuration codeparrot-valid-29167601d8e69487\n"
     ]
    }
   ],
   "source": [
    "#hide_output\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "def create_dataloaders(dataset_name):\n",
    "    train_data = load_dataset(dataset_name+'-train', split=\"train\",\n",
    "                              streaming=True)\n",
    "    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer,\n",
    "                                    seed=args.seed)\n",
    "    valid_data = load_dataset(dataset_name+'-valid', split=\"validation\",\n",
    "                              streaming=True)\n",
    "    \n",
    "    train_dataset = ConstantLengthDataset(tokenizer, train_data,\n",
    "                                          seq_length=args.seq_length)\n",
    "    valid_dataset = ConstantLengthDataset(tokenizer, valid_data,\n",
    "                                          seq_length=args.seq_length)\n",
    "    \n",
    "    train_dataloader=DataLoader(train_dataset, batch_size=args.train_batch_size)\n",
    "    eval_dataloader=DataLoader(valid_dataset, batch_size=args.valid_batch_size)\n",
    "    return train_dataloader, eval_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_grouped_params(model, no_decay=[\"bias\", \"LayerNorm.weight\"]):\n",
    "    params_with_wd, params_without_wd = [], []\n",
    "    for n, p in model.named_parameters():\n",
    "        if any(nd in n for nd in no_decay):\n",
    "            params_without_wd.append(p)\n",
    "        else:\n",
    "            params_with_wd.append(p)\n",
    "    return [{'params': params_with_wd, 'weight_decay': args.weight_decay},\n",
    "            {'params': params_without_wd, 'weight_decay': 0.0}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate():\n",
    "    model.eval()\n",
    "    losses = []\n",
    "    for step, batch in enumerate(eval_dataloader):\n",
    "        with torch.no_grad():\n",
    "            outputs = model(batch, labels=batch)\n",
    "        loss = outputs.loss.repeat(args.valid_batch_size)\n",
    "        losses.append(accelerator.gather(loss))\n",
    "        if args.max_eval_steps > 0 and step >= args.max_eval_steps: break\n",
    "    loss = torch.mean(torch.cat(losses))\n",
    "    try:\n",
    "\t\tperplexity = torch.exp(loss)\n",
    "\texcept OverflowError:\n",
    "\t\tperplexity = torch.tensor(float(\"inf\"))\n",
    "    return loss.item(), perplexity.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(args.seed)\n",
    "\n",
    "# Accelerator\n",
    "accelerator = Accelerator()\n",
    "samples_per_step = accelerator.state.num_processes * args.train_batch_size\n",
    "\n",
    "# Logging\n",
    "logger, tb_writer, run_name = setup_logging(project_name.split(\"/\")[1])\n",
    "logger.info(accelerator.state)\n",
    "\n",
    "# Load model and tokenizer\n",
    "if accelerator.is_main_process:\n",
    "    hf_repo = Repository(\"./\", clone_from=project_name, revision=run_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(\"./\", gradient_checkpointing=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"./\")\n",
    "\n",
    "# Load dataset and dataloader\n",
    "train_dataloader, eval_dataloader = create_dataloaders(dataset_name)\n",
    "\n",
    "# Prepare the optimizer and learning rate scheduler\n",
    "optimizer = AdamW(get_grouped_params(model), lr=args.learning_rate)\n",
    "lr_scheduler = get_scheduler(name=args.lr_scheduler_type, optimizer=optimizer,\n",
    "                             num_warmup_steps=args.num_warmup_steps,\n",
    "                             num_training_steps=args.max_train_steps,)\n",
    "def get_lr():\n",
    "    return optimizer.param_groups[0]['lr']\n",
    "\n",
    "# Prepare everything with our `accelerator` (order of args is not important)\n",
    "model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n",
    "    model, optimizer, train_dataloader, eval_dataloader)\n",
    "\n",
    "# Train model\n",
    "model.train()\n",
    "completed_steps = 0\n",
    "for step, batch in enumerate(train_dataloader, start=1):\n",
    "    loss = model(batch, labels=batch).loss\n",
    "    log_metrics(step, {'lr': get_lr(), 'samples': step*samples_per_step,\n",
    "                       'steps': completed_steps, 'loss/train': loss.item()})\n",
    "    loss = loss / args.gradient_accumulation_steps\n",
    "    accelerator.backward(loss)\n",
    "    if step % args.gradient_accumulation_steps == 0:\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        completed_steps += 1\n",
    "    if step % args.save_checkpoint_steps == 0:\n",
    "        logger.info('Evaluating and saving model checkpoint')\n",
    "        eval_loss, perplexity = evaluate()\n",
    "        log_metrics(step, {'loss/eval': eval_loss, 'perplexity': perplexity})\n",
    "        accelerator.wait_for_everyone()\n",
    "        unwrapped_model = accelerator.unwrap_model(model)\n",
    "        if accelerator.is_main_process:\n",
    "            unwrapped_model.save_pretrained(\"./\")\n",
    "            hf_repo.push_to_hub(commit_message=f'step {step}')\n",
    "        model.train()\n",
    "    if completed_steps >= args.max_train_steps:\n",
    "        break\n",
    "\n",
    "# Evaluate and save the last checkpoint\n",
    "logger.info('Evaluating and saving model after training')\n",
    "eval_loss, perplexity = evaluate()\n",
    "log_metrics(step, {'loss/eval': eval_loss, 'perplexity': perplexity})\n",
    "accelerator.wait_for_everyone()\n",
    "unwrapped_model = accelerator.unwrap_model(model)\n",
    "if accelerator.is_main_process:\n",
    "    unwrapped_model.save_pretrained(\"./\")\n",
    "    hf_repo.push_to_hub(commit_message=f'final model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Training Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results and Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-10-20 18:29:01.107727: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-10-20 18:29:01.107759: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n"
     ]
    }
   ],
   "source": [
    "#hide_output\n",
    "from transformers import pipeline, set_seed\n",
    "\n",
    "model_ckpt = 'transformersbook/codeparrot-small'\n",
    "generation = pipeline('text-generation', model=model_ckpt, device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from transformers import set_seed \n",
    "\n",
    "def first_block(string):\n",
    "    return re.split('\\nclass|\\ndef|\\n#|\\n@|\\nprint|\\nif', string)[0].rstrip()\n",
    "\n",
    "def complete_code(pipe, prompt, max_length=64, num_completions=4, seed=1):\n",
    "    set_seed(seed)\n",
    "    gen_kwargs = {\"temperature\":0.4, \"top_p\":0.95, \"top_k\":0, \"num_beams\":1,\n",
    "                  \"do_sample\":True,}\n",
    "    code_gens = generation(prompt, num_return_sequences=num_completions, \n",
    "                            max_length=max_length, **gen_kwargs)\n",
    "    code_strings = []\n",
    "    for code_gen in code_gens:\n",
    "        generated_code = first_block(code_gen['generated_text'][len(prompt):])\n",
    "        code_strings.append(generated_code)\n",
    "    print(('\\n'+'='*80 + '\\n').join(code_strings))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    return math.sqrt(a * b)\n",
      "================================================================================\n",
      "\n",
      "    return a * b / 2.0\n",
      "================================================================================\n",
      "\n",
      "    return a * b\n",
      "================================================================================\n",
      "\n",
      "    return a * b / a\n"
     ]
    }
   ],
   "source": [
    "prompt = '''def area_of_rectangle(a: float, b: float):\n",
    "    \"\"\"Return the area of the rectangle.\"\"\"'''\n",
    "complete_code(generation, prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    if not html:\n",
      "        return []\n",
      "    return [url for url in re.findall(r'<a href=\"(/[^/]+/[^\"]+?)\">', html)]\n",
      "================================================================================\n",
      "\n",
      "    return [url for url in re.findall(r'<a href=\"(.*?)\"', html)\n",
      "            if url]\n",
      "================================================================================\n",
      "\n",
      "    return [url for url in re.findall(r'<a href=\"(/.*)\",', html)]\n",
      "================================================================================\n",
      "\n",
      "    return re.findall(r'<a href=\"(.*?)\" class=\"url\"[^>]*>', html)\n"
     ]
    }
   ],
   "source": [
    "prompt = '''def get_urls_from_html(html):\n",
    "    \"\"\"Get all embedded URLs in a HTML string.\"\"\"'''\n",
    "complete_code(generation, prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://github.com/huggingface/transformers | /allenai | /facebook |\n",
      "/asteroid-team | /google | /amazon | /speechbrain | /microsoft | /grammarly |\n",
      "/models | /inference-api | /distilbert-base-uncased |\n",
      "/dbmdz/bert-large-cased-finetuned-conll03-english |\n",
      "https://huggingface.co/transformers | https://arxiv.org/abs/1811.06031 |\n",
      "https://arxiv.org/abs/1803.10631 | https://transformer.huggingface.co/ | /coref\n",
      "| https://medium.com/huggingface/distilbert-8cf3380435b5\n"
     ]
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "def get_urls_from_html(html):\n",
    "    return [url for url in re.findall(r'<a href=\"(.*?)\"', html) if url]\n",
    "\n",
    "print(\" | \".join(get_urls_from_html(requests.get('https://hf.co/').text)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **NOTE**: In the following code block, a large GPT-2 checkpoint is loaded into memory. On platforms like Colab and Kaggle, this can cause the instance to crash due to insufficient RAM or GPU memory. You can still run the example if you replace the large model with the small one by using `model_ckpt = \"transformersbook/codeparrot-small\"`.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    return np.mean(a)\n",
      "================================================================================\n",
      "\n",
      "    return np.mean(a)\n",
      "================================================================================\n",
      "\n",
      "    return np.mean(a)\n",
      "================================================================================\n",
      "\n",
      "    return np.mean(a)\n"
     ]
    }
   ],
   "source": [
    "model_ckpt = 'transformersbook/codeparrot'\n",
    "generation = pipeline('text-generation', model=model_ckpt, device=0)\n",
    "\n",
    "prompt = '''# a function in native python:\n",
    "def mean(a):\n",
    "    return sum(a)/len(a)\n",
    "\n",
    "# the same function using numpy:\n",
    "import numpy as np\n",
    "def mean(a):'''\n",
    "complete_code(generation, prompt, max_length=64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "reg = DummyRegressor()\n",
      "\n",
      "forest = RandomForestClassifier(n_estimators=20)\n",
      "\n",
      "forest.fit(X, y)\n",
      "================================================================================\n",
      "\n",
      "clf = ExtraTreesClassifier(n_estimators=100, max_features='sqrt')\n",
      "clf.fit(X, y)\n",
      "================================================================================\n",
      "\n",
      "clf = RandomForestClassifier(n_estimators=20, n_jobs=n_jobs, random_state=1)\n",
      "clf.fit(X, y)\n",
      "================================================================================\n",
      "\n",
      "clf = RandomForestClassifier(n_estimators=20)\n",
      "clf.fit(X, y)\n"
     ]
    }
   ],
   "source": [
    "prompt = '''X = np.random.randn(100, 100)\n",
    "y = np.random.randint(0, 1, 100)\n",
    "\n",
    "# fit random forest classifier with 20 estimators'''\n",
    "complete_code(generation, prompt, max_length=96)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}