training_medium_news.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "BuqP5gVvxMZP",
    "outputId": "336f0d1f-4163-4af7-c2ef-093813b1ac40"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading package lists... Done\n",
      "Building dependency tree       \n",
      "Reading state information... Done\n",
      "unzip is already the newest version (6.0-21ubuntu1).\n",
      "wget is already the newest version (1.19.4-1ubuntu2.2).\n",
      "Suggested packages:\n",
      "  spell\n",
      "The following NEW packages will be installed:\n",
      "  nano\n",
      "0 upgraded, 1 newly installed, 0 to remove and 8 not upgraded.\n",
      "Need to get 231 kB of archives.\n",
      "After this operation, 778 kB of additional disk space will be used.\n",
      "Get:1 http://archive.ubuntu.com/ubuntu bionic/main amd64 nano amd64 2.9.3-2 [231 kB]\n",
      "Fetched 231 kB in 1s (353 kB/s)\n",
      "Selecting previously unselected package nano.\n",
      "(Reading database ... 131183 files and directories currently installed.)\n",
      "Preparing to unpack .../nano_2.9.3-2_amd64.deb ...\n",
      "Unpacking nano (2.9.3-2) ...\n",
      "Setting up nano (2.9.3-2) ...\n",
      "update-alternatives: using /bin/nano to provide /usr/bin/editor (editor) in auto mode\n",
      "update-alternatives: using /bin/nano to provide /usr/bin/pico (pico) in auto mode\n",
      "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
      "Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.6/dist-packages (0.4)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (0.24.2)\n",
      "Requirement already satisfied: python-dateutil>=2.5.0 in /usr/local/lib/python3.6/dist-packages (from pandas) (2.5.3)\n",
      "Requirement already satisfied: numpy>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from pandas) (1.16.5)\n",
      "Requirement already satisfied: pytz>=2011k in /usr/local/lib/python3.6/dist-packages (from pandas) (2018.9)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.6/dist-packages (from python-dateutil>=2.5.0->pandas) (1.12.0)\n",
      "Collecting torch-encoding\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/4b/66/8eff2dd732336d866ce08bc383f4747d60ba988b96a5a02619ced0c8e61c/torch_encoding-1.0.1-py2.py3-none-any.whl (84kB)\n",
      "\u001b[K     |████████████████████████████████| 92kB 4.1MB/s \n",
      "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (4.28.1)\n",
      "Requirement already satisfied: nose in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (1.3.7)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (2.21.0)\n",
      "Requirement already satisfied: torch>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (1.1.0)\n",
      "Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (1.3.1)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (4.3.0)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch-encoding) (1.16.5)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->torch-encoding) (2.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->torch-encoding) (2019.6.16)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->torch-encoding) (3.0.4)\n",
      "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->torch-encoding) (1.24.3)\n",
      "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from Pillow->torch-encoding) (0.46)\n",
      "Installing collected packages: torch-encoding\n",
      "Successfully installed torch-encoding-1.0.1\n",
      "Collecting torch\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/30/57/d5cceb0799c06733eefce80c395459f28970ebb9e896846ce96ab579a3f1/torch-1.2.0-cp36-cp36m-manylinux1_x86_64.whl (748.8MB)\n",
      "\u001b[K     |████████████████████████████████| 748.9MB 28kB/s \n",
      "\u001b[?25hCollecting torchvision\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/06/e6/a564eba563f7ff53aa7318ff6aaa5bd8385cbda39ed55ba471e95af27d19/torchvision-0.4.0-cp36-cp36m-manylinux1_x86_64.whl (8.8MB)\n",
      "\u001b[K     |████████████████████████████████| 8.8MB 1.0MB/s \n",
      "\u001b[?25hRequirement already satisfied, skipping upgrade: numpy in /usr/local/lib/python3.6/dist-packages (from torch) (1.16.5)\n",
      "Requirement already satisfied, skipping upgrade: six in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.12.0)\n",
      "Requirement already satisfied, skipping upgrade: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision) (4.3.0)\n",
      "Requirement already satisfied, skipping upgrade: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision) (0.46)\n",
      "Installing collected packages: torch, torchvision\n",
      "  Found existing installation: torch 1.1.0\n",
      "    Uninstalling torch-1.1.0:\n",
      "      Successfully uninstalled torch-1.1.0\n",
      "  Found existing installation: torchvision 0.3.0\n",
      "    Uninstalling torchvision-0.3.0:\n",
      "      Successfully uninstalled torchvision-0.3.0\n",
      "Successfully installed torch-1.2.0 torchvision-0.4.0\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (4.28.1)\n",
      "Requirement already satisfied: boto3 in /usr/local/lib/python3.6/dist-packages (1.9.220)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (2.21.0)\n",
      "Collecting regex\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/6f/a6/99eeb5904ab763db87af4bd71d9b1dfdd9792681240657a4c0a599c10a81/regex-2019.08.19.tar.gz (654kB)\n",
      "\u001b[K     |████████████████████████████████| 655kB 3.5MB/s \n",
      "\u001b[?25hCollecting sentencepiece\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/14/3d/efb655a670b98f62ec32d66954e1109f403db4d937c50d779a75b9763a29/sentencepiece-0.1.83-cp36-cp36m-manylinux1_x86_64.whl (1.0MB)\n",
      "\u001b[K     |████████████████████████████████| 1.0MB 38.1MB/s \n",
      "\u001b[?25hCollecting sacremoses\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/df/24/0b86f494d3a5c7531f6d0c77d39fd8f9d42e651244505d3d737e31db9a4d/sacremoses-0.0.33.tar.gz (802kB)\n",
      "\u001b[K     |████████████████████████████████| 808kB 24.1MB/s \n",
      "\u001b[?25hRequirement already satisfied: botocore<1.13.0,>=1.12.220 in /usr/local/lib/python3.6/dist-packages (from boto3) (1.12.220)\n",
      "Requirement already satisfied: s3transfer<0.3.0,>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from boto3) (0.2.1)\n",
      "Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /usr/local/lib/python3.6/dist-packages (from boto3) (0.9.4)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests) (3.0.4)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests) (2.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests) (2019.6.16)\n",
      "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests) (1.24.3)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses) (1.12.0)\n",
      "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses) (7.0)\n",
      "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses) (0.13.2)\n",
      "Requirement already satisfied: python-dateutil<3.0.0,>=2.1; python_version >= \"2.7\" in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.220->boto3) (2.5.3)\n",
      "Requirement already satisfied: docutils<0.16,>=0.10 in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.220->boto3) (0.15.2)\n",
      "Building wheels for collected packages: regex, sacremoses\n",
      "  Building wheel for regex (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for regex: filename=regex-2019.8.19-cp36-cp36m-linux_x86_64.whl size=609236 sha256=77526bcd80e32269a2f62ba37a7236a23375e07e2c1a0720f0d1667eb0113f04\n",
      "  Stored in directory: /root/.cache/pip/wheels/90/04/07/b5010fb816721eb3d6dd64ed5cc8111ca23f97fdab8619b5be\n",
      "  Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for sacremoses: filename=sacremoses-0.0.33-cp36-none-any.whl size=833106 sha256=8e6211efae254899c3a2f558fa85ec5b566438088efc61517e715bdb4c3f1729\n",
      "  Stored in directory: /root/.cache/pip/wheels/70/87/56/e40575cca30d12fee8875d523b8878b7aba866a9f03b2fd983\n",
      "Successfully built regex sacremoses\n",
      "Installing collected packages: regex, sentencepiece, sacremoses\n",
      "Successfully installed regex-2019.8.19 sacremoses-0.0.33 sentencepiece-0.1.83\n",
      "Collecting pytorch_transformers\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a3/b7/d3d18008a67e0b968d1ab93ad444fc05699403fa662f634b2f2c318a508b/pytorch_transformers-1.2.0-py3-none-any.whl (176kB)\n",
      "\u001b[K     |████████████████████████████████| 184kB 3.5MB/s \n",
      "\u001b[?25hRequirement already satisfied: sentencepiece in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (0.1.83)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (1.16.5)\n",
      "Requirement already satisfied: torch>=1.0.0 in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (1.2.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (2.21.0)\n",
      "Requirement already satisfied: regex in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (2019.8.19)\n",
      "Requirement already satisfied: sacremoses in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (0.0.33)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (4.28.1)\n",
      "Requirement already satisfied: boto3 in /usr/local/lib/python3.6/dist-packages (from pytorch_transformers) (1.9.220)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->pytorch_transformers) (2.8)\n",
      "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->pytorch_transformers) (1.24.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->pytorch_transformers) (2019.6.16)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->pytorch_transformers) (3.0.4)\n",
      "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->pytorch_transformers) (7.0)\n",
      "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->pytorch_transformers) (0.13.2)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->pytorch_transformers) (1.12.0)\n",
      "Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /usr/local/lib/python3.6/dist-packages (from boto3->pytorch_transformers) (0.9.4)\n",
      "Requirement already satisfied: s3transfer<0.3.0,>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from boto3->pytorch_transformers) (0.2.1)\n",
      "Requirement already satisfied: botocore<1.13.0,>=1.12.220 in /usr/local/lib/python3.6/dist-packages (from boto3->pytorch_transformers) (1.12.220)\n",
      "Requirement already satisfied: docutils<0.16,>=0.10 in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.220->boto3->pytorch_transformers) (0.15.2)\n",
      "Requirement already satisfied: python-dateutil<3.0.0,>=2.1; python_version >= \"2.7\" in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.220->boto3->pytorch_transformers) (2.5.3)\n",
      "Installing collected packages: pytorch-transformers\n",
      "Successfully installed pytorch-transformers-1.2.0\n"
     ]
    }
   ],
   "source": [
    "!apt install wget unzip nano\n",
    "!pip install googledrivedownloader\n",
    "!pip install pandas\n",
    "!pip install torch-encoding\n",
    "!pip install torch torchvision -U\n",
    "!pip install tqdm boto3 requests regex sentencepiece sacremoses\n",
    "!pip install pytorch_transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qDUPKBJvz2y9"
   },
   "outputs": [],
   "source": [
    "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "\n",
    "import pandas as pd \n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "vVyN1pb95B0v",
    "outputId": "18718b9a-2cdd-4c62-abf8-e81f77e54510"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading 1bkgxYlj5BYXFJZ8L9fnJH2pJ8Cnx9_Ge into ./medium_dataset_news.zip... Done.\n",
      "Unzipping...Done.\n"
     ]
    }
   ],
   "source": [
    "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
    "  \n",
    "gdd.download_file_from_google_drive(file_id='1bkgxYlj5BYXFJZ8L9fnJH2pJ8Cnx9_Ge',\n",
    "                                    dest_path='./medium_dataset_news.zip',\n",
    "                                    unzip=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 128
    },
    "colab_type": "code",
    "id": "yrHgLygE0c0M",
    "outputId": "4951677a-c413-40a2-a6c2-f33c85394f74"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22592, 3)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>url</th>\n",
       "      <th>tokens</th>\n",
       "      <th>category</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>https://www.huffingtonpost.com/entry/will-smit...</td>\n",
       "      <td>[720, 34, 6158, 38, 15513, 3, 198, 298, 1425, ...</td>\n",
       "      <td>entertainment</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>https://www.huffingtonpost.com/entry/hugh-gran...</td>\n",
       "      <td>[720, 34, 6158, 38, 15513, 3, 198, 298, 1425, ...</td>\n",
       "      <td>entertainment</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 url  ...       category\n",
       "0  https://www.huffingtonpost.com/entry/will-smit...  ...  entertainment\n",
       "1  https://www.huffingtonpost.com/entry/hugh-gran...  ...  entertainment\n",
       "\n",
       "[2 rows x 3 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('medium_news_training.csv')\n",
    "df.tokens = df.tokens.apply(eval)\n",
    "print(df.shape)\n",
    "df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "khFaF1xU0iTx",
    "outputId": "148b6dbc-bf03-4eba-b99f-643b2d3a54d3"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://github.com/huggingface/pytorch-transformers/archive/master.zip\" to /root/.cache/torch/hub/master.zip\n",
      "100%|██████████| 1042301/1042301 [00:00<00:00, 4037010.63B/s]\n",
      "100%|██████████| 456318/456318 [00:00<00:00, 3184374.44B/s]\n",
      "Using cache found in /root/.cache/torch/hub/huggingface_pytorch-transformers_master\n",
      "100%|██████████| 293/293 [00:00<00:00, 69227.75B/s]\n",
      "100%|██████████| 1520013706/1520013706 [00:30<00:00, 49837185.02B/s]\n"
     ]
    }
   ],
   "source": [
    "tokenizer = torch.hub.load('huggingface/pytorch-transformers:1.2.0', 'tokenizer', 'gpt2-medium')\n",
    "model = torch.hub.load('huggingface/pytorch-transformers:1.2.0', 'modelWithLMHead', 'gpt2-medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QR_UQmaa98rE"
   },
   "outputs": [],
   "source": [
    "def sample_text(seed, tokens_count, temperature=1, topK=30):\n",
    "  model.cuda()\n",
    "  model.eval()\n",
    "  \n",
    "  tokens = tokenizer.encode(seed)\n",
    "  last_past = None\n",
    "  \n",
    "  with torch.no_grad():\n",
    "    for i in range(tokens_count):\n",
    "      \n",
    "      if last_past == None:\n",
    "        logits, past = model(torch.tensor([tokens]).cuda())\n",
    "      else:\n",
    "        logits, past = model(torch.tensor([tokens[-1:]]).cuda(), past=last_past)\n",
    "        \n",
    "      logit = logits[0,-1]\n",
    "\n",
    "      topk_scores, topk_indeces = torch.topk(logit, topK)\n",
    "      \n",
    "      topk_scores = topk_scores.div(temperature).softmax(0)\n",
    "      topk_index = torch.multinomial(topk_scores, 1)\n",
    "      \n",
    "      sampled_token = topk_indeces[topk_index].item()\n",
    "            \n",
    "      tokens.append(sampled_token)\n",
    "      \n",
    "      last_past = past    \n",
    "    \n",
    "  print(tokenizer.decode(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 544
    },
    "colab_type": "code",
    "id": "crnB6XH1-AyO",
    "outputId": "a2aff7db-a6ff-4cff-c265-e820c8a7c411"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " $CATEGORY$\n",
      "entertainment\n",
      "$CATEGORY$\n",
      "entertainment/animation/tv\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment/animation\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment/animation\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment/animation\n",
      "\n",
      "$CATEGORY$\n",
      "\n",
      "entertainment/animation\n",
      "\n",
      "CPU times: user 6.3 s, sys: 1.54 s, total: 7.84 s\n",
      "Wall time: 7.98 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "sample_text(f\"\"\"$CATEGORY$\n",
    "entertainment\n",
    "$CATEGORY$\n",
    "\"\"\", tokens_count = 100, temperature=1, topK=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "zufvHRMP__aC",
    "outputId": "f0e48200-81cc-4121-ed4c-63a753e0e158"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 500])"
      ]
     },
     "execution_count": 9,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = TensorDataset(torch.tensor(df.tokens.values.tolist()))\n",
    "loader = DataLoader(dataset, batch_size=1, shuffle=True)\n",
    "next(iter(loader))[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bP5xWJZT_wTp"
   },
   "outputs": [],
   "source": [
    "def compute_total_loss(tokens):\n",
    "  loss, _, _ = model(tokens, labels=tokens)\n",
    "\n",
    "  optimizer.zero_grad()\n",
    "  loss.backward()\n",
    "  optimizer.step()\n",
    "\n",
    "  return loss.item()\n",
    "\n",
    "def compute_aggregated_loss(tokens):\n",
    "  last_past = None\n",
    "  aggregated_loss = None\n",
    "\n",
    "  for i in range(4):\n",
    "    sampled = tokens[:, i*256:(i+1)*256]\n",
    "    loss, _, past = model(sampled, labels=sampled, past=last_past)\n",
    "\n",
    "    if aggregated_loss is None:\n",
    "      aggregated_loss = loss\n",
    "    else:\n",
    "      aggregated_loss += loss\n",
    "\n",
    "  optimizer.zero_grad()\n",
    "  aggregated_loss.backward()\n",
    "  optimizer.step()\n",
    "\n",
    "  return aggregated_loss.item() / 4\n",
    "  \n",
    "def compute_loss(tokens):\n",
    "  last_past = None\n",
    "  partial_loss = None\n",
    "\n",
    "  for i in range(4):\n",
    "    sampled = tokens[:, i*256:(i+1)*256]\n",
    "    loss, _, past = model(sampled, labels=sampled, past=last_past)\n",
    "    last_past = past[0].detach(), past[1].detach()\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if partial_loss is None:\n",
    "      partial_loss = loss.item()\n",
    "    else:\n",
    "      partial_loss += loss.item()\n",
    "      \n",
    "\n",
    "  return partial_loss / 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Whjr6TQm-K4q"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "#aggregate = False\n",
    "\n",
    "#if (aggregate):\n",
    "  #for param in model.parameters():\n",
    "    #param.requires_grad = False\n",
    "\n",
    "  #for param in model.lm_head.parameters():\n",
    "    #param.requires_grad = True\n",
    "    \n",
    "  #optimizer = optim.Adam(model.lm_head.parameters(), lr=0.00001)\n",
    "#else:\n",
    "  #for param in model.parameters():\n",
    "    #param.requires_grad = True\n",
    "    \n",
    "  #optimizer = optim.Adam(model.parameters(), lr=0.00001)\n",
    "  \n",
    "for param in model.parameters():\n",
    "  param.requires_grad = False\n",
    "\n",
    "for param in model.lm_head.parameters():\n",
    "  param.requires_grad = True\n",
    "\n",
    "optimizer = optim.Adam(model.lm_head.parameters(), lr=0.00001)\n",
    "  \n",
    "print_every = 100\n",
    "\n",
    "model.cuda()\n",
    "\n",
    "epochs = 1\n",
    "\n",
    "last_loss = None\n",
    "\n",
    "for epoch in range(1,epochs+1):\n",
    "  total_loss = 0\n",
    "  \n",
    "  for batch, tokens in enumerate(loader):\n",
    "    model.train()\n",
    "    \n",
    "    tokens = tokens[0].cuda()\n",
    "    total_loss += compute_total_loss(tokens)\n",
    "    \n",
    "    #if (aggregate): \n",
    "      #total_loss += compute_aggregated_loss(tokens)\n",
    "    #else:\n",
    "      #total_loss += compute_loss(tokens)\n",
    "\n",
    "    if (batch+1) % print_every == 0 or (batch+1) == len(loader):\n",
    "      loss_normalized = total_loss/(batch+1)\n",
    "      \n",
    "      print(epoch, (len(loader), batch+1), loss_normalized) \n",
    "      sample_text(f\"\"\"$CATEGORY$\n",
    "      entertainment\n",
    "      $CATEGORY$\n",
    "      \"\"\", tokens_count = 150, temperature=0.1, topK=100)\n",
    "      \n",
    "      if last_loss == None:\n",
    "        last_loss = loss_normalized\n",
    "            \n",
    "      if last_loss >= loss_normalized:\n",
    "        print('saving the model')\n",
    "        torch.save(model.state_dict(), 'medium_fake_news_model.pt')\n",
    "        last_loss = loss_normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EK7HMKYuNgcg"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "training-medium-news.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}