scripts/wikipedia_train_from_scratch_lm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train RoBERTa from Scratch\n",
    "\n",
    "This notebook is based on a [tutorial](https://huggingface.co/blog/how-to-train) from the Hugging Face Transformers Library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G-kkz81OY6xH"
   },
   "source": [
    "## Train a tokenizer\n",
    "\n",
    "We choose to train a byte-level Byte-pair encoding tokenizer (the same as GPT-2), with the same special tokens as RoBERTa. Let’s arbitrarily pick its size to be 52,000.\n",
    "\n",
    "We recommend training a byte-level BPE (rather than let’s say, a WordPiece tokenizer like BERT) because it will start building its vocabulary from an alphabet of single bytes, so all words will be decomposable into tokens (no more `<unk>` tokens!).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jzmr6jY3R2sj",
    "outputId": "770d9a69-0b7a-46d2-f09f-c6457cd95de5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mounted at /content/drive\n",
      "/content/drive/MyDrive\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "%cd drive/MyDrive/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5duRggBRZKvP",
    "outputId": "b8ba0942-1222-48da-bc29-15563f50cd6e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/huggingface/transformers\n",
      "  Cloning https://github.com/huggingface/transformers to /tmp/pip-req-build-2kwmka7j\n",
      "  Running command git clone -q https://github.com/huggingface/transformers /tmp/pip-req-build-2kwmka7j\n",
      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
      "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
      "Collecting tokenizers==0.9.4\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/0f/1c/e789a8b12e28be5bc1ce2156cf87cb522b379be9cadc7ad8091a4cc107c4/tokenizers-0.9.4-cp36-cp36m-manylinux2010_x86_64.whl (2.9MB)\n",
      "\u001b[K     |████████████████████████████████| 2.9MB 5.2MB/s \n",
      "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (2.23.0)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (20.8)\n",
      "Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (0.8)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (3.0.12)\n",
      "Collecting sacremoses\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7d/34/09d19aff26edcc8eb2a01bed8e98f13a1537005d31e95233fd48216eed10/sacremoses-0.0.43.tar.gz (883kB)\n",
      "\u001b[K     |████████████████████████████████| 890kB 46.5MB/s \n",
      "\u001b[?25hRequirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (4.41.1)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (1.19.4)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.6/dist-packages (from transformers==4.2.0.dev0) (2019.12.20)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==4.2.0.dev0) (1.24.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==4.2.0.dev0) (2020.12.5)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==4.2.0.dev0) (3.0.4)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==4.2.0.dev0) (2.10)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from packaging->transformers==4.2.0.dev0) (2.4.7)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==4.2.0.dev0) (1.15.0)\n",
      "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==4.2.0.dev0) (7.1.2)\n",
      "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==4.2.0.dev0) (1.0.0)\n",
      "Building wheels for collected packages: transformers\n",
      "  Building wheel for transformers (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for transformers: filename=transformers-4.2.0.dev0-cp36-none-any.whl size=1527284 sha256=ccb5d2bd601a0949412c0555fe9efcc3a911ee411aa8b4d6c6b760cfebf0a8d9\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-xu32q26l/wheels/70/d3/52/b3fa4f8b8ef04167ac62e5bb2accb62ae764db2a378247490e\n",
      "Successfully built transformers\n",
      "Building wheels for collected packages: sacremoses\n",
      "  Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for sacremoses: filename=sacremoses-0.0.43-cp36-none-any.whl size=893261 sha256=502f977d15b0d8f39b4cdbc1685d537fab7854430ba3f9cb3c69105bb130cf25\n",
      "  Stored in directory: /root/.cache/pip/wheels/29/3c/fd/7ce5c3f0666dab31a50123635e6fb5e19ceb42ce38d4e58f45\n",
      "Successfully built sacremoses\n",
      "Installing collected packages: tokenizers, sacremoses, transformers\n",
      "Successfully installed sacremoses-0.0.43 tokenizers-0.9.4 transformers-4.2.0.dev0\n",
      "tokenizers                    0.9.4          \n",
      "transformers                  4.2.0.dev0     \n",
      "Collecting textacy\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/65/99/054efc5dea92c84a850639c490541de6cba29bc148debc3c73848c5e64c2/textacy-0.10.1-py3-none-any.whl (183kB)\n",
      "\u001b[K     |████████████████████████████████| 184kB 6.5MB/s \n",
      "\u001b[?25hRequirement already satisfied: spacy<3.0.0,>=2.2.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (2.2.4)\n",
      "Requirement already satisfied: pyemd>=0.5.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (0.5.1)\n",
      "Requirement already satisfied: numpy>=1.17.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (1.19.4)\n",
      "Requirement already satisfied: tqdm>=4.19.6 in /usr/local/lib/python3.6/dist-packages (from textacy) (4.41.1)\n",
      "Requirement already satisfied: requests>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (2.23.0)\n",
      "Requirement already satisfied: scikit-learn<0.24.0,>=0.19.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (0.22.2.post1)\n",
      "Requirement already satisfied: cachetools>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from textacy) (4.2.0)\n",
      "Collecting pyphen>=0.9.4\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7c/5a/5bc036e01389bc6a6667a932bac3e388de6e7fa5777a6ff50e652f60ec79/Pyphen-0.10.0-py3-none-any.whl (1.9MB)\n",
      "\u001b[K     |████████████████████████████████| 1.9MB 6.7MB/s \n",
      "\u001b[?25hRequirement already satisfied: srsly>=0.0.5 in /usr/local/lib/python3.6/dist-packages (from textacy) (1.0.5)\n",
      "Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (2.5)\n",
      "Collecting cytoolz>=0.8.0\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/58/67/1c60da8ba831bfefedb64c78b9f6820bdf58972797c95644ee3191daf27a/cytoolz-0.11.0.tar.gz (477kB)\n",
      "\u001b[K     |████████████████████████████████| 481kB 26.5MB/s \n",
      "\u001b[?25hRequirement already satisfied: scipy>=0.17.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (1.4.1)\n",
      "Requirement already satisfied: joblib>=0.13.0 in /usr/local/lib/python3.6/dist-packages (from textacy) (1.0.0)\n",
      "Collecting jellyfish>=0.7.0\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/6c/09/927ae35fc5a9f70abb6cc2c27ee88fc48549f7bc4786c1d4b177c22e997d/jellyfish-0.8.2-cp36-cp36m-manylinux2014_x86_64.whl (93kB)\n",
      "\u001b[K     |████████████████████████████████| 102kB 10.3MB/s \n",
      "\u001b[?25hRequirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (0.4.1)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (51.0.0)\n",
      "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (0.8.0)\n",
      "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (1.1.3)\n",
      "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (1.0.0)\n",
      "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (7.4.0)\n",
      "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (3.0.5)\n",
      "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (1.0.5)\n",
      "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy<3.0.0,>=2.2.0->textacy) (2.0.5)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.10.0->textacy) (1.24.3)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.10.0->textacy) (2.10)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.10.0->textacy) (3.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.10.0->textacy) (2020.12.5)\n",
      "Requirement already satisfied: decorator>=4.3.0 in /usr/local/lib/python3.6/dist-packages (from networkx>=2.0->textacy) (4.4.2)\n",
      "Requirement already satisfied: toolz>=0.8.0 in /usr/local/lib/python3.6/dist-packages (from cytoolz>=0.8.0->textacy) (0.11.1)\n",
      "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy<3.0.0,>=2.2.0->textacy) (3.3.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy<3.0.0,>=2.2.0->textacy) (3.4.0)\n",
      "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy<3.0.0,>=2.2.0->textacy) (3.7.4.3)\n",
      "Building wheels for collected packages: cytoolz\n",
      "  Building wheel for cytoolz (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for cytoolz: filename=cytoolz-0.11.0-cp36-cp36m-linux_x86_64.whl size=1225555 sha256=09d6f852067f93e3e4f09027b6c60faf46e20c1d062162c6483b86dc93bd0954\n",
      "  Stored in directory: /root/.cache/pip/wheels/a1/32/3c/9c9926b510647cacdde744b2c7acdf1ccd5896fbb7f8d5df0c\n",
      "Successfully built cytoolz\n",
      "Installing collected packages: pyphen, cytoolz, jellyfish, textacy\n",
      "Successfully installed cytoolz-0.11.0 jellyfish-0.8.2 pyphen-0.10.0 textacy-0.10.1\n"
     ]
    }
   ],
   "source": [
    "# We won't need TensorFlow here\n",
    "# !pip install -y tensorflow\n",
    "# Install `transformers` from master\n",
    "!pip install git+https://github.com/huggingface/transformers\n",
    "!pip list | grep -E 'transformers|tokenizers'\n",
    "!pip install textacy\n",
    "# transformers version at notebook update --- 2.11.0\n",
    "# tokenizers version at notebook update --- 0.8.0rc1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sIbyx59Q8w_J"
   },
   "outputs": [],
   "source": [
    "# load wikipedia file and make corpus. This cell takes many hours to run because of the size of the wikipedia dump.\n",
    "# for smaller wikipedia dumps, look here: https://dumps.wikimedia.org/enwiki/latest/\n",
    "import sys\n",
    "from gensim.corpora import WikiCorpus\n",
    "\n",
    "def make_corpus(in_f, out_f):\n",
    "\n",
    "    \"\"\"Convert wikipedia xml dump file to text corpus\"\"\"\n",
    "\n",
    "    output = open(out_f, 'w')\n",
    "    print('starting corpus')\n",
    "    wiki = WikiCorpus(in_f)\n",
    "    print('loaded corpus')\n",
    "\n",
    "    i = 0\n",
    "    for text in wiki.get_texts():\n",
    "    output.write(bytes(' '.join(text), 'utf-8').decode('utf-8') + '\\n')\n",
    "    i = i + 1\n",
    "    if (i % 100 == 0):\n",
    "        print('Processed ' + str(i) + ' articles')\n",
    "    output.close()\n",
    "    print('Processing complete!')\n",
    " \n",
    "# !wget -c https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2\n",
    "# make_corpus('enwiki-latest-pages-articles.xml.bz2', 'wiki_en_full.txt')\n",
    "\n",
    "# !wget -c https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles10.xml-p4045403p5399366.bz2\n",
    "# make_corpus('enwiki-latest-pages-articles10.xml-p4045403p5399366.bz2', 'wiki_en_10.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IMnymRDLe0hi",
    "outputId": "97b00e3f-6ada-44ed-c5b4-827d7d5097b8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7min 47s, sys: 7.45 s, total: 7min 54s\n",
      "Wall time: 2min 11s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "from pathlib import Path\n",
    "\n",
    "from tokenizers import ByteLevelBPETokenizer\n",
    "\n",
    "paths = ['wiki_en_10.txt']\n",
    "\n",
    "# Initialize a tokenizer\n",
    "tokenizer = ByteLevelBPETokenizer()\n",
    "\n",
    "# Customize training\n",
    "tokenizer.train(files=paths, vocab_size=52_000, min_frequency=2, special_tokens=[\n",
    "    \"<s>\",\n",
    "    \"<pad>\",\n",
    "    \"</s>\",\n",
    "    \"<unk>\",\n",
    "    \"<mask>\",\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EIS-irI0f32P",
    "outputId": "cc068b8b-2fa1-4d33-9731-7763cf67f469"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘wikipedia_model_3e’: File exists\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['wikipedia_model_3e/vocab.json', 'wikipedia_model_3e/merges.txt']"
      ]
     },
     "execution_count": 6,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "!mkdir wikipedia_model_3e\n",
    "tokenizer.save_model(\"wikipedia_model_3e\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "t7L6wh8Iyj25",
    "outputId": "ff8cdb3e-1d52-4a46-bf4a-4b91e24c069d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/content/drive/MyDrive/lm_training_experiments\n"
     ]
    }
   ],
   "source": [
    "%cd lm_training_experiments/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "tKVWB8WShT-z"
   },
   "outputs": [],
   "source": [
    "from tokenizers.implementations import ByteLevelBPETokenizer\n",
    "from tokenizers.processors import BertProcessing\n",
    "\n",
    "\n",
    "tokenizer = ByteLevelBPETokenizer(\n",
    "    \"./wikipedia_model_3e/vocab.json\",\n",
    "    \"./wikipedia_model_3e/merges.txt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "hO5M3vrAhcuj"
   },
   "outputs": [],
   "source": [
    "tokenizer._tokenizer.post_processor = BertProcessing(\n",
    "    (\"</s>\", tokenizer.token_to_id(\"</s>\")),\n",
    "    (\"<s>\", tokenizer.token_to_id(\"<s>\")),\n",
    ")\n",
    "tokenizer.enable_truncation(max_length=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "X8ya5_7rhjKS",
    "outputId": "56a6aaf6-2c04-45d1-aae1-e8f3386edc22"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<s>', 'I', 'Ġam', 'Ġa', 'Ġfireman', '.', '</s>']"
      ]
     },
     "execution_count": 6,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.encode(\"I am a fireman.\").tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WQpUC_CDhnWW"
   },
   "source": [
    "## Train a language model from scratch\n",
    "\n",
    "**Update:** This section follows along the [`run_language_modeling.py`](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_language_modeling.py) script, using our new [`Trainer`](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py) directly. Feel free to pick the approach you like best.\n",
    "\n",
    "> We’ll train a RoBERTa-like model, which is a BERT-like with a couple of changes (check the [documentation](https://huggingface.co/transformers/model_doc/roberta.html) for more details).\n",
    "\n",
    "As the model is BERT-like, we’ll train it on a task of *Masked language modeling*, i.e. the predict how to fill arbitrary tokens that we randomly mask in the dataset. This is taken care of by the example script.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kD140sFjh0LQ",
    "outputId": "41fafdb5-4e00-4619-eb4b-f53a8cc908ed"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue Dec 29 01:18:41 2020       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 460.27.04    Driver Version: 418.67       CUDA Version: 10.1     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |\n",
      "| N/A   39C    P0    24W / 300W |      0MiB / 16130MiB |      0%      Default |\n",
      "|                               |                      |                 ERR! |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "# Check that we have a GPU\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VNZZs-r6iKAV",
    "outputId": "3bcaf01e-b4e5-4f13-c97a-26bbf29b3f42"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check that PyTorch sees it\n",
    "import torch\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u0qQzgrBi1OX"
   },
   "source": [
    "### We'll define the following config for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "LTXXutqeDzPi"
   },
   "outputs": [],
   "source": [
    "from transformers import RobertaConfig\n",
    "\n",
    "config = RobertaConfig(\n",
    "    vocab_size=52_000,\n",
    "    max_position_embeddings=514,\n",
    "    num_attention_heads=12,\n",
    "    num_hidden_layers=6,\n",
    "    type_vocab_size=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yAwQ82JiE5pi"
   },
   "source": [
    "Now let's re-create our tokenizer in transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "4keFBUjQFOD1"
   },
   "outputs": [],
   "source": [
    "from transformers import RobertaTokenizerFast\n",
    "\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"./wikipedia_model_3e\", max_len=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6yNCw-3hFv9h"
   },
   "source": [
    "Finally let's initialize our model.\n",
    "\n",
    "**Important:**\n",
    "\n",
    "As we are training from scratch, we only initialize from a config, not from an existing pretrained model or checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "BzMqR-dzF4Ro"
   },
   "outputs": [],
   "source": [
    "from transformers import RobertaForMaskedLM\n",
    "\n",
    "model = RobertaForMaskedLM(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jU6JhBSTKiaM",
    "outputId": "64fc0516-0e2e-4657-867b-f7e34c30e8ce"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "83504416"
      ]
     },
     "execution_count": 10,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.num_parameters()\n",
    "# => 84 million parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jBtUHRMliOLM"
   },
   "source": [
    "### Now let's build our training Dataset\n",
    "\n",
    "We'll build our dataset by applying our tokenizer to our text file.\n",
    "\n",
    "Here, as we only have one text file, we don't even need to customize our `Dataset`. We'll just use the `LineByLineDataset` out-of-the-box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GlvP_A-THEEl",
    "outputId": "ae2d1ef0-0dbd-4997-dfea-8ff1c922829c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/transformers/data/datasets/language_modeling.py:128: FutureWarning: This dataset will be removed from the library soon, preprocessing should be handled with the 🤗 Datasets library. You can have a look at this example script for pointers: https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_mlm.py\n",
      "  FutureWarning,\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "from transformers import LineByLineTextDataset\n",
    "\n",
    "dataset = LineByLineTextDataset(\n",
    "    tokenizer=tokenizer,\n",
    "    file_path=\"./wiki_en_10.txt\",\n",
    "    block_size=128,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hDLs73HcIHk5"
   },
   "source": [
    "Like in the [`run_language_modeling.py`](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_language_modeling.py) script, we need to define a data_collator.\n",
    "\n",
    "This is just a small helper that will help us batch different samples of the dataset together into an object that PyTorch knows how to perform backprop on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zTgWPa9Dipk2"
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm=True, mlm_probability=0.15\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ri2BIQKqjfHm"
   },
   "source": [
    "### Finally, we are all set to initialize our Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YpvnFFmZJD-N"
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./wikipedia_model_3e\",\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=64,\n",
    "    save_steps=10_000\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=dataset  \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o6sASa36Nf-N"
   },
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 568
    },
    "id": "VmaHZXzmkNtJ",
    "outputId": "c4aa08a8-0144-4b27-809a-26681031352c"
   },
   "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='7371' max='7371' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7371/7371 48:28, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>7.880714</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>7.322539</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>7.120128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>6.984567</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>6.899908</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>6.826775</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>6.777635</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>6.726898</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>6.669105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>6.635973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5500</td>\n",
       "      <td>6.592266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>6.582412</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6500</td>\n",
       "      <td>6.547765</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7000</td>\n",
       "      <td>6.536538</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 29min 6s, sys: 19min 26s, total: 48min 33s\n",
      "Wall time: 48min 28s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=7371, training_loss=6.847868527474647, metrics={'train_runtime': 2908.7329, 'train_samples_per_second': 2.534, 'total_flos': 30249366474276864, 'epoch': 3.0})"
      ]
     },
     "execution_count": 19,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_ZkooHz1-_2h"
   },
   "source": [
    "#### 🎉 Save final model (+ tokenizer + config) to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "QDNgPls7_l13"
   },
   "outputs": [],
   "source": [
    "trainer.save_model(\"./wikipedia_model_3e\")\n",
    "# %cd lm_training_experiments\n",
    "\n",
    "# if model is pretrained, use the following line to import the model directly\n",
    "# model = RobertaForMaskedLM(config=config).from_pretrained('./wikipedia_model_3e')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d0caceCy_p1-"
   },
   "source": [
    "## Check that the LM actually trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ltXgXyCbAJLY",
    "outputId": "71c999eb-7c2a-48af-903f-728c097846c3"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaModel were not initialized from the model checkpoint at ./wikipedia_model_3e and are newly initialized: ['roberta.pooler.dense.weight', 'roberta.pooler.dense.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 pipeline\n",
    "\n",
    "unmasker = pipeline(\n",
    "    \"fill-mask\",\n",
    "    model=\"./wikipedia_model_3e\",\n",
    "    tokenizer=\"./wikipedia_model_3e\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Knowledge Graph Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UIvgZ3S6AO0z"
   },
   "outputs": [],
   "source": [
    "# The sun <mask>.\n",
    "# =>\n",
    "from run_training_kg_experiments import *\n",
    "\n",
    "run_experiments(tokenizer, model, unmasker, \"Roberta3e\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "wikipedia_esperberto_modification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}