01_train_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d378ba87-e270-4606-a977-25c03afba013",
   "metadata": {},
   "source": [
    "# A simple NER\n",
    "For this project we preprocess the conllu dataset, train a custom model for token classification (roberta embeddings + lstm layer) for which we try to find the optimal hyperparameters and finally we train and evaluate the final model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "595ddc33-5273-48ed-ba91-c0d97b2c72e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#===============================================================================\n",
    "#\n",
    "#           FILE: 01_train_models.ipynb\n",
    "#         AUTHOR: Bianca Ciobanica\n",
    "#          EMAIL: bianca.ciobanica@student.uclouvain.be\n",
    "#\n",
    "#           BUGS: \n",
    "#        VERSION: 3.11.4\n",
    "#        CREATED: 20-05-2024 \n",
    "#\n",
    "#===============================================================================\n",
    "#    DESCRIPTION:  \n",
    "#    \n",
    "#   DEPENDENCIES:  torch, transformers, accelerate, evaluate, datasets, tqdm\n",
    "#                  polars, ray, numpy\n",
    "#\n",
    "#          USAGE: jupyter notebook 01_train_models.ipynb \n",
    "#==============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f50c6b4-07cb-4efb-8e7d-9743dfe8fd20",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "edcccbc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import cuda\n",
    "from torch.nn import CrossEntropyLoss, LSTM, Module, Linear, Dropout, LayerNorm\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import AdamW\n",
    "\n",
    "from transformers.utils import logging\n",
    "from transformers import RobertaTokenizerFast, RobertaModel, RobertaConfig, RobertaForTokenClassification, DataCollatorForTokenClassification, get_scheduler\n",
    "from transformers.modeling_outputs import TokenClassifierOutput\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from functools import partial\n",
    "import os\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "from ray import tune\n",
    "from ray import train\n",
    "from ray.train import Checkpoint, get_checkpoint\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "import ray.cloudpickle as pickle\n",
    "\n",
    "import numpy as np\n",
    "from accelerate import Accelerator\n",
    "from datasets import Dataset\n",
    "import evaluate\n",
    "import polars as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a99713f7-ef86-469f-a66f-1d7797acb8e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cuda is available :  True\n"
     ]
    }
   ],
   "source": [
    "logging.set_verbosity_info()\n",
    "\n",
    "print(\"Cuda is available : \", torch.cuda.is_available())\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd09ebda-a6e8-4443-811d-ebe0a395c54f",
   "metadata": {},
   "source": [
    "### File paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "db2d24b9-c5be-48b7-810a-6a80c7898ec5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DATASET\n",
    "train_file = \"./corpus/dataset/train.conllu\"\n",
    "dev_file = \"./corpus/dataset/val.conllu\"\n",
    "test_file = \"./corpus/dataset/test.conllu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8f6bef4c-7390-4eba-95dc-2bec301eba56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MODEL output dir\n",
    "os.makedirs('./models', exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "892ff865-e333-4009-b2f0-c86aa7e459a7",
   "metadata": {},
   "source": [
    "## Data preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b732fcdf-d743-496d-a2f1-961f67765967",
   "metadata": {},
   "source": [
    "### Corpus Loading\n",
    "<b>Reference:</b> [Introduction to the CoNLL-2003 Shared Task: Language-Independent Named Entity Recognition](https://aclanthology.org/W03-0419) (Tjong Kim Sang & De Meulder, CoNLL 2003)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77c7c1a6-47a4-4618-b45b-812e1464c6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_conllu_file(input_file, encoding):\n",
    "    data = None\n",
    "\n",
    "    X = []\n",
    "    Y = []\n",
    "    tagset = {}\n",
    "\n",
    "    with open(input_file, \"r\", encoding=\"utf-8\") as fin:\n",
    "        current_sentence_tokens = []\n",
    "        current_sentence_labels = []\n",
    "        \n",
    "        line = fin.readline()\n",
    "        while line:\n",
    "            line = line.rstrip()\n",
    "            if not line:\n",
    "                X.append(current_sentence_tokens)\n",
    "                Y.append(current_sentence_labels)\n",
    "                current_sentence_tokens = []\n",
    "                current_sentence_labels = []\n",
    "\n",
    "                line = fin.readline()\n",
    "                continue\n",
    "\n",
    "            data = line.split()\n",
    "            \n",
    "            token = data[1]\n",
    "            label = data[2]\n",
    "\n",
    "            if label not in tagset:\n",
    "                tagset[label] = len(tagset)\n",
    "            current_sentence_tokens.append(token)\n",
    "            current_sentence_labels.append(label)\n",
    "\n",
    "            line = fin.readline()\n",
    "\n",
    "    return X, Y, tagset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4ce474-37ac-408c-82fb-329b6d9a2513",
   "metadata": {},
   "source": [
    "Here we split train and test and also remap id to labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "dd9045c0-a55b-42bd-a01c-bf23b5091d26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tagset  |  label id\n",
      "     O      0\n",
      " B-ORG      1\n",
      " I-ORG      2\n",
      "B-MISC      3\n",
      "I-MISC      4\n",
      " B-PER      5\n",
      " I-PER      6\n",
      " B-LOC      7\n",
      " I-LOC      8\n"
     ]
    }
   ],
   "source": [
    "train_x, train_y, tagset = load_conllu_file(train_file, \"utf-8\")\n",
    "test_x, test_y, _ = load_conllu_file(test_file, \"utf-8\")\n",
    "dev_x, dev_y, _ = load_conllu_file(dev_file, \"utf-8\")\n",
    "\n",
    "tagset_label2id = {'O': 0,'B-ORG': 1, 'I-ORG': 2, 'B-MISC': 3, 'I-MISC': 4, 'B-PER': 5, 'I-PER': 6, 'B-LOC': 7, 'I-LOC': 8}\n",
    "tagset_id2label = dict(zip(tagset_label2id.values(), tagset_label2id.keys()))\n",
    "\n",
    "# check if no errros\n",
    "# print(tagset.keys() == tagset_custom.keys())\n",
    "\n",
    "sys.stdout.write(\"Tagset  |  label id\\n\")\n",
    "for label, label_id in tagset_label2id.items():\n",
    "    sys.stdout.write(\"%6s      %i\\n\" % (label ,label_id))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "455e9eef-8f75-474d-91df-ede9cca2f6e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total size :  22134\n",
      "train size :  0.68\n",
      "test size  :  0.17\n",
      "dev size   :  0.16\n"
     ]
    }
   ],
   "source": [
    "total_len = len(dev_x) + len(train_x) + len(test_x)\n",
    "train_len = len(train_x) / total_len\n",
    "test_len = len(test_x) / total_len\n",
    "dev_len = len(dev_x) / total_len\n",
    "\n",
    "print(\"total size : \", total_len)\n",
    "print(\"train size : \", round(train_len, 2)) # 70 %\n",
    "print(\"test size  : \", round(test_len, 2)) # 15 %\n",
    "print(\"dev size   : \", round(dev_len, 2)) # 16 %"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bd31b8b2-6fbc-4d40-927a-20c7d2070b01",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 'O', 1: 'B-ORG', 2: 'I-ORG', 3: 'B-MISC', 4: 'I-MISC', 5: 'B-PER', 6: 'I-PER', 7: 'B-LOC', 8: 'I-LOC'}\n"
     ]
    }
   ],
   "source": [
    "print(tagset_id2label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e422945c-bf62-4fc7-a6ae-699cbecfc627",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# example\n",
    "#index = 0\n",
    "#sentence = train_x[index]\n",
    "#tags = train_y[index]\n",
    "#sys.stdout.write(\"\\nTokens: [%d]\\n  %s\\n\\n\" % (len(sentence), str(sentence)))\n",
    "#sys.stdout.write(\"Labels: [%d]\\n  %s\\n\" % (len(tags), str(tags)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399de707-a25f-46d9-9d00-1441564d220f",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Dataset creation\n",
    "We create a dataset object with train and test split which will be used for our training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "560aff0f-e035-4737-a001-284eea45473a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_set = Dataset.from_dict({\"tokens\": train_x, \"labels\" : train_y})\n",
    "test_set = Dataset.from_dict({\"tokens\": test_x, \"labels\" : test_y})\n",
    "dev_set = Dataset.from_dict({\"tokens\": dev_x, \"labels\" : dev_y})\n",
    "\n",
    "#dataset = DatasetDict({\"train\" : train_set, \"test\" : test_set})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fde52800-6054-43de-b3cf-ea380960df9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Longest sequence in train :  113\n",
      "Longest sequence in test  :  124\n"
     ]
    }
   ],
   "source": [
    "train_max_len = len(max(train_x, key = len))\n",
    "test_max_len = len(max(test_x, key = len))\n",
    "print(\"Longest sequence in train : \", train_max_len) \n",
    "print(\"Longest sequence in test  : \", test_max_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16b02d45-7e54-4b58-9128-2adf14181e07",
   "metadata": {},
   "source": [
    "### Tokenizer initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "542a21a0-27f0-447e-936b-c4843eee75df",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loading file vocab.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/vocab.json\n",
      "loading file merges.txt from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/merges.txt\n",
      "loading file tokenizer.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/tokenizer.json\n",
      "loading file added_tokens.json from cache at None\n",
      "loading file special_tokens_map.json from cache at None\n",
      "loading file tokenizer_config.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/tokenizer_config.json\n",
      "loading configuration file config.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/config.json\n",
      "Model config RobertaConfig {\n",
      "  \"_name_or_path\": \"roberta-base\",\n",
      "  \"architectures\": [\n",
      "    \"RobertaForMaskedLM\"\n",
      "  ],\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"bos_token_id\": 0,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"eos_token_id\": 2,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-05,\n",
      "  \"max_position_embeddings\": 514,\n",
      "  \"model_type\": \"roberta\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 1,\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.41.1\",\n",
      "  \"type_vocab_size\": 1,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50265\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base', add_prefix_space=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d5d0ff2e-d9f6-45ad-b55d-edde8a87ab78",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_func(row):\n",
    "    return tokenizer(row,\n",
    "                    is_split_into_words=True,\n",
    "                     add_special_tokens=False\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a34dc513-4bcc-4d95-8951-b2f8d15a8b16",
   "metadata": {},
   "source": [
    "### Align tokens with new labels\n",
    "We use `RobertaTokenizerFast` for tokenizing our inputs then use the method `words_ids()` which gives the position the original token. Thus, if the original token is separated into pieces after using the RobertaTokenizer, its index will repeat. We create a for loop that if next word is current word and current word has a **B-XXX** label (meaning the original token is split in pieces), we add **I-XXX** label. Given that we mapped the ids for the labels in pairs (uneven id == **B-XXX** and even id == **I-XXX**) we are able to create the new desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "25ec8279-2838-4b6b-8dd0-67cab0e8e678",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens    : Only      France    and       Britain   backed    Fischler  's        proposal  .         \n",
      "Labels    :O         B-LOC     O         B-LOC     O         B-PER     O         O         O         \n",
      "Label ids :1         5         1         5         1         3         1         1         1         \n",
      "\n",
      "Original number of tokens: 9\n",
      "Number of input_ids: 12 \n",
      "\n",
      "Input ID: 4041       Original Token: Only            Tokenized Token:  Only\n",
      "Input ID: 1470       Original Token: France          Tokenized Token:  France\n",
      "Input ID: 8          Original Token: and             Tokenized Token:  and\n",
      "Input ID: 1444       Original Token: Britain         Tokenized Token:  Britain\n",
      "Input ID: 4094       Original Token: backed          Tokenized Token:  backed\n",
      "Input ID: 274        Original Token: Fischler        Tokenized Token:  F\n",
      "Input ID: 13239      Original Token: 's              Tokenized Token: isch\n",
      "Input ID: 1371       Original Token: proposal        Tokenized Token: ler\n",
      "Input ID: 128        Original Token: .               Tokenized Token:  '\n",
      "Input ID: 29         Original Token: None            Tokenized Token: s\n",
      "Input ID: 2570       Original Token: None            Tokenized Token:  proposal\n",
      "Input ID: 479        Original Token: None            Tokenized Token: .\n"
     ]
    }
   ],
   "source": [
    "def print_example(words, labels):\n",
    "    line1 = \"Tokens    : \"\n",
    "    line2 = \"Labels    :\"\n",
    "    line3 = \"Label ids :\"\n",
    "    for word, label in zip(words, labels):\n",
    "        #full_label = label_names[label]\n",
    "        max_length = max(len(words), len(labels))\n",
    "        line1 += word + \" \" * (max_length - len(word) + 1)\n",
    "        line2 += label + \" \" * (max_length - len(label) + 1)\n",
    "        line3 += str(tagset[label]) + \" \" * (max_length - len(str(tagset[label])) + 1)\n",
    "    print(line1)\n",
    "    print(line2)\n",
    "    print(line3)\n",
    "def helper_visualizing():\n",
    "    index = 12  # example\n",
    "    \n",
    "    print_example(train_x[index], train_y[index])\n",
    "    \n",
    "    inputs = tokenizer(train_x[index], is_split_into_words=True, add_special_tokens=False)\n",
    "    print(\"\\nOriginal number of tokens:\", len(train_x[index]))\n",
    "    print(\"Number of input_ids:\", len(inputs[\"input_ids\"]), \"\\n\")\n",
    "    \n",
    "    \n",
    "    \n",
    "    for i, token_id in enumerate(inputs['input_ids']):\n",
    "        if i < len(train_x[index]):\n",
    "            token = tokenizer.decode([token_id])\n",
    "            original_token = train_x[index][i]\n",
    "            print(f\"Input ID: {token_id:<10} Original Token: {original_token:<15} Tokenized Token: {token}\")\n",
    "        else:\n",
    "            print(f\"Input ID: {token_id:<10} Original Token: None            Tokenized Token: {tokenizer.decode([token_id])}\")\n",
    "    \n",
    "helper_visualizing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "83c9ab89-cbc3-4bfa-b051-760cf3290b36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def align_labels_with_tokens(labels, word_ids):\n",
    "    # word id is the original position in the input sequence\n",
    "    new_labels = []\n",
    "    current_word = None\n",
    "    \n",
    "    for word_id in word_ids:\n",
    "        # if not same word\n",
    "        if word_id != current_word:    \n",
    "            current_word = word_id\n",
    "            label = -100 if word_id is None else labels[word_id]\n",
    "            new_labels.append(label)\n",
    "            \n",
    "        elif word_id is None:\n",
    "            new_labels.append(-100)\n",
    "            \n",
    "       # same word as previous (piece of token)\n",
    "        else:\n",
    "            label = labels[word_id]\n",
    "            # if the label is B-XXX we change it to I-XXX\n",
    "            if label % 2 == 1:\n",
    "                label += 1\n",
    "            new_labels.append(label)\n",
    "\n",
    "    return new_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b5a3bd34-ceb3-4612-9401-64939b9f5510",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Check to see if new labels are correct\n",
      "\n",
      "Original tokens: ['BRUSSELS', '1996-08-22']\n",
      "New Tokens :  ['BR', 'USS', 'ELS', '1996', '-', '08', '-', '22']\n",
      "position of new tokens in original tokens:  [0, 0, 0, 1, 1, 1, 1, 1]\n",
      "Original labels: ['B-LOC', 'O']\n",
      "Original label ids :  [7, 0]\n",
      "\n",
      "\n",
      "\n",
      "new labels : \n",
      "Tokens    : BR       USS      ELS      1996     -        08       -        22       \n",
      "Labels    :B-LOC    I-LOC    I-LOC    O        O        O        O        O        \n",
      "Label ids :5        8        8        1        1        1        1        1        \n"
     ]
    }
   ],
   "source": [
    "# test alignment function for new labels\n",
    "def test_align():\n",
    "    print(\"Check to see if new labels are correct\\n\")\n",
    "    index = 2\n",
    "\n",
    "    inputs = tokenizer(train_x[index], is_split_into_words=True, add_special_tokens=False)\n",
    "    labels = list(map(lambda label : tagset_label2id[label], train_y[index]))\n",
    "    new_tokens = \" \".join(tokenizer.convert_ids_to_tokens(inputs['input_ids'])).replace(\"Ġ\", \" \").split()\n",
    "    \n",
    "    print(\"Original tokens:\", train_x[index])\n",
    "    print(\"New Tokens : \",  new_tokens)\n",
    "    word_ids = inputs.word_ids()\n",
    "    print(\"position of new tokens in original tokens: \", word_ids)\n",
    "    print(\"Original labels:\", train_y[index])\n",
    "    print(\"Original label ids : \", labels)\n",
    "    \n",
    "    print()\n",
    "    print()\n",
    "    \n",
    "    new_labels_id = align_labels_with_tokens(labels, word_ids)\n",
    "    new_labels = [tagset_id2label[label_id] for label_id in new_labels_id]\n",
    "    print()\n",
    "    print(\"new labels : \")\n",
    "    print_example(new_tokens, new_labels)\n",
    "    \n",
    "test_align()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "04b4552f-4afe-48c5-b4f1-d274c0caa54c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_and_align_labels(dataset):\n",
    "    \n",
    "    all_tokens = dataset.pop(\"tokens\")\n",
    "    all_labels = dataset.pop(\"labels\")\n",
    "\n",
    "    tokenized_dataset = tokenize_func(all_tokens)\n",
    "    max_len = len(max(tokenized_dataset['input_ids'], key=len))\n",
    "\n",
    "    new_labels = []\n",
    "    all_labels_id = []\n",
    "    \n",
    "    # go through each labels\n",
    "    for i, labels in enumerate(all_labels):\n",
    "        # conver labels to their id\n",
    "        labels_id = [tagset_label2id[label] for label in labels]\n",
    "        token_ids = tokenized_dataset.word_ids(i)\n",
    "\n",
    "        new_labels.append(align_labels_with_tokens(labels_id, token_ids))\n",
    "        all_labels_id.append(labels_id)\n",
    "\n",
    "    #tokenized_dataset[\"new_labels\"] = new_labels\n",
    "    tokenized_dataset[\"labels\"] = new_labels\n",
    "    \n",
    "    \n",
    "    return tokenized_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d0a1ff-f57b-478c-a5e6-8edea4749286",
   "metadata": {},
   "source": [
    "### Final train and test set structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4d68cd5f-a6ec-47b3-a3cf-bf2310210fa6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "709be44a93c54f0a96f86f137b02f106",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/14986 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "974b1bde23fd40c19f6abeb0a1e3a77b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/3465 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "760489c173904b5f842622402cc9b1d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/3683 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# map train and test with tokenizing func and add new labels column\n",
    "train_set = train_set.map(\n",
    "    tokenize_and_align_labels,\n",
    "    batched=True,\n",
    ")\n",
    "dev_set = dev_set.map(\n",
    "    tokenize_and_align_labels,\n",
    "    batched=True,\n",
    ")\n",
    "test_set = test_set.map(\n",
    "    tokenize_and_align_labels,\n",
    "    batched=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f97964a3-8dec-4fa9-b5f4-ba3a4496b2df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21ace2ce08cd4e428924bbdff0085560",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/14986 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a411ffcc2e041469555108d49857f94",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/3465 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6dafb004680945da9423c1cfd44a5500",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/3683 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# save to disk for tuning with ray_gridsearch.py\n",
    "train_set.save_to_disk('corpus/dataset/train')\n",
    "dev_set.save_to_disk('corpus/dataset/dev')\n",
    "test_set.save_to_disk('corpus/dataset/test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "03fb1ebd-a37d-41a8-97d4-f03b866eef74",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_train_max_size = len(max(train_set['input_ids'], key=len))\n",
    "input_test_max_size = len(max(test_set['input_ids'], key=len))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7868ac44-9f51-4311-8aa2-232dc78968a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(train_set.num_rows)\n",
    "#print(test_set.num_rows / (train_set.num_rows + test_set.num_rows))\n",
    "#\n",
    "#print(train_set.features)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b41f370-5792-42ee-9e93-2f6a087d9d29",
   "metadata": {},
   "source": [
    "`train_set` has 14986 elements (80 %)  and `test_set` has 3683 elements (20%). Both have `input_ids`, `attention_mask`, `new_labels` , `new_labels_id` as features that will serve for our model.\n",
    "\n",
    "The tagset contains 9 labels."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0d964ae-ed5d-4a2b-beb3-4b31bb420d06",
   "metadata": {},
   "source": [
    "## Compute metrics for evaluation\n",
    "We use [seqeval][1] for sequence labeling evaluation. It takes the lists of labels as strings, not integers, so we will need to fully decode the predictions and labels before passing them to the metric.\n",
    "\n",
    "[1]: https://github.com/chakki-works/seqeval?tab=readme-ov-file\n",
    "\n",
    "We also create a post processing function in order to convert ids to their original labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "248ccf2a-05b2-4ef7-85f7-abc2b71df34f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_ids_2_labels(model_predictions, true_labels):\n",
    "\n",
    "    # go through each batch and ignore when id = -100 (padding token)\n",
    "    true_labels = [\n",
    "        [tagset_id2label[label_id] \n",
    "        for label_id in batch if label_id != -100 ] \n",
    "        for batch in true_labels ]\n",
    "    \n",
    "    model_predictions = [\n",
    "        [tagset_id2label[pred] \n",
    "         for (pred, lab) in zip(prediction, label) if lab != -100 ]\n",
    "        for prediction, label in zip(model_predictions, true_labels) ]\n",
    "    \n",
    "    return model_predictions, true_labels # both are same length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "759503e4-7086-4171-b36e-1c314f186e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(model_predictions):\n",
    "    \"\"\" input : ŷ,  logits (log prob of model's prediction)\n",
    "                y, true label of current prediction\n",
    "\n",
    "                converts prediction ids to labels in string for computing scores\n",
    "                \n",
    "        returns dict with overall precision, recall, f1, accuracy\n",
    "    \"\"\"\n",
    "    logits, true_labels = model_predictions\n",
    "\n",
    "    predictions = np.argmax(logits, axis=1) # take the predicted tag with highest score\n",
    "\n",
    "    return metric.compute(predictions=predictions, references=true_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0490cd00-6d1b-4d8a-86b2-f21d29c80f1c",
   "metadata": {},
   "source": [
    "## Model initialization\n",
    "\n",
    "First, we setup the [model][2]. We tried the base version (125M params) of `RobertaForTokenClassification`, a RoBERTa[1] Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks.\n",
    "\n",
    "The model returns `loss`, `scores`, `hidden_states`, `attentions`.\n",
    "\n",
    "[1]: Liu, Yinhan, et al. (2019). RoBERTa: A Robustly Optimized BERT Pretraining Approach. \n",
    "\n",
    "[2]: https://huggingface.co/FacebookAI/roberta-large\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bf627d6f-e242-4697-9e67-56f6df3d4b93",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cbianc/Documents/M2/Q2/CL_seminar/tagging/venv/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "loading configuration file config.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/config.json\n",
      "Model config RobertaConfig {\n",
      "  \"architectures\": [\n",
      "    \"RobertaForMaskedLM\"\n",
      "  ],\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"bos_token_id\": 0,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"eos_token_id\": 2,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-05,\n",
      "  \"max_position_embeddings\": 514,\n",
      "  \"model_type\": \"roberta\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 1,\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"return_dict\": false,\n",
      "  \"transformers_version\": \"4.41.1\",\n",
      "  \"type_vocab_size\": 1,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50265\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "config = RobertaConfig.from_pretrained(\"roberta-base\", return_dict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "23a73999-c616-4f92-b377-eb9ef1f957b2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loading configuration file config.json from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/config.json\n",
      "Model config RobertaConfig {\n",
      "  \"architectures\": [\n",
      "    \"RobertaForMaskedLM\"\n",
      "  ],\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"bos_token_id\": 0,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"eos_token_id\": 2,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"id2label\": {\n",
      "    \"0\": \"O\",\n",
      "    \"1\": \"B-ORG\",\n",
      "    \"2\": \"I-ORG\",\n",
      "    \"3\": \"B-MISC\",\n",
      "    \"4\": \"I-MISC\",\n",
      "    \"5\": \"B-PER\",\n",
      "    \"6\": \"I-PER\",\n",
      "    \"7\": \"B-LOC\",\n",
      "    \"8\": \"I-LOC\"\n",
      "  },\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"label2id\": {\n",
      "    \"B-LOC\": 7,\n",
      "    \"B-MISC\": 3,\n",
      "    \"B-ORG\": 1,\n",
      "    \"B-PER\": 5,\n",
      "    \"I-LOC\": 8,\n",
      "    \"I-MISC\": 4,\n",
      "    \"I-ORG\": 2,\n",
      "    \"I-PER\": 6,\n",
      "    \"O\": 0\n",
      "  },\n",
      "  \"layer_norm_eps\": 1e-05,\n",
      "  \"max_position_embeddings\": 514,\n",
      "  \"model_type\": \"roberta\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 1,\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.41.1\",\n",
      "  \"type_vocab_size\": 1,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50265\n",
      "}\n",
      "\n",
      "loading weights file model.safetensors from cache at /home/cbianc/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b/model.safetensors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config of the RobertaForTokenClassification model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at roberta-base were not used when initializing RobertaForTokenClassification: ['lm_head.bias', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.layer_norm.bias', 'lm_head.layer_norm.weight']\n",
      "- This IS expected if you are initializing RobertaForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of RobertaForTokenClassification were not initialized from the model checkpoint at roberta-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RobertaForTokenClassification(\n",
       "  (roberta): RobertaModel(\n",
       "    (embeddings): RobertaEmbeddings(\n",
       "      (word_embeddings): Embedding(50265, 768, padding_idx=1)\n",
       "      (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "      (token_type_embeddings): Embedding(1, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): RobertaEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x RobertaLayer(\n",
       "          (attention): RobertaAttention(\n",
       "            (self): RobertaSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): RobertaSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): RobertaIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): RobertaOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=9, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Config of the RobertaForTokenClassification model\")\n",
    "robertaClass_model = RobertaForTokenClassification.from_pretrained('roberta-base', num_labels=len(tagset), label2id=tagset_label2id, id2label=tagset_id2label)\n",
    "robertaClass_model.to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "676d8fd4-e04d-4550-976b-857dd9d6f05a",
   "metadata": {},
   "source": [
    "### LSTM layer\n",
    "We then modified the structure in order to add an [LSTM][1] layer with 2 hidden layers (stacked LSTM).\n",
    "\n",
    "simple lstm unit\n",
    "\n",
    "![lstm unit](./pictures/lstm_unit.png)\n",
    "\n",
    "forget gate, input gate, output gate\n",
    "\n",
    "![lstm_structure](./pictures/lstm_structure.png)\n",
    "\n",
    "[1]: http://web.stanford.edu/~jurafsky/slp3/9.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a8e3188-87cc-4c2f-a15b-6a42c2f3b1db",
   "metadata": {},
   "source": [
    "#### Parameters\n",
    "https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html\n",
    "\n",
    "- **input_size**: The number of expected features in the input x.\n",
    "- **hidden_size**: The number of features in the hidden state h.\n",
    "- **num_layers**: Number of recurrent layers. E.g., setting `num_layers=2` would mean stacking two LSTMs together to form a stacked LSTM, with the second LSTM taking in outputs of the first LSTM and computing the final results. Default: 1.\n",
    "- **bias**: If `False`, then the layer does not use bias weights `b_ih` and `b_hh`. Default: `True`.\n",
    "- **batch_first**: If `True`, then the input and output tensors are provided as `(batch, seq, feature)` instead of `(seq, batch, feature)`. Note that this does not apply to hidden or cell states. See the Inputs/Outputs sections below for details. Default: `False`.\n",
    "- **dropout**: If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to `dropout`. Default: 0.\n",
    "- **bidirectional**: If `True`, becomes a bidirectional LSTM. Default: `False`.\n",
    "- **proj_size**: If > 0, will use LSTM with projections of corresponding size. Default: 0.\n",
    "\n",
    "#### Input\n",
    "\n",
    "**Inputs**: `input`, `(h_0, c_0)`\n",
    "\n",
    "- `input`: tensor of shape `(L, Hin)` for unbatched input, `(L, N, Hin)` when `batch_first=False` or `(N, L, Hin)` when `batch_first=True` containing the features of the input sequence. The input can also be a packed variable length sequence. See `torch.nn.utils.rnn.pack_padded_sequence()` or `torch.nn.utils.rnn.pack_sequence()` for details.\n",
    "- `h_0`: tensor of shape `(D * num_layers, Hout)` for unbatched input or `(D * num_layers, N, Hout)` containing the initial hidden state for each element in the input sequence. Defaults to zeros if `(h_0, c_0)` is not provided.\n",
    "- `c_0`: tensor of shape `(D * num_layers, Hcell)` for unbatched input or `(D * num_layers, N, Hcell)` containing the initial cell state for each element in the input sequence. Defaults to zeros if `(h_0, c_0)` is not provided.\n",
    "\n",
    "**Note**: <span style=\"color:red;\">Here our input is the last hidden layer (embeddings) of the `RobertaModel`.</span> We get the outputs of the `roberta` model ([BaseModelOutputWithPoolingAndCrossAttentions][1]) and the first element is the last hidden layer. This is what we need to pass to the LSTM as input.\n",
    "\n",
    "[1]: https://github.com/huggingface/transformers/blob/v4.41.2/src/transformers/models/roberta/modeling_roberta.py#L679\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8ce9c2e-1734-4d40-8f35-6f0c0b2d729f",
   "metadata": {},
   "source": [
    "#### Stucture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3e58846c-5387-4e71-821b-6f683a1c26ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRobertaTokenClassWithLSTM(Module):\n",
    "    def __init__(self, config, num_labels, hidden_dim, n_hidden_layers, lstm_dropout):\n",
    "        super().__init__()\n",
    "        self.num_labels = num_labels # number of classes to predict\n",
    "\n",
    "        # layer 1\n",
    "        self.roberta = RobertaModel(config, add_pooling_layer=False)\n",
    "\n",
    "        # layer 2\n",
    "        self.dropout = Dropout(0.3) # same shape as input\n",
    "        \n",
    "        # layer 3\n",
    "        self.lstm = LSTM(config.hidden_size, hidden_size=hidden_dim, num_layers=n_hidden_layers, batch_first=True) # takes last hidden layer output from roberta\n",
    "\n",
    "        self.lnorm = LayerNorm(hidden_dim)\n",
    "        \n",
    "        self.lstm_dropout = Dropout(lstm_dropout) # same shape as input\n",
    "        \n",
    "        self.fc = Linear(hidden_dim, num_labels)\n",
    "\n",
    "        # Initialize weights and apply final processing\n",
    "        self.roberta.post_init()\n",
    "       \n",
    "    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,\n",
    "                position_ids=None, head_mask=None, inputs_embeds=None, output_attentions=None,\n",
    "                output_hidden_states=None, labels=None, return_dict=True):\n",
    "        \n",
    "        outputs = self.roberta(\n",
    "            input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "            position_ids=position_ids,\n",
    "            head_mask=head_mask,\n",
    "            inputs_embeds=inputs_embeds,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "\n",
    "        sequence_output = outputs[0] # output of last hidden layer\n",
    "        sequence_output = self.dropout(sequence_output)\n",
    "        sequence_output,_ = self.lstm(sequence_output) # LSTM Outputs: output, (h_n, c_n) -> we only need output for token classification\n",
    "        sequence_output = self.lstm_dropout(sequence_output)\n",
    "        sequence_output = self.lnorm(sequence_output)\n",
    "\n",
    "        logits = self.fc(sequence_output)\n",
    "      \n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            labels = labels.to(logits.device)\n",
    "            loss_fct = CrossEntropyLoss()\n",
    "            # Only keep active parts of the loss\n",
    "            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))\n",
    "            \n",
    "        return (loss, logits)\n",
    "        "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "67efc7dc-a922-4e45-8cc0-1a155078c72f",
   "metadata": {},
   "source": [
    "#### Diagram\n",
    "![model diagram](./pictures/ner_lstm_diagram.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "6abe3e5b-a535-488e-b459-1c8eb1c6e42d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config of the RoBERTA + LSTM \n",
      " CustomRobertaTokenClassWithLSTM(\n",
      "  (roberta): RobertaModel(\n",
      "    (embeddings): RobertaEmbeddings(\n",
      "      (word_embeddings): Embedding(50265, 768, padding_idx=1)\n",
      "      (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
      "      (token_type_embeddings): Embedding(1, 768)\n",
      "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "    (encoder): RobertaEncoder(\n",
      "      (layer): ModuleList(\n",
      "        (0-11): 12 x RobertaLayer(\n",
      "          (attention): RobertaAttention(\n",
      "            (self): RobertaSelfAttention(\n",
      "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "            (output): RobertaSelfOutput(\n",
      "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "          )\n",
      "          (intermediate): RobertaIntermediate(\n",
      "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
      "            (intermediate_act_fn): GELUActivation()\n",
      "          )\n",
      "          (output): RobertaOutput(\n",
      "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
      "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.3, inplace=False)\n",
      "  (lstm): LSTM(768, 256, batch_first=True)\n",
      "  (lnorm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "  (lstm_dropout): Dropout(p=0.2, inplace=False)\n",
      "  (fc): Linear(in_features=256, out_features=9, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "custom_model = CustomRobertaTokenClassWithLSTM(config, num_labels=len(tagset), hidden_dim=256, n_hidden_layers=1, lstm_dropout=0.2)\n",
    "custom_model = custom_model.to(device)\n",
    "\n",
    "print(\"Config of the RoBERTA + LSTM \\n\",custom_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "5c26672e-cfd9-4d11-820c-531db5288ec2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#print(\"Model parameters:\")\n",
    "#for name, param in custom_model.named_parameters():\n",
    "#    print(f\"Parameter name: {name}, Size: {param.size()}, Requires grad: {param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f5c1d7-3d45-4120-a21d-7bdfb06b493c",
   "metadata": {},
   "source": [
    "## Check before training\n",
    "The output layer of our model is a probability distribution over our classes (we take the log of softmax activation function) as we want to predict the y possible tag $c$ among the 9 others classes $C$. The initial loss is then defined by \n",
    "$-ln(1/$ number of classes$) = -2.19$ (The natural log penalizes more the model). We use the [cross entropy loss][1] function ($l_n$) for our model:\n",
    "$$ l_n = -w_{y_n} \\log \\left( \\frac{\\exp(x_{n,y_n})}{\\sum_{c=1}^{C} \\exp(x_{n,c})} \\right) $$\n",
    "\n",
    "Given that before training weights are initialized at random, the probability distribution should be uniform \n",
    "(& the mean is $1/C$). Thus at the beginning the loss should be close to $ln(1/9)$. Let's compute this and check for our case.\n",
    "\n",
    "We give one batch to the model (input ids, true labels, attention mask as we padded) and get the outputs (loss and logits) and then compare the mean of the loss $log(1/9)$ for each tokens with the initial loss -2.19 . We also check the shape of the logits (shape of ŷ) which is (*batch_size*, *sequence_length*, *num_labels*).\n",
    "\n",
    "[1]: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "739c173a-abce-4ecc-a19d-6fdc7383409f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "_________________________________________\n",
      "Small check on Roberta4TokenClass before training\n",
      "\n",
      "ID 2 token  | Target class\n",
      "\n",
      "EU          B-ORG  1\n",
      "rejects     O  0\n",
      "German      B-MISC  3\n",
      "call        O  0\n",
      "to          O  0\n",
      "boycott     O  0\n",
      "British     B-MISC  3\n",
      "lamb        O  0\n",
      ".           O  0\n",
      "\n",
      "We pad inputs and labels to longest sequence in batch, here we used longest sequence in train\n",
      "\n",
      "Model's loss :  2.3993735313415527 ~ initial loss 2.19 \n",
      "\n",
      "True predictions :  [['B-PER', 'B-PER', 'B-PER', 'B-PER', 'B-PER', 'B-PER', 'B-PER', 'B-PER', 'B-PER']]\n",
      "True labels :  [['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']]\n",
      "Correct shape : \n",
      "batch size   : 1 \n",
      "input_size   : 163 \n",
      "target size  : 9\n",
      "\n",
      "Logit tensor shape :  torch.Size([1, 163, 9])\n",
      "\n",
      "_________________________________________\n",
      "Small check on Roberta+LSTM before training\n",
      "\n",
      "ID 2 token  | Target class\n",
      "\n",
      "EU          B-ORG  1\n",
      "rejects     O  0\n",
      "German      B-MISC  3\n",
      "call        O  0\n",
      "to          O  0\n",
      "boycott     O  0\n",
      "British     B-MISC  3\n",
      "lamb        O  0\n",
      ".           O  0\n",
      "\n",
      "We pad inputs and labels to longest sequence in batch, here we used longest sequence in train\n",
      "\n",
      "Model's loss :  2.6245949268341064 ~ initial loss 2.19 \n",
      "\n",
      "True predictions :  [['B-MISC', 'I-MISC', 'B-PER', 'I-MISC', 'B-LOC', 'I-ORG', 'B-LOC', 'B-PER', 'I-MISC']]\n",
      "True labels :  [['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']]\n",
      "Correct shape : \n",
      "batch size   : 1 \n",
      "input_size   : 163 \n",
      "target size  : 9\n",
      "\n",
      "Logit tensor shape :  torch.Size([1, 163, 9])\n"
     ]
    }
   ],
   "source": [
    "#### Check before trainig ####\n",
    "def check_dimensions(model, name=None):\n",
    "\n",
    "    print('\\n_________________________________________')\n",
    "    print(f'Small check on {name} before training\\n')\n",
    "    batch = [train_set[0]]\n",
    "    batch_size = len(batch)\n",
    "    \n",
    "    ids = train_set[0][\"input_ids\"] + [1] * input_train_max_size\n",
    "    targets = train_set[0][\"labels\"] + [-100] * input_train_max_size\n",
    "    mask = train_set[0][\"attention_mask\"] + [0] * input_train_max_size\n",
    "    \n",
    "    input_ids_size = len(ids)\n",
    "    \n",
    "    print(\"ID 2 token  | Target class\\n\")\n",
    "    for token, label in zip(tokenizer.convert_ids_to_tokens(ids), targets):\n",
    "        if label == -100:\n",
    "            continue\n",
    "        print('{0:10}  {1}  {2}'.format(token.replace(\"Ġ\", \"\"), tagset_id2label[label], label))\n",
    "    \n",
    "    print(\"\\nWe pad inputs and labels to longest sequence in batch, here we used longest sequence in train\")\n",
    "    \n",
    "    # create batch of size 1\n",
    "    ids_tensor = torch.tensor(ids).to(device).unsqueeze(0)\n",
    "    targets_tensor = torch.tensor(targets).to(device).unsqueeze(0)\n",
    "    mask_tensor = torch.tensor(mask).to(device).unsqueeze(0)\n",
    "    \n",
    "    outputs = model(input_ids=ids_tensor, labels=targets_tensor, attention_mask=mask_tensor) # First elem is loss and second the logits\n",
    "    initial_loss = outputs[0]\n",
    "    \n",
    "    print(\"\\nModel's loss : \", initial_loss.item(), \"~ initial loss 2.19 \\n\")\n",
    "\n",
    "    logits = outputs[1]\n",
    "    predictions =  logits.argmax(dim=-1)\n",
    "    true_preds, true_labels = convert_ids_2_labels(predictions.cpu().numpy(), [targets])\n",
    "    print(\"True predictions : \", true_preds)\n",
    "    print(\"True labels : \", true_labels)\n",
    "    \n",
    "    #print(compute_metrics((training_logits.argmax(dim=-1).cpu().numpy(), [batch[0]['labels']])))\n",
    "    # check the shape\n",
    "    print(\"Correct shape : \")\n",
    "    print(\"batch size   :\", batch_size, \"\\ninput_size   :\", input_ids_size, \"\\ntarget size  :\", len(tagset))\n",
    "    print(\"\\nLogit tensor shape : \", logits.shape)\n",
    "\n",
    "check_dimensions(robertaClass_model, \"Roberta4TokenClass\") # roberta4TokenClassif\n",
    "check_dimensions(custom_model, \"Roberta+LSTM\") # robertaWithLSTM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70e77d0-d262-4533-9a25-6299d9c0dc07",
   "metadata": {},
   "source": [
    "## Dataloaders and hyperparameters\n",
    "We setup the [dataloaders][1] so we can iterate over batches. Dataloaders also enables dynamic padding for sequences with different lengths which is more efficient than doing it on the whole corpus beforehand.\n",
    "\n",
    "Here we use the `collate_fn` argument to pass `data_collator` that will be called to transform the list of samples into a batch. This will apply padding also labels and not only tokens.\n",
    "\n",
    "[1]: <https://pytorch.org/tutorials/beginner/basics/data_tutorial.html>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "102efdbd-6d36-4db3-9920-ede3d9c1d77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "55252492-36fd-4322-94aa-2661f4b5365f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_train_dev(batch_size):\n",
    "    \n",
    "    train_dataloader = DataLoader(\n",
    "        train_set,\n",
    "        collate_fn=data_collator,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        pin_memory=True\n",
    "    )\n",
    "    test_dataloader = DataLoader(\n",
    "        test_set,\n",
    "        collate_fn=data_collator,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        pin_memory=True\n",
    "    )\n",
    "\n",
    "    return train_dataloader, test_dataloader\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6359ab11-04ab-40d2-987a-f4cf10e78d98",
   "metadata": {},
   "source": [
    "We use the [AdamW][1] optimizer for the gradient descent and finally the scheduler to manage\n",
    "the learning during the training process.\n",
    "\n",
    "[1]: https://pytorch.org/docs/stable/generated/torch.optim.Adam.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "861b1c07-68c8-4f9b-8a8b-273d1bd6971b",
   "metadata": {},
   "source": [
    "## Training loop\n",
    "In this section we do the magic with some maths. We build the training and evaluation. For each epoch, a first loop iterates over each batch of the train dataloader. The purpose of the training loop is to minimize the loss of the model through backpropagation using an adapive gradient descent.\n",
    "\n",
    "**Training loop:**\n",
    "\n",
    "(1) Get the outputs of the model and the loss. The shape is *batch_size*, *sequence_len*, *num_output_classes* \\\n",
    "(2) We do backpropagation of the gradient \\\n",
    "(3) Then we update the parameters (i.e. weights and biases) \\\n",
    "(4) Then we update the learning rate \\\n",
    "(5) Finally, we reset the gradients to prevent an accumulation\n",
    "\n",
    "Once the training loop has been perform, the model mode changes to evaluation to verify if the model learns something through epochs. This evaluation loop iterates over each batch of the `test_dataloader`.\n",
    "\n",
    "**Evaluation loop:**\n",
    "\n",
    "(1) Each batch passes through the model and we collect the outputs \\\n",
    "(2) From those outputs we extract the logits on the CPU in the form of numpy arrays in order to **postprocess** them\n",
    "\n",
    "Once both lists are completed, the **compute_metrics** computes the metrics for classification. It provides the **exact-match** score, **f1-score**, **accuracy**, **precision**, **recall** for this epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "4ad4f9eb-1bc9-414e-bd46-698dfc9d6f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, n_epochs, optimizer, lr_scheduler, train_dataloader):\n",
    "    \n",
    "    for epoch in range(n_epochs):\n",
    "        model.train()\n",
    "        for batch in train_dataloader:\n",
    "            batch = accelerator.prepare(batch).to(device)\n",
    "            \n",
    "            outputs = model(**batch) \n",
    "            loss = outputs[0]\n",
    "            \n",
    "            #print(f\"Loss: {loss.item()}\")\n",
    "            \n",
    "            accelerator.backward(loss)\n",
    "            #torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad() \n",
    "            progress_bar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "66de5d84-33a6-4cc6-b4a3-2c8d4407a0eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, test_loader):\n",
    "    metric = evaluate.load('seqeval')\n",
    "\n",
    "    ### Evaluation ###\n",
    "    model.eval()\n",
    "    for batch in test_loader:\n",
    "        batch = accelerator.prepare(batch).to(device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "\n",
    "        predictions = outputs[1].argmax(dim=-1) \n",
    "        labels = batch[\"labels\"] \n",
    "\n",
    "        predictions_gathered = accelerator.gather(predictions).detach().cpu().clone().numpy()\n",
    "        labels_gathered = accelerator.gather(labels).detach().cpu().clone().numpy()\n",
    "\n",
    "        true_predictions, true_labels = convert_ids_2_labels(predictions_gathered, labels_gathered) \n",
    "        metric.add_batch(predictions=true_predictions, references=true_labels) \n",
    "\n",
    "    metrics = metric.compute()\n",
    "    print(metrics)\n",
    "    print({key: round(metrics[f\"overall_{key}\"],3) for key in [\"precision\", \"recall\", \"f1\", \"accuracy\"]})\n",
    "\n",
    "    # save\n",
    "    accelerator.wait_for_everyone()\n",
    "    unwrapped_model = accelerator.unwrap_model(model)\n",
    "    torch.save(unwrapped_model.state_dict(), f\"./models/final_model_lstm\")\n",
    "\n",
    "    if accelerator.is_main_process:\n",
    "        tokenizer.save_pretrained(f\"./models/final_lstm_tokenizer\")\n",
    "\n",
    "    return {'metrics': metrics, 'model_name' : \"RobertaLSTMForTokenClass\"}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b83ce9a-34ed-4be9-a6a7-1adbc083af35",
   "metadata": {},
   "source": [
    "## Optimal hyperparameters\n",
    "After we searched for optimal hyperparamters using a separate script `ray_gridsearch.py` we load the best config and train our final model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "2d030ce9-9804-4e67-8976-f67e841f8cc2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e25feab63ff4b20ac56b6f7b5088a2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4685 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "### Initialize optimal hyperparameters\n",
    "dropout = 0.745\n",
    "num_lstm_layers = 3\n",
    "hidden_dim = 256\n",
    "n_epochs = 5\n",
    "batch_size = 16\n",
    "\n",
    "### Train and test set\n",
    "train_loader, test_loader = load_train_dev(batch_size) \n",
    "\n",
    "n_training_steps = n_epochs * len(train_loader)  \n",
    "progress_bar = tqdm(range(n_training_steps))\n",
    "\n",
    "### accelerator ###\n",
    "accelerator = Accelerator(mixed_precision='fp16')\n",
    "\n",
    "### Model, Otpimizer\n",
    "final_model = CustomRobertaTokenClassWithLSTM(config, len(tagset), hidden_dim, num_lstm_layers, dropout).to(device)\n",
    "final_optimizer = AdamW(final_model.parameters(), lr=1e-5)\n",
    "final_lr_scheduler = get_scheduler(optimizer=final_optimizer, name=\"linear\", num_warmup_steps=0, num_training_steps=n_training_steps)\n",
    "\n",
    "### Run train\n",
    "train_model(final_model, n_epochs, final_optimizer, final_lr_scheduler, train_loader) # ready for the spaghetti"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "cc548782-2977-4a78-8e36-a0eab9364ad2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'LOC': {'precision': 0.5437956204379562, 'recall': 0.6252997601918465, 'f1': 0.5817066369213608, 'number': 1668}, 'MISC': {'precision': 0.19047619047619047, 'recall': 0.10256410256410256, 'f1': 0.13333333333333333, 'number': 702}, 'ORG': {'precision': 0.3971915747241725, 'recall': 0.4768211920529801, 'f1': 0.4333789329685362, 'number': 1661}, 'PER': {'precision': 0.3840676457498887, 'recall': 0.533704390847248, 'f1': 0.4466873706004141, 'number': 1617}, 'overall_precision': 0.4237417775738106, 'overall_recall': 0.4904390934844193, 'overall_f1': 0.45465736561345915, 'overall_accuracy': 0.8895520931594184}\n",
      "{'precision': 0.424, 'recall': 0.49, 'f1': 0.455, 'accuracy': 0.89}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tokenizer config file saved in ./models/final_lstm_tokenizer/tokenizer_config.json\n",
      "Special tokens file saved in ./models/final_lstm_tokenizer/special_tokens_map.json\n"
     ]
    }
   ],
   "source": [
    "### Run eval on test set\n",
    "results = test_model(final_model, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77e0c3ac-e3ec-4092-9482-a71eb86018e9",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "Here we create a df with overall metrics for each model and for each epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0cd84557-e4c3-4527-a278-301aaa669ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_eval_2_csv(results):\n",
    "    results_df = [{\n",
    "        'model_name': results['model_name'],\n",
    "        'overall_f1': results['metrics']['overall_f1'],\n",
    "        'overall_recall': results['metrics']['overall_recall'],\n",
    "        'overall_precision': results['metrics']['overall_precision'],\n",
    "        'overall_accuracy': results['metrics']['overall_accuracy'],\n",
    "    }]\n",
    "    \n",
    "    \n",
    "    eval_df = pl.DataFrame(results_df)\n",
    "    \n",
    "    eval_df.write_csv(\"evaluations_df.csv\")\n",
    "\n",
    "    return eval_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "015d09f1-58ff-49aa-943d-f18f585e1d20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scores after training loop\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>model_name</th><th>overall_f1</th><th>overall_recall</th><th>overall_precision</th><th>overall_accuracy</th></tr><tr><td>str</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>&quot;RobertaLSTMForTokenClass&quot;</td><td>0.454657</td><td>0.490439</td><td>0.423742</td><td>0.889552</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 5)\n",
       "┌──────────────────────────┬────────────┬────────────────┬───────────────────┬──────────────────┐\n",
       "│ model_name               ┆ overall_f1 ┆ overall_recall ┆ overall_precision ┆ overall_accuracy │\n",
       "│ ---                      ┆ ---        ┆ ---            ┆ ---               ┆ ---              │\n",
       "│ str                      ┆ f64        ┆ f64            ┆ f64               ┆ f64              │\n",
       "╞══════════════════════════╪════════════╪════════════════╪═══════════════════╪══════════════════╡\n",
       "│ RobertaLSTMForTokenClass ┆ 0.454657   ┆ 0.490439       ┆ 0.423742          ┆ 0.889552         │\n",
       "└──────────────────────────┴────────────┴────────────────┴───────────────────┴──────────────────┘"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_results = write_eval_2_csv(results)\n",
    "print(\"Scores after training loop\")\n",
    "final_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4535d976-00f7-432a-8ebd-e04ba07b1694",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Results\n",
    "These are the final results after training our custom Roberta + LSTM layer\n",
    "\n",
    "My scores after eval loop\n",
    "\n",
    "| model_name               | overall_f1 | overall_recall | overall_precision | overall_accuracy |\n",
    "|--------------------------|------------|----------------|-------------------|------------------|\n",
    "| RobertaLSTMForTokenClass | 0.469755   | 0.499115       | 0.443658          | 0.892496         |\n",
    "\n",
    "\n",
    "It seems like the model is not able to capture all true positive instances as we have a very low f1. This could be due to imbalanced classes. Ideally, we should have checked the proportion for each set of tags and resample if needed in order to avoid unbalanced classes as unbalanced data can lead to underpredicted or conversely overpredicted classes. We could also find more optimal hyperparameters and perform cross-validation to increase the generalization ability of the model."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}