ALL_TOKENS.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Relation extraction with BERT\n",
    "\n",
    "---\n",
    "\n",
    "The goal of this notebook is to show how to use [BERT](https://arxiv.org/abs/1810.04805)\n",
    "to [extract relation](https://en.wikipedia.org/wiki/Relationship_extraction) from text.\n",
    "\n",
    "Used libraries:\n",
    "- [PyTorch](https://pytorch.org/)\n",
    "- [PyTorch-Lightning](https://pytorch-lightning.readthedocs.io/en/latest/)\n",
    "- [Transformers](https://huggingface.co/transformers/index.html)\n",
    "\n",
    "Used datasets:\n",
    "- SemEval 2010 Task 8 - [paper](https://arxiv.org/pdf/1911.10422.pdf) - [download](https://github.com/sahitya0000/Relation-Classification/blob/master/corpus/SemEval2010_task8_all_data.zip?raw=true)\n",
    "- Google IISc Distant Supervision (GIDS) - [paper](https://arxiv.org/pdf/1804.06987.pdf) - [download](https://drive.google.com/open?id=1gTNAbv8My2QDmP-OHLFtJFlzPDoCG4aI)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## High level overview\n",
    "\n",
    "We will experiment with two architectures: single-classifier & duo-classifier\n",
    "\n",
    "![Single-classifier architecture](https://raw.githubusercontent.com/dthung1602/bert-relation-extraction/notebook/images/high_level.png \"Single-classifier architecture\")\n",
    "\n",
    "![Duo-classifier architecture](https://raw.githubusercontent.com/dthung1602/bert-relation-extraction/notebook/images/high_level_bin.png \"Duo-classifier architecture\")\n",
    "\n",
    "The classifiers are implemented as follows:\n",
    "\n",
    "![All hidden](https://raw.githubusercontent.com/dthung1602/bert-relation-extraction/notebook/images/all_hidden.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "## Install dependencies\n",
    "\n",
    "This project uses [Python 3.7+](https://www.python.org/downloads/release/python-378/)"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": true
   },
   "cell_type": "code",
   "source": [
    "!pip install requests==2.23.0 numpy==1.18.5 pandas==1.0.3 \\\n",
    "    scikit-learn==0.23.1 pytorch-lightning==0.8.4 torch==1.5.1 \\\n",
    "    transformers==3.0.2 sklearn==0.0 tqdm==4.45.0 neptune-client==0.4.119 \\\n",
    "    matplotlib==3.1.0 scikit-plot==0.3.7"
   ],
   "execution_count": 1,
   "outputs": [
    {
     "output_type": "stream",
     "text": "Requirement already satisfied: requests==2.23.0 in /opt/conda/lib/python3.7/site-packages (2.23.0)\nRequirement already satisfied: numpy==1.18.5 in /opt/conda/lib/python3.7/site-packages (1.18.5)\nRequirement already satisfied: pandas==1.0.3 in /opt/conda/lib/python3.7/site-packages (1.0.3)\nRequirement already satisfied: scikit-learn==0.23.1 in /opt/conda/lib/python3.7/site-packages (0.23.1)\nCollecting pytorch-lightning==0.8.4\n  Downloading pytorch_lightning-0.8.4-py3-none-any.whl (304 kB)\n\u001B[K     |████████████████████████████████| 304 kB 2.9 MB/s eta 0:00:01\n\u001B[?25hRequirement already satisfied: torch==1.5.1 in /opt/conda/lib/python3.7/site-packages (1.5.1)\nCollecting transformers==3.0.2\n  Downloading transformers-3.0.2-py3-none-any.whl (769 kB)\n\u001B[K     |████████████████████████████████| 769 kB 9.3 MB/s eta 0:00:01\n\u001B[?25hRequirement already satisfied: sklearn==0.0 in /opt/conda/lib/python3.7/site-packages (0.0)\nRequirement already satisfied: tqdm==4.45.0 in /opt/conda/lib/python3.7/site-packages (4.45.0)\nCollecting neptune-client==0.4.119\n  Downloading neptune-client-0.4.119.tar.gz (90 kB)\n\u001B[K     |████████████████████████████████| 90 kB 4.5 MB/s  eta 0:00:01\n\u001B[?25hCollecting matplotlib==3.1.0\n  Downloading matplotlib-3.1.0-cp37-cp37m-manylinux1_x86_64.whl (13.1 MB)\n\u001B[K     |████████████████████████████████| 13.1 MB 15.3 MB/s eta 0:00:01\n\u001B[?25hRequirement already satisfied: scikit-plot==0.3.7 in /opt/conda/lib/python3.7/site-packages (0.3.7)\nRequirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests==2.23.0) (2020.6.20)\nRequirement already satisfied: idna<3,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests==2.23.0) (2.9)\nRequirement already satisfied: chardet<4,>=3.0.2 in /opt/conda/lib/python3.7/site-packages (from requests==2.23.0) (3.0.4)\nRequirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests==2.23.0) (1.24.3)\nRequirement already satisfied: python-dateutil>=2.6.1 in /opt/conda/lib/python3.7/site-packages (from pandas==1.0.3) (2.8.1)\nRequirement already satisfied: pytz>=2017.2 in /opt/conda/lib/python3.7/site-packages (from pandas==1.0.3) (2019.3)\nRequirement already satisfied: joblib>=0.11 in /opt/conda/lib/python3.7/site-packages (from scikit-learn==0.23.1) (0.14.1)\nRequirement already satisfied: scipy>=0.19.1 in /opt/conda/lib/python3.7/site-packages (from scikit-learn==0.23.1) (1.4.1)\nRequirement already satisfied: threadpoolctl>=2.0.0 in /opt/conda/lib/python3.7/site-packages (from scikit-learn==0.23.1) (2.1.0)\nRequirement already satisfied: tensorboard>=1.14 in /opt/conda/lib/python3.7/site-packages (from pytorch-lightning==0.8.4) (2.2.2)\nRequirement already satisfied: future>=0.17.1 in /opt/conda/lib/python3.7/site-packages (from pytorch-lightning==0.8.4) (0.18.2)\nRequirement already satisfied: PyYAML>=5.1 in /opt/conda/lib/python3.7/site-packages (from pytorch-lightning==0.8.4) (5.3.1)\nRequirement already satisfied: sentencepiece!=0.1.92 in /opt/conda/lib/python3.7/site-packages (from transformers==3.0.2) (0.1.91)\nRequirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers==3.0.2) (2020.4.4)\nRequirement already satisfied: sacremoses in /opt/conda/lib/python3.7/site-packages (from transformers==3.0.2) (0.0.43)\nRequirement already satisfied: packaging in /opt/conda/lib/python3.7/site-packages (from transformers==3.0.2) (20.1)\nRequirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers==3.0.2) (3.0.10)\nCollecting tokenizers==0.8.1.rc1\n  Downloading tokenizers-0.8.1rc1-cp37-cp37m-manylinux1_x86_64.whl (3.0 MB)\n\u001B[K     |████████████████████████████████| 3.0 MB 39.4 MB/s eta 0:00:01\n\u001B[?25hCollecting bravado\n  Downloading bravado-10.6.2-py2.py3-none-any.whl (37 kB)\nRequirement already satisfied: click>=7.0 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (7.1.1)\nCollecting py3nvml\n  Downloading py3nvml-0.2.6-py3-none-any.whl (55 kB)\n\u001B[K     |████████████████████████████████| 55 kB 2.5 MB/s  eta 0:00:01\n\u001B[?25hRequirement already satisfied: oauthlib>=2.1.0 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (3.0.1)\nRequirement already satisfied: Pillow>=1.1.6 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (7.2.0)\nRequirement already satisfied: PyJWT in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (1.7.1)\nRequirement already satisfied: requests-oauthlib>=1.0.0 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (1.2.0)\nRequirement already satisfied: six>=1.12.0 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (1.14.0)\nRequirement already satisfied: websocket-client>=0.35.0 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (0.57.0)\nRequirement already satisfied: GitPython>=2.0.8 in /opt/conda/lib/python3.7/site-packages (from neptune-client==0.4.119) (3.1.1)\nRequirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.7/site-packages (from matplotlib==3.1.0) (0.10.0)\nRequirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.7/site-packages (from matplotlib==3.1.0) (1.2.0)\nRequirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /opt/conda/lib/python3.7/site-packages (from matplotlib==3.1.0) (2.4.7)\nRequirement already satisfied: google-auth<2,>=1.6.3 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (1.14.0)\nRequirement already satisfied: markdown>=2.6.8 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (3.2.1)\nRequirement already satisfied: protobuf>=3.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (3.12.2)\nRequirement already satisfied: setuptools>=41.0.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (46.1.3.post20200325)\nRequirement already satisfied: werkzeug>=0.11.15 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (1.0.1)\nRequirement already satisfied: absl-py>=0.4 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (0.9.0)\nRequirement already satisfied: wheel>=0.26; python_version >= \"3\" in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (0.34.2)\nRequirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (1.7.0)\nRequirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (0.4.1)\nRequirement already satisfied: grpcio>=1.24.3 in /opt/conda/lib/python3.7/site-packages (from tensorboard>=1.14->pytorch-lightning==0.8.4) (1.30.0)\nCollecting msgpack-python\n  Downloading msgpack-python-0.5.6.tar.gz (138 kB)\n\u001B[K     |████████████████████████████████| 138 kB 27.4 MB/s eta 0:00:01\n\u001B[?25hRequirement already satisfied: simplejson in /opt/conda/lib/python3.7/site-packages (from bravado->neptune-client==0.4.119) (3.17.0)\nCollecting bravado-core>=5.16.1\n  Downloading bravado_core-5.17.0-py2.py3-none-any.whl (67 kB)\n\u001B[K     |████████████████████████████████| 67 kB 3.7 MB/s  eta 0:00:01\n\u001B[?25hCollecting monotonic\n  Downloading monotonic-1.5-py2.py3-none-any.whl (5.3 kB)\nRequirement already satisfied: typing-extensions in /opt/conda/lib/python3.7/site-packages (from bravado->neptune-client==0.4.119) (3.7.4.1)\nCollecting xmltodict\n  Downloading xmltodict-0.12.0-py2.py3-none-any.whl (9.2 kB)\nRequirement already satisfied: gitdb<5,>=4.0.1 in /opt/conda/lib/python3.7/site-packages (from GitPython>=2.0.8->neptune-client==0.4.119) (4.0.4)\nRequirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch-lightning==0.8.4) (3.1.1)\nRequirement already satisfied: rsa<4.1,>=3.1.4 in /opt/conda/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch-lightning==0.8.4) (4.0)\nRequirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch-lightning==0.8.4) (0.2.7)\nRequirement already satisfied: msgpack>=0.5.2 in /opt/conda/lib/python3.7/site-packages (from bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (1.0.0)\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "Collecting swagger-spec-validator>=2.0.1\n  Downloading swagger_spec_validator-2.7.3-py2.py3-none-any.whl (27 kB)\nCollecting jsonref\n  Downloading jsonref-0.2-py3-none-any.whl (9.3 kB)\nRequirement already satisfied: jsonschema[format]>=2.5.1 in /opt/conda/lib/python3.7/site-packages (from bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (3.2.0)\nRequirement already satisfied: smmap<4,>=3.0.1 in /opt/conda/lib/python3.7/site-packages (from gitdb<5,>=4.0.1->GitPython>=2.0.8->neptune-client==0.4.119) (3.0.2)\nRequirement already satisfied: pyasn1>=0.1.3 in /opt/conda/lib/python3.7/site-packages (from rsa<4.1,>=3.1.4->google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch-lightning==0.8.4) (0.4.8)\nRequirement already satisfied: importlib-metadata; python_version < \"3.8\" in /opt/conda/lib/python3.7/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (1.6.0)\nRequirement already satisfied: attrs>=17.4.0 in /opt/conda/lib/python3.7/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (19.3.0)\nRequirement already satisfied: pyrsistent>=0.14.0 in /opt/conda/lib/python3.7/site-packages (from jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (0.16.0)\nCollecting strict-rfc3339; extra == \"format\"\n  Downloading strict-rfc3339-0.7.tar.gz (17 kB)\nCollecting jsonpointer>1.13; extra == \"format\"\n  Downloading jsonpointer-2.0-py2.py3-none-any.whl (7.6 kB)\nCollecting webcolors; extra == \"format\"\n  Downloading webcolors-1.11.1-py3-none-any.whl (9.9 kB)\nCollecting rfc3987; extra == \"format\"\n  Downloading rfc3987-1.3.8-py2.py3-none-any.whl (13 kB)\nRequirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata; python_version < \"3.8\"->jsonschema[format]>=2.5.1->bravado-core>=5.16.1->bravado->neptune-client==0.4.119) (3.1.0)\nBuilding wheels for collected packages: neptune-client, msgpack-python, strict-rfc3339\n  Building wheel for neptune-client (setup.py) ... \u001B[?25ldone\n\u001B[?25h  Created wheel for neptune-client: filename=neptune_client-0.4.119-py2.py3-none-any.whl size=150019 sha256=24eefc6cd2cbc90df1983e37059960aa343e75afbfb03ea737ac951abf70570b\n  Stored in directory: /root/.cache/pip/wheels/64/3b/7e/69f84d99e2109788f757ef707b3ea51921f16891e42929eb31\n  Building wheel for msgpack-python (setup.py) ... \u001B[?25ldone\n\u001B[?25h  Created wheel for msgpack-python: filename=msgpack_python-0.5.6-cp37-cp37m-linux_x86_64.whl size=302583 sha256=c3c9ae7d1d425337a36a4bea31906e7c450e5103f4c043ce3b21046891ec8bd5\n  Stored in directory: /root/.cache/pip/wheels/f8/6c/02/92ebc97f3b99ad5bfc675be2c513f9cb3504fdbe338314f377\n  Building wheel for strict-rfc3339 (setup.py) ... \u001B[?25ldone\n\u001B[?25h  Created wheel for strict-rfc3339: filename=strict_rfc3339-0.7-py3-none-any.whl size=18119 sha256=de8d4e5e937346f769c65931e9cccda619852ed7aafce68c8b25236b9729b063\n  Stored in directory: /root/.cache/pip/wheels/f3/1d/9f/2a74caecb81b8beb9a4fbe1754203d4b7cf42ef5d39e0d2311\nSuccessfully built neptune-client msgpack-python strict-rfc3339\n\u001B[31mERROR: plotnine 0.7.0 has requirement matplotlib>=3.1.1, but you'll have matplotlib 3.1.0 which is incompatible.\u001B[0m\n\u001B[31mERROR: pandas-profiling 2.6.0 has requirement matplotlib>=3.2.0, but you'll have matplotlib 3.1.0 which is incompatible.\u001B[0m\n\u001B[31mERROR: osmnx 0.15.1 has requirement geopandas>=0.7, but you'll have geopandas 0.6.3 which is incompatible.\u001B[0m\n\u001B[31mERROR: osmnx 0.15.1 has requirement matplotlib>=3.2, but you'll have matplotlib 3.1.0 which is incompatible.\u001B[0m\n\u001B[31mERROR: mizani 0.7.1 has requirement matplotlib>=3.1.1, but you'll have matplotlib 3.1.0 which is incompatible.\u001B[0m\n\u001B[31mERROR: hypertools 0.6.2 has requirement scikit-learn<0.22,>=0.19.1, but you'll have scikit-learn 0.23.1 which is incompatible.\u001B[0m\n\u001B[31mERROR: allennlp 1.0.0 has requirement transformers<2.12,>=2.9, but you'll have transformers 3.0.2 which is incompatible.\u001B[0m\nInstalling collected packages: pytorch-lightning, tokenizers, transformers, msgpack-python, swagger-spec-validator, jsonref, bravado-core, monotonic, bravado, xmltodict, py3nvml, neptune-client, matplotlib, strict-rfc3339, jsonpointer, webcolors, rfc3987\n  Attempting uninstall: tokenizers\n    Found existing installation: tokenizers 0.7.0\n    Uninstalling tokenizers-0.7.0:\n      Successfully uninstalled tokenizers-0.7.0\n  Attempting uninstall: transformers\n    Found existing installation: transformers 2.11.0\n    Uninstalling transformers-2.11.0:\n      Successfully uninstalled transformers-2.11.0\n  Attempting uninstall: matplotlib\n    Found existing installation: matplotlib 3.2.1\n    Uninstalling matplotlib-3.2.1:\n      Successfully uninstalled matplotlib-3.2.1\nSuccessfully installed bravado-10.6.2 bravado-core-5.17.0 jsonpointer-2.0 jsonref-0.2 matplotlib-3.1.0 monotonic-1.5 msgpack-python-0.5.6 neptune-client-0.4.119 py3nvml-0.2.6 pytorch-lightning-0.8.4 rfc3987-1.3.8 strict-rfc3339-0.7 swagger-spec-validator-2.7.3 tokenizers-0.8.1rc1 transformers-3.0.2 webcolors-1.11.1 xmltodict-0.12.0\n\u001B[33mWARNING: You are using pip version 20.1.1; however, version 20.2.3 is available.\nYou should consider upgrading via the '/opt/conda/bin/python3.7 -m pip install --upgrade pip' command.\u001B[0m\n",
     "name": "stdout"
    }
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Import needed modules"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% \n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "import gc\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "from abc import ABC, abstractmethod\n",
    "from collections import OrderedDict\n",
    "from random import randint\n",
    "from typing import Iterable, Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from matplotlib.figure import Figure\n",
    "from pandas import DataFrame\n",
    "from pytorch_lightning import LightningModule, seed_everything\n",
    "from pytorch_lightning import Trainer as LightningTrainer\n",
    "from pytorch_lightning.logging.neptune import NeptuneLogger\n",
    "from sklearn.metrics import *\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import label_binarize\n",
    "from sklearn.utils import column_or_1d\n",
    "from torch import Tensor, nn\n",
    "from torch.nn import functional as F\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data import DataLoader, IterableDataset\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import *"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Define constants"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% \n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "# --- Random seed ---\n",
    "SEED = 2020\n",
    "seed_everything(SEED)\n",
    "\n",
    "# --- Directory ---\n",
    "ROOT_DIR = os.path.abspath(\".\")\n",
    "PROCESSED_DATA_DIR = os.path.join(ROOT_DIR, \"data/processed\") \n",
    "METADATA_FILE_NAME = os.path.join(PROCESSED_DATA_DIR, \"metadata.json\")\n",
    "CHECKPOINT_DIR = os.path.join(ROOT_DIR, \"checkpoint\")\n",
    "\n",
    "KAGGLE_ENV = bool(os.getenv(\"KAGGLE_URL_BASE\"))\n",
    "if KAGGLE_ENV:\n",
    "    # in Kaggle environment\n",
    "    # 2 datasets should already been added to the notebook\n",
    "    RAW_DATA_DIR = os.path.join(ROOT_DIR, \"../input\")\n",
    "else:\n",
    "    # in local environment\n",
    "    RAW_DATA_DIR =  os.path.join(ROOT_DIR, \"data/raw\")\n",
    "\n",
    "# --- Datasets ---\n",
    "DATASET_MAPPING = {\n",
    "    \"SemEval2010Task8\": {\n",
    "        \"dir\": os.path.join(RAW_DATA_DIR,\"semeval2010-task-8\"),\n",
    "        \"keep_test_order\": True,\n",
    "        \"precision_recall_curve_baseline_img\": None,\n",
    "    },\n",
    "    \"GIDS\": {\n",
    "        \"dir\": os.path.join(RAW_DATA_DIR,\"gids-dataset\"),\n",
    "        \"keep_test_order\": False,\n",
    "        \"precision_recall_curve_baseline_img\": os.path.join(RAW_DATA_DIR,\"gids-dataset/GIDS_precision_recall_curve.png\"),\n",
    "    }\n",
    "}\n",
    "\n",
    "# change this variable to switch dataset in later tasks\n",
    "DATASET_NAME = list(DATASET_MAPPING.keys())[1]\n",
    "\n",
    "# --- Subject & object markup ---\n",
    "SUB_START_CHAR = \"[\"\n",
    "SUB_END_CHAR = \"]\"\n",
    "OBJ_START_CHAR = \"{\"\n",
    "OBJ_END_CHAR = \"}\"\n",
    "\n",
    "# --- BERT variants ---\n",
    "# See https://huggingface.co/transformers/pretrained_models.html for the full list\n",
    "AVAILABLE_PRETRAINED_MODELS = [\n",
    "    \"distilbert-base-uncased\", # 0\n",
    "    \"distilbert-base-cased\",   # 1\n",
    "    \"bert-base-uncased\",       # 2\n",
    "    \"distilgpt2\",              # 3\n",
    "    \"gpt2\",                    # 4\n",
    "    \"distilroberta-base\",      # 5\n",
    "    \"roberta-base\",            # 6\n",
    "    \"albert-base-v1\",          # 7\n",
    "    \"albert-base-v2\",          # 8\n",
    "    \"bert-large-uncased\",      # 9\n",
    "]\n",
    "\n",
    "# change this variable to switch pretrained language model\n",
    "PRETRAINED_MODEL = AVAILABLE_PRETRAINED_MODELS[2]\n",
    "\n",
    "# if e1 is not related to e2, should \"e2 not related to e1\" be added to the training set\n",
    "ADD_REVERSE_RELATIONSHIP = True\n",
    "\n",
    "# --- Neptune logger ---\n",
    "# Create a free account at https://neptune.ai/,\n",
    "# then get the API token and create a project\n",
    "NEPTUNE_API_TOKEN = \" INSERT YOUR API TOKEN HERE \"\n",
    "NEPTUNE_PROJECT_NAME = \" INSERT YOUR PROJECT NAME HERE \""
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "## Preprocess\n",
    "\n",
    "![Preprocessing](https://raw.githubusercontent.com/dthung1602/bert-relation-extraction/notebook/images/preprocess.png)\n",
    "\n",
    "First, we define a custom label encoder. What this label encoder offers but `sklearn.preprocessing.LabelEncoder` fails\n",
    "to provide:\n",
    "- Order preservation: labels will be encoded in order they appear in the dataset. Labels appears earlier will have\n",
    "  smaller id. We need this to ensure the `no relation` class is always encoded as `0`\n",
    "- Multiple fit: `sklearn.preprocessing.LabelEncoder` forgets what is fit in the last time `fit` is called while our\n",
    "  encoder keeps adding new labels to existing ones. This is useful when we process large dataset in batches."
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "class OrdinalLabelEncoder:\n",
    "    def __init__(self, init_labels=None):\n",
    "        if init_labels is None:\n",
    "            init_labels = []\n",
    "        self.mapping = OrderedDict({l: i for i, l in enumerate(init_labels)})\n",
    "\n",
    "    @property\n",
    "    def classes_(self):\n",
    "        return list(self.mapping.keys())\n",
    "\n",
    "    def fit_transform(self, y):\n",
    "        return self.fit(y).transform(y)\n",
    "\n",
    "    def fit(self, y):\n",
    "        y = column_or_1d(y, warn=True)\n",
    "        new_classes = pd.Series(y).unique()\n",
    "        for cls in new_classes:\n",
    "            if cls not in self.mapping:\n",
    "                self.mapping[cls] = len(self.mapping)\n",
    "        return self\n",
    "\n",
    "    def transform(self, y):\n",
    "        y = column_or_1d(y, warn=True)\n",
    "        return [self.mapping[value] for value in y]\n",
    "    "
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "Abstract preprocessor class:"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": true
   },
   "cell_type": "code",
   "source": [
    "class AbstractPreprocessor(ABC):\n",
    "    DATASET_NAME = \"\"\n",
    "    VAL_DATA_PROPORTION = 0.2\n",
    "    NO_RELATION_LABEL = \"\"\n",
    "\n",
    "    def __init__(self, tokenizer: PreTrainedTokenizer):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.SUB_START_ID, self.SUB_END_ID, self.OBJ_START_ID, self.OBJ_END_ID \\\n",
    "            = tokenizer.convert_tokens_to_ids([SUB_START_CHAR, SUB_END_CHAR, OBJ_START_CHAR, OBJ_END_CHAR])\n",
    "        self.label_encoder = OrdinalLabelEncoder([self.NO_RELATION_LABEL])\n",
    "\n",
    "    def preprocess_data(self, reprocess: bool):\n",
    "        print(f\"\\n---> Preprocessing {self.DATASET_NAME} dataset <---\")\n",
    "        \n",
    "        # create processed data dir\n",
    "        if not os.path.exists(PROCESSED_DATA_DIR):\n",
    "            print(\"Creating processed data directory \" + PROCESSED_DATA_DIR)\n",
    "            os.makedirs(PROCESSED_DATA_DIR)\n",
    "\n",
    "        # stop preprocessing if file existed\n",
    "        json_file_names = [self.get_dataset_file_name(k) for k in (\"train\", \"val\", \"test\")]\n",
    "        existed_files = [fn for fn in json_file_names if os.path.exists(fn)]\n",
    "        if existed_files:\n",
    "            file_text = \"- \" + \"\\n- \".join(existed_files)\n",
    "            if not reprocess:\n",
    "                print(\"The following files already exist:\")\n",
    "                print(file_text)\n",
    "                print(\"Preprocessing is skipped. See option --reprocess.\")\n",
    "                return\n",
    "            else:\n",
    "                print(\"The following files will be overwritten:\")\n",
    "                print(file_text)\n",
    "\n",
    "        train_data, val_data, test_data = self._preprocess_data()\n",
    "\n",
    "        print(\"Saving to json files\")\n",
    "        self._write_data_to_file(train_data, \"train\")\n",
    "        self._write_data_to_file(val_data, \"val\")\n",
    "        self._write_data_to_file(test_data, \"test\")\n",
    "\n",
    "        self._save_metadata({\n",
    "            \"train_size\": len(train_data),\n",
    "            \"val_size\": len(val_data),\n",
    "            \"test_size\": len(test_data),\n",
    "            \"no_relation_label\": self.NO_RELATION_LABEL,\n",
    "            **self._get_label_mapping()\n",
    "        })\n",
    "\n",
    "        self._create_secondary_data_files()\n",
    "\n",
    "        print(\"---> Done ! <---\")\n",
    "\n",
    "    @abstractmethod\n",
    "    def _preprocess_data(self) -> Tuple[DataFrame, DataFrame, DataFrame]:\n",
    "        pass\n",
    "\n",
    "    def _create_secondary_data_files(self):\n",
    "        \"\"\"\n",
    "        From the primary data file, create a data file with binary labels\n",
    "        and a data file with only sentences classified as \"related\"\n",
    "        \"\"\"\n",
    "\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            root_metadata = json.load(f)\n",
    "            metadata = root_metadata[self.DATASET_NAME]\n",
    "\n",
    "        related_only_count = {\n",
    "            \"train\": 0,\n",
    "            \"val\": 0,\n",
    "            \"test\": 0,\n",
    "        }\n",
    "\n",
    "        for key in [\"train\", \"test\", \"val\"]:\n",
    "            print(f\"Creating secondary files for {key} data\")\n",
    "\n",
    "            origin_file = open(self.get_dataset_file_name(key))\n",
    "            bin_file = open(self.get_dataset_file_name(f\"{key}_binary\"), \"w\")\n",
    "            related_file = open(self.get_dataset_file_name(f\"{key}_related_only\"), \"w\")\n",
    "\n",
    "            total = metadata[f\"{key}_size\"]\n",
    "\n",
    "            for line in tqdm(origin_file, total=total):\n",
    "                data = json.loads(line)\n",
    "                if data[\"label\"] != 0:\n",
    "                    related_only_count[key] += 1\n",
    "                    data[\"label\"] -= 1 # label in \"related_only\" files is 1 less than the original label\n",
    "                    related_file.write(json.dumps(data) + \"\\n\")\n",
    "                    data[\"label\"] = 1 # in binary dataset, all \"related\" classes have label 1\n",
    "                    bin_file.write(json.dumps(data) + \"\\n\")\n",
    "                else:\n",
    "                    bin_file.write(json.dumps(data) + \"\\n\")\n",
    "\n",
    "            origin_file.close()\n",
    "            bin_file.close()\n",
    "            related_file.close()\n",
    "\n",
    "        print(\"Updating metadata.json\")\n",
    "        for key in [\"train\", \"test\", \"val\"]:\n",
    "            metadata[f\"{key}_related_only_size\"] = related_only_count[key]\n",
    "        root_metadata[self.DATASET_NAME] = metadata\n",
    "        with open(METADATA_FILE_NAME, \"w\") as f:\n",
    "            json.dump(root_metadata, f, indent=4)\n",
    "\n",
    "    def _find_sub_obj_pos(self, input_ids_list: Iterable) -> DataFrame:\n",
    "        \"\"\"\n",
    "        Find subject and object position in a sentence\n",
    "        \"\"\"\n",
    "        sub_start_pos = [self._index(s, self.SUB_START_ID) + 1 for s in input_ids_list]\n",
    "        sub_end_pos = [self._index(s, self.SUB_END_ID, sub_start_pos[i]) for i, s in enumerate(input_ids_list)]\n",
    "        obj_start_pos = [self._index(s, self.OBJ_START_ID) + 1 for s in input_ids_list]\n",
    "        obj_end_pos = [self._index(s, self.OBJ_END_ID, obj_start_pos[i]) for i, s in enumerate(input_ids_list)]\n",
    "        return DataFrame({\n",
    "            \"sub_start_pos\": sub_start_pos,\n",
    "            \"sub_end_pos\": sub_end_pos,\n",
    "            \"obj_start_pos\": obj_start_pos,\n",
    "            \"obj_end_pos\": obj_end_pos,\n",
    "        })\n",
    "\n",
    "    @staticmethod\n",
    "    def _index(lst: list, ele: int, start: int = 0) -> int:\n",
    "        \"\"\"\n",
    "        Find an element in a list. Returns -1 if not found instead of raising an exception.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return lst.index(ele, start)\n",
    "        except ValueError:\n",
    "            return -1\n",
    "\n",
    "    def _clean_data(self, raw_sentences: list, labels: list) -> DataFrame:\n",
    "        if not raw_sentences:\n",
    "            return DataFrame()\n",
    "\n",
    "        tokens = self.tokenizer(raw_sentences, truncation=True, padding=\"max_length\")\n",
    "        data = DataFrame(tokens.data)\n",
    "        data[\"label\"] = self.label_encoder.fit_transform(labels)\n",
    "        sub_obj_position = self._find_sub_obj_pos(data[\"input_ids\"])\n",
    "        data = pd.concat([data, sub_obj_position], axis=1)\n",
    "        data = self._remove_invalid_sentences(data)\n",
    "        return data\n",
    "\n",
    "    def _remove_invalid_sentences(self, data: DataFrame) -> DataFrame:\n",
    "        \"\"\"\n",
    "        Remove sentences without subject/object or whose subject/object\n",
    "        is beyond the maximum length the model supports\n",
    "        \"\"\"\n",
    "        seq_max_len = self.tokenizer.model_max_length\n",
    "        return data.loc[\n",
    "            (data[\"sub_end_pos\"] < seq_max_len)\n",
    "            & (data[\"obj_end_pos\"] < seq_max_len)\n",
    "            & (data[\"sub_end_pos\"] > -1)\n",
    "            & (data[\"obj_end_pos\"] > -1)\n",
    "        ]\n",
    "\n",
    "    def _get_label_mapping(self):\n",
    "        \"\"\"\n",
    "        Returns a mapping from id to label and vise versa from the label encoder\n",
    "        \"\"\"\n",
    "        # all labels\n",
    "        id_to_label = dict(enumerate(self.label_encoder.classes_))\n",
    "        label_to_id = {v: k for k, v in id_to_label.items()}\n",
    "\n",
    "        # for the related_only dataset\n",
    "        # ignore id 0, which represent no relation\n",
    "        id_to_label_related_only = {k - 1: v for k, v in id_to_label.items() if k != 0}\n",
    "        label_to_id_related_only = {v: k for k, v in id_to_label_related_only.items()}\n",
    "\n",
    "        return {\n",
    "            \"id_to_label\": id_to_label,\n",
    "            \"label_to_id\": label_to_id,\n",
    "            \"id_to_label_related_only\": id_to_label_related_only,\n",
    "            \"label_to_id_related_only\": label_to_id_related_only,            \n",
    "        }\n",
    "\n",
    "    def _write_data_to_file(self, dataframe: DataFrame, subset: str):\n",
    "        \"\"\"Write data in a dataframe to train/val/test file\"\"\"\n",
    "        lines = \"\"\n",
    "        for _, row in dataframe.iterrows():\n",
    "            lines += row.to_json() + \"\\n\"\n",
    "        with open(self.get_dataset_file_name(subset), \"w\") as file:\n",
    "            file.write(lines)\n",
    "\n",
    "    def _save_metadata(self, metadata: dict):\n",
    "        \"\"\"Save metadata to metadata.json\"\"\"\n",
    "        # create metadata file\n",
    "        if not os.path.exists(METADATA_FILE_NAME):\n",
    "            print(f\"Create metadata file at {METADATA_FILE_NAME}\")\n",
    "            with open(METADATA_FILE_NAME, \"w\") as f:\n",
    "                f.write(\"{}\\n\")\n",
    "\n",
    "        # add metadata\n",
    "        print(\"Saving metadata\")\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            root_metadata = json.load(f)\n",
    "        with open(METADATA_FILE_NAME, \"w\") as f:\n",
    "            root_metadata[self.DATASET_NAME] = metadata\n",
    "            json.dump(root_metadata, f, indent=4)\n",
    "\n",
    "    @classmethod\n",
    "    def get_dataset_file_name(cls, key: str) -> str:\n",
    "        return os.path.join(PROCESSED_DATA_DIR, f\"{cls.DATASET_NAME.lower()}_{key}.json\")\n"
   ],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Concrete preprocessor for each dataset:"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class SemEval2010Task8Preprocessor(AbstractPreprocessor):\n",
    "    DATASET_NAME = \"SemEval2010Task8\"\n",
    "    NO_RELATION_LABEL = \"Other\"\n",
    "    RAW_TRAIN_FILE_NAME = os.path.join(DATASET_MAPPING[\"SemEval2010Task8\"][\"dir\"],\n",
    "                                       \"SemEval2010_task8_training/TRAIN_FILE.TXT\")\n",
    "    RAW_TEST_FILE_NAME = os.path.join(DATASET_MAPPING[\"SemEval2010Task8\"][\"dir\"],\n",
    "                                      \"SemEval2010_task8_testing_keys/TEST_FILE_FULL.TXT\")\n",
    "    RAW_TRAIN_DATA_SIZE = 8000\n",
    "    RAW_TEST_DATA_SIZE = 2717\n",
    "\n",
    "    def _preprocess_data(self):\n",
    "        print(\"Processing training data\")\n",
    "        train_data = self._process_file(\n",
    "            self.RAW_TRAIN_FILE_NAME,\n",
    "            self.RAW_TRAIN_DATA_SIZE,\n",
    "            ADD_REVERSE_RELATIONSHIP,\n",
    "        )\n",
    "\n",
    "        print(\"Processing test data\")\n",
    "        test_data = self._process_file(\n",
    "            self.RAW_TEST_FILE_NAME,\n",
    "            self.RAW_TEST_DATA_SIZE,\n",
    "            False,\n",
    "        )\n",
    "\n",
    "        print(\"Splitting train & validate data\")\n",
    "        train_data, val_data = train_test_split(train_data, shuffle=True, random_state=SEED)\n",
    "\n",
    "        return train_data, val_data, test_data\n",
    "\n",
    "    def _process_file(self, file_name: str, dataset_size: int, add_reverse: bool) -> DataFrame:\n",
    "        raw_sentences = []\n",
    "        labels = []\n",
    "        with open(file_name) as f:\n",
    "            for _ in tqdm(range(dataset_size)):\n",
    "                sent = f.readline()\n",
    "                label, sub, obj = self._process_label(f.readline())\n",
    "                labels.append(label)\n",
    "                raw_sentences.append(self._process_sentence(sent, sub, obj))\n",
    "                if label == \"Other\" and add_reverse:\n",
    "                    labels.append(label)\n",
    "                    raw_sentences.append(self._process_sentence(sent, obj, sub))\n",
    "                f.readline()\n",
    "                f.readline()\n",
    "\n",
    "        return self._clean_data(raw_sentences, labels)\n",
    "\n",
    "    @staticmethod\n",
    "    def _process_sentence(sentence: str, sub: int, obj: int) -> str:\n",
    "        return sentence.split(\"\\t\")[1][1:-2] \\\n",
    "            .replace(f\"<e{sub}>\", SUB_START_CHAR) \\\n",
    "            .replace(f\"</e{sub}>\", SUB_END_CHAR) \\\n",
    "            .replace(f\"<e{obj}>\", OBJ_START_CHAR) \\\n",
    "            .replace(f\"</e{obj}>\", OBJ_END_CHAR)\n",
    "\n",
    "    @staticmethod\n",
    "    def _process_label(label: str) -> Tuple[str, int, int]:\n",
    "        label = label.strip()\n",
    "        if label == \"Other\":\n",
    "            return label, 1, 2\n",
    "        nums = list(filter(str.isdigit, label))\n",
    "        return label, int(nums[0]), int(nums[1])\n",
    "\n",
    "\n",
    "class GIDSPreprocessor(AbstractPreprocessor):\n",
    "    DATASET_NAME = \"GIDS\"\n",
    "    RAW_TRAIN_FILE_NAME = os.path.join(DATASET_MAPPING[\"GIDS\"][\"dir\"], \"train.tsv\")\n",
    "    RAW_VAL_FILE_NAME = os.path.join(DATASET_MAPPING[\"GIDS\"][\"dir\"], \"val.tsv\")\n",
    "    RAW_TEST_FILE_NAME = os.path.join(DATASET_MAPPING[\"GIDS\"][\"dir\"], \"test.tsv\")\n",
    "    TRAIN_SIZE = 11297\n",
    "    VAL_SIZE = 1864\n",
    "    TEST_SIZE = 5663\n",
    "    NO_RELATION_LABEL = \"NA\"\n",
    "\n",
    "    def _process_file(self, file_name: str, add_reverse: bool) -> DataFrame:\n",
    "        \"\"\"\n",
    "        Process a file in batches\n",
    "        Return the total data size\n",
    "        \"\"\"\n",
    "        with open(file_name) as in_file:\n",
    "            lines = in_file.readlines()\n",
    "            \n",
    "        raw_sentences = []\n",
    "        labels = []\n",
    "\n",
    "        for line in tqdm(lines):\n",
    "            _, _, sub, obj, label, sent = line.split(\"\\t\")\n",
    "            sent = sent.replace(\"###END###\", \"\")\n",
    "\n",
    "            # add subject markup\n",
    "            new_sub = SUB_START_CHAR + \" \" + sub.replace(\"_\", \" \") + \" \" + SUB_END_CHAR\n",
    "            new_obj = OBJ_START_CHAR + \" \" +  obj.replace(\"_\", \" \") + \" \" + OBJ_END_CHAR\n",
    "            sent = sent.replace(sub, new_sub).replace(obj, new_obj)\n",
    "            raw_sentences.append(sent)\n",
    "            labels.append(label)\n",
    "\n",
    "            if add_reverse and label == self.NO_RELATION_LABEL:\n",
    "                new_sub = OBJ_START_CHAR + \" \" + sub.replace(\"_\", \" \") + \" \" + OBJ_END_CHAR\n",
    "                new_obj = SUB_START_CHAR + \" \" + obj.replace(\"_\", \" \") + \" \" + SUB_END_CHAR\n",
    "                sent = sent.replace(sub, new_sub).replace(obj, new_obj)\n",
    "                raw_sentences.append(sent)\n",
    "                labels.append(label)\n",
    "\n",
    "        return self._clean_data(raw_sentences, labels)\n",
    "\n",
    "    def _preprocess_data(self):\n",
    "        \n",
    "        print(\"Process train dataset\")\n",
    "        train_data = self._process_file(\n",
    "            self.RAW_TRAIN_FILE_NAME,\n",
    "            ADD_REVERSE_RELATIONSHIP,\n",
    "        )\n",
    "\n",
    "        print(\"Process val dataset\")\n",
    "        val_data = self._process_file(\n",
    "            self.RAW_VAL_FILE_NAME,\n",
    "            False,\n",
    "        )\n",
    "        \n",
    "        print(\"Process test dataset\")\n",
    "        test_data = self._process_file(\n",
    "            self.RAW_TEST_FILE_NAME,\n",
    "            False,\n",
    "        )\n",
    "\n",
    "        return train_data, val_data, test_data"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Factory method to create preprocessors:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_preprocessor_class(dataset_name: str = DATASET_NAME):\n",
    "    return globals()[f\"{dataset_name}Preprocessor\"]\n",
    "        \n",
    "def get_preprocessor(dataset_name: str = DATASET_NAME)-> AbstractPreprocessor:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL, use_fast=True)\n",
    "    # some tokenizer, like GPTTokenizer, doesn't have pad_token\n",
    "    # in this case, we use eos token as pad token\n",
    "    if tokenizer.pad_token is None:\n",
    "        tokenizer.pad_token = tokenizer.eos_token\n",
    "        \n",
    "    preprocessors_class = get_preprocessor_class(dataset_name)\n",
    "    return preprocessors_class(tokenizer)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Preprocess data:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true
   },
   "cell_type": "code",
   "source": [
    "preprocessor = get_preprocessor()\n",
    "preprocessor.preprocess_data(reprocess=True)"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "---> Preprocessing GIDS dataset <---\n",
      "The following files will be overwritten:\n",
      "- /media/dthung1602/WORKING/bert-relation-extraction/data/processed/gids_train.json\n",
      "- /media/dthung1602/WORKING/bert-relation-extraction/data/processed/gids_val.json\n",
      "- /media/dthung1602/WORKING/bert-relation-extraction/data/processed/gids_test.json\n",
      "Process train dataset\n",
      "\n",
      "Process val dataset\n",
      "\n",
      "Process test dataset\n",
      "\n",
      "Saving to json files\n",
      "Saving metadata\n",
      "Creating secondary files for train data\n",
      "\n",
      "Creating secondary files for test data\n",
      "\n",
      "Creating secondary files for val data\n",
      "\n",
      "Updating metadata.json\n",
      "---> Done ! <---\n"
     ]
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=11297.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "81d4e2f97f064b6985947569655601fb"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=1864.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "be7a5d2b167e4efeaaafd5798fbcfcaa"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=5663.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "98ca783355254c6099bd00bef2ec2057"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=11280.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ce1ae36758144730a71a64a55b43bb04"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=5653.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5a0267500f26474190797864742ae376"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, max=1864.0), HTML(value='')))",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f2d2cd2f3570405caf5cd8450268e870"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "## Dataset\n",
    "\n",
    "We adopt the \"smart batching\" technique from [here](https://towardsdatascience.com/divide-hugging-face-transformers-training-time-by-2-or-more-21bf7129db9q-21bf7129db9e)\n",
    "\n",
    "Here's a brief diagram:\n",
    "\n",
    "![Smart batching](https://raw.githubusercontent.com/dthung1602/bert-relation-extraction/notebook/images/uniform_length_padding.png)"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "class GenericDataset(IterableDataset):\n",
    "    \"\"\"A generic dataset for train/val/test data for both SemEval and GIDS dataset\"\"\"\n",
    "\n",
    "    def __init__(self, dataset_name: str, subset: str, batch_size: int, label_transform: str):\n",
    "        assert subset in [\"train\", \"val\", \"test\"]\n",
    "        assert label_transform in [\"none\", \"binary\", \"related_only\"]\n",
    "\n",
    "        file_name = subset if label_transform == \"none\" \\\n",
    "            else f\"{subset}_{label_transform}\"\n",
    "\n",
    "        preprocessor_class = get_preprocessor_class()\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            metadata = json.load(f)[dataset_name]\n",
    "\n",
    "        size = metadata[f\"{subset}_related_only_size\"] \\\n",
    "            if label_transform is \"related_only\" \\\n",
    "            else metadata[f\"{subset}_size\"]\n",
    "\n",
    "        self.subset = subset\n",
    "        self.batch_size = batch_size\n",
    "        self.length = math.ceil(size / batch_size)\n",
    "        self.file = open(preprocessor_class.get_dataset_file_name(file_name))\n",
    "\n",
    "        self.keep_test_order = self.subset == \"test\" and DATASET_MAPPING[dataset_name][\"keep_test_order\"]\n",
    "\n",
    "    def __del__(self):\n",
    "        if self.file:\n",
    "            self.file.close()\n",
    "\n",
    "    def __iter__(self):\n",
    "        \"\"\"\n",
    "        Implement \"smart batching\"\n",
    "        \"\"\"\n",
    "\n",
    "        data = [json.loads(line) for line in self.file]\n",
    "        if not self.keep_test_order:\n",
    "            data = sorted(data, key=lambda x: sum(x[\"attention_mask\"]))\n",
    "\n",
    "        new_data = []\n",
    "\n",
    "        while len(data) > 0:\n",
    "            if self.keep_test_order or len(data) < self.batch_size:\n",
    "                idx = 0\n",
    "            else:\n",
    "                idx = randint(0, len(data) - self.batch_size)\n",
    "            batch = data[idx:idx + self.batch_size]\n",
    "            max_len = max([sum(b[\"attention_mask\"]) for b in batch])\n",
    "\n",
    "            for b in batch:\n",
    "                input_data = {}\n",
    "                for k, v in b.items():\n",
    "                    if k != \"label\":\n",
    "                        if isinstance(v, list):\n",
    "                            input_data[k] = torch.tensor(v[:max_len])\n",
    "                        else:\n",
    "                            input_data[k] = torch.tensor(v)\n",
    "                label = torch.tensor(b[\"label\"])\n",
    "                new_data.append((input_data, label))\n",
    "\n",
    "            del data[idx:idx + self.batch_size]\n",
    "\n",
    "        yield from new_data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def as_batches(self):\n",
    "        input_data = []\n",
    "        label = []\n",
    "        \n",
    "        def create_batch():\n",
    "            return (\n",
    "                {k: torch.stack([x[k] for x in input_data]).cuda() for k in input_data[0].keys()},\n",
    "                torch.tensor(label).cuda()\n",
    "            )\n",
    "        \n",
    "        for ip, l in self:\n",
    "            input_data.append(ip)\n",
    "            label.append(l)\n",
    "            if len(input_data) == self.batch_size:\n",
    "                yield create_batch()\n",
    "                input_data.clear()\n",
    "                label.clear()\n",
    "\n",
    "        yield create_batch()"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "## Classifiers"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "class BaseClassifier(LightningModule, ABC):\n",
    "    \"\"\"\n",
    "    Base class of all classifiers\n",
    "    \"\"\"\n",
    "    \n",
    "    dataset_label_transform = None\n",
    "    num_classes = None\n",
    "\n",
    "    @abstractmethod\n",
    "    def loss_function(self, logits: Tensor, label: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Calculate the loss of the model\n",
    "        It MUST take care of the last activation layer\n",
    "        \"\"\"\n",
    "        pass\n",
    "    \n",
    "    @abstractmethod\n",
    "    def log_metrics(self, epoch_type: str, logits: Tensor, label: Tensor) -> dict:\n",
    "        pass\n",
    "\n",
    "    def __init__(self, pretrained_language_model, dataset_name, batch_size, learning_rate, decay_lr_speed,\n",
    "                 activation_function, weight_decay, linear_size, dropout_p):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        self.language_model = AutoModel.from_pretrained(pretrained_language_model)\n",
    "        config = self.language_model.config\n",
    "        self.max_seq_len = config.max_position_embeddings\n",
    "        self.hidden_size = config.hidden_size\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout_p)\n",
    "        self.linear = nn.Linear(self.hidden_size * self.max_seq_len, linear_size)\n",
    "        self.activation_function = getattr(nn, activation_function)()\n",
    "        self.linear_output = nn.Linear(linear_size, self.num_classes)\n",
    "\n",
    "    def forward(self, sub_start_pos, sub_end_pos,\n",
    "                obj_start_pos, obj_end_pos, *args, **kwargs) -> Tensor:\n",
    "        language_model_output = self.language_model(*args, **kwargs)\n",
    "        if isinstance(language_model_output, tuple):\n",
    "            language_model_output = language_model_output[0]\n",
    "\n",
    "        bz, seq_len, _ = language_model_output.shape\n",
    "        x = torch.zeros([bz, self.max_seq_len, self.hidden_size]).cuda()\n",
    "        x[:,:seq_len,:] = language_model_output\n",
    "\n",
    "        x = x.reshape([bz, -1])\n",
    "        x = self.dropout(x)\n",
    "        x = self.linear(x)\n",
    "        x = self.activation_function(x)\n",
    "        x = self.dropout(x)\n",
    "        logits = self.linear_output(x)\n",
    "\n",
    "        return logits\n",
    "   \n",
    "    def train_dataloader(self) -> DataLoader:\n",
    "        return self.__get_dataloader(\"train\")\n",
    "\n",
    "    def val_dataloader(self) -> DataLoader:\n",
    "        return self.__get_dataloader(\"val\")\n",
    "\n",
    "    def test_dataloader(self) -> DataLoader:\n",
    "        return self.__get_dataloader(\"test\")\n",
    "\n",
    "    def __get_dataloader(self, subset: str) -> DataLoader:\n",
    "        batch_size = self.hparams.batch_size\n",
    "        dataset = GenericDataset(\n",
    "            self.hparams.dataset_name,\n",
    "            subset, \n",
    "            batch_size, \n",
    "            self.dataset_label_transform\n",
    "        )\n",
    "        return DataLoader(\n",
    "            dataset,\n",
    "            batch_size=batch_size,\n",
    "            num_workers=1\n",
    "        )\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = AdamW(\n",
    "            [p for p in self.parameters() if p.requires_grad],\n",
    "            lr=self.hparams.learning_rate,\n",
    "            weight_decay=self.hparams.weight_decay\n",
    "        )\n",
    "        scheduler = LambdaLR(optimizer, lambda epoch: self.hparams.decay_lr_speed[epoch])\n",
    "        return [optimizer], [scheduler]\n",
    "    \n",
    "    def training_step(self, batch: Tuple[dict, Tensor], batch_nb: int) -> dict:\n",
    "        input_data, label = batch\n",
    "        logits = self(**input_data)\n",
    "\n",
    "        loss = self.loss_function(logits, label)\n",
    "        log = {\"train_loss\": loss}\n",
    "\n",
    "        return {\"loss\": loss, \"log\": log}\n",
    "\n",
    "    def __eval_step(self, batch:  Tuple[dict, Tensor]) -> dict:\n",
    "        input_data, label = batch\n",
    "        logits = self(**input_data)\n",
    "\n",
    "        return {\n",
    "            \"logits\": logits,\n",
    "            \"label\": label,\n",
    "        }\n",
    "    \n",
    "    def validation_step(self, batch: Tuple[dict, Tensor], batch_nb: int) -> dict:\n",
    "        return self.__eval_step(batch)\n",
    "    \n",
    "    def test_step(self, batch: Tuple[dict, Tensor], batch_nb: int) -> dict:\n",
    "        return self.__eval_step(batch)\n",
    "\n",
    "    def __eval_epoch_end(self, epoch_type: str, outputs: Iterable[dict]) -> dict:\n",
    "        assert epoch_type in [\"test\", \"val\"]\n",
    "        \n",
    "        logits = torch.cat([x[\"logits\"] for x in outputs]).cpu()\n",
    "        label = torch.cat([x[\"label\"] for x in outputs]).cpu()\n",
    "        \n",
    "        logs = self.log_metrics(epoch_type, logits, label)\n",
    "        \n",
    "        return {\"progress_bar\": logs}\n",
    "    \n",
    "    def validation_epoch_end(self, outputs: Iterable[dict]) -> dict:\n",
    "        return self.__eval_epoch_end(\"val\", outputs)\n",
    "\n",
    "    def test_epoch_end(self, outputs: Iterable[dict]) -> dict:\n",
    "        return self.__eval_epoch_end(\"test\", outputs)\n",
    "\n",
    "    def numeric_labels_to_text(self, label):\n",
    "        \"\"\"Revert labels from number to text\"\"\"\n",
    "        if self.dataset_label_transform == \"binary\":\n",
    "            label = [\"Positive\" if x else \"Negative\" for x in label]\n",
    "        else:\n",
    "            with open(METADATA_FILE_NAME) as f:\n",
    "                meta = json.load(f)[self.hparams.dataset_name]\n",
    "            if self.dataset_label_transform == \"none\":\n",
    "                mapping = meta[\"id_to_label\"]\n",
    "            else:\n",
    "                mapping = meta[\"id_to_label_related_only\"]\n",
    "            label = [mapping[str(int(x))] for x in label]\n",
    "        return label\n",
    "\n",
    "    @staticmethod\n",
    "    def plot_confusion_matrix(predicted_label, label) -> Figure:\n",
    "        result = confusion_matrix(label, predicted_label)\n",
    "        display = ConfusionMatrixDisplay(result)\n",
    "        fig, ax = plt.subplots(figsize=(16, 12))\n",
    "        display.plot(cmap=plt.cm.get_cmap(\"Blues\"), ax=ax, xticks_rotation='vertical')\n",
    "        return fig\n",
    "    \n",
    "    def log_confusion_matrix(self, prefix: str, predicted_label: Tensor, label: Tensor):\n",
    "        predicted_label = self.numeric_labels_to_text(predicted_label)\n",
    "        label = self.numeric_labels_to_text(label)\n",
    "        fig = self.plot_confusion_matrix(predicted_label, label)\n",
    "        self.logger.experiment.log_image(f\"{prefix}_confusion_matrix\", fig)\n",
    "\n",
    "\n",
    "class MulticlassClassifier(BaseClassifier, ABC):\n",
    "    \"\"\"\n",
    "    Base class for multiclass classifiers\n",
    "    \"\"\"\n",
    "\n",
    "    def loss_function(self, logits: Tensor, label: Tensor)-> Tensor:\n",
    "        return F.cross_entropy(logits, label)\n",
    "\n",
    "    @staticmethod\n",
    "    def logits_to_label(logits: Tensor) -> Tensor:\n",
    "        return torch.argmax(logits, dim=-1)\n",
    "\n",
    "    def log_metrics(self, epoch_type: str, logits: Tensor, label: Tensor) -> dict:\n",
    "        predicted_label = self.logits_to_label(logits)\n",
    "        self.log_confusion_matrix(epoch_type, predicted_label, label)\n",
    "\n",
    "        logs = {\n",
    "            f\"{epoch_type}_avg_loss\": float(self.loss_function(logits, label)),\n",
    "            f\"{epoch_type}_acc\": accuracy_score(label, predicted_label),\n",
    "            f\"{epoch_type}_pre_weighted\": precision_score(label, predicted_label, average=\"weighted\"),\n",
    "            f\"{epoch_type}_rec_weighted\": recall_score(label, predicted_label, average=\"weighted\"),\n",
    "            f\"{epoch_type}_f1_weighted\": f1_score(label, predicted_label, average=\"weighted\"),\n",
    "            f\"{epoch_type}_pre_macro\": precision_score(label, predicted_label, average=\"macro\"),\n",
    "            f\"{epoch_type}_rec_macro\": recall_score(label, predicted_label, average=\"macro\"),\n",
    "            f\"{epoch_type}_f1_macro\": f1_score(label, predicted_label, average=\"macro\"),\n",
    "        }\n",
    "\n",
    "        for k, v in logs.items():\n",
    "            self.logger.experiment.log_metric(k, v)\n",
    "\n",
    "        return logs\n",
    "\n",
    "\n",
    "class StandardClassifier(MulticlassClassifier):\n",
    "    \"\"\"\n",
    "    A classifier that can recognize the \"not related\" as well as other relations\n",
    "    \"\"\"\n",
    "    \n",
    "    dataset_label_transform = \"none\"\n",
    "\n",
    "    def __init__(self, dataset_name, **kwargs):\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            self.num_classes = len(json.load(f)[dataset_name][\"label_to_id\"])\n",
    "        self.test_proposed_answer = None\n",
    "        super().__init__(dataset_name=dataset_name, **kwargs)\n",
    "\n",
    "\n",
    "    def log_metrics(self, epoch_type: str, logits: Tensor, label: Tensor)-> dict:\n",
    "        if epoch_type == \"test\":\n",
    "            self.test_proposed_answer = self.logits_to_label(logits).tolist()\n",
    "        self.__log_precision_recall_curve(epoch_type, logits, label)\n",
    "        return super().log_metrics(epoch_type, logits, label)\n",
    "\n",
    "    def __log_precision_recall_curve(self, epoch_type: str, logits: Tensor, label: Tensor):\n",
    "        \"\"\"\n",
    "        Log the micro-averaged precision recall curve\n",
    "        Ref: https://scikit-learn.org/stable/auto_examples/model_selection/plot_precision_recall.html\n",
    "        \"\"\"\n",
    "\n",
    "        label = torch.tensor(label_binarize(label, classes=list(range(self.num_classes)))).flatten()\n",
    "        logits = logits.flatten()\n",
    "\n",
    "        pre, rec, thresholds = precision_recall_curve(label, logits)\n",
    "        f1s = 2 * pre * rec / (pre + rec)\n",
    "        ix = np.argmax(f1s)\n",
    "\n",
    "        fig, ax = plt.subplots(figsize=(10, 10))\n",
    "        \n",
    "        # render the baseline curves as background for comparison\n",
    "        background = DATASET_MAPPING[self.hparams.dataset_name][\"precision_recall_curve_baseline_img\"]\n",
    "        if background:\n",
    "            img = plt.imread(background)\n",
    "            ax.imshow(img, extent=[0, 1, 0, 1])\n",
    "        \n",
    "        no_skill = len(label[label == 1]) / len(label)\n",
    "        ax.plot(rec, pre, label=\"Our proposed model\", color=\"blue\")\n",
    "        ax.set_xlabel(\"Recall\")\n",
    "        ax.set_ylabel(\"Precision\")\n",
    "        ax.legend()\n",
    "\n",
    "        self.logger.experiment.log_image(f\"{epoch_type}_pre_rec_curve\", fig)\n",
    "        self.logger.experiment.log_metric(\n",
    "            f\"{epoch_type}_average_precision_score_micro\", \n",
    "            average_precision_score(label, logits, average=\"micro\")\n",
    "        )\n",
    "\n",
    "\n",
    "class BinaryClassifier(BaseClassifier):\n",
    "    \"\"\"\n",
    "    A binary classifier that picks out \"not-related\" sentences\n",
    "    \"\"\"\n",
    "    \n",
    "    dataset_label_transform = \"binary\"\n",
    "    num_classes = 1\n",
    "\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.thresholds = {}\n",
    "\n",
    "    def forward(self, *args, **kwargs):\n",
    "        return super().forward(*args, **kwargs).flatten()\n",
    "\n",
    "    @staticmethod\n",
    "    def yhat_to_label(y_hat: Tensor, threshold: float) -> Tensor:\n",
    "        return (y_hat > threshold).long()\n",
    "\n",
    "    def loss_function(self, logits: Tensor, label: Tensor) -> Tensor:\n",
    "        return F.binary_cross_entropy_with_logits(logits, label.float())\n",
    "    \n",
    "    def log_metrics(self, epoch_type: str, logits: Tensor, label: Tensor) -> dict:\n",
    "        y_hat = torch.sigmoid(logits)\n",
    "        \n",
    "        if epoch_type == \"val\":\n",
    "            self.__find_thresholds(y_hat, label)\n",
    "        \n",
    "        self.__log_output_distribution(epoch_type, y_hat, label)\n",
    "        \n",
    "        logs = {\n",
    "            f\"{epoch_type}_avg_loss\": float(self.loss_function(logits, label)),\n",
    "            f\"{epoch_type}_roc_auc\": self.__roc_auc_score(label, y_hat),\n",
    "        }\n",
    "        \n",
    "        for criteria, threshold in self.thresholds.items():\n",
    "            prefix = f\"{epoch_type}_{criteria}\"\n",
    "            predicted_label = self.yhat_to_label(y_hat, threshold)\n",
    "            self.log_confusion_matrix(prefix, predicted_label, label)\n",
    "            \n",
    "            logs[f\"{prefix}_acc\"] = accuracy_score(label, predicted_label)\n",
    "            logs[f\"{prefix}_pre\"] = precision_score(label, predicted_label, average=\"binary\")\n",
    "            logs[f\"{prefix}_rec\"] = recall_score(label, predicted_label, average=\"binary\")\n",
    "            logs[f\"{prefix}_f1\"] = f1_score(label, predicted_label, average=\"binary\")\n",
    "\n",
    "        for k, v in logs.items():\n",
    "            self.logger.experiment.log_metric(k, v)\n",
    "            \n",
    "        return logs\n",
    "    \n",
    "    @staticmethod\n",
    "    def __roc_auc_score(label: Tensor, y_hat: Tensor) -> float:\n",
    "        try:\n",
    "            return roc_auc_score(label, y_hat)\n",
    "        except ValueError:\n",
    "            return 0\n",
    "    \n",
    "    def __find_thresholds(self, y_hat: Tensor, label: Tensor):\n",
    "        \"\"\"\n",
    "        Find 3 classification thresholds based on 3 criteria:\n",
    "        - The one that yields highest accuracy\n",
    "        - The \"best point\" in the ROC curve\n",
    "        - The one that yields highest f1\n",
    "        The results are logged and stored in self.threshold\n",
    "        \"\"\"\n",
    "        # best accuracy\n",
    "        best_acc = 0\n",
    "        best_acc_threshold = None\n",
    "        for y in y_hat:\n",
    "            y_predicted = self.yhat_to_label(y_hat, threshold=y)\n",
    "            acc = accuracy_score(label, y_predicted)\n",
    "            if best_acc < acc:\n",
    "                best_acc = acc\n",
    "                best_acc_threshold = y\n",
    "        self.thresholds[\"best_acc\"] = best_acc_threshold\n",
    "\n",
    "        # ROC curve\n",
    "        # https://machinelearningmastery.com/threshold-moving-for-imbalanced-classification/\n",
    "        fpr, tpr, thresholds = roc_curve(label, y_hat)\n",
    "        gmeans = tpr * (1 - fpr)\n",
    "        ix = np.argmax(gmeans)\n",
    "        self.thresholds[\"best_roc\"] = thresholds[ix]\n",
    "        fig, ax = plt.subplots(figsize=(16, 12))\n",
    "        ax.plot([0,1], [0,1], linestyle=\"--\", label=\"No Skill\")\n",
    "        ax.plot(fpr, tpr, marker=\".\", label=\"Logistic\")\n",
    "        ax.scatter(fpr[ix], tpr[ix], marker=\"o\", color=\"black\", label=\"Best\")\n",
    "        ax.set_xlabel(\"False Positive Rate\")\n",
    "        ax.set_ylabel(\"True Positive Rate\")\n",
    "        ax.legend()\n",
    "        self.logger.experiment.log_image(\"roc_curve\", fig)\n",
    "\n",
    "        # precision recall curve\n",
    "        # https://machinelearningmastery.com/roc-curves-and-precision-recall-curves-for-classification-in-python/\n",
    "        pre, rec, thresholds = precision_recall_curve(label, y_hat)\n",
    "        f1s = 2 * pre * rec / (pre + rec)\n",
    "        ix = np.argmax(f1s)\n",
    "        self.thresholds[\"best_f1\"] = thresholds[ix]\n",
    "\n",
    "        fig, ax = plt.subplots(figsize=(16, 12))\n",
    "        no_skill = len(label[label == 1]) / len(label)\n",
    "        ax.plot([0, 1], [no_skill, no_skill], linestyle=\"--\", label=\"No Skill\")\n",
    "        ax.plot(rec, pre, marker=\".\", label=\"Logistic\")\n",
    "        ax.scatter(rec[ix], pre[ix], marker=\"o\", color=\"black\", label=\"Best F1\")\n",
    "        ax.set_xlabel(\"Recall\")\n",
    "        ax.set_ylabel(\"Precision\")\n",
    "        ax.legend()\n",
    "        self.logger.experiment.log_image(\"pre_rec_curve\", fig)\n",
    "\n",
    "        # log thresholds\n",
    "        for k, v in self.thresholds.items():\n",
    "            self.logger.experiment.log_metric(f\"threshold_{k}\", v)\n",
    "\n",
    "    def __log_output_distribution(self, epoch_type: str, y_hat: Tensor, label: Tensor):\n",
    "        \"\"\"\n",
    "        Log the distribution of the model output and 3 thresholds with log scale and linear scale\n",
    "        \"\"\"\n",
    "        y_neg = y_hat[label == 0].numpy()\n",
    "        y_pos = y_hat[label == 1].numpy()\n",
    "\n",
    "        for scale in [\"linear\", \"log\"]:\n",
    "            fig, ax = plt.subplots(figsize=(16, 12))\n",
    "            ax.set_yscale(scale)\n",
    "            ax.hist([y_neg, y_pos], stacked=True, bins=50, label=[\"No relation\", \"Related\"])\n",
    "            ylim = ax.get_ylim()\n",
    "            for k, v in self.thresholds.items():\n",
    "                ax.plot([v, v], ylim, linestyle=\"--\", label=f\"{k} threshold\")\n",
    "            ax.legend()\n",
    "            self.logger.experiment.log_image(f\"{epoch_type}_distribution_{scale}_scale\", fig)\n",
    "        \n",
    "\n",
    "class RelationClassifier(MulticlassClassifier):\n",
    "    \"\"\"\n",
    "    A classifier that recognizes relations except for \"not-related\"\n",
    "    \"\"\"\n",
    "\n",
    "    dataset_label_transform = \"related_only\"\n",
    "\n",
    "    def __init__(self, dataset_name, **kwargs):\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            self.num_classes = len(json.load(f)[dataset_name][\"label_to_id_related_only\"])\n",
    "        super().__init__(dataset_name=dataset_name, **kwargs)"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## The official scorer\n",
    "\n",
    "Some datasets comes with official scorers. We will define them in this session."
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true
   },
   "cell_type": "code",
   "source": [
    "class AbstractScorer(ABC):\n",
    "    def __init__(self, experiment_no: int, logger):\n",
    "        self.experiment_no = experiment_no\n",
    "        self.logger = logger\n",
    "\n",
    "    @abstractmethod\n",
    "    def score(self, proposed_answer: dict):\n",
    "        pass\n",
    "\n",
    "class SemEval2010Task8Scorer(AbstractScorer):\n",
    "    RESULT_FILE = \"semeval2010_task8_official_score_{}_{}.txt\"\n",
    "    PROPOSED_ANSWER_FILE = \"semeval2010_task8_proposed_answer.txt\"\n",
    "    SCORER = os.path.join(DATASET_MAPPING[\"SemEval2010Task8\"][\"dir\"], \"SemEval2010_task8_scorer-v1.2/semeval2010_task8_scorer-v1.2.pl\")\n",
    "    FORMAT_CHECKER = os.path.join(DATASET_MAPPING[\"SemEval2010Task8\"][\"dir\"], \"SemEval2010_task8_scorer-v1.2/semeval2010_task8_format_checker.pl\")\n",
    "    ANSWER_KEY = os.path.join(DATASET_MAPPING[\"SemEval2010Task8\"][\"dir\"], \"SemEval2010_task8_testing_keys/TEST_FILE_KEY.TXT\")\n",
    "\n",
    "    def score(self, proposed_answer: dict):\n",
    "        # write test_result to file\n",
    "        with open(METADATA_FILE_NAME) as f:\n",
    "            metadata = json.load(f)\n",
    "            id_to_label = {int(k): v for k, v in metadata[DATASET_NAME][\"id_to_label\"].items()}\n",
    "\n",
    "        for criteria, answer in proposed_answer.items():\n",
    "            result_file = self.RESULT_FILE.format(self.experiment_no, criteria)\n",
    "            i = 8001\n",
    "            with open(self.PROPOSED_ANSWER_FILE, \"w\") as f:\n",
    "                for r in answer:\n",
    "                    f.write(f\"{i}\\t{id_to_label[r]}\\n\")\n",
    "                    i += 1\n",
    "\n",
    "            # call the official scorer\n",
    "            os.system(f\"perl {self.FORMAT_CHECKER} {self.PROPOSED_ANSWER_FILE}\")\n",
    "            os.system(f\"perl {self.SCORER} {self.PROPOSED_ANSWER_FILE} {self.ANSWER_KEY} > {result_file}\")\n",
    "\n",
    "            # log the official score\n",
    "            with open(result_file) as f:\n",
    "                result = f.read()\n",
    "                print(f\">>> Classifier with criteria: {criteria} <<<\")\n",
    "                print(result)\n",
    "                print(\"\\n\\n\")\n",
    "            self.logger.experiment.log_artifact(result_file)\n",
    "\n",
    "def get_official_scorer(experiment_no: int, logger, dataset_name: str = DATASET_NAME) -> AbstractScorer:\n",
    "    cls = globals().get(dataset_name + \"Scorer\")\n",
    "    if cls:\n",
    "        return cls(experiment_no, logger)"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Claiming back memory & disk space\n",
    "\n",
    "See [this](https://stackoverflow.com/a/61707643/7342188) and [this](https://stackoverflow.com/a/57860310/7342188)"
   ]
  },
  {
   "metadata": {
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "1 / 0"
   ],
   "execution_count": 21,
   "outputs": [
    {
     "output_type": "error",
     "ename": "ZeroDivisionError",
     "evalue": "division by zero",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mZeroDivisionError\u001B[0m                         Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-21-bc757c3fda29>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m\u001B[0m\n\u001B[0;32m----> 1\u001B[0;31m \u001B[0;36m1\u001B[0m \u001B[0;34m/\u001B[0m \u001B[0;36m0\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m",
      "\u001B[0;31mZeroDivisionError\u001B[0m: division by zero"
     ]
    }
   ]
  },
  {
   "metadata": {
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "trainer = classifier = rel_trainer = rel_classifier = bin_trainer = bin_classifier = None\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ],
   "execution_count": 23,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Training standard classifier"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true
   },
   "cell_type": "code",
   "source": [
    "GPUS = 1\n",
    "MIN_EPOCHS = MAX_EPOCHS = 3\n",
    "\n",
    "BATCH_SIZE = 16\n",
    "LEARNING_RATE = 2e-05\n",
    "LEARNING_RATE_DECAY_SPEED = [1, 1, 0.75, 0.5, 0.25, 0.1, 0.075, 0.05, 0.025, 0.01]\n",
    "\n",
    "LINEAR_SIZE = 1024\n",
    "\n",
    "DROPOUT_P = 0.2\n",
    "ACTIVATION_FUNCTION = \"PReLU\"\n",
    "WEIGHT_DECAY = 0.01 # default = 0.01"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true
   },
   "cell_type": "code",
   "source": [
    "logger = NeptuneLogger(\n",
    "    api_key=NEPTUNE_API_TOKEN,\n",
    "    project_name=NEPTUNE_PROJECT_NAME,\n",
    "    close_after_fit=False,\n",
    ")\n",
    "\n",
    "try:\n",
    "    for i in range(4):\n",
    "        print(f\"--------- EXPERIMENT {i} ---------\")\n",
    "\n",
    "        classifier = trainer = None\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        trainer = LightningTrainer(\n",
    "            gpus=GPUS,\n",
    "            min_epochs=MIN_EPOCHS,\n",
    "            max_epochs=MAX_EPOCHS,\n",
    "            default_root_dir=CHECKPOINT_DIR,\n",
    "            reload_dataloaders_every_epoch=True, # needed as we loop over a file,\n",
    "            deterministic=False,\n",
    "            checkpoint_callback=False,\n",
    "            logger=logger\n",
    "        )\n",
    "\n",
    "        classifier = StandardClassifier(\n",
    "            pretrained_language_model=PRETRAINED_MODEL,\n",
    "            dataset_name=DATASET_NAME,\n",
    "            batch_size=BATCH_SIZE,\n",
    "            learning_rate=LEARNING_RATE,\n",
    "            decay_lr_speed=LEARNING_RATE_DECAY_SPEED,\n",
    "            dropout_p=DROPOUT_P,\n",
    "            activation_function=ACTIVATION_FUNCTION,\n",
    "            weight_decay=WEIGHT_DECAY,\n",
    "            linear_size=LINEAR_SIZE,\n",
    "        )\n",
    "\n",
    "        trainer.fit(classifier)\n",
    "        trainer.test(classifier)\n",
    "\n",
    "        scorer = get_official_scorer(i, logger)\n",
    "        if scorer:\n",
    "            scorer.score({\n",
    "                \"standard\": classifier.test_proposed_answer,\n",
    "            })\n",
    "        else:\n",
    "            print(\"No official scorer found\")\n",
    "\n",
    "except Exception as e:\n",
    "    logger.experiment.stop(str(e))\n",
    "    raise e\n",
    "\n",
    "else:\n",
    "    logger.experiment.stop()"
   ],
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "text": "https://ui.neptune.ai/hung/bert-relation-extraction/e/BERT-301\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "NeptuneLogger will work in online mode\n",
     "name": "stderr"
    },
    {
     "output_type": "stream",
     "text": "--------- EXPERIMENT 0 ---------\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "GPU available: True, used: True\nTPU available: False, using: 0 TPU cores\nCUDA_VISIBLE_DEVICES: [0]\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=440473133.0, style=ProgressStyle(descri…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "afe13149cfb44f88b23aa102231338a0"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "\n  | Name                | Type      | Params\n--------------------------------------------------\n0 | language_model      | BertModel | 109 M \n1 | dropout             | Dropout   | 0     \n2 | linear              | Linear    | 787 K \n3 | activation_function | Tanh      | 0     \n4 | linear_output       | Linear    | 5 K   \n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:209: RuntimeWarning: invalid value encountered in true_divide\n/opt/conda/lib/python3.7/site-packages/sklearn/metrics/_classification.py:1221: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n  _warn_prf(average, modifier, msg_start, len(result))\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dcc7c10088c24363bdd08862b61878df"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "stream",
     "text": "\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Testing', layout=Layout(flex='2'), max=…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8e2dc76c4c6e482589adbb5eeac6c994"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "--------------------------------------------------------------------------------\nTEST RESULTS\n{'test_acc': 0.7910843799752344,\n 'test_avg_loss': 0.6946395635604858,\n 'test_f1_macro': 0.796556195878245,\n 'test_f1_weighted': 0.7810788346092403,\n 'test_pre_macro': 0.7940858117023544,\n 'test_pre_weighted': 0.7839553351383417,\n 'test_rec_macro': 0.8104579170607679,\n 'test_rec_weighted': 0.7910843799752344}\n--------------------------------------------------------------------------------\n\nNo official scorer found\n--------- EXPERIMENT 1 ---------\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\nGPU available: True, used: True\nTPU available: False, using: 0 TPU cores\nCUDA_VISIBLE_DEVICES: [0]\n\n  | Name                | Type      | Params\n--------------------------------------------------\n0 | language_model      | BertModel | 109 M \n1 | dropout             | Dropout   | 0     \n2 | linear              | Linear    | 787 K \n3 | activation_function | Tanh      | 0     \n4 | linear_output       | Linear    | 5 K   \n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:209: RuntimeWarning: invalid value encountered in true_divide\n/opt/conda/lib/python3.7/site-packages/sklearn/metrics/_classification.py:1221: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n  _warn_prf(average, modifier, msg_start, len(result))\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f6dc4af7884a498fb8eb0eeb57eedee4"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:209: RuntimeWarning: invalid value encountered in true_divide\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "stream",
     "text": "\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Testing', layout=Layout(flex='2'), max=…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "947e60d6fcf340af8e837ddffec695b7"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "--------------------------------------------------------------------------------\nTEST RESULTS\n{'test_acc': 0.7778170882717141,\n 'test_avg_loss': 0.7352004647254944,\n 'test_f1_macro': 0.7832999070829824,\n 'test_f1_weighted': 0.766742538258222,\n 'test_pre_macro': 0.7824102061635398,\n 'test_pre_weighted': 0.7707975130483338,\n 'test_rec_macro': 0.7978208607064214,\n 'test_rec_weighted': 0.7778170882717141}\n--------------------------------------------------------------------------------\n\nNo official scorer found\n--------- EXPERIMENT 2 ---------\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\nGPU available: True, used: True\nTPU available: False, using: 0 TPU cores\nCUDA_VISIBLE_DEVICES: [0]\n\n  | Name                | Type      | Params\n--------------------------------------------------\n0 | language_model      | BertModel | 109 M \n1 | dropout             | Dropout   | 0     \n2 | linear              | Linear    | 787 K \n3 | activation_function | Tanh      | 0     \n4 | linear_output       | Linear    | 5 K   \n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:209: RuntimeWarning: invalid value encountered in true_divide\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/sklearn/metrics/_classification.py:1221: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n  _warn_prf(average, modifier, msg_start, len(result))\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e01f02887e6e40509d08a7b47ea81981"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "stream",
     "text": "\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Testing', layout=Layout(flex='2'), max=…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "08a23dc76eb74bce9e16acf86e49a660"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n",
     "name": "stderr"
    },
    {
     "output_type": "stream",
     "text": "--------------------------------------------------------------------------------\nTEST RESULTS\n{'test_acc': 0.7820626216168406,\n 'test_avg_loss': 0.7605665922164917,\n 'test_f1_macro': 0.7927035659274893,\n 'test_f1_weighted': 0.7764553912950387,\n 'test_pre_macro': 0.7867127297287052,\n 'test_pre_weighted': 0.7745966466502907,\n 'test_rec_macro': 0.802214598326939,\n 'test_rec_weighted': 0.7820626216168406}\n--------------------------------------------------------------------------------\n\nNo official scorer found\n--------- EXPERIMENT 3 ---------\n",
     "name": "stdout"
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\nGPU available: True, used: True\nTPU available: False, using: 0 TPU cores\nCUDA_VISIBLE_DEVICES: [0]\n\n  | Name                | Type      | Params\n--------------------------------------------------\n0 | language_model      | BertModel | 109 M \n1 | dropout             | Dropout   | 0     \n2 | linear              | Linear    | 787 K \n3 | activation_function | Tanh      | 0     \n4 | linear_output       | Linear    | 5 K   \n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": ""
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:212: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:110: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n/opt/conda/lib/python3.7/site-packages/sklearn/metrics/_classification.py:1221: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n  _warn_prf(average, modifier, msg_start, len(result))\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-processing data loading (e.g. batch size > 1), this can lead to unintended side effects since the samples will be duplicated.\n  warnings.warn(*args, **kwargs)\n/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f5c2ca0c7ff94bb6aac6f4fe0e3b89a4"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "text": "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 2 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n  warnings.warn(*args, **kwargs)\n",
     "name": "stderr"
    }
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "cell_type": "markdown",
   "source": [
    "## Training binary classifier"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "GPUS = 1\n",
    "\n",
    "BIN_MIN_EPOCHS = BIN_MAX_EPOCHS = 2\n",
    "\n",
    "BIN_BATCH_SIZE = 32\n",
    "BIN_LEARNING_RATE = 2e-05\n",
    "BIN_LEARNING_RATE_DECAY_SPEED = [1, 1, 0.75, 0.5, 0.5, 0.25, 0.25, 0.1, 0.075, 0.05, 0.025, 0.01]\n",
    "\n",
    "BIN_LINEAR_SIZE = 256\n",
    "\n",
    "BIN_DROPOUT_P = 0.2\n",
    "BIN_ACTIVATION_FUNCTION = \"PReLU\"\n",
    "BIN_WEIGHT_DECAY = 0.01 # default = 0.01\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "_kg_hide-output": true,
    "_kg_hide-input": false,
    "scrolled": true
   },
   "cell_type": "code",
   "source": [
    "bin_logger = NeptuneLogger(\n",
    "    api_key=NEPTUNE_API_TOKEN,\n",
    "    project_name=NEPTUNE_PROJECT_NAME,\n",
    "    close_after_fit=False,\n",
    ")\n",
    "\n",
    "try:\n",
    "    for i in range(4):\n",
    "        print(f\"--------- EXPERIMENT {i} ---------\")\n",
    "    \n",
    "        bin_classifier = bin_trainer = None\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "    \n",
    "        bin_trainer = LightningTrainer(\n",
    "            gpus=GPUS,\n",
    "            min_epochs=BIN_MIN_EPOCHS,\n",
    "            max_epochs=BIN_MAX_EPOCHS,\n",
    "            default_root_dir=CHECKPOINT_DIR,\n",
    "            reload_dataloaders_every_epoch=True, # needed as we loop over a file,\n",
    "            deterministic=False,\n",
    "            checkpoint_callback=False,\n",
    "            logger=bin_logger,\n",
    "        )\n",
    "    \n",
    "        bin_classifier = BinaryClassifier(\n",
    "            pretrained_language_model=PRETRAINED_MODEL,\n",
    "            dataset_name=DATASET_NAME,\n",
    "            batch_size=BIN_BATCH_SIZE,\n",
    "            learning_rate=BIN_LEARNING_RATE,\n",
    "            decay_lr_speed=BIN_LEARNING_RATE_DECAY_SPEED,\n",
    "            linear_size=BIN_LINEAR_SIZE,\n",
    "            dropout_p=BIN_DROPOUT_P,\n",
    "            activation_function=BIN_ACTIVATION_FUNCTION,\n",
    "            weight_decay=BIN_WEIGHT_DECAY,\n",
    "        )\n",
    "\n",
    "        bin_trainer.fit(bin_classifier)\n",
    "        bin_trainer.test(bin_classifier)\n",
    "\n",
    "except Exception as e:\n",
    "    bin_logger.experiment.stop(str(e))\n",
    "    raise e\n",
    "\n",
    "else:\n",
    "    bin_logger.experiment.stop()\n",
    "    "
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Train relation classifier"
   ]
  },
  {
   "metadata": {
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "GPUS = 1\n",
    "REL_MIN_EPOCHS = REL_MAX_EPOCHS = 4\n",
    "\n",
    "REL_BATCH_SIZE = 32\n",
    "REL_LEARNING_RATE = 2e-05\n",
    "REL_LEARNING_RATE_DECAY_SPEED = [1, 1, 0.75, 0.5, 0.25, 0.1, 0.075, 0.05, 0.025, 0.01]\n",
    "\n",
    "REL_LINEAR_SIZE = 512\n",
    "\n",
    "REL_DROPOUT_P = 0.1\n",
    "REL_ACTIVATION_FUNCTION = \"PReLU\"\n",
    "REL_WEIGHT_DECAY = 0.01 # default = 0.01"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "rel_logger = NeptuneLogger(\n",
    "    api_key=NEPTUNE_API_TOKEN,\n",
    "    project_name=NEPTUNE_PROJECT_NAME,\n",
    "    close_after_fit=False,\n",
    ")\n",
    "\n",
    "try:\n",
    "    for i in range(4):\n",
    "        print(f\"--------- EXPERIMENT {i} ---------\")\n",
    "    \n",
    "        rel_classifier = rel_trainer = None\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "    \n",
    "        rel_trainer = LightningTrainer(\n",
    "            gpus=GPUS,\n",
    "            min_epochs=REL_MIN_EPOCHS,\n",
    "            max_epochs=REL_MAX_EPOCHS,\n",
    "            default_root_dir=CHECKPOINT_DIR,\n",
    "            reload_dataloaders_every_epoch=True, # needed as we loop over a file,\n",
    "            deterministic=False,\n",
    "            checkpoint_callback=False,\n",
    "            logger=rel_logger\n",
    "        )\n",
    "    \n",
    "        rel_classifier = RelationClassifier(\n",
    "            pretrained_language_model=PRETRAINED_MODEL,\n",
    "            dataset_name=DATASET_NAME,\n",
    "            batch_size=REL_BATCH_SIZE,\n",
    "            learning_rate=REL_LEARNING_RATE,\n",
    "            decay_lr_speed=REL_LEARNING_RATE_DECAY_SPEED,\n",
    "            dropout_p=REL_DROPOUT_P,\n",
    "            activation_function=REL_ACTIVATION_FUNCTION,\n",
    "            weight_decay=REL_WEIGHT_DECAY,\n",
    "            linear_size=REL_LINEAR_SIZE,\n",
    "        )\n",
    "    \n",
    "        rel_trainer.fit(rel_classifier)\n",
    "        rel_trainer.test(rel_classifier)\n",
    "\n",
    "except Exception as e:\n",
    "    rel_logger.experiment.stop(str(e))\n",
    "    raise e\n",
    "\n",
    "else:\n",
    "    rel_logger.experiment.stop()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Train 2 classifiers independently then test together"
   ]
  },
  {
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "trusted": true,
    "scrolled": false
   },
   "cell_type": "code",
   "source": [
    "def test_together(experiment_no: int, logger, b_classifier: BinaryClassifier, r_classifier: RelationClassifier, dataset_name: str = DATASET_NAME,\n",
    "                  bin_batch_size = BIN_BATCH_SIZE, batch_size: int = REL_BATCH_SIZE):\n",
    "    \n",
    "    b_classifier.freeze()\n",
    "    r_classifier.freeze()\n",
    "\n",
    "    true_answer = []\n",
    "\n",
    "    # run binary classifier\n",
    "    print(\"Running binary classifier\")\n",
    "    dataset = GenericDataset(dataset_name, subset=\"test\", batch_size=bin_batch_size, label_transform=\"none\")\n",
    "    binary_classify_results = { criteria: [] for criteria in b_classifier.thresholds.keys() }\n",
    "\n",
    "    for input_data, true_label  in tqdm(dataset.as_batches(), total=len(dataset)):\n",
    "        # append true answers\n",
    "        true_answer += true_label.tolist()\n",
    "\n",
    "        # run bin classifier\n",
    "        logits = b_classifier(**input_data)\n",
    "        y_hat = torch.sigmoid(logits)\n",
    "        for criteria, threshold in b_classifier.thresholds.items():\n",
    "            label = b_classifier.yhat_to_label(y_hat, threshold)\n",
    "            binary_classify_results[criteria] += label.tolist()\n",
    "\n",
    "    # run relation classifier\n",
    "    print(\"Running relation classifier\")\n",
    "    dataset = GenericDataset(dataset_name, subset=\"test\", batch_size=batch_size, label_transform=\"none\")\n",
    "    relation_classify_result = []\n",
    "\n",
    "    for input_data, true_label  in tqdm(dataset.as_batches(), total=len(dataset)):\n",
    "        logits = r_classifier(**input_data)\n",
    "        label = r_classifier.logits_to_label(logits) + 1\n",
    "        relation_classify_result += label.tolist()\n",
    "\n",
    "    # combine results\n",
    "    print(\"Combining results\")\n",
    "    proposed_answer = {}\n",
    "    for criteria in b_classifier.thresholds.keys():\n",
    "        results = zip(relation_classify_result, binary_classify_results[criteria])\n",
    "        final_label = [relation_result if bin_result else 0 for relation_result, bin_result in results]\n",
    "        proposed_answer[criteria] = final_label\n",
    "\n",
    "    # log metric\n",
    "    final_metrics = {}\n",
    "    for criteria in b_classifier.thresholds.keys():\n",
    "        pa = proposed_answer[criteria]\n",
    "        \n",
    "        final_metrics.update({\n",
    "            f\"test_combined_{criteria}_acc\": accuracy_score(true_answer, pa),\n",
    "            f\"test_combined_{criteria}_pre_micro\": precision_score(true_answer, pa, average=\"micro\"),\n",
    "            f\"test_combined_{criteria}_rec_micro\": recall_score(true_answer, pa, average=\"micro\"),\n",
    "            f\"test_combined_{criteria}_f1_micro\": f1_score(true_answer, pa, average=\"micro\"),\n",
    "            f\"test_combined_{criteria}_pre_macro\": precision_score(true_answer, pa, average=\"macro\"),\n",
    "            f\"test_combined_{criteria}_rec_macro\": recall_score(true_answer, pa, average=\"macro\"),\n",
    "            f\"test_combined_{criteria}_f1_macro\": f1_score(true_answer, pa, average=\"macro\"),\n",
    "        })\n",
    "        \n",
    "        fig = BaseClassifier.plot_confusion_matrix(pa, true_answer)\n",
    "        logger.experiment.log_image(f\"test_combined_{criteria}_confusion_matrix\", fig)\n",
    "\n",
    "    for k, v in final_metrics.items():\n",
    "        print(f\"{k}: {v * 100}\")\n",
    "\n",
    "    for k, v in final_metrics.items():\n",
    "        logger.experiment.log_metric(k, v)\n",
    "    \n",
    "    # run the official scorer\n",
    "    scorer = get_official_scorer(experiment_no, logger)\n",
    "    if scorer:\n",
    "        scorer.score(proposed_answer)\n",
    "    else:\n",
    "        print(\"No official scorer found\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {
    "trusted": true,
    "scrolled": true
   },
   "cell_type": "code",
   "source": [
    "combine_logger = NeptuneLogger(\n",
    "    api_key=NEPTUNE_API_TOKEN,\n",
    "    project_name=NEPTUNE_PROJECT_NAME,\n",
    "    close_after_fit=False,\n",
    ")\n",
    "\n",
    "try:\n",
    "    for i in range(4):\n",
    "        print(f\"--------- EXPERIMENT {i} ---------\")\n",
    "    \n",
    "        # clean up\n",
    "        bin_classifier = bin_trainer = rel_classifier = rel_trainer = None\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "    \n",
    "        # relation classifier\n",
    "        rel_trainer = LightningTrainer(\n",
    "            gpus=GPUS,\n",
    "            min_epochs=REL_MIN_EPOCHS,\n",
    "            max_epochs=REL_MAX_EPOCHS,\n",
    "            default_root_dir=CHECKPOINT_DIR,\n",
    "            reload_dataloaders_every_epoch=True, # needed as we loop over a file,\n",
    "            deterministic=False,\n",
    "            checkpoint_callback=False,\n",
    "            logger=combine_logger\n",
    "        )\n",
    "        rel_classifier = RelationClassifier(\n",
    "            pretrained_language_model=PRETRAINED_MODEL,\n",
    "            dataset_name=DATASET_NAME,\n",
    "            batch_size=REL_BATCH_SIZE,\n",
    "            learning_rate=REL_LEARNING_RATE,\n",
    "            decay_lr_speed=REL_LEARNING_RATE_DECAY_SPEED,\n",
    "            dropout_p=REL_DROPOUT_P,\n",
    "            activation_function=REL_ACTIVATION_FUNCTION,\n",
    "            weight_decay=REL_WEIGHT_DECAY,\n",
    "            linear_size=REL_LINEAR_SIZE,\n",
    "        )\n",
    "        rel_trainer.fit(rel_classifier)\n",
    "    \n",
    "        # binary classifier\n",
    "        bin_trainer = LightningTrainer(\n",
    "            gpus=GPUS,\n",
    "            min_epochs=BIN_MIN_EPOCHS,\n",
    "            max_epochs=BIN_MAX_EPOCHS,\n",
    "            default_root_dir=CHECKPOINT_DIR,\n",
    "            reload_dataloaders_every_epoch=True, # needed as we loop over a file,\n",
    "            deterministic=False,\n",
    "            checkpoint_callback=False,\n",
    "            logger=combine_logger,\n",
    "        )\n",
    "        bin_classifier = BinaryClassifier(\n",
    "            pretrained_language_model=PRETRAINED_MODEL,\n",
    "            dataset_name=DATASET_NAME,\n",
    "            batch_size=BIN_BATCH_SIZE,\n",
    "            learning_rate=BIN_LEARNING_RATE,\n",
    "            decay_lr_speed=BIN_LEARNING_RATE_DECAY_SPEED,\n",
    "            dropout_p=BIN_DROPOUT_P,\n",
    "            activation_function=BIN_ACTIVATION_FUNCTION,\n",
    "            weight_decay=BIN_WEIGHT_DECAY,\n",
    "            linear_size=BIN_LINEAR_SIZE,\n",
    "        )\n",
    "        bin_trainer.fit(bin_classifier)\n",
    "        \n",
    "        # test together\n",
    "        test_together(i, combine_logger, bin_classifier, rel_classifier)\n",
    "\n",
    "except Exception as e:\n",
    "    combine_logger.experiment.stop(str(e))\n",
    "    raise e\n",
    "\n",
    "else:\n",
    "    combine_logger.experiment.stop()"
   ],
   "execution_count": null,
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "language": "python",
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.6",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}