train_all_5_BERTs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import conllu\n",
    "import tqdm\n",
    "import torch\n",
    "import json\n",
    "from sklearn.metrics import confusion_matrix, classification_report\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "from torch.nn import functional\n",
    "from functools import partial\n",
    "from transformers import BertModel\n",
    "from torch import nn\n",
    "from transformers import BertTokenizer\n",
    "from sklearn.metrics import confusion_matrix, accuracy_score,f1_score, precision_score, recall_score, classification_report\n",
    "from utils.get_conllu_dataset import CoNLLDataset\n",
    "from utils.get_parl_corpus_token_data import ParlamentaryCorpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    gpu = True\n",
    "    print(\"Using GPU\")\n",
    "else:\n",
    "    gpu = False\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Conllu stuff\n",
    "def filter_tags(x):\n",
    "    return x        \n",
    "\n",
    "def convert_to_list_dict(path, file):\n",
    "    path = path.format(file)\n",
    "    with open(path, encoding=\"UTF-8\") as infile:\n",
    "        lst = []\n",
    "        tokens = list(conllu.parse_incr(infile))\n",
    "        for sent in tokens:\n",
    "                dic = {\n",
    "                \"idx\": sent.metadata[\"sent_id\"],\n",
    "                \"text\": sent.metadata[\"text\"].lower(),\n",
    "                \"tokens\": [token[\"form\"].lower() for token in sent],\n",
    "                \"lemmas\": [token[\"lemma\"] for token in sent],\n",
    "                \"pos_tags\": [token[\"upos\"] for token in sent],\n",
    "                \"ner_tags\": [filter_tags(token[\"misc\"].get(\"name\", \"O\")) for token in sent],\n",
    "            }\n",
    "                lst.append(dic) \n",
    "        print(\"Converting {} to list of dictionaries\\n     {} elements converted..\".format(file, len(lst)))\n",
    "    return lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bert(nn.Module):\n",
    "    def __init__(self, ner_vocab, model_path=None, freeze=False):\n",
    "        super().__init__()\n",
    "        self._bert = BertModel.from_pretrained(\n",
    "            model_path\n",
    "        )\n",
    "        hidden_size = self._bert.config.hidden_size\n",
    "        self._linear = nn.Linear(hidden_size, len(ner_vocab))\n",
    "\n",
    "        if freeze:\n",
    "            for param in self._bert.parameters():\n",
    "                param.requires_grad = False #Freezing bert layer\n",
    "\n",
    "    def forward(self, batch, mask):\n",
    "        b = self._bert(\n",
    "            input_ids=batch[\"input_ids\"], attention_mask=batch[\"attention_mask\"]\n",
    "        )                                                                   # Notes to self\n",
    "        pooler = b.last_hidden_state[:, mask].diagonal().permute(2, 0, 1) #https://pytorch.org/docs/stable/generated/torch.permute.html\n",
    "        return self._linear(pooler)                                     #https://pytorch.org/docs/stable/generated/torch.diagonal.html "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function borrowed from https://github.com/ltgoslo/NorBERT/\n",
    "\n",
    "def collate_fn(batch, gpu=False):\n",
    "    longest_y = max([y.size(0) for X, y in batch])\n",
    "    x = [X for X, y in batch]\n",
    "    y = torch.stack(\n",
    "        [functional.pad(y, (0, longest_y - y.size(0)), value=-1) for X, y in batch]) #https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html\n",
    "    if gpu:\n",
    "        y = y.to(\"cuda\")\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/ltgoslo/NorBERT/\n",
    "\n",
    "def build_mask(tokenizer, ids):\n",
    "    tok_sents = [tokenizer.convert_ids_to_tokens(i) for i in ids]\n",
    "    mask = []\n",
    "    for sentence in tok_sents:\n",
    "        current = []\n",
    "        for n, token in enumerate(sentence):\n",
    "            if token in tokenizer.all_special_tokens[1:] or token.startswith(\"##\"): # ## masked\n",
    "                continue\n",
    "            else:\n",
    "                current.append(n)\n",
    "        mask.append(current)\n",
    "\n",
    "    mask = tokenizer.pad({\"input_ids\": mask}, return_tensors=\"pt\")[\"input_ids\"]\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(input_data, tokenizer, model, gpu=False):\n",
    "    input_data = tokenizer(\n",
    "        input_data, is_split_into_words=True, return_tensors=\"pt\", padding=True\n",
    "    )\n",
    "    if gpu:\n",
    "        input_data = input_data.to(\"cuda\")\n",
    "    batch_mask = build_mask(tokenizer, input_data[\"input_ids\"])\n",
    "    y_pred = model(input_data, batch_mask).permute(0, 2, 1).argmax(dim=1)\n",
    "    return y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(device, train_dataset, dev_loader, tokenizer, model, optimiser, scheduler, criterion, no_ne_index, MAX_EPOCHS, OUTPUT_PATH, PATIENCE):\n",
    "    \n",
    "    accuracy_lst = []\n",
    "    loss_history_train, loss_history_dev, val_accuracy, learn_rate_history = {}, {}, {}, {}\n",
    "    inpatience = 0\n",
    "    \n",
    "    for epoch in range(MAX_EPOCHS):\n",
    "        print(\"Shuffling dataset..\")\n",
    "        train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, collate_fn=partial(collate_fn, gpu=gpu))\n",
    "        print(\"Learn rate: \", end=\"\")\n",
    "        print(optimiser.param_groups[0]['lr'])\n",
    "        # TRAIN\n",
    "        train_loss_accumulation = 0\n",
    "        model.train()\n",
    "        train_iter = tqdm.tqdm(train_loader)\n",
    "        for x, y in train_iter:\n",
    "            x = tokenizer(\n",
    "                x, is_split_into_words=True, return_tensors=\"pt\", padding=True\n",
    "            )\n",
    "            if gpu:\n",
    "                x = x.to(device)\n",
    "            batch_mask = build_mask(tokenizer, x[\"input_ids\"])\n",
    "            optimiser.zero_grad()\n",
    "            y_pred = model(x, batch_mask).permute(0, 2, 1)\n",
    "            loss = criterion(y_pred, y)\n",
    "            train_loss_accumulation+=loss.item()\n",
    "            loss.backward()\n",
    "            optimiser.step()\n",
    "            train_iter.set_postfix_str(f\"Train loss: {loss.item()}\")\n",
    "        \n",
    "                \n",
    "        # EVAL\n",
    "        val_loss_accumulation = 0        \n",
    "        model.eval()\n",
    "        dev_iter = tqdm.tqdm(dev_loader)\n",
    "        correct, total = 0, 0\n",
    "        for x, y in dev_iter:\n",
    "            y_pred = predict(x, tokenizer, model, gpu=gpu)\n",
    "            correct += (((y_pred == y).logical_and(y != no_ne_index)).nonzero(as_tuple=False).size(0))\n",
    "            total += ((y != no_ne_index).logical_and(y != -1).nonzero(as_tuple=False).size(0))\n",
    "            x = tokenizer(\n",
    "                x, is_split_into_words=True, return_tensors=\"pt\", padding=True\n",
    "            )\n",
    "            if gpu:\n",
    "                x = x.to(device)\n",
    "            batch_mask = build_mask(tokenizer, x[\"input_ids\"])\n",
    "            y_pred = model(x, batch_mask).permute(0, 2, 1)\n",
    "            dev_loss = criterion(y_pred, y)\n",
    "            val_loss_accumulation+=dev_loss.item()\n",
    "            dev_iter.set_postfix_str(f\"Validation loss: {loss.item()}\")\n",
    "\n",
    "        scheduler.step(dev_loss)\n",
    "    \n",
    "        print(\"Epoch number: {}\".format(epoch))\n",
    "        accuracy = correct / total\n",
    "        accuracy_lst.append(accuracy)\n",
    "        current_acc_best = max(accuracy_lst)\n",
    "        print(accuracy)\n",
    "        print(f\"List of accuracies: {accuracy_lst}\")\n",
    "        print(f\"Best accuracy so far: {current_acc_best}\")\n",
    "        print(f\"Validation accuracy = {correct} / {total} = {accuracy}\")\n",
    "        if ((accuracy < current_acc_best) and (inpatience <= PATIENCE)): \n",
    "            print(f\"No improvement. Inpatience counter increased..\")\n",
    "            print(f\"Patience set at {PATIENCE}\")\n",
    "            inpatience +=1\n",
    "            print(f\"Inpatience counter reached {inpatience}\")\n",
    "            if inpatience == PATIENCE:\n",
    "                print(f\"Patience tolerance reached. Early stopping at epoch {epoch}!\")\n",
    "                break\n",
    "        else:\n",
    "            current_acc_best = accuracy\n",
    "            print(\"New best!\")\n",
    "            print(\"Setting inpatience counter to 0\")\n",
    "            inpatience = 0\n",
    "            torch.save(model.state_dict(), OUTPUT_PATH)\n",
    "\n",
    "        loss_history_train[\"epoch: {}\".format(epoch)] = train_loss_accumulation\n",
    "        loss_history_dev[\"epoch: {}\".format(epoch)] = val_loss_accumulation\n",
    "        val_accuracy[\"epoch: {}\".format(epoch)] = accuracy\n",
    "        learn_rate_history[\"epoch: {}\".format(epoch)] = optimiser.param_groups[0]['lr']\n",
    "        \n",
    "\n",
    "    return loss_history_train, loss_history_dev, val_accuracy, epoch, learn_rate_history\n",
    "\n",
    "# \"\"\" TODO: legg til tekst fil med loss og andre stats\"\"\"\n",
    "# INTERGRATE TODO PATIENCE\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_test(test_set, ner_vocab, tokenizer, model):\n",
    "    model.eval()\n",
    "    predicted_labels = []\n",
    "    test_set = tqdm.tqdm(test_set)\n",
    "    for x, y in test_set:\n",
    "        y_pred = predict(x, tokenizer, model, gpu=gpu)\n",
    "        predicted = [ner_vocab[element] for element in y_pred[0]]\n",
    "        predicted_labels += predicted\n",
    "    return predicted_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_freeze(freeze, model):\n",
    "    if freeze:\n",
    "        lr=0.001\n",
    "        optimiser = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    else:\n",
    "        lr = 2e-5\n",
    "        optimiser = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    return optimiser, lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_info_json(model_name, loss_history_train, loss_history_dev, val_accuracy, learn_rate_history, MAX_EPOCHS, stop_epoch, lr, seed, PATIENCE):\n",
    "    info = {\n",
    "        \"Seed\": seed,\n",
    "        \"Max epochs set\": MAX_EPOCHS,\n",
    "        \"Train stop at epoch (begins at 0)\": stop_epoch,\n",
    "        \"Learning rate\": lr,\n",
    "        \"Validation accuracy history\": val_accuracy,\n",
    "        \"Train loss history\": loss_history_train,\n",
    "        \"Val loss history\": loss_history_dev,\n",
    "        \"Learn rate history\": learn_rate_history,\n",
    "        \"Patience set\": PATIENCE\n",
    "    }\n",
    "\n",
    "    file_name = \"{}_seed_{}.json\".format(model_name, seed)\n",
    "    with open(file_name, \"w\") as write:\n",
    "        json.dump(info, write, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"all_conllu/{0}.conllu\"\n",
    "file_list = [\"no_bokmaal-ud-dev\", \"no_bokmaal-ud-test\", \"no_bokmaal-ud-train\", \"no_nynorsk-ud-dev\", \"no_nynorsk-ud-test\", \"no_nynorsk-ud-train\"]\n",
    "\n",
    "dev_split_no = convert_to_list_dict(path, file_list[0])\n",
    "test_split_no = convert_to_list_dict(path, file_list[1])\n",
    "train_split_no = convert_to_list_dict(path, file_list[2])\n",
    "\n",
    "dev_split_ny = convert_to_list_dict(path, file_list[3])\n",
    "test_split_ny = convert_to_list_dict(path, file_list[4])\n",
    "train_split_ny = convert_to_list_dict(path, file_list[5])\n",
    "\n",
    "print(\"Combining train, dev and test sets..\")\n",
    "dev_split = dev_split_no + dev_split_ny\n",
    "test_split = test_split_no + test_split_ny\n",
    "train_split = train_split_no + train_split_ny\n",
    "print(\"Success!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assumes that it is the same vocab in dev and test\n",
    "from itertools import chain\n",
    "def get_labels(train_split, add_UNK=True):\n",
    "    label_vocab = [set(y[\"ner_tags\"]) for y in train_split]\n",
    "    label_vocab = list(chain(*[d for d in label_vocab]))\n",
    "    label_vocab = list(set(sorted(label_vocab)))\n",
    "    \n",
    "    if add_UNK==True:\n",
    "        label_vocab.append(\"@UNK\")\n",
    "\n",
    "    tmp1, tmp2 = [], []\n",
    "    for label in label_vocab:\n",
    "        if \"-\" in label:\n",
    "            tmp1.append(label)\n",
    "        else:\n",
    "            tmp2.append(label)\n",
    "    tmp1.sort(key=lambda x: (x.split(\"-\")[1], x.split(\"-\")[0]))\n",
    "    tmp1 = [x for x in tmp1]\n",
    "    tmp1.extend(tmp2)\n",
    "    label_vocab = tmp1\n",
    "    return label_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# ner_vocab =     ['B-GPE_LOC', 'I-DRV', 'I-LOC', 'B-PER', 'I-PER', 'B-PROD', \n",
    "#                 'I-GPE_ORG', 'B-GPE_ORG', 'B-EVT', 'B-DRV', 'I-PROD', 'B-ORG', 'B-MISC',\n",
    "#                 'I-MISC', 'I-GPE_LOC', 'B-LOC', 'I-ORG', 'I-EVT', 'O', '@UNK']\n",
    "\n",
    "\n",
    "ner_vocab = get_labels(train_split)\n",
    "\n",
    "x_train_tokens = [x[\"tokens\"] for x in train_split]\n",
    "y_train_labels = [y[\"ner_tags\"] for y in train_split]\n",
    "train_dataset = CoNLLDataset(x_train_tokens, y_train_labels, ner_vocab)\n",
    "\n",
    "x_dev_tokens = [x[\"tokens\"] for x in dev_split]\n",
    "y_dev_labels = [y[\"ner_tags\"] for y in dev_split]\n",
    "dev_dataset = CoNLLDataset(x_dev_tokens, y_dev_labels, ner_vocab)\n",
    "\n",
    "x_test_tokens = [x[\"tokens\"] for x in test_split]\n",
    "y_test_labels = [y[\"ner_tags\"] for y in test_split]\n",
    "test_dataset = CoNLLDataset(x_test_tokens, y_test_labels, ner_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepping data for datao to predict\n",
    "dev_loader = DataLoader(\n",
    "        dev_dataset, batch_size=64, shuffle=False, collate_fn=partial(collate_fn, gpu=gpu))\n",
    "test_loader = DataLoader(\n",
    "        test_dataset, batch_size=1, shuffle=False, collate_fn=partial(collate_fn, gpu=gpu))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Training block params\"\"\"\n",
    "\n",
    "\n",
    "\"\"\" SET BERT JSON NAME AND MODEL TYPE \"\"\"\n",
    "model_name_type = {\n",
    "                    \"norBERT\":\"ltgoslo/norbert\",\n",
    "                    \"norBERT2\":\"ltgoslo/norbert2\",\n",
    "                    \"mBERT\":\"bert-base-multilingual-cased\",\n",
    "                    \"NB-BERT\":\"NbAiLab/nb-bert-base\",\n",
    "                    \"scandi-ner\":\"saattrupdan/nbailab-base-ner-scandi\",\n",
    "                    }\n",
    "\n",
    "\n",
    "\"\"\" SET NUMBER OF EPOCHS \"\"\"\n",
    "MAX_EPOCHS = 1\n",
    "\n",
    "\"\"\" SET PATIENCE IF  \"\"\"\n",
    "PATIENCE = 10\n",
    "\n",
    "\"\"\" SEEDS \"\"\"\n",
    "seeds = [8, 37, 42, 101, 1024]\n",
    "\n",
    "\"\"\" SET FREEZE CONDITION \"\"\"\n",
    "freeze = False\n",
    "\n",
    "\"\"\" SET IF MODEL SHOULD TRAIN \"\"\"\n",
    "do_training = True\n",
    "\n",
    "\"\"\" LABEL THAT IS NOT NE \"\"\"\n",
    "no_ne_index = train_dataset.ner_indexer[\"O\"] # For train parameter\n",
    "        \n",
    "\"\"\" OUTPUT PATH NAMES \"\"\"\n",
    "output_path = \"Provide a path\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, model_type in model_name_type.items():\n",
    "    tokenizer = BertTokenizer.from_pretrained(model_type, do_basic_tokenize=False)\n",
    "    model_name = name\n",
    "    if do_training == True:\n",
    "        for seed in seeds:\n",
    "            # Setting deterministic behaviour \n",
    "            torch.backends.cudnn.deterministic = True\n",
    "            torch.manual_seed(seed)\n",
    "            torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "            # Instaniating model\n",
    "            bert = Bert(ner_vocab, model_type, freeze).to(device)\n",
    "\n",
    "            # Some paramters for train\n",
    "            criterion = nn.CrossEntropyLoss(ignore_index=-1)\n",
    "            optimiser, lr = is_freeze(freeze, bert)\n",
    "            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimiser, mode='min',\n",
    "                    factor=0.93, patience=1, threshold=1e-7, threshold_mode='abs')\n",
    "\n",
    "            OUTPUT_PATH = r\"{}_adam_lr_{}__maxEpochs_{}__seed_{}\".format(output_path, lr, MAX_EPOCHS, seed)\n",
    "\n",
    "            loss_history_train, loss_history_dev, val_accuracy, stop_epoch, learn_rate_history = train(\n",
    "                                                device, train_dataset, dev_loader, \n",
    "                                                tokenizer, bert, optimiser, scheduler, criterion,\n",
    "                                                no_ne_index, MAX_EPOCHS, OUTPUT_PATH, PATIENCE\n",
    "                                                )\n",
    "\n",
    "            info = model_info_json(model_name, loss_history_train, loss_history_dev, val_accuracy, learn_rate_history,\n",
    "                                MAX_EPOCHS, stop_epoch, lr, seed, PATIENCE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 64-bit",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "1e5f1d1d80de0e037045f44f922af48a35ae55538dc259ba83b72568cdb122e4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}