6-Bert-fine-tuning-NER-CONLL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT - Fine-tuning NER\n",
    "\n",
    "\n",
    "## Préparation des données CoNLL-2003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm, trange\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from torchtext import data\n",
    "from torchtext import datasets\n",
    "from torchtext.datasets import SequenceTaggingDataset\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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>Sentence #</th>\n",
       "      <th>text</th>\n",
       "      <th>tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>204557</th>\n",
       "      <td>14985</td>\n",
       "      <td>Plymouth</td>\n",
       "      <td>B-ORG</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204558</th>\n",
       "      <td>14985</td>\n",
       "      <td>2</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204559</th>\n",
       "      <td>14985</td>\n",
       "      <td>Preston</td>\n",
       "      <td>B-ORG</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204560</th>\n",
       "      <td>14985</td>\n",
       "      <td>1</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204561</th>\n",
       "      <td>14986</td>\n",
       "      <td>Division</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204562</th>\n",
       "      <td>14986</td>\n",
       "      <td>three</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204563</th>\n",
       "      <td>14987</td>\n",
       "      <td>Swansea</td>\n",
       "      <td>B-ORG</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204564</th>\n",
       "      <td>14987</td>\n",
       "      <td>1</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204565</th>\n",
       "      <td>14987</td>\n",
       "      <td>Lincoln</td>\n",
       "      <td>B-ORG</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>204566</th>\n",
       "      <td>14987</td>\n",
       "      <td>2</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        Sentence #      text    tag\n",
       "204557       14985  Plymouth  B-ORG\n",
       "204558       14985         2      O\n",
       "204559       14985   Preston  B-ORG\n",
       "204560       14985         1      O\n",
       "204561       14986  Division      O\n",
       "204562       14986     three      O\n",
       "204563       14987   Swansea  B-ORG\n",
       "204564       14987         1      O\n",
       "204565       14987   Lincoln  B-ORG\n",
       "204566       14987         2      O"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv(\"data_ner/data_train.csv\", encoding=\"latin1\").fillna(method=\"ffill\")\n",
    "test = pd.read_csv(\"data_ner/data_test.csv\", encoding=\"latin1\").fillna(method=\"ffill\")\n",
    "data.tail(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On récupère les phrases et les tags avec la fonction suivante."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SentenceGetter(object):\n",
    "\n",
    "    def __init__(self, data):\n",
    "        self.n_sent = 1\n",
    "        self.data = data\n",
    "        self.empty = False\n",
    "        agg_func = lambda s: [(w,  t) for w, t in zip(s[\"text\"].values.tolist(),\n",
    "                                                          \n",
    "                                                           s[\"tag\"].values.tolist())]\n",
    "        self.grouped = self.data.groupby(\"Sentence #\").apply(agg_func)\n",
    "        self.sentences = [s for s in self.grouped]\n",
    "\n",
    "    def get_next(self):\n",
    "        try:\n",
    "            s = self.grouped[\"Sentence: {}\".format(self.n_sent)]\n",
    "            self.n_sent += 1\n",
    "            return s\n",
    "        except:\n",
    "            return None\n",
    "      \n",
    "getter = SentenceGetter(data)\n",
    "get_test = SentenceGetter(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']\n"
     ]
    }
   ],
   "source": [
    "sentences = [[word[0] for word in sentence] for sentence in getter.sentences]\n",
    "sentences_test = [[word[0] for word in sentence] for sentence in get_test.sentences]\n",
    "print(sentences[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']\n"
     ]
    }
   ],
   "source": [
    "labels = [[s[1] for s in sentence] for sentence in getter.sentences]\n",
    "labels_test = [[s[1] for s in sentence] for sentence in get_test.sentences]\n",
    "print(labels[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On attribut un indice pour chaque tag et on ajoute le tag PAD."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_values = list(set(data[\"tag\"].values))\n",
    "tag_values.append(\"PAD\")\n",
    "tag2idx = {t: i for i, t in enumerate(tag_values)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'I-MISC': 0,\n",
       " 'B-ORG': 1,\n",
       " 'I-LOC': 2,\n",
       " 'B-PER': 3,\n",
       " 'B-LOC': 4,\n",
       " 'I-ORG': 5,\n",
       " 'O': 6,\n",
       " 'B-MISC': 7,\n",
       " 'I-PER': 8,\n",
       " 'PAD': 9}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tag2idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'1.4.0'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n",
    "from transformers import BertTokenizer, BertConfig\n",
    "\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ici, nous corrigeons certaines configurations. \n",
    "Nous limiterons la longueur de notre séquence à 512 tokens et nous utiliserons une taille de batch de 32 comme suggéré par l'article de Bert. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LEN = 512\n",
    "bs = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Maintenant, nous tokenisons toutes les phrases. Étant donné que le tokenizer BERT est basé sur un tokenizer Wordpiece, il divisera les tokens en tokens de sous-mots. Par exemple, \"gunships\" sera divisé en deux tokens \"guns\" et \"##hips\". \n",
    "La fonction suivante nous permet de résoudre ce problème."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_and_preserve_labels(sentence, text_labels):\n",
    "    tokenized_sentence = []\n",
    "    labels = []\n",
    "\n",
    "    for word, label in zip(sentence, text_labels):\n",
    "\n",
    "        # Tokenize the mot et compte le nombre de sous-mots \n",
    "        tokenized_word = tokenizer.tokenize(word)\n",
    "        n_subwords = len(tokenized_word)\n",
    "        tokenized_sentence.extend(tokenized_word)\n",
    "\n",
    "        # Ajoute le même tag pour tous les sous-mots d'un mot\n",
    "        labels.extend([label] * n_subwords)\n",
    "\n",
    "    return tokenized_sentence, labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_texts_and_labels = [\n",
    "    tokenize_and_preserve_labels(sent, labs)\n",
    "    for sent, labs in zip(sentences, labels)\n",
    "]\n",
    "# pour le test\n",
    "tokenized_texts_and_labels_test = [\n",
    "    tokenize_and_preserve_labels(sent, labs)\n",
    "    for sent, labs in zip(sentences_test, labels_test)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_texts = [token_label_pair[0] for token_label_pair in tokenized_texts_and_labels]\n",
    "labels = [token_label_pair[1] for token_label_pair in tokenized_texts_and_labels]\n",
    "\n",
    "# pour le test\n",
    "tokenized_texts_test = [token_label_pair[0] for token_label_pair in tokenized_texts_and_labels_test]\n",
    "labels_test = [token_label_pair[1] for token_label_pair in tokenized_texts_and_labels_test]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ensuite, nous coupons et remplissons les séquences de tokens et de labels à la longueur souhaitée en ajoutant des tokens `PAD`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ids = pad_sequences([tokenizer.convert_tokens_to_ids(txt) for txt in tokenized_texts],\n",
    "                          maxlen=MAX_LEN, dtype=\"long\", value=0.0,\n",
    "                          truncating=\"post\", padding=\"post\")\n",
    "\n",
    "#test\n",
    "input_ids_test = pad_sequences([tokenizer.convert_tokens_to_ids(txt) for txt in tokenized_texts_test],\n",
    "                          maxlen=MAX_LEN, dtype=\"long\", value=0.0,\n",
    "                          truncating=\"post\", padding=\"post\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = pad_sequences([[tag2idx.get(l) for l in lab] for lab in labels],\n",
    "                     maxlen=MAX_LEN, value=tag2idx[\"PAD\"], padding=\"post\",\n",
    "                     dtype=\"long\", truncating=\"post\")\n",
    "\n",
    "#test\n",
    "tags_test = pad_sequences([[tag2idx.get(l) for l in lab] for lab in labels_test],\n",
    "                     maxlen=MAX_LEN, value=tag2idx[\"PAD\"], padding=\"post\",\n",
    "                     dtype=\"long\", truncating=\"post\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Le modèle de Bert prend en charge des masques d'attention ou attention_mask. Nous créons donc ici le masque pour ignorer les éléments PAD dans les séquences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_masks = [[float(i != 0.0) for i in ii] for ii in input_ids]\n",
    "\n",
    "# test\n",
    "attention_masks_test = [[float(i != 0.0) for i in ii] for ii in input_ids_test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_inputs, val_inputs, tr_tags, val_tags = train_test_split(input_ids, tags,\n",
    "                                                            random_state=2018, test_size=0.1)\n",
    "tr_masks, val_masks, _, _ = train_test_split(attention_masks, input_ids,\n",
    "                                             random_state=2018, test_size=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On convertit nos inputs, tags et masques en tensor torch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train\n",
    "tr_inputs = torch.tensor(tr_inputs)\n",
    "tr_tags = torch.tensor(tr_tags)\n",
    "tr_masks = torch.tensor(tr_masks)\n",
    "# valid\n",
    "val_inputs = torch.tensor(val_inputs)\n",
    "val_tags = torch.tensor(val_tags)\n",
    "val_masks = torch.tensor(val_masks)\n",
    "# test\n",
    "test_inputs = torch.tensor(input_ids_test)\n",
    "test_tags = torch.tensor(tags_test)\n",
    "test_masks = torch.tensor(attention_masks_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "La dernière étape consiste à définir les itérateurs. Nous mélangons les données au moment de l'apprentissage avec le RandomSampler et au moment du test, nous les transmettons simplement séquentiellement avec le SequentialSampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = TensorDataset(tr_inputs, tr_masks, tr_tags)\n",
    "train_sampler = RandomSampler(train_data)\n",
    "train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=bs)\n",
    "\n",
    "valid_data = TensorDataset(val_inputs, val_masks, val_tags)\n",
    "valid_sampler = SequentialSampler(valid_data)\n",
    "valid_dataloader = DataLoader(valid_data, sampler=valid_sampler, batch_size=bs)\n",
    "\n",
    "test_data = TensorDataset(test_inputs, test_masks, test_tags)\n",
    "test_sampler = SequentialSampler(test_data)\n",
    "test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=bs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construction du modèle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Le package transformer fournit une classe BertForTokenClassification pour les prédictions au niveau du token. \n",
    "\n",
    "BertForTokenClassification est un modèle de réglage fin qui enveloppe BertModel et ajoute un classificateur au niveau du token au-dessus du BertModel. Le classificateur au niveau du token est une couche linéaire qui prend comme entrée le dernier état caché de la séquence. Nous chargeons le modèle pré-formé et fournissons le nombre de tags possibles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "from transformers import BertForTokenClassification, AdamW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BertForTokenClassification.from_pretrained(\n",
    "    'bert-base-uncased',\n",
    "    num_labels=len(tag2idx),\n",
    "    output_attentions = False,\n",
    "    output_hidden_states = False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "n_gpu = torch.cuda.device_count()\n",
    "#model.cuda();\n",
    "device = 'cpu'\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - Optimiseur et scheduler : \n",
    " \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_optimizer = list(model.named_parameters())\n",
    "no_decay = ['bias', 'gamma', 'beta']\n",
    "optimizer_grouped_parameters = [\n",
    "    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n",
    "        'weight_decay_rate': 0.01},\n",
    "    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],\n",
    "        'weight_decay_rate': 0.0}\n",
    "]\n",
    "\n",
    "optimizer = AdamW(\n",
    "    optimizer_grouped_parameters,\n",
    "    lr=3e-5,\n",
    "    eps=1e-8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import get_linear_schedule_with_warmup\n",
    "\n",
    "epochs = 10\n",
    "max_grad_norm = 1.0\n",
    "\n",
    "\n",
    "total_steps = len(train_dataloader) * epochs\n",
    "\n",
    "scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=total_steps\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "TAG_PAD_IDX = tag2idx.get('PAD')\n",
    "criterion = nn.CrossEntropyLoss(ignore_index = TAG_PAD_IDX)\n",
    "criterion = criterion.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Entraînement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score, classification_report\n",
    "def train(model, iterator, optimizer, criterion):\n",
    "    tag_pad_idx = 0\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    total_loss = 0\n",
    "    f1 = 0\n",
    "    \n",
    "    eval_loss, eval_accuracy = 0, 0\n",
    "    nb_eval_steps, nb_eval_examples = 0, 0\n",
    "    predictions , true_labels = [], []\n",
    "    \n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        \n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "        b_input_ids, b_input_mask, b_labels = batch \n",
    "        model.zero_grad()\n",
    "        outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)\n",
    "        loss = outputs[0]    \n",
    "        loss.backward()\n",
    "        total_loss += loss.item()\n",
    "        torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "  \n",
    "        logits = outputs[1].detach().cpu().numpy()\n",
    "        label_ids = b_labels.to('cpu').numpy()\n",
    "        predictions.extend([list(p) for p in np.argmax(logits, axis=2)])\n",
    "        true_labels.extend(label_ids)\n",
    "        pred_tags = [tag_values[p_i] for p, l in zip(predictions, true_labels)\n",
    "                                 for p_i, l_i in zip(p, l) if tag_values[l_i] != \"PAD\"]\n",
    "        valid_tags = [tag_values[l_i] for l in true_labels\n",
    "                                  for l_i in l if tag_values[l_i] != \"PAD\"]\n",
    "    \n",
    "        f1s = f1_score(pred_tags, valid_tags, average=\"micro\")\n",
    "        f1 = f1 + f1s\n",
    "    f1_s = f1/len(train_dataloader)\n",
    "    \n",
    "    avg_train_loss = total_loss / len(train_dataloader)\n",
    "    return avg_train_loss, f1_s\n",
    "\n",
    "def evaluate(model, iterator, criterion):\n",
    "    model.eval()\n",
    "    eval_loss, eval_accuracy, f1 = 0, 0, 0\n",
    "    nb_eval_steps, nb_eval_examples = 0, 0\n",
    "    predictions , true_labels = [], []\n",
    "    preds = []\n",
    "    labels = []\n",
    "    for batch in valid_dataloader:\n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "        b_input_ids, b_input_mask, b_labels = batch\n",
    "        with torch.no_grad():\n",
    "            outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)\n",
    "        logits = outputs[1].detach().cpu().numpy()\n",
    "        label_ids = b_labels.to('cpu').numpy()\n",
    "        eval_loss += outputs[0].mean().item()\n",
    "        predictions.extend([list(p) for p in np.argmax(logits, axis=2)])\n",
    "        true_labels.extend(label_ids)\n",
    "        \n",
    "        pred_tags = [tag_values[p_i] for p, l in zip(predictions, true_labels)\n",
    "                                 for p_i, l_i in zip(p, l) if tag_values[l_i] != \"PAD\"]\n",
    "        valid_tags = [tag_values[l_i] for l in true_labels\n",
    "                                  for l_i in l if tag_values[l_i] != \"PAD\"]\n",
    "    \n",
    "        f1s = f1_score(pred_tags, valid_tags, average=\"micro\")\n",
    "        f1 = f1 + f1s\n",
    "        preds.append(pred_tags)\n",
    "        labels.append(valid_tags)\n",
    "        \n",
    "    f1_s = f1 /  len(valid_dataloader)\n",
    "    eval_loss = eval_loss / len(valid_dataloader)\n",
    "    \n",
    "    return eval_loss, f1_s, preds, labels\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs\n",
    "\n",
    "N_EPOCHS = 5\n",
    "\n",
    "t_loss = []\n",
    "t_f1 = []\n",
    "v_loss = []\n",
    "v_f1 = []\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "\n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss, train_f1 = train(model, train_dataloader, optimizer, criterion)\n",
    "    t_loss.append(train_loss)\n",
    "    t_f1.append(train_f1) \n",
    "    \n",
    "    valid_loss, valid_f1,_,_ = evaluate(model, valid_dataloader, criterion)\n",
    "    v_loss.append(valid_loss)\n",
    "    v_f1.append(valid_f1)  \n",
    "    \n",
    "    end_time = time.time()\n",
    "\n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), 'tut5-model_ner.pt')\n",
    "    if epoch%1 == 0:\n",
    "        print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "        print(f'\\tTrain Loss: {train_loss:.3f} | Train F1 score: {train_f1*100:.2f}%')\n",
    "        print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. F1 score: {valid_f1*100:.2f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt \n",
    "import numpy as np \n",
    "sns.set()\n",
    "\n",
    "x = np.linspace(0, N_EPOCHS,N_EPOCHS)\n",
    "\n",
    "plt.plot(x,t_loss)\n",
    "plt.plot(x,v_loss)\n",
    "plt.title(\"Loss\")\n",
    "plt.legend([\"Train loss\", \"Valid loss\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, N_EPOCHS,N_EPOCHS)\n",
    "\n",
    "plt.plot(x,t_f1)\n",
    "plt.plot(x,v_f1)\n",
    "plt.title(\"F1 score\")\n",
    "plt.legend([\"Train F1\", \"Valid F1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load('tut5-model_ner.pt'))\n",
    "\n",
    "test_loss, test_f1, preds, labels = evaluate(model, test_dataloader, criterion)\n",
    "print(f'Test Loss: {test_loss:.3f} |  Test F1 score: {test_f1*100:.2f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict =  [item for sublist in preds for item in sublist]\n",
    "true =  [item for sublist in labels for item in sublist]\n",
    "\n",
    "print(classification_report(\n",
    "                y_true=true,\n",
    "                y_pred=predict\n",
    "            ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ent_manquantes = []\n",
    "for i in tag2idx:\n",
    "    if i not in true:\n",
    "        ent_manquantes.append(i)\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "confusion = confusion_matrix(true, predict)\n",
    "confusion = confusion_matrix(true, predict)\n",
    "l = [i for i in tag2idx if i not in ent_manquantes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "confusion_df =pd.DataFrame(confusion)\n",
    "\n",
    "confusion_df.columns=[i for i in TAG.vocab.itos][1:]\n",
    "s = pd.Series([i for i in TAG.vocab.itos][1:])\n",
    "confusion_df = confusion_df.set_index([s])\n",
    "\n",
    "confusion_df['LOC'] = confusion_df['B-LOC'] + confusion_df['I-LOC']\n",
    "confusion_df['PER'] = confusion_df['B-PER'] + confusion_df['I-PER']\n",
    "confusion_df['ORG'] = confusion_df['B-ORG'] + confusion_df['I-ORG']\n",
    "confusion_df['MISC'] = confusion_df['B-MISC'] + confusion_df['I-MISC']\n",
    "\n",
    "\n",
    "confusion_df = confusion_df.drop(columns=[ i for i in TAG.vocab.itos if i != 'O' and i != '<pad>'])\n",
    "\n",
    "confusion_df.loc['LOC'] = confusion_df.loc['B-LOC'] + confusion_df.loc['I-LOC']\n",
    "confusion_df.loc['PER'] = confusion_df.loc['B-PER'] + confusion_df.loc['I-PER']\n",
    "confusion_df.loc['ORG'] = confusion_df.loc['B-ORG'] + confusion_df.loc['I-ORG']\n",
    "confusion_df.loc['MISC'] = confusion_df.loc['B-MISC'] + confusion_df.loc['I-MISC']\n",
    "\n",
    "confusion_df = confusion_df.drop([i for i in TAG.vocab.itos if i != 'O' and i != '<pad>'])\n",
    "\n",
    "confusion_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm = confusion_df.to_numpy()\n",
    "\n",
    "TP = np.diag(cm)\n",
    "FP = np.sum(cm, axis=0) - TP\n",
    "FN = np.sum(cm, axis=1) - TP\n",
    "\n",
    "num_classes = 4\n",
    "TN = []\n",
    "for i in range(num_classes):\n",
    "    temp = np.delete(cm, i, 0)    \n",
    "    temp = np.delete(temp, i, 1) \n",
    "    TN.append(sum(sum(temp)))\n",
    "\n",
    "precision = [0]*len(TP)\n",
    "recall = [0]*len(TP)\n",
    "f1 = [0]*len(TP)\n",
    "for i in range(len(TP)):\n",
    "    if TP[i]+FP[i] != 0:\n",
    "        precision[i] = TP[i]/(TP[i]+FP[i])\n",
    "    else:\n",
    "        precision[i] = 0\n",
    "\n",
    "    if TP[i]+FN[i] != 0:\n",
    "        recall[i] = TP[i]/(TP[i]+FN[i])\n",
    "    else:\n",
    "        recall[i] = 0\n",
    "\n",
    "    if precision[i]+recall[i] != 0:\n",
    "        f1[i] = (2*precision[i]*recall[i])/(precision[i]+recall[i])\n",
    "    else:\n",
    "        f1[i] = 0\n",
    "\n",
    "\n",
    "prf_df =pd.DataFrame()\n",
    "prf_df['Precision'] = precision\n",
    "prf_df['Recall'] = recall\n",
    "prf_df['F1-score'] = f1\n",
    "\n",
    "s = pd.Series([i for i in confusion_df.index])\n",
    "prf_df = prf_df.set_index([s])\n",
    "\n",
    "prf_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inférence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sentence = 'The will deliver a speech about the conflict in North Korea tomorrow in New York with my friend Mary Kate.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_sentence = tokenizer.encode(test_sentence)\n",
    "input_ids = torch.tensor([tokenized_sentence])#.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    output = model(input_ids)\n",
    "label_indices = np.argmax(output[0].numpy(), axis=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = tokenizer.convert_ids_to_tokens(input_ids.to('cpu').numpy()[0])\n",
    "new_tokens, new_labels = [], []\n",
    "for token, label_idx in zip(tokens, label_indices[0]):\n",
    "    if token.startswith(\"##\"):\n",
    "        new_tokens[-1] = new_tokens[-1] + token[2:]\n",
    "    else:\n",
    "        new_labels.append(tag_values[label_idx])\n",
    "        new_tokens.append(token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for token, label in zip(new_tokens, new_labels):\n",
    "    print(\"{}\\t{}\".format(label, token))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}