Train BERT.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": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T14:36:23.896252Z",
     "start_time": "2020-05-24T14:36:23.881984Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from os.path import expanduser\n",
    "home = expanduser(\"~\")\n",
    "os.environ[\"JAVA_HOME\"] = f\"{home}/.sdkman/candidates/java/11.0.7.hs-adpt\"  #Set right JAVA version\n",
    "data_home = \"/ssd2/arthur/MsMarcoTREC/\"\n",
    "def path(x):\n",
    "    return os.path.join(data_home, x)\n",
    "\n",
    "try:\n",
    "    import pyserini\n",
    "except:\n",
    "    !pip install pyserini==0.9.2.0 # install pyserini\n",
    "try:\n",
    "    import tqdm\n",
    "except:\n",
    "    !pip install tqdm # Good for progress bars!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:53:32.427832Z",
     "start_time": "2020-05-20T20:53:31.030388Z"
    }
   },
   "outputs": [],
   "source": [
    "import jnius_config\n",
    "jnius_config.add_options('-Xmx16G') # Adjust to your machine. Probably less than 16G.\n",
    "from pyserini.search import pysearch\n",
    "import subprocess\n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:53:36.703105Z",
     "start_time": "2020-05-20T20:53:36.159207Z"
    }
   },
   "outputs": [],
   "source": [
    "# Anserini uses this \"SimpleSearcher\" object for interfacing with the index.\n",
    "index_path = path(\"lucene-index.msmarco-doc.pos+docvectors+rawdocs\")\n",
    "searcher = pysearch.SimpleSearcher(index_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting Anserini top-10\n",
    "We will use pyserini 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. 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",
    "obs.: Potentinaly, it would be faster to use Anserini's `batch_search()` method, since it works in multiple threads. However, the lack of feedback (i.e. How may queries have been processed already) and higher memory footprint could cause issues."
   ]
  },
  {
   "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": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:53:45.095002Z",
     "start_time": "2020-05-20T20:53:44.394558Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load the relevant query-document pairs\n",
    "relevant_docs = defaultdict(lambda:[])\n",
    "for file in [path(\"qrels/msmarco-doctrain-qrels.tsv\"), 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)                            "
   ]
  },
  {
   "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": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:55:22.572882Z",
     "start_time": "2020-05-20T20:55:11.163769Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Already found file /ssd2/arthur/MsMarcoTREC/queries/msmarco-doctrain-queries.tsv. Cowardly refusing to run this again. Will only load querytexts.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "407e17cda787446db192d2b2710eff96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Retrieval batches', max=37.0, style=ProgressStyle(descrip…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Already found file /ssd2/arthur/MsMarcoTREC/queries/msmarco-docdev-queries.tsv. Cowardly refusing to run this again. Will only load querytexts.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a4832441572e46a9897bfd39b9ba0c5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Retrieval batches', max=1.0, style=ProgressStyle(descript…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pattern = re.compile('([^\\s\\w]|_)+')\n",
    "\n",
    "anserini_top_10 = defaultdict(lambda:[])\n",
    "searcher.set_bm25_similarity(0.9, 0.4)\n",
    "pairs_per_split = defaultdict(lambda: [])\n",
    "threads = 42 # Number of Threads to use when retrieving\n",
    "k = 10       # Number of documents to retrieve \n",
    "neg_samples = 2 # Number of negatives samples to use\n",
    "batch_size = 10000 # Batch size for each retrieval step on Anserini\n",
    "\n",
    "query_texts = dict()\n",
    "for split in [\"train\", \"dev\"]:\n",
    "    file_path = path(f\"queries/msmarco-doc{split}-queries.tsv\")\n",
    "    run_search=True\n",
    "    if os.path.isfile(f\"{split}_triples.pkl\"):\n",
    "        print(f\"Already found file {file_path}. Cowardly refusing to run this again. Will only load querytexts.\")\n",
    "        pairs_per_split[split] = pickle.load(open(path(f\"{split}_triples.pkl\"), 'rb'))\n",
    "        run_search = False\n",
    "    number_of_queries = int(subprocess.run(f\"wc -l {file_path}\".split(), capture_output=True).stdout.split()[0])\n",
    "    number_of_batches = math.ceil(number_of_queries/batch_size)\n",
    "    pbar = tqdm(total=number_of_batches, desc=\"Retrieval batches\")\n",
    "    queries = []\n",
    "    query_ids = []\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": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:56:58.030400Z",
     "start_time": "2020-05-20T20:56:57.979290Z"
    },
    "code_folding": [],
    "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",
    "    - There is an extra class here that does that, if you want to follow that path.\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": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:57:00.212022Z",
     "start_time": "2020-05-20T20:57:00.162305Z"
    }
   },
   "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": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T20:57:05.559885Z",
     "start_time": "2020-05-20T20:57:00.890040Z"
    }
   },
   "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": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T21:29:57.630269Z",
     "start_time": "2020-05-20T21:29:57.600264Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from transformers import DistilBertModel, BertModel\n",
    "\n",
    "class BertRelevanceRanker(nn.Module):\n",
    "    def __init__(self, model=\"distilbert-base-uncased\"):\n",
    "        \"\"\"Creates an instance of Bert Relevance Ranker. \n",
    "        It feeds two senteces into a pre-trained BERT model, extracts the [CLS] token and feeds it into a one-layer FFNN\"\"\"\n",
    "        super().__init__()\n",
    "        self.distil = False\n",
    "        self.loss_fct = nn.CrossEntropyLoss()\n",
    "        if \"distil\" in model:\n",
    "            self.distil = True\n",
    "            self.bert = DistilBertModel.from_pretrained(model)\n",
    "        else:\n",
    "            self.bert = BertModel.from_pretrained(model)\n",
    "        self.config = self.bert.config\n",
    "        self.linear1 = nn.Linear(self.bert.config.dim, self.bert.config.dim)\n",
    "        self.linear2 = nn.Linear(self.bert.config.dim, 2)\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "    \n",
    "    def forward(self, input_ids, attention_mask, token_type_ids=None, labels=None):\n",
    "        if not self.distil and token_type_ids is None:\n",
    "            raise ValueError(\"Model is not distilBERT and it did not received token_type_ids!\")\n",
    "        if not self.distil:\n",
    "            outputs = self.bert(input_ids, attention_mask, token_type_ids)\n",
    "        else:\n",
    "            pooled_output = self.bert(input_ids, attention_mask)[0][:, 0]\n",
    "        pooled_output = self.linear1(pooled_output)\n",
    "        pooled_output = nn.ReLU()(pooled_output)\n",
    "        pooled_output = self.dropout(pooled_output)\n",
    "        logits = self.linear2(pooled_output)\n",
    "        outputs = (logits,)\n",
    "        if labels is not None:\n",
    "            loss = self.loss_fct(logits.view(-1, 2), labels.view(-1))\n",
    "            outputs = (loss, ) + outputs\n",
    "        return outputs\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-20T21:57:24.855038Z",
     "start_time": "2020-05-20T21:57:21.130878Z"
    }
   },
   "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. Usefull in a Notebook.\n",
    "except:\n",
    "    pass\n",
    "\n",
    "# model = DistilBertForSequenceClassification.from_pretrained(\"distilbert-base-uncased\")\n",
    "model = BertRelevanceRanker()\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",
    "    \n",
    "    # For a 1080Ti, 16 samples fit on a GPU confortably for a DistilBert model. A bert-base, not more than 8. 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 DataLoader is a nice device for generating batches for you easily.\n",
    "# It receives any object that implementes __getitem__(self, idx) and __len__(self)\n",
    "\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": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-21T04:35:09.400749Z",
     "start_time": "2020-05-20T21:57:26.628857Z"
    },
    "code_folding": [],
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "257afcc80db44ec0bae6ffc582f5de11",
       "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": "5a8c5c00e3154d56af65068049b450c2",
       "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": [
      "Training loss: 0.7429962754249573 Learning Rate: 1.189286903572618e-08\n",
      "Training loss: 0.6137410998344421 Learning Rate: 1.1904761904761907e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66c1c77f3e614ebd8d50f258a0158f17",
       "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.5293512055623834\n",
      "  Acuracy Dev = 0.7571878279118573\n",
      "  F1 Dev = 0.6457657216337028\n",
      "  ROC Dev = 0.8253530962920068\n",
      "Training loss: 0.5174795389175415 Learning Rate: 2.3797630940488087e-05\n",
      "Training loss: 0.3816486895084381 Learning Rate: 3.5690499976214264e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "49d2ff1b6c0642079171f7eff4c7f487",
       "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.5495730855000485\n",
      "  Acuracy Dev = 0.778174186778594\n",
      "  F1 Dev = 0.6443711728685821\n",
      "  ROC Dev = 0.8438980489934637\n",
      "Training loss: 0.4404618442058563 Learning Rate: 4.7583369011940445e-05\n",
      "Training loss: 0.43097397685050964 Learning Rate: 4.894708466137037e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67275bddc5044dd88a6456619ee61a09",
       "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.5561074157707073\n",
      "  Acuracy Dev = 0.784218258132214\n",
      "  F1 Dev = 0.6384415219073072\n",
      "  ROC Dev = 0.8563399891928818\n",
      "Training loss: 0.369981974363327 Learning Rate: 4.762565476851191e-05\n",
      "Training loss: 0.42483997344970703 Learning Rate: 4.630422487565345e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f3bb28efba441d6be4911615d846ee2",
       "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.57188751809242\n",
      "  Acuracy Dev = 0.7889611752360965\n",
      "  F1 Dev = 0.6810454199441766\n",
      "  ROC Dev = 0.8639991786767734\n",
      "Training loss: 0.3996042013168335 Learning Rate: 4.498279498279498e-05\n",
      "Training loss: 0.41216525435447693 Learning Rate: 4.366136508993652e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "932aa76e2f1f4b5ba1283e28f8fbe427",
       "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.5806967078935027\n",
      "  Acuracy Dev = 0.794501573976915\n",
      "  F1 Dev = 0.6897731592953998\n",
      "  ROC Dev = 0.8672993079445701\n",
      "Training loss: 0.36422398686408997 Learning Rate: 4.233993519707805e-05\n",
      "Training loss: 0.3106476366519928 Learning Rate: 4.101850530421959e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e6d0afbc11241a5b8ed496edcc16889",
       "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.5771497711912323\n",
      "  Acuracy Dev = 0.791311647429171\n",
      "  F1 Dev = 0.6897154268597104\n",
      "  ROC Dev = 0.8676886167711565\n",
      "Training loss: 0.3256611227989197 Learning Rate: 3.9697075411361125e-05\n",
      "Training loss: 0.3615630567073822 Learning Rate: 3.837564551850266e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7700e35d87bd4bbd9f619e4f084ee887",
       "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.5832131863608856\n",
      "  Acuracy Dev = 0.7948793284365163\n",
      "  F1 Dev = 0.6965915440491711\n",
      "  ROC Dev = 0.8700944340962947\n",
      "Training loss: 0.322876513004303 Learning Rate: 3.705421562564419e-05\n",
      "Training loss: 0.4105748236179352 Learning Rate: 3.5732785732785736e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6672f23ea264a169659aa886191572a",
       "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.5883950206287615\n",
      "  Acuracy Dev = 0.7997061909758657\n",
      "  F1 Dev = 0.6954301761552208\n",
      "  ROC Dev = 0.8730471911409979\n",
      "Training loss: 0.31920647621154785 Learning Rate: 3.4411355839927267e-05\n",
      "Training loss: 0.3777822256088257 Learning Rate: 3.3089925947068804e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d942fb8159d483c99cbc3e34feb5fed",
       "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.5856603889092339\n",
      "  Acuracy Dev = 0.8007555089192026\n",
      "  F1 Dev = 0.675551910327387\n",
      "  ROC Dev = 0.8757322753280555\n",
      "Training loss: 0.49660611152648926 Learning Rate: 3.176849605421034e-05\n",
      "Training loss: 0.4242243766784668 Learning Rate: 3.0447066161351874e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eff7504c396f402d9a27b26d12c71350",
       "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.597058873022513\n",
      "  Acuracy Dev = 0.8061699895068206\n",
      "  F1 Dev = 0.6968622817382172\n",
      "  ROC Dev = 0.8781652721702214\n",
      "Training loss: 0.4124586284160614 Learning Rate: 2.912563626849341e-05\n",
      "Training loss: 0.27975550293922424 Learning Rate: 2.7804206375634945e-05\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ff24a46f93d4ba4b54249c50ade5d5d",
       "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": [
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "\r",
      "Training loss: 0.26966866850852966 Learning Rate: 2.7776456347884916e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5be88cd9a21d4edba61e5032537ec1cd",
       "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.5963447782572543\n",
      "  Acuracy Dev = 0.8050786988457502\n",
      "  F1 Dev = 0.7003097573567372\n",
      "  ROC Dev = 0.879107783460562\n",
      "Training loss: 0.2813105285167694 Learning Rate: 2.6455026455026456e-05\n",
      "Training loss: 0.3830776512622833 Learning Rate: 2.513359656216799e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21cc4ed1ee0e487fb6514bb0f499d8c1",
       "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.5958801274717004\n",
      "  Acuracy Dev = 0.8047429171038825\n",
      "  F1 Dev = 0.7001804588811549\n",
      "  ROC Dev = 0.8808074605233469\n",
      "Training loss: 0.3111289143562317 Learning Rate: 2.3812166669309527e-05\n",
      "Training loss: 0.38978663086891174 Learning Rate: 2.249073677645106e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "033a56b60ff241f8b2fe9ae76731dd17",
       "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.5990996255714716\n",
      "  Acuracy Dev = 0.80700944386149\n",
      "  F1 Dev = 0.7012345679012346\n",
      "  ROC Dev = 0.8803928430754654\n",
      "Training loss: 0.3295662999153137 Learning Rate: 2.1169306883592597e-05\n",
      "Training loss: 0.26243922114372253 Learning Rate: 1.984787699073413e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f4425ce11df847a990d4c34c12975ad0",
       "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.5991149488934595\n",
      "  Acuracy Dev = 0.8084784889821616\n",
      "  F1 Dev = 0.6884686283880658\n",
      "  ROC Dev = 0.8823239998837508\n",
      "Training loss: 0.29984936118125916 Learning Rate: 1.8526447097875668e-05\n",
      "Training loss: 0.38034573197364807 Learning Rate: 1.7205017205017205e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83649a33348241c887b295d24fea477c",
       "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.6085644484307055\n",
      "  Acuracy Dev = 0.8130535152151102\n",
      "  F1 Dev = 0.7064329027155286\n",
      "  ROC Dev = 0.8832567761742633\n",
      "Training loss: 0.3424564301967621 Learning Rate: 1.588358731215874e-05\n",
      "Training loss: 0.33130455017089844 Learning Rate: 1.4562157419300276e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b213af1583cf467ab77667a0c46b37e0",
       "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.6064592677699424\n",
      "  Acuracy Dev = 0.8112906610703043\n",
      "  F1 Dev = 0.7086195722618276\n",
      "  ROC Dev = 0.8849555376079928\n",
      "Training loss: 0.3085682988166809 Learning Rate: 1.3240727526441813e-05\n",
      "Training loss: 0.45851755142211914 Learning Rate: 1.1919297633583346e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bcf6a127141041b6a4e60c0c68b322b2",
       "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.6096129126893163\n",
      "  Acuracy Dev = 0.8131794333683106\n",
      "  F1 Dev = 0.7112927288058637\n",
      "  ROC Dev = 0.8856695697207766\n",
      "Training loss: 0.38201770186424255 Learning Rate: 1.0597867740724883e-05\n",
      "Training loss: 0.29469871520996094 Learning Rate: 9.276437847866419e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "837ac55f799f4ed4a912f2d14abbb4f0",
       "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.6080821099123832\n",
      "  Acuracy Dev = 0.8105771248688353\n",
      "  F1 Dev = 0.7191836226743824\n",
      "  ROC Dev = 0.8863449947902688\n",
      "Training loss: 0.28664761781692505 Learning Rate: 7.955007955007956e-06\n",
      "Training loss: 0.25723132491111755 Learning Rate: 6.633578062149491e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "baf1f8a036794f318a39713afbb36838",
       "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.6087906985691338\n",
      "  Acuracy Dev = 0.8128436516264428\n",
      "  F1 Dev = 0.709454616537434\n",
      "  ROC Dev = 0.8863522920763777\n",
      "Training loss: 0.22093525528907776 Learning Rate: 5.312148169291026e-06\n",
      "Training loss: 0.24985454976558685 Learning Rate: 3.9907182764325615e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b979245d0e5d4ce881b1ffffe7699ea1",
       "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.6079977362388271\n",
      "  Acuracy Dev = 0.8115844700944386\n",
      "  F1 Dev = 0.7138759640512462\n",
      "  ROC Dev = 0.8870989261328032\n",
      "Training loss: 0.24013452231884003 Learning Rate: 2.6692883835740977e-06\n",
      "Training loss: 0.4052871763706207 Learning Rate: 1.3478584907156336e-06\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b5b2628e162445daa6c5e79ecd216022",
       "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.6083502009552072\n",
      "  Acuracy Dev = 0.8121301154249738\n",
      "  F1 Dev = 0.7121913580246914\n",
      "  ROC Dev = 0.8872791108323864\n",
      "Training loss: 0.2829817235469818 Learning Rate: 2.6428597857169283e-08\n",
      "\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. Avoid 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 Rate: {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.config.save_pretrained(output_dir)\n",
    "            torch.save(model_to_save.state_dict(), output_dir+\"/pytorch_model.bin\")\n",
    "            \n",
    "\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.config.save_pretrained(output_dir)\n",
    "torch.save(model_to_save.state_dict(), output_dir+\"/pytorch_model.bin\")"
   ]
  },
  {
   "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
}