Train BERT-Terrier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT Training script\n",
    "\n",
    "In order to train a BERT model, we first need to generate positive and negative samples. To make things more realistic, we will first retrieve a top-10 results for each trainint query using Anserini, and then, randomly pick a few that are not relevant (2) as \"negative sampling\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T14:10:11.096423Z",
     "start_time": "2020-05-20T14:10:11.069331Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"JAVA_HOME\"] = \"/home/arthur/.sdkman/candidates/java/8.0.242-open\" # Make sure this points to the right place for your Java Home!\n",
    "data_home = \"/ssd2/arthur/MsMarcoTREC/\"\n",
    "def get_path(x):\n",
    "    return os.path.join(data_home, x)\n",
    "\n",
    "import tqdm\n",
    "import pyterrier as pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T14:11:34.861721Z",
     "start_time": "2020-05-20T14:11:34.799789Z"
    }
   },
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import random\n",
    "import pickle\n",
    "import sys\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import os\n",
    "from collections import defaultdict\n",
    "import math\n",
    "import subprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T14:12:47.122173Z",
     "start_time": "2020-05-20T14:12:47.041313Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IndexingType.CLASSIC\n",
      "IndexingType.CLASSIC\n"
     ]
    }
   ],
   "source": [
    "# Terrier uses this Indexer object for interfacing with the index.\n",
    "index_path = get_path(\"terrier-index\")\n",
    "try:\n",
    "    pt.init(mem=16384)\n",
    "except:\n",
    "    pass\n",
    "index = pt.TRECCollectionIndexer(index_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting Terrier top-10\n",
    "We will use Terrier to retrieve the top-10 results using BM25. It doesn't need to be perfect, so, we won't care about fine-tunning it. Default settings should be enough\n",
    "\n",
    "The way this works is by:\n",
    "1. Submit queries to Terrier\n",
    "<!-- 1. submiting each query as a new search on Anserini, with the `SimpleSearcher.search()` method -->\n",
    "2. For each query, find $neg\\_samples$ negative samples from the top-$k$ results from BM25.\n",
    "3. Store these and the positive samples in a list\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2020-04-06T14:09:22.934Z"
    }
   },
   "source": [
    "### Loading all relevant docs from the qrels file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T14:11:37.915699Z",
     "start_time": "2020-05-20T14:11:36.753216Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load the relevant query-document pairs\n",
    "relevant_docs = defaultdict(lambda:[])\n",
    "for file in [get_path(\"qrels/msmarco-doctrain-qrels.tsv\"), get_path(\"qrels/msmarco-docdev-qrels.tsv\")]:\n",
    "    for line in open(file):\n",
    "        query_id, _, doc_id, rel = line.split()\n",
    "        assert rel == \"1\"\n",
    "        relevant_docs[query_id].append(doc_id)\n",
    "train_qrels = pt.Utils.parse_qrels(get_path(\"qrels/msmarco-doctrain-qrels.tsv\"))\n",
    "dev_qrels = pt.Utils.parse_qrels(get_path(\"qrels/msmarco-docdev-qrels.tsv\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the top-10 using BM25 and create a training set based on this\n",
    "Some notes:\n",
    "\n",
    "- If it finds the `.pkl` file created in the end of the loop, it won't re-compute everything.\n",
    "- Each query is sanitized before being submitted to Anserini. (lines 22-32)\n",
    "- The code will \"batch\" a number of queries to be submitted at once to Anserini, and will run these in parallel. This is much faster than one at a time, and more efficient than all of the queries at once.\n",
    "- We store the end results in a pickle file, that is a list with the triples `query_id, doc_id, label`. \n",
    "- Should take about 1.5h to finish.\n",
    "- Each element in the output list is: `[query_id, document_id, label]` where `label` is `1` for relevant and `0` for non-relevant\n",
    " \n",
    "### **PAY ATTENTION TO YOUR MACHINE**\n",
    "this notebook was ran at DeepIR, with 56 threads and 128GB of memory. Make sure to pick a fair number of threads, and a batchsize that fits confortably on memory. BE MINDFULL OF FAIR USAGE OF THE MACHINE. Check if someone else is using the machine, and chose a fair number of threads/batch size. In this configuration, 42 threads and batch_size 10000, this took about 6 minutes to finish. YMMV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T14:58:13.660688Z",
     "start_time": "2020-05-20T14:58:13.656278Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T15:00:33.481734Z",
     "start_time": "2020-05-20T15:00:33.475480Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'querying.processes': 'terrierql:TerrierQLParser,parsecontrols:TerrierQLToControls,parseql:TerrierQLToMatchingQueryTerms,matchopql:MatchingOpQLParser,applypipeline:ApplyTermPipeline,localmatching:LocalManager$ApplyLocalMatching,qe:QueryExpansion,labels:org.terrier.learning.LabelDecorator,filters:LocalManager$PostFilterProcess',\n",
       " 'querying.postfilters': 'decorate:SimpleDecorate,site:SiteFilter,scope:Scope',\n",
       " 'querying.default.controls': 'wmodel:DPH,parsecontrols:on,parseql:on,applypipeline:on,terrierql:on,localmatching:on,filters:on,decorate:on',\n",
       " 'querying.allowed.controls': 'scope,qe,qemodel,start,end,site,scope,applypipeline',\n",
       " 'termpipelines': 'Stopwords,PorterStemmer'}"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pt.BatchRetrieve.default_properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T15:42:51.123173Z",
     "start_time": "2020-05-20T15:42:49.917174Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qid</th>\n",
       "      <th>docid</th>\n",
       "      <th>docno</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>query</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1515835</td>\n",
       "      <td>D2983801</td>\n",
       "      <td>0</td>\n",
       "      <td>16.953258</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>1730877</td>\n",
       "      <td>D2983802</td>\n",
       "      <td>1</td>\n",
       "      <td>16.825129</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>731948</td>\n",
       "      <td>D1509608</td>\n",
       "      <td>2</td>\n",
       "      <td>16.737326</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>627921</td>\n",
       "      <td>D2983799</td>\n",
       "      <td>3</td>\n",
       "      <td>16.684804</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>1113515</td>\n",
       "      <td>D3551182</td>\n",
       "      <td>4</td>\n",
       "      <td>16.615293</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>96</th>\n",
       "      <td>1</td>\n",
       "      <td>2632375</td>\n",
       "      <td>D3047962</td>\n",
       "      <td>96</td>\n",
       "      <td>15.425559</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>97</th>\n",
       "      <td>1</td>\n",
       "      <td>3071388</td>\n",
       "      <td>D1164288</td>\n",
       "      <td>97</td>\n",
       "      <td>15.419781</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>98</th>\n",
       "      <td>1</td>\n",
       "      <td>1386134</td>\n",
       "      <td>D1703652</td>\n",
       "      <td>98</td>\n",
       "      <td>15.406517</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>99</th>\n",
       "      <td>1</td>\n",
       "      <td>1955030</td>\n",
       "      <td>D2325165</td>\n",
       "      <td>99</td>\n",
       "      <td>15.405913</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>100</th>\n",
       "      <td>1</td>\n",
       "      <td>797353</td>\n",
       "      <td>D3003707</td>\n",
       "      <td>100</td>\n",
       "      <td>15.401334</td>\n",
       "      <td>hello world 2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>101 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    qid    docid     docno  rank      score          query\n",
       "0     1  1515835  D2983801     0  16.953258  hello world 2\n",
       "1     1  1730877  D2983802     1  16.825129  hello world 2\n",
       "2     1   731948  D1509608     2  16.737326  hello world 2\n",
       "3     1   627921  D2983799     3  16.684804  hello world 2\n",
       "4     1  1113515  D3551182     4  16.615293  hello world 2\n",
       "..   ..      ...       ...   ...        ...            ...\n",
       "96    1  2632375  D3047962    96  15.425559  hello world 2\n",
       "97    1  3071388  D1164288    97  15.419781  hello world 2\n",
       "98    1  1386134  D1703652    98  15.406517  hello world 2\n",
       "99    1  1955030  D2325165    99  15.405913  hello world 2\n",
       "100   1   797353  D3003707   100  15.401334  hello world 2\n",
       "\n",
       "[101 rows x 6 columns]"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever = pt.BatchRetrieve(index, controls = {\"wmodel\":\"BM25\", \"end\":\"99\"})\n",
    "retriever.transform(\"hello world 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:45:59.927297Z",
     "start_time": "2020-05-20T15:49:49.013869Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 367013 for train\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-87-49bc3f306885>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0mnumber_of_queries\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mqueries_df\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"Found {number_of_queries} for {split}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m     \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mretriever\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mqueries_df\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m     \u001b[0mresults\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mget_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"results/terrier-{split}.csv\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/bert4IR/lib/python3.7/site-packages/pyterrier/batchretrieve.py\u001b[0m in \u001b[0;36mtransform\u001b[0;34m(self, queries)\u001b[0m\n\u001b[1;32m    182\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    183\u001b[0m             \u001b[0;31m# now ask Terrier to run the request\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 184\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmanager\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrunSearchRequest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    185\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msrq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetResults\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    186\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "pattern = re.compile('([^\\s\\w]|_)+')\n",
    "\n",
    "terrier_top10 = defaultdict(lambda:[])\n",
    "pairs_per_split = defaultdict(lambda: [])\n",
    "k = 10       # Number of documents to retrieve \n",
    "neg_samples = 2 # Number of negatives samples to use\n",
    "\n",
    "query_texts = dict()\n",
    "retriever = pt.BatchRetrieve(index, {\"wmodel\":\"BM25\", \"end\":f\"{k-1}\"})\n",
    "\n",
    "for split in [\"train\", \"dev\"]:\n",
    "    file_path = get_path(f\"queries/msmarco-doc{split}-queries.tsv\")\n",
    "    run_search=True\n",
    "    queries_df = pt.Utils.parse_singleline_topics_file(file_path)\n",
    "    number_of_queries = len(queries_df)\n",
    "    print(f\"Found {number_of_queries} for {split}\")\n",
    "    results = retriever.transform(queries_df)\n",
    "    results.to_csv(get_path(f\"results/terrier-{split}.csv\"))\n",
    "    \n",
    "#    # Run Terrier retrieve step\n",
    "#     for idx, line in enumerate(open(file_path, encoding=\"utf-8\")):\n",
    "#         query_id, query = line.strip().split(\"\\t\")\n",
    "#         query_ids.append(query_id)\n",
    "#         query = unicodedata.normalize(\"NFKD\", query) # Force queries into UTF-8\n",
    "#         query = pattern.sub(' ',query) # Remove non-ascii characters. It clears up most of the issues we may find on the query datasets\n",
    "#         query_texts[query_id] = query\n",
    "#         if run_search is False:\n",
    "#             continue\n",
    "#         queries.append(query)\n",
    "#         if len(queries) == batch_size or idx == number_of_queries-1:\n",
    "#             results = searcher.batch_search(queries, query_ids, k=k, threads=threads)\n",
    "#             pbar.update()\n",
    "#             for query, query_id in zip(queries, query_ids):\n",
    "#                 retrieved_docs_ids = [hit.docid for hit in results[query_id]]\n",
    "#                 relevant_docs_for_query = relevant_docs[query_id]\n",
    "#                 retrieved_non_relevant_documents = set(retrieved_docs_ids).difference(set(relevant_docs_for_query))\n",
    "                   \n",
    "#                 if len(retrieved_non_relevant_documents) < 2:\n",
    "#                     print(f\"query {query} has less than 2 retrieved docs.\")\n",
    "#                     continue\n",
    "#                 random_negative_samples = random.sample(retrieved_non_relevant_documents, neg_samples)\n",
    "#                 pairs_per_split[split] += [(query_id, doc_id, 1) for doc_id in relevant_docs_for_query]\n",
    "#                 pairs_per_split[split] += [(query_id, doc_id, 0) for doc_id in random_negative_samples]\n",
    "#             queries = []\n",
    "#             query_ids = []\n",
    "#     pickle.dump(pairs_per_split[split], open(path(f\"{split}_triples.pkl\"), 'wb'))\n",
    "#     pbar.close()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Creation\n",
    "\n",
    "This dataset is too big to fit in memory. Therefore, it's a good idea to leave it in disk, and retrieve as needed.\n",
    "\n",
    "To do so, we will create three files: \n",
    "- `msmarco_samples.txt`: file with every sample already tokenized and in the right format to be used as input to BERT.\n",
    "- `msmarco_offset.pkl`: pickle file with a dictionary with the file address for each of these samples. Will make it WAY faster to retrieve data from disk. \n",
    "- `msmarco_index.pkl`: A pickle file with a dictionary mapping each sample id (`queryid_docid`) to it's numbered position on the previous file. This will enable us to find a sample by index, and not only ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T08:32:12.069465Z",
     "start_time": "2020-04-16T08:32:11.232129Z"
    },
    "code_folding": [
     37,
     88
    ],
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "import torch\n",
    "\n",
    "# This is our main Dataset class.\n",
    "class MsMarcoDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 samples,\n",
    "                 tokenizer,\n",
    "                 searcher,\n",
    "                 split,\n",
    "                 tokenizer_batch=8000):\n",
    "        '''Initialize a Dataset object. \n",
    "        Arguments:\n",
    "            samples: A list of samples. Each sample should be a tuple with (query_id, doc_id, <label>), where label is optional\n",
    "            tokenizer: A tokenizer object from Hugging Face's Tokenizer lib. (need to implement encode_batch())\n",
    "            searcher: A PySerini Simple Searcher object. Should implement the .doc() method\n",
    "            split: A strong indicating if we are in a train, dev or test dataset.\n",
    "            tokenizer_batch: How many samples to be tokenized at once by the tokenizer object.\n",
    "            The biggest bottleneck is the searcher, not the tokenizer.\n",
    "        '''\n",
    "        self.searcher = searcher\n",
    "        self.split = split\n",
    "        # If we already have the data pre-computed, we shouldn't need to re-compute it.\n",
    "        self.split = split\n",
    "        if (os.path.isfile(path(f\"{split}_msmarco_samples.tsv\"))\n",
    "                and os.path.isfile(path(f\"{split}_msmarco_offset.pkl\"))\n",
    "                and os.path.isfile(path(f\"{split}_msmarco_index.pkl\"))):\n",
    "            print(\"Already found every meaningful file. Cowardly refusing to re-compute.\")\n",
    "            self.samples_offset_dict = pickle.load(open(path(f\"{split}_msmarco_offset.pkl\"), 'rb'))\n",
    "            self.index_dict = pickle.load(open(path(f\"{split}_msmarco_index.pkl\"), 'rb'))\n",
    "            return\n",
    "        self.tokenizer = tokenizer\n",
    "        print(\"Loading and tokenizing dataset...\")\n",
    "        self.samples_offset_dict = dict()\n",
    "        self.index_dict = dict()\n",
    "\n",
    "        self.samples_file = open(path(f\"{split}_msmarco_samples.tsv\"),'w',encoding=\"utf-8\")\n",
    "        self.processed_samples = 0\n",
    "        query_batch = []\n",
    "        doc_batch = []\n",
    "        sample_ids_batch = []\n",
    "        labels_batch = []\n",
    "        number_of_batches = math.ceil(len(samples) // tokenizer_batch)\n",
    "        # A progress bar to display how far we are.\n",
    "        batch_pbar = tqdm(total=number_of_batches, desc=\"Tokenizer batches\")\n",
    "        for i, sample in enumerate(samples):\n",
    "            if split==\"train\" or split == \"dev\":\n",
    "                label = sample[2]\n",
    "                labels_batch.append(label)\n",
    "            query_batch.append(query_texts[sample[0]])\n",
    "            doc_batch.append(self._get_document_content_from_id(sample[1]))\n",
    "            sample_ids_batch.append(f\"{sample[0]}_{sample[1]}\")\n",
    "            #If we hit the number of samples for this batch OR this is the last sample\n",
    "            if len(query_batch) == tokenizer_batch or i == len(samples) - 1:\n",
    "                self._tokenize_and_dump_batch(doc_batch, query_batch, labels_batch, sample_ids_batch)\n",
    "                batch_pbar.update()\n",
    "                query_batch = []\n",
    "                doc_batch = []\n",
    "                sample_ids_batch = []\n",
    "                if split == \"train\" or split == \"dev\":\n",
    "                    labels_batch = []\n",
    "        batch_pbar.close()\n",
    "        # Dump files in disk, so we don't need to go over it again.\n",
    "        self.samples_file.close()\n",
    "        pickle.dump(self.index_dict, open(path(f\"{self.split}_msmarco_index.pkl\"), 'wb'))\n",
    "        pickle.dump(self.samples_offset_dict, open(path(f\"{self.split}_msmarco_offset.pkl\"), 'wb'))\n",
    "\n",
    "    def _tokenize_and_dump_batch(self, doc_batch, query_batch, labels_batch,\n",
    "                                 sample_ids_batch):\n",
    "        '''tokenizes and dumps the samples in the current batch\n",
    "        It also store the positions from the current file into the samples_offset_dict.\n",
    "        '''\n",
    "        # Use the tokenizer object\n",
    "        tokens = self.tokenizer.encode_batch(list(zip(query_batch, doc_batch)))\n",
    "        for idx, (sample_id, token) in enumerate(zip(sample_ids_batch, tokens)):\n",
    "            #BERT supports up to 512 tokens. If we have more than that, we need to remove some tokens from the document\n",
    "            if len(token.ids) >= 512:\n",
    "                token_ids = token.ids[:511]\n",
    "                token_ids.append(tokenizer.token_to_id(\"[SEP]\"))\n",
    "                segment_ids = token.type_ids[:512]\n",
    "            # With less tokens, we need to \"pad\" the vectors up to 512.\n",
    "            else:\n",
    "                padding = [0] * (512 - len(token.ids))\n",
    "                token_ids = token.ids + padding\n",
    "                segment_ids = token.type_ids + padding\n",
    "            # How far in the file are we? This is where we need to go to find the documents later.\n",
    "            file_location = self.samples_file.tell()\n",
    "            # If we have labels\n",
    "            if self.split==\"train\" or split == \"dev\":\n",
    "                self.samples_file.write(f\"{sample_id}\\t{token_ids}\\t{segment_ids}\\t{labels_batch[idx]}\\n\")\n",
    "            else:\n",
    "                self.samples_file.write(f\"{sample_id}\\t{token_ids}\\t{segment_ids}\\n\")\n",
    "            self.samples_offset_dict[sample_id] = file_location\n",
    "            self.index_dict[self.processed_samples] = sample_id\n",
    "            self.processed_samples += 1\n",
    "\n",
    "    def _get_document_content_from_id(self, doc_id):\n",
    "        '''Get the raw text value from the doc_id\n",
    "        There is probably an easier way to do that, but this works.\n",
    "        '''\n",
    "        doc_text = self.searcher.doc(doc_id).lucene_document().getField(\"raw\").stringValue()\n",
    "        return doc_text[7:-8]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        '''Returns a sample with index idx\n",
    "        DistilBERT does not take into account segment_ids. (indicator if the token comes from the query or the document) \n",
    "        However, for the sake of completness, we are including it here, together with the attention mask\n",
    "        position_ids, with the positional encoder, is not needed. It's created for you inside the model.\n",
    "        '''\n",
    "        if isinstance(idx, int):\n",
    "            idx = self.index_dict[idx]\n",
    "        with open(path(f\"{self.split}_msmarco_samples.tsv\"), 'r', encoding=\"utf-8\") as inf:\n",
    "            inf.seek(self.samples_offset_dict[idx])\n",
    "            line = inf.readline().split(\"\\t\")\n",
    "            try:\n",
    "                sample_id = line[0]\n",
    "                input_ids = eval(line[1])\n",
    "                token_type_ids = eval(line[2])\n",
    "                input_mask = [1] * 512\n",
    "            except:\n",
    "                print(line, idx)\n",
    "                raise IndexError\n",
    "            # If it's a training dataset, we also have a label tag.\n",
    "            if split==\"train\" or split == \"dev\":\n",
    "                label = int(line[3])\n",
    "                return (torch.tensor(input_ids, dtype=torch.long),\n",
    "                        torch.tensor(input_mask, dtype=torch.long),\n",
    "                        torch.tensor(token_type_ids, dtype=torch.long),\n",
    "                        torch.tensor([label], dtype=torch.long))\n",
    "            return (torch.tensor(input_ids, dtype=torch.long),\n",
    "                    torch.tensor(input_mask, dtype=torch.long),\n",
    "                    torch.tensor(token_type_ids, dtype=torch.long))\n",
    "    def __len__(self):\n",
    "        return len(self.samples_offset_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training script\n",
    "For actually training our model, we need to do the following:\n",
    "1. Create a DataLoader object for train and one for dev. This will help with batching and such.\n",
    "2. Load a BERT pre-trained model. For this example, we are using DistilBert. Because it's smaller and faster.\n",
    "    - For ease of use, we will use the `DistilBertForSequenceClassification` model. It's ready for computing whether two senteces are related.\n",
    "    - Also note that, for this model, weirdly enough, $1$ is NOT RELEVANT and  $0$ is RELEVANT\n",
    "    - Alternativelly, we can use the default `DistilBert` and extract the `[CLS]` token embedding and feed it to a shallow NN using PyTorch or even Sklearn and a linear regression.\n",
    "3. Create a training loop that for every $X$ samples will check the results on the dev dataset.\n",
    "4. Store breakpoints every $N$ steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T08:32:15.412034Z",
     "start_time": "2020-04-16T08:32:14.212203Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import DistilBertForSequenceClassification\n",
    "from torch.utils.data import DataLoader\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "\n",
    "tokenizer = BertWordPieceTokenizer(\"/ssd2/arthur/bert-axioms/tokenizer/bert-base-uncased-vocab.txt\", lowercase=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T08:32:20.792503Z",
     "start_time": "2020-04-16T08:32:19.314639Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Already found every meaningful file. Cowardly refusing to re-compute.\n",
      "Already found every meaningful file. Cowardly refusing to re-compute.\n"
     ]
    }
   ],
   "source": [
    "train_dataset = MsMarcoDataset(pairs_per_split[\"train\"], tokenizer, searcher, split = \"train\")\n",
    "dev_dataset = MsMarcoDataset(pairs_per_split[\"dev\"], tokenizer, searcher, split = \"dev\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We NEED to use GPUs for this. If you don't have access to some GPUs you can try Google Colab OR if you are a MSc from WIS, get in touch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T08:32:47.633978Z",
     "start_time": "2020-04-16T08:32:41.000963Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running on 5 GPUS, on 80-sized batches\n",
      "*********Total optmization steps: 42042*********\n"
     ]
    }
   ],
   "source": [
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "# With these configurations, on DeepIR, it takes ~3h/batch to train, with ~2batches/s\n",
    "GPUS_TO_USE = [2,4,5,6,7] # If you have multiple GPUs, pick the ones you want to use.\n",
    "number_of_cpus = 24 # Number of CPUS to use when loading your dataset.\n",
    "n_epochs = 2 # How may passes over the whole dataset to complete\n",
    "weight_decay = 0.0 # Some papers define a weight decay, meaning, the weights on some layers will decay slower overtime. By default, we don't do this.\n",
    "lr = 0.00005 # Learning rate for the fine-tunning.\n",
    "warmup_proportion = 0.1 # Percentage of training steps to perform before we start to decrease the learning rate.\n",
    "steps_to_print = 1000 # How many steps to wait before printing loss\n",
    "steps_to_eval = 2000 # How many steps to wait before running an eval step\n",
    "\n",
    "# This is our base model\n",
    "try:\n",
    "    del model\n",
    "    torch.cuda.empty_cache() # Make sure we have a clean slate\n",
    "except:\n",
    "    pass\n",
    "model = DistilBertForSequenceClassification.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    # Asssign the model to GPUs, specifying to use Data parallelism.\n",
    "    model = torch.nn.DataParallel(model, device_ids=GPUS_TO_USE)\n",
    "    # The main model should be on the first GPU\n",
    "    device = torch.device(f\"cuda:{GPUS_TO_USE[0]}\") \n",
    "    model.to(device)\n",
    "    # For a 1080Ti, 16 samples fit on a GPU confortably. So, the train batch size will be 16*the number of GPUS\n",
    "    train_batch_size = len(GPUS_TO_USE) * 16\n",
    "    print(f\"running on {len(GPUS_TO_USE)} GPUS, on {train_batch_size}-sized batches\")\n",
    "else:\n",
    "    print(\"Are you sure about it? We will try to run this in CPU, but it's a BAD idea...\")\n",
    "    device = torch.device(\"cpu\")\n",
    "    train_batch_size = 16\n",
    "    model.to(device)\n",
    "\n",
    "# A data loader is a nice device for generating batches for you easily.\n",
    "# It receives any object that implementes __getitem__(self, idx) and __len__(self)\n",
    "train_data_loader = DataLoader(train_dataset, batch_size=train_batch_size, num_workers=number_of_cpus,shuffle=True)\n",
    "dev_data_loader = DataLoader(dev_dataset, batch_size=32, num_workers=number_of_cpus,shuffle=True)\n",
    "\n",
    "#how many optimization steps to run, given the NUMBER OF BATCHES. (The len of the dataloader is the number of batches).\n",
    "num_train_optimization_steps = len(train_data_loader) * n_epochs\n",
    "\n",
    "#which layers will not have a linear weigth decay when training\n",
    "no_decay = ['bias', 'LayerNorm.weight']\n",
    "\n",
    "#all parameters to be optimized by our fine tunning.\n",
    "optimizer_grouped_parameters = [\n",
    "    {'params': [p for n, p in model.named_parameters() if not any( nd in n for nd in no_decay)], 'weight_decay': weight_decay},\n",
    "    {'params': [p for n, p in model.named_parameters() if any( nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "]\n",
    "\n",
    "#We use the AdamW optmizer here.\n",
    "optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=1e-8) \n",
    "\n",
    "# How many steps to wait before we start to decrease the learning rate\n",
    "warmup_steps = num_train_optimization_steps * warmup_proportion \n",
    "# A scheduler to take care of the above.\n",
    "scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=num_train_optimization_steps)\n",
    "print(f\"*********Total optmization steps: {num_train_optimization_steps}*********\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T15:36:29.760219Z",
     "start_time": "2020-04-16T08:33:00.028286Z"
    },
    "code_folding": [],
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8c3c1b6a115a4d7e8fba091e6d34c0d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Epochs', max=2.0, style=ProgressStyle(description_width='…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d36943663e1b4f8d850ab330adbb9360",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Batches', max=21021.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.6017579445571333\n",
      "Train accuracy: 0.5875\n",
      "Training loss: 0.6824558973312378\n",
      "Learning rate: 1.189286903572618e-08\n",
      "Train ROC: 0.7813390313390314\n",
      "Train accuracy: 0.725\n",
      "Training loss: 0.5217803120613098\n",
      "Learning rate: 5.9583273868988165e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "487b9119a40b4ce9af97c57c5f73dcac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.4734283789961621\n",
      "  Acuracy Dev = 0.7344386149003148\n",
      "  F1 Dev = 0.527800582132995\n",
      "  ROC Dev = 0.7823920346899583\n",
      "Train ROC: 0.8581818181818182\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.4277583062648773\n",
      "Learning rate: 1.1904761904761907e-05\n",
      "Train ROC: 0.8382913806254768\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.42712002992630005\n",
      "Learning rate: 1.7851196422624997e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0eb1983f28f4479aa78db8219f784e24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5346556823798122\n",
      "  Acuracy Dev = 0.7642812172088143\n",
      "  F1 Dev = 0.6432926829268293\n",
      "  ROC Dev = 0.8302346199964222\n",
      "Train ROC: 0.8413696715583509\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.450906366109848\n",
      "Learning rate: 2.3797630940488087e-05\n",
      "Train ROC: 0.8988095238095238\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.37316620349884033\n",
      "Learning rate: 2.9744065458351178e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ce5dc2b50084641bbceee2ad6efca9d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5421112003521867\n",
      "  Acuracy Dev = 0.7692759706190976\n",
      "  F1 Dev = 0.6513604363544111\n",
      "  ROC Dev = 0.8383614184643108\n",
      "Train ROC: 0.7827208252740168\n",
      "Train accuracy: 0.725\n",
      "Training loss: 0.5373059511184692\n",
      "Learning rate: 3.5690499976214264e-05\n",
      "Train ROC: 0.8839031339031339\n",
      "Train accuracy: 0.7875\n",
      "Training loss: 0.39931368827819824\n",
      "Learning rate: 4.1636934494077355e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a847ff1b0ab542479a38dae02426a8bf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5487017845350912\n",
      "  Acuracy Dev = 0.7785099685204617\n",
      "  F1 Dev = 0.6388831862040649\n",
      "  ROC Dev = 0.8454836767280532\n",
      "Train ROC: 0.7886108714408974\n",
      "Train accuracy: 0.8375\n",
      "Training loss: 0.4352118670940399\n",
      "Learning rate: 4.7583369011940445e-05\n",
      "Train ROC: 0.7998575498575499\n",
      "Train accuracy: 0.7875\n",
      "Training loss: 0.5039113759994507\n",
      "Learning rate: 4.9607799607799605e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0e0af3b7fea4ff3afaf4e5ca7f882cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5533506899824174\n",
      "  Acuracy Dev = 0.7814480587618048\n",
      "  F1 Dev = 0.6433317350503459\n",
      "  ROC Dev = 0.8502054766730774\n",
      "Train ROC: 0.8254985754985754\n",
      "Train accuracy: 0.7625\n",
      "Training loss: 0.48662814497947693\n",
      "Learning rate: 4.894708466137037e-05\n",
      "Train ROC: 0.8504542278127183\n",
      "Train accuracy: 0.775\n",
      "Training loss: 0.4507910907268524\n",
      "Learning rate: 4.828636971494114e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d7a1f0f28e74ba19f6a97b7622d19cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5661902327233137\n",
      "  Acuracy Dev = 0.7871983210912906\n",
      "  F1 Dev = 0.6679329316216925\n",
      "  ROC Dev = 0.8568459000104849\n",
      "Train ROC: 0.8676761026991442\n",
      "Train accuracy: 0.7625\n",
      "Training loss: 0.4336676597595215\n",
      "Learning rate: 4.762565476851191e-05\n",
      "Train ROC: 0.9045424621461489\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.41981038451194763\n",
      "Learning rate: 4.696493982208268e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a91038bc509b44c2beacbe834e43a2c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5595527075761935\n",
      "  Acuracy Dev = 0.7861070304302203\n",
      "  F1 Dev = 0.6433370660694289\n",
      "  ROC Dev = 0.8604986653775032\n",
      "Train ROC: 0.8386666666666668\n",
      "Train accuracy: 0.7875\n",
      "Training loss: 0.4822412431240082\n",
      "Learning rate: 4.630422487565345e-05\n",
      "Train ROC: 0.8916363636363637\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.3781545162200928\n",
      "Learning rate: 4.5643509929224216e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "636706cd78c64f4aa62393db8ce9797d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5685093263044734\n",
      "  Acuracy Dev = 0.7894648478488983\n",
      "  F1 Dev = 0.6661341853035144\n",
      "  ROC Dev = 0.8612058401817764\n",
      "Train ROC: 0.8610928242264647\n",
      "Train accuracy: 0.8\n",
      "Training loss: 0.4607420563697815\n",
      "Learning rate: 4.498279498279498e-05\n",
      "Train ROC: 0.8792727272727273\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.40605759620666504\n",
      "Learning rate: 4.4322080036365746e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbbfa3362bd24424bdde1873e1e4f1a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5709956259110867\n",
      "  Acuracy Dev = 0.7858551941238195\n",
      "  F1 Dev = 0.6880655416972365\n",
      "  ROC Dev = 0.8635486376525725\n",
      "Train ROC: 0.8839031339031338\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.39896130561828613\n",
      "Learning rate: 4.366136508993652e-05\n",
      "Train ROC: 0.925595238095238\n",
      "Train accuracy: 0.9\n",
      "Training loss: 0.30391794443130493\n",
      "Learning rate: 4.300065014350729e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60aaaf1141364364afe2049bc8db2971",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5805320459245902\n",
      "  Acuracy Dev = 0.7950052465897167\n",
      "  F1 Dev = 0.6871236386931455\n",
      "  ROC Dev = 0.8677017962802882\n",
      "Train ROC: 0.8891369047619048\n",
      "Train accuracy: 0.7875\n",
      "Training loss: 0.38106560707092285\n",
      "Learning rate: 4.233993519707805e-05\n",
      "Train ROC: 0.8736299161831076\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.43353912234306335\n",
      "Learning rate: 4.167922025064882e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e1914df460b42a18f6b9fc85d80b663",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5843719330785806\n",
      "  Acuracy Dev = 0.7976495278069256\n",
      "  F1 Dev = 0.6896285328011331\n",
      "  ROC Dev = 0.8690233542170219\n",
      "Train ROC: 0.9032882011605416\n",
      "Train accuracy: 0.8625\n",
      "Training loss: 0.4046548008918762\n",
      "Learning rate: 4.101850530421959e-05\n",
      "Train ROC: 0.8981191222570533\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.3345804214477539\n",
      "Learning rate: 4.035779035779036e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bad4f6e48cb4da196e5781d810283af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5732463267039056\n",
      "  Acuracy Dev = 0.7840923399790136\n",
      "  F1 Dev = 0.6979448032883148\n",
      "  ROC Dev = 0.8687817549867656\n",
      "Train ROC: 0.9106666666666666\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.39230263233184814\n",
      "Learning rate: 3.9697075411361125e-05\n",
      "Train ROC: 0.9059065934065933\n",
      "Train accuracy: 0.8\n",
      "Training loss: 0.39423346519470215\n",
      "Learning rate: 3.9036360464931894e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83a06754e8b34042a1f1a46d66c9b44e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5861583283817675\n",
      "  Acuracy Dev = 0.8003357817418678\n",
      "  F1 Dev = 0.6820399705902012\n",
      "  ROC Dev = 0.8744936155128205\n",
      "Train ROC: 0.8058181818181818\n",
      "Train accuracy: 0.775\n",
      "Training loss: 0.47148066759109497\n",
      "Learning rate: 3.837564551850266e-05\n",
      "Train ROC: 0.8763636363636363\n",
      "Train accuracy: 0.8\n",
      "Training loss: 0.4132147431373596\n",
      "Learning rate: 3.7714930572073424e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bad9ef4321dd46b98677601adaa75b63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5907510539538373\n",
      "  Acuracy Dev = 0.8018887722980063\n",
      "  F1 Dev = 0.6939834024896266\n",
      "  ROC Dev = 0.8763013368041515\n",
      "Train ROC: 0.8833010960670535\n",
      "Train accuracy: 0.7875\n",
      "Training loss: 0.4441492259502411\n",
      "Learning rate: 3.705421562564419e-05\n",
      "Train ROC: 0.8772321428571428\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.3992480933666229\n",
      "Learning rate: 3.639350067921497e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "91ade5aa0d8b4cc681be26c938e2835f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5875612980804374\n",
      "  Acuracy Dev = 0.7983630640083945\n",
      "  F1 Dev = 0.6980894922071392\n",
      "  ROC Dev = 0.8744116726579988\n",
      "Train ROC: 0.8689927583936801\n",
      "Train accuracy: 0.8\n",
      "Training loss: 0.43431875109672546\n",
      "Learning rate: 3.5732785732785736e-05\n",
      "Train ROC: 0.9059561128526645\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.3638037443161011\n",
      "Learning rate: 3.50720707863565e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d22c5d3a6a54884b7c4dfc53ba2539c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5895932620272767\n",
      "  Acuracy Dev = 0.8019727177334732\n",
      "  F1 Dev = 0.6881692002643754\n",
      "  ROC Dev = 0.8770030138703295\n",
      "Train ROC: 0.9283865401207938\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.38181766867637634\n",
      "Learning rate: 3.4411355839927267e-05\n",
      "Train ROC: 0.9494301994301995\n",
      "Train accuracy: 0.9\n",
      "Training loss: 0.2794446647167206\n",
      "Learning rate: 3.3750640893498035e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89385ce930664020acca52ce17d4754f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5908281271775208\n",
      "  Acuracy Dev = 0.8036096537250788\n",
      "  F1 Dev = 0.6819386853375025\n",
      "  ROC Dev = 0.8776348415199899\n",
      "Train ROC: 0.8875\n",
      "Train accuracy: 0.8375\n",
      "Training loss: 0.3492681384086609\n",
      "Learning rate: 3.3089925947068804e-05\n",
      "Train ROC: 0.9163636363636364\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.3378193974494934\n",
      "Learning rate: 3.242921100063957e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "566d3a0174f6478f8ce9aefc125bf546",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5926705793325872\n",
      "  Acuracy Dev = 0.8010073452256034\n",
      "  F1 Dev = 0.7050332856342936\n",
      "  ROC Dev = 0.8772818130083159\n",
      "Train ROC: 0.8832417582417582\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.39545726776123047\n",
      "Learning rate: 3.176849605421034e-05\n",
      "Train ROC: 0.8398437500000001\n",
      "Train accuracy: 0.725\n",
      "Training loss: 0.4513436257839203\n",
      "Learning rate: 3.110778110778111e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18146f22f0cd489fa59c2b3f3137104d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5897693029618618\n",
      "  Acuracy Dev = 0.8036935991605456\n",
      "  F1 Dev = 0.6719045948789898\n",
      "  ROC Dev = 0.8792138735543703\n",
      "Train ROC: 0.9107505070993915\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.37642183899879456\n",
      "Learning rate: 3.0447066161351874e-05\n",
      "Train ROC: 0.8763440860215054\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.3669411838054657\n",
      "Learning rate: 2.9786351214922643e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "921e5ca3f4964655971a22fb1ea2dde8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5984480191476851\n",
      "  Acuracy Dev = 0.8065897166841552\n",
      "  F1 Dev = 0.7008180755746007\n",
      "  ROC Dev = 0.8799395075989522\n",
      "Train ROC: 0.9187979539641944\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.3746066093444824\n",
      "Learning rate: 2.912563626849341e-05\n",
      "Train ROC: 0.9016927083333333\n",
      "Train accuracy: 0.8625\n",
      "Training loss: 0.3908333480358124\n",
      "Learning rate: 2.846492132206418e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f515fc4cd0c74adeb8e6193fb189e3e8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5912168546520766\n",
      "  Acuracy Dev = 0.802392444910808\n",
      "  F1 Dev = 0.693169968717414\n",
      "  ROC Dev = 0.8808847776677455\n",
      "Train ROC: 0.9014675052410901\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.3724188506603241\n",
      "Learning rate: 2.7804206375634945e-05\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c00a78dc7d44f599362df3861ab1cb4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Batches', max=21021.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.9216666666666667\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.3073306977748871\n",
      "Learning rate: 2.7776456347884916e-05\n",
      "Train ROC: 0.9636617749825297\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.23995323479175568\n",
      "Learning rate: 2.7115741401455684e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "757230c1edd749f68eeb173225692fc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.598748947823059\n",
      "  Acuracy Dev = 0.8066736621196222\n",
      "  F1 Dev = 0.7016839378238342\n",
      "  ROC Dev = 0.8803806782894426\n",
      "Train ROC: 0.8902340597255851\n",
      "Train accuracy: 0.8625\n",
      "Training loss: 0.37392547726631165\n",
      "Learning rate: 2.6455026455026456e-05\n",
      "Train ROC: 0.8761755485893418\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.3874090909957886\n",
      "Learning rate: 2.5794311508597225e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41752d04dfeb458d860051b1a886fa19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.8923636363636364\n",
      "Train accuracy: 0.8375\n",
      "Training loss: 0.3760998845100403\n",
      "Learning rate: 2.4472881615738758e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9e45ef4076f4c1aad787595751ddf12",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6014140481669834\n",
      "  Acuracy Dev = 0.808604407135362\n",
      "  F1 Dev = 0.7019218198457315\n",
      "  ROC Dev = 0.8818576672311645\n",
      "Train ROC: 0.9086538461538461\n",
      "Train accuracy: 0.8125\n",
      "Training loss: 0.3471875488758087\n",
      "Learning rate: 2.3812166669309527e-05\n",
      "Train ROC: 0.9189378057302586\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.33646488189697266\n",
      "Learning rate: 2.3151451722880292e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c53a6f76fa4e4dd690c6856582e35a2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5964450211995388\n",
      "  Acuracy Dev = 0.804784889821616\n",
      "  F1 Dev = 0.7022978941304487\n",
      "  ROC Dev = 0.8792647008763546\n",
      "Train ROC: 0.9585454545454545\n",
      "Train accuracy: 0.9125\n",
      "Training loss: 0.2510432004928589\n",
      "Learning rate: 2.249073677645106e-05\n",
      "Train ROC: 0.9397321428571428\n",
      "Train accuracy: 0.8875\n",
      "Training loss: 0.2831451892852783\n",
      "Learning rate: 2.183002183002183e-05\n",
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.5944969588656094\n",
      "  Acuracy Dev = 0.8054564533053515\n",
      "  F1 Dev = 0.6881517863150104\n",
      "  ROC Dev = 0.8814327519292304\n",
      "Train ROC: 0.9425287356321839\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.3084748685359955\n",
      "Learning rate: 2.1169306883592597e-05\n",
      "Train ROC: 0.9017857142857144\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.35808029770851135\n",
      "Learning rate: 2.0508591937163366e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d8bc6dc9945c46ecbbef8dc7f1f3ebe0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2adf6bb6b3af480da1ee6a2dc9197ed3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6040351162609724\n",
      "  Acuracy Dev = 0.8090661070304302\n",
      "  F1 Dev = 0.7109727428680347\n",
      "  ROC Dev = 0.8843723809937099\n",
      "Train ROC: 0.9475890985324947\n",
      "Train accuracy: 0.8875\n",
      "Training loss: 0.2689463794231415\n",
      "Learning rate: 1.8526447097875668e-05\n",
      "Train ROC: 0.91015625\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.3051362633705139\n",
      "Learning rate: 1.7865732151446436e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "134dc654922b42b19aad4bbf9631cc54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.605562752790373\n",
      "  Acuracy Dev = 0.809485834207765\n",
      "  F1 Dev = 0.7148690244362083\n",
      "  ROC Dev = 0.8858565483086762\n",
      "Train ROC: 0.9546703296703296\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.26548585295677185\n",
      "Learning rate: 1.7205017205017205e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1051084f8b246419f7c8cfec16eaa34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.9007686932215233\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.3778061866760254\n",
      "Learning rate: 1.5222872365729507e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dfe7087bf75f44e4ae4ccf0e60af5d36",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6049565812285839\n",
      "  Acuracy Dev = 0.8114165792235047\n",
      "  F1 Dev = 0.6987192382485079\n",
      "  ROC Dev = 0.8856391200999794\n",
      "Train ROC: 0.9308333333333334\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.3253900408744812\n",
      "Learning rate: 1.4562157419300276e-05\n",
      "Train ROC: 0.9716024340770791\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.22869662940502167\n",
      "Learning rate: 1.3901442472871042e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "080ab67049704321a43fe5dec5d48bff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6049439806644868\n",
      "  Acuracy Dev = 0.8080167890870934\n",
      "  F1 Dev = 0.7188345217605114\n",
      "  ROC Dev = 0.88657091337748\n",
      "Train ROC: 0.8993710691823898\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.3934527039527893\n",
      "Learning rate: 1.3240727526441813e-05\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd50759e085c4182807ea67de6b8209d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6062855956509797\n",
      "  Acuracy Dev = 0.8112486883525708\n",
      "  F1 Dev = 0.7080438875543725\n",
      "  ROC Dev = 0.887348135786378\n",
      "Train ROC: 0.9282990083905416\n",
      "Train accuracy: 0.8625\n",
      "Training loss: 0.31258562207221985\n",
      "Learning rate: 1.1919297633583346e-05\n",
      "Train ROC: 0.880184331797235\n",
      "Train accuracy: 0.775\n",
      "Training loss: 0.4340452253818512\n",
      "Learning rate: 1.1258582687154115e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ace983222d04681b32855eadf0e9232",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6085985056226089\n",
      "  Acuracy Dev = 0.812507869884575\n",
      "  F1 Dev = 0.7108924988673871\n",
      "  ROC Dev = 0.8874412128490645\n",
      "Train ROC: 0.889397406559878\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.37193986773490906\n",
      "Learning rate: 1.0597867740724883e-05\n",
      "Train ROC: 0.9560931899641578\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.2541200816631317\n",
      "Learning rate: 9.937152794295652e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2c7d924e801448ebed811cdfa671a7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6081039524874794\n",
      "  Acuracy Dev = 0.810912906610703\n",
      "  F1 Dev = 0.717749514441451\n",
      "  ROC Dev = 0.8877547939998727\n",
      "Train ROC: 0.8971354166666667\n",
      "Train accuracy: 0.8\n",
      "Training loss: 0.3900657892227173\n",
      "Learning rate: 9.276437847866419e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.98\n",
      "Train accuracy: 0.8625\n",
      "Training loss: 0.24718227982521057\n",
      "Learning rate: 8.615722901437187e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd2d5a39f1914c309d800f504685d412",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6097731075456313\n",
      "  Acuracy Dev = 0.8122560335781742\n",
      "  F1 Dev = 0.7176314626601856\n",
      "  ROC Dev = 0.8886919938489388\n",
      "Train ROC: 0.9105454545454545\n",
      "Train accuracy: 0.825\n",
      "Training loss: 0.37262165546417236\n",
      "Learning rate: 7.955007955007956e-06\n",
      "Train ROC: 0.9341692789968652\n",
      "Train accuracy: 0.8375\n",
      "Training loss: 0.29048022627830505\n",
      "Learning rate: 7.294293008578723e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f48698e6ec5649f1be4b9970998fd15f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6086857166213097\n",
      "  Acuracy Dev = 0.8128436516264428\n",
      "  F1 Dev = 0.7088855520010445\n",
      "  ROC Dev = 0.8885818646815047\n",
      "Train ROC: 0.9749784296807593\n",
      "Train accuracy: 0.9125\n",
      "Training loss: 0.20650334656238556\n",
      "Learning rate: 6.633578062149491e-06\n",
      "Train ROC: 0.9628127112914131\n",
      "Train accuracy: 0.9125\n",
      "Training loss: 0.25391659140586853\n",
      "Learning rate: 5.9728631157202585e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e57789028d004709bf74bf945f598f87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.9587242026266416\n",
      "Train accuracy: 0.8875\n",
      "Training loss: 0.30974555015563965\n",
      "Learning rate: 4.651433222861795e-06\n",
      "Train ROC: 0.9170909090909091\n",
      "Train accuracy: 0.85\n",
      "Training loss: 0.33698007464408875\n",
      "Learning rate: 3.33000333000333e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eeef9458a3ff4f8ba05dfd788a4fdf08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train ROC: 0.9287272727272726\n",
      "Train accuracy: 0.875\n",
      "Training loss: 0.3170314431190491\n",
      "Learning rate: 2.0085734371448653e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac622c4b7dcb448293d27a2700715fb4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Dev batch', max=745.0, style=ProgressStyle(description_wi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "***** Eval results *****\n",
      "  AP Dev = 0.6107921747560067\n",
      "  Acuracy Dev = 0.8133473242392445\n",
      "  F1 Dev = 0.715864800971184\n",
      "  ROC Dev = 0.8894114087507152\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "import numpy as np\n",
    "import datetime\n",
    "\n",
    "try:\n",
    "    from sklearn.metrics import f1_score, average_precision_score, accuracy_score, roc_auc_score\n",
    "except:\n",
    "    !pip install sklearn\n",
    "    from sklearn.metrics import f1_score, average_precision_score, accuracy_score, roc_auc_score\n",
    "\n",
    "global_step = 0 # Number of steps performed so far\n",
    "tr_loss = 0.0 # Training loss\n",
    "model.zero_grad() # Initialize gradients to 0\n",
    "\n",
    "for _ in tqdm(range(n_epochs), desc=\"Epochs\"):\n",
    "    for step, batch in tqdm(enumerate(train_data_loader), desc=\"Batches\", total=len(train_data_loader)):\n",
    "        model.train()\n",
    "        # get the batch inpute\n",
    "        inputs = {\n",
    "            'input_ids': batch[0].to(device),\n",
    "            'attention_mask': batch[1].to(device),\n",
    "            'labels': batch[3].to(device)\n",
    "        }\n",
    "        # Run through the network.\n",
    "        \n",
    "        with warnings.catch_warnings():\n",
    "            # There is a very annoying warning here when we are using multiple GPUS,\n",
    "            # As described here: https://github.com/huggingface/transformers/issues/852.\n",
    "            # We can safely ignore this.\n",
    "            warnings.simplefilter(\"ignore\")\n",
    "            outputs = model(**inputs)\n",
    "        loss = outputs[0]\n",
    "\n",
    "        loss = loss.sum()/len(model.device_ids) # Average over all GPUS.\n",
    "        # Clipping gradients. Avoud gradient explosion, if the gradient is too large.\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "        # Backward pass on the network\n",
    "        loss.backward()\n",
    "        tr_loss += loss.item()\n",
    "        # Run the optimizer with the gradients\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        model.zero_grad()\n",
    "        if step % steps_to_print == 0:\n",
    "            # Logits is the actual output from the network. \n",
    "            # This is the probability of being relevant or not.\n",
    "            # You can check its shape (Should be a vector sized 2) with logits.shape()\n",
    "            logits = outputs[1]\n",
    "            # Send the logits to the CPU and in numpy form. Easier to check what is going on.\n",
    "            preds = logits.detach().cpu().numpy()\n",
    "            \n",
    "            # Bring the labels to CPU too.\n",
    "            tqdm.write(f\"Training loss: {loss.item()} Learning Ragte: {scheduler.get_last_lr()[0]}\")\n",
    "        global_step += 1\n",
    "        \n",
    "        # Run an evluation step over the eval dataset. Let's see how we are going.\n",
    "        if global_step%steps_to_eval == 0:\n",
    "            eval_loss = 0.0\n",
    "            nb_eval_steps = 0\n",
    "            preds = None\n",
    "            out_label_ids = None\n",
    "            for batch in tqdm(dev_data_loader, desc=\"Dev batch\"):\n",
    "                model.eval()\n",
    "                with torch.no_grad(): # Avoid upgrading gradients here\n",
    "                    inputs = {'input_ids': batch[0].to(device),\n",
    "                      'attention_mask': batch[1].to(device),\n",
    "                      'labels': batch[3].to(device)}\n",
    "                    with warnings.catch_warnings():\n",
    "                        warnings.simplefilter(\"ignore\")\n",
    "                        outputs = model(**inputs)\n",
    "                    tmp_eval_loss, logits = outputs[:2] # Logits is the actual output. Probabilities between 0 and 1.\n",
    "                    eval_loss += tmp_eval_loss.mean().item()\n",
    "                    nb_eval_steps += 1\n",
    "                    # Concatenate all outputs to evaluate in the end.\n",
    "                    if preds is None:\n",
    "                        preds = logits.detach().cpu().numpy() # PRedictions into numpy mode\n",
    "                        out_label_ids = inputs['labels'].detach().cpu().numpy().flatten() # Labels assigned by model\n",
    "                    else:\n",
    "                        batch_predictions = logits.detach().cpu().numpy()\n",
    "                        preds = np.append(preds, batch_predictions, axis=0)\n",
    "                        out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy().flatten(), axis=0)\n",
    "                eval_loss = eval_loss / nb_eval_steps\n",
    "            results = {}\n",
    "            results[\"ROC Dev\"] = roc_auc_score(out_label_ids, preds[:, 1])\n",
    "            preds = np.argmax(preds, axis=1)\n",
    "            results[\"Acuracy Dev\"] = accuracy_score(out_label_ids, preds)\n",
    "            results[\"F1 Dev\"] = f1_score(out_label_ids, preds)\n",
    "            results[\"AP Dev\"] = average_precision_score(out_label_ids, preds)\n",
    "            tqdm.write(\"***** Eval results *****\")\n",
    "            for key in sorted(results.keys()):\n",
    "                tqdm.write(f\"  {key} = {str(results[key])}\")\n",
    "            output_dir = path(f\"checkpoints/checkpoint-{global_step}\")\n",
    "            if not os.path.isdir(output_dir):\n",
    "                os.makedirs(path(output_dir))\n",
    "#             print(f\"Saving model checkpoint to {output_dir}\")\n",
    "            model_to_save = model.module if hasattr(model, 'module') else model\n",
    "            model_to_save.save_pretrained(output_dir)\n",
    "\n",
    "# Save final model \n",
    "output_dir = path(f\"models/distilBERT-{str(datetime.date.today())}\")\n",
    "if not os.path.isdir(output_dir):\n",
    "    os.makedirs(path(output_dir))\n",
    "model_to_save = model.module if hasattr(model, 'module') else model\n",
    "model_to_save.save_pretrained(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-16T19:43:37.911307Z",
     "start_time": "2020-04-16T19:43:37.889677Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): DistilBertForSequenceClassification(\n",
       "    (distilbert): DistilBertModel(\n",
       "      (embeddings): Embeddings(\n",
       "        (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "        (position_embeddings): Embedding(512, 768)\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (transformer): Transformer(\n",
       "        (layer): ModuleList(\n",
       "          (0): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "          (1): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "          (2): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "          (3): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "          (4): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "          (5): TransformerBlock(\n",
       "            (attention): MultiHeadSelfAttention(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
       "            )\n",
       "            (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (ffn): FFN(\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            )\n",
       "            (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pre_classifier): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       "    (dropout): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7 (Bert4IR)",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}