WSD_with_BERT.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x5I3y0a98T_H"
      },
      "source": [
        "# Libraries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJ0wHYFs8bjo"
      },
      "source": [
        "## Install modules if necessary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZRHidmor8iDt",
        "outputId": "684115b4-61fb-4290-cc49-6b3f622971d1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: transformers in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (4.12.5)\n",
            "Requirement already satisfied: tqdm>=4.27 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (4.62.3)\n",
            "Requirement already satisfied: numpy>=1.17 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (1.21.4)\n",
            "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (0.10.3)\n",
            "Requirement already satisfied: packaging>=20.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (21.3)\n",
            "Requirement already satisfied: filelock in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (3.4.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.1.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (0.2.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (2021.11.10)\n",
            "Requirement already satisfied: requests in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (2.26.0)\n",
            "Requirement already satisfied: sacremoses in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (0.0.46)\n",
            "Requirement already satisfied: pyyaml>=5.1 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from huggingface-hub<1.0,>=0.1.0->transformers) (4.0.1)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from packaging>=20.0->transformers) (3.0.6)\n",
            "Requirement already satisfied: colorama in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from tqdm>=4.27->transformers) (0.4.4)\n",
            "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from requests->transformers) (3.3)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from requests->transformers) (2.0.9)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from requests->transformers) (2021.10.8)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from requests->transformers) (1.26.7)\n",
            "Requirement already satisfied: six in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from sacremoses->transformers) (1.16.0)\n",
            "Requirement already satisfied: click in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from sacremoses->transformers) (8.0.3)\n",
            "Requirement already satisfied: joblib in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from sacremoses->transformers) (1.1.0)\n",
            "Requirement already satisfied: xmltodict in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (0.12.0)\n",
            "Requirement already satisfied: pandas in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (1.3.4)\n",
            "Requirement already satisfied: pytz>=2017.3 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from pandas) (2021.3)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from pandas) (2.8.2)\n",
            "Requirement already satisfied: numpy>=1.17.3 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from pandas) (1.21.4)\n",
            "Requirement already satisfied: six>=1.5 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from python-dateutil>=2.7.3->pandas) (1.16.0)\n",
            "Requirement already satisfied: numpy in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (1.21.4)\n",
            "Looking in links: https://download.pytorch.org/whl/cu102/torch_stable.html\n",
            "Requirement already satisfied: torch==1.10.0+cu102 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (1.10.0+cu102)\n",
            "Requirement already satisfied: torchvision==0.11.1+cu102 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (0.11.1+cu102)\n",
            "Requirement already satisfied: torchaudio===0.10.0+cu102 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (0.10.0+cu102)\n",
            "Requirement already satisfied: typing-extensions in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from torch==1.10.0+cu102) (4.0.1)\n",
            "Requirement already satisfied: numpy in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from torchvision==0.11.1+cu102) (1.21.4)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from torchvision==0.11.1+cu102) (8.4.0)\n",
            "Requirement already satisfied: nltk in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (3.6.5)\n",
            "Requirement already satisfied: regex>=2021.8.3 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from nltk) (2021.11.10)\n",
            "Requirement already satisfied: tqdm in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from nltk) (4.62.3)\n",
            "Requirement already satisfied: click in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from nltk) (8.0.3)\n",
            "Requirement already satisfied: joblib in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from nltk) (1.1.0)\n",
            "Requirement already satisfied: colorama in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from click->nltk) (0.4.4)\n",
            "Requirement already satisfied: sklearn in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (0.0)\n",
            "Requirement already satisfied: scikit-learn in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from sklearn) (1.0.1)\n",
            "Requirement already satisfied: joblib>=0.11 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from scikit-learn->sklearn) (1.1.0)\n",
            "Requirement already satisfied: scipy>=1.1.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from scikit-learn->sklearn) (1.7.3)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from scikit-learn->sklearn) (3.0.0)\n",
            "Requirement already satisfied: numpy>=1.14.6 in c:\\users\\hms17\\miniconda3\\envs\\wsd_env\\lib\\site-packages (from scikit-learn->sklearn) (1.21.4)\n"
          ]
        }
      ],
      "source": [
        "!pip install transformers\n",
        "!pip install xmltodict\n",
        "!pip install pandas\n",
        "!pip install numpy\n",
        "!pip3 install torch==1.10.0+cu102 torchvision==0.11.1+cu102 torchaudio===0.10.0+cu102 -f https://download.pytorch.org/whl/cu102/torch_stable.html\n",
        "!pip install nltk\n",
        "!pip install sklearn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LGHkvHoK8Pqm"
      },
      "source": [
        "## Import Modules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vkp0g8yO75EP",
        "outputId": "c5801c56-4634-456e-810a-1cb29c7d3271"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[nltk_data] Downloading package wordnet to\n",
            "[nltk_data]     C:\\Users\\hms17\\AppData\\Roaming\\nltk_data...\n",
            "[nltk_data]   Package wordnet is already up-to-date!\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import xml.etree.ElementTree as ET\n",
        "import torch\n",
        "from torch import nn\n",
        "import nltk\n",
        "from nltk.corpus import wordnet as wn\n",
        "from sklearn.metrics import f1_score\n",
        "nltk.download('wordnet')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "qytmMQ628IVN"
      },
      "outputs": [],
      "source": [
        "import transformers as trf\n",
        "import random, time, datetime, warnings, re\n",
        "from os.path import exists as pathExists\n",
        "from os import rename as renameFile\n",
        "from collections import defaultdict\n",
        "warnings.filterwarnings('ignore') "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "dev = torch.device('cuda:0')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hm7VNWZU8wTM"
      },
      "source": [
        "# Data Preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbFaiOfG85dd"
      },
      "source": [
        "## Loading Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "DccQvagqSal3"
      },
      "outputs": [],
      "source": [
        "# This function reads a dataset .xml file and its gold.key.txt file, and returns a list of dictionaries each containing a piece of data \n",
        "# with its proper label and other necessary components.\n",
        "def loadDataset(path):\n",
        "  data = ET.parse(path + '.data.xml')\n",
        "  with open(path + '.gold.key.txt') as file:\n",
        "      labels =  [line.rstrip() for line in file.readlines()]\n",
        "  root = data.getroot()\n",
        "  dataset = []\n",
        "  for doc in root:\n",
        "    for raw_sent in doc:\n",
        "      whole_sentence = []\n",
        "      instances = [(i, x) for i, x in enumerate(raw_sent) if x.tag == 'instance']\n",
        "      for term in raw_sent:\n",
        "        whole_sentence.append(term.text.lower())\n",
        "      whole_sentence = ' '.join(whole_sentence)\n",
        "      for idx, inst in instances:\n",
        "        gold_label = labels.pop(0).split()\n",
        "        assert(gold_label[0] == inst.attrib['id'])\n",
        "        all_senses = [lemma for sense in wn.synsets(inst.text) for lemma in sense.lemmas()]\n",
        "        try:\n",
        "          label = torch.tensor([ [x.key() for x in all_senses].index(gold_label[1]) ])\n",
        "        except ValueError:\n",
        "          continue\n",
        "\n",
        "        if any(c in set(\".-\\\\/~()\") for c in inst.text): continue\n",
        "        dataset.append({\n",
        "            'sentence': whole_sentence,\n",
        "            'idx': whole_sentence.split().index(inst.text.lower()),\n",
        "            'polyseme': inst.text.lower(),\n",
        "            'lemma' : inst.attrib['lemma'],\n",
        "            'senses': all_senses,\n",
        "            'label': label\n",
        "        })\n",
        "  return dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "xxIXyO1K0wJx"
      },
      "outputs": [],
      "source": [
        "# Should take about a minute to load all the data (the training set is especially large)\n",
        "training = loadDataset('Datasets/Training/SemCor/semcor')\n",
        "validation = loadDataset('Datasets/Validation/semeval2007/semeval2007')\n",
        "testing = {\n",
        "    'SE2': loadDataset('Datasets/Testing/senseval2/senseval2'),\n",
        "    'SE3': loadDataset('Datasets/Testing/senseval3/senseval3'),\n",
        "    'SE13': loadDataset('Datasets/Testing/semeval2013/semeval2013'),\n",
        "    'SE15': loadDataset('Datasets/Testing/semeval2015/semeval2015')\n",
        "}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BEcV86t9-G_S"
      },
      "source": [
        "## DistilBERT Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E1BjCVRy-MQq",
        "outputId": "90ddb5fb-0c25-457f-c775-09c5e2c7c5b0"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_projector.bias', 'vocab_transform.bias', 'vocab_layer_norm.bias', 'vocab_transform.weight', 'vocab_layer_norm.weight']\n",
            "- This IS expected if you are initializing DistilBertModel 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 DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        }
      ],
      "source": [
        "tokenizer = trf.DistilBertTokenizer.from_pretrained('distilbert-base-uncased', truncation=True, padding=True, do_lower_case=True)\n",
        "config = trf.DistilBertConfig.from_pretrained('distilbert-base-uncased', output_hidden_states=True)\n",
        "bert_model = trf.DistilBertModel.from_pretrained('distilbert-base-uncased', config=config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9mjvadBVjwfq"
      },
      "source": [
        "# Classifier"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Qgb4ROeDn-v_"
      },
      "outputs": [],
      "source": [
        "def getEncodedIndeces(text, idx):\n",
        "  encodings = [(x, tokenizer(x, add_special_tokens=False)['input_ids']) for x in text.split()]\n",
        "  start, end = 0, 0\n",
        "  for i, encoding in enumerate(encodings):\n",
        "    if i == idx:\n",
        "      start += 1\n",
        "      end = start + len(encoding[1])\n",
        "      return (start, end)\n",
        "    else:\n",
        "      start += len(encoding[1])\n",
        "\n",
        "cached_output = ('', None)\n",
        "def getFeatureVec(data):\n",
        "  global cached_output\n",
        "  text = data['sentence']\n",
        "  if text == cached_output[0]:\n",
        "    hidden_states = cached_output[1]\n",
        "  else:\n",
        "    encoded_input = tokenizer(text, return_tensors='pt')\n",
        "    output = bert_model(**encoded_input)\n",
        "    hidden_states = output[1][0][0]\n",
        "    cached_output = (text, hidden_states.clone().detach())\n",
        "  s, e = getEncodedIndeces(text, data['idx'])\n",
        "  avg = torch.mean(hidden_states[s:e], 0)\n",
        "  return avg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([768])"
            ]
          },
          "execution_count": 79,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "getFeatureVec(training[1]).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "def evaluate(model, dataset):\n",
        "  loss_function = nn.CrossEntropyLoss().to(dev)\n",
        "  loss = 0.0\n",
        "  f1 = 0\n",
        "  predictions = defaultdict(list)\n",
        "  labels = defaultdict(list)\n",
        "  model.eval().cuda()\n",
        "  with torch.no_grad():\n",
        "    for i, data in enumerate(dataset):\n",
        "      feature_vec = getFeatureVec(data).to(dev)\n",
        "      label = data['label'].to(dev)\n",
        "      outputs = torch.reshape(model(feature_vec.to(dev), data), (1, -1))\n",
        "      loss = loss_function(outputs, label)\n",
        "      predictions[data['polyseme']].append(outputs)\n",
        "      labels[data['polyseme']].append(label)\n",
        "      \n",
        "      loss += loss.item()\n",
        "\n",
        "      if i % 100 == 99:\n",
        "        print(\"{}/{} complete.\".format(i+1, len(dataset)))\n",
        "    \n",
        "    all_f1_scores = []\n",
        "    f1_weights = []\n",
        "    for polyseme in labels.keys():\n",
        "      preds_tensor = torch.stack(predictions[polyseme])\n",
        "      preds_tensor = torch.argmax(preds_tensor.reshape((-1, preds_tensor.shape[-1])), dim=1)\n",
        "      labels_tensor = torch.stack(labels[polyseme]).reshape((-1,))\n",
        "      all_f1_scores.append(f1_score(labels_tensor.cpu(), preds_tensor.cpu(), average='weighted'))\n",
        "      f1_weights.append(len(labels_tensor))\n",
        "    \n",
        "    avgLoss = loss / len(dataset)\n",
        "    f1 = np.average(all_f1_scores, weights=f1_weights) * 100\n",
        "    return avgLoss, f1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Code for sanity-checking that getEncodedIndeces extracts the correct tokens.\n",
        "\n",
        "# for data in training:\n",
        "#     f = getFeatureVec(data)\n",
        "#     # print(data['sentence'], \"| Definition of:\", data['polyseme'])\n",
        "#     # data['senses'][data['label']].synset().definition()\n",
        "#     encodings = {x:tokenizer(x, add_special_tokens=False)['input_ids'] for x in data['sentence'].split()}\n",
        "#     idxs = getEncodedIndeces(data['sentence'], data['idx'])\n",
        "#     # print(encodings)\n",
        "#     extracted_token = tokenizer(data['sentence'])['input_ids'][idxs[0]:idxs[1]]\n",
        "#     # print(extracted_token)\n",
        "#     try:\n",
        "#         assert(encodings[data['polyseme'].lower()] == extracted_token)\n",
        "#     except:\n",
        "#         print(encodings)\n",
        "#         print(data['polyseme'], encodings[data['polyseme'].lower()])\n",
        "#         print(extracted_token)\n",
        "#         break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## MLP Neural Net Structure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "-cmmtwiNjwFQ"
      },
      "outputs": [],
      "source": [
        "class MLP_Classifier(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.H = 768\n",
        "        self.L1 = nn.Linear(self.H, self.H)\n",
        "        self.ReLU = nn.ReLU()\n",
        "        self.L2 = nn.ModuleDict({})\n",
        "\n",
        "    def forward(self, features, data):\n",
        "        x = self.L1(features)\n",
        "        x = self.ReLU(x)\n",
        "\n",
        "        polyseme, numSenses = \"polyseme_\" + data['polyseme'], len(data['senses'])\n",
        "        if polyseme not in self.L2: \n",
        "          self.L2.update({polyseme: nn.Linear(self.H, numSenses).to(dev)})\n",
        "        \n",
        "        x = self.L2[polyseme](x)\n",
        "        \n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load model if it is present"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [],
      "source": [
        "def loadModel(path):\n",
        "    if pathExists(path):\n",
        "        print(\"Saved model found. Loading saved model\")\n",
        "        \n",
        "        # This code makes sure the saved model is loaded properly, including the ModuleDict which stores the nn.Linear() layer for each polyseme\n",
        "        combined_state_dict = torch.load(path)\n",
        "        L1_state_dict = {k:v for k,v in combined_state_dict.items() if k in ['L1.weight', 'L1.bias']}\n",
        "        L2_state_dict = {k:v for k,v in combined_state_dict.items() if k not in ['L1.weight', 'L1.bias']}\n",
        "\n",
        "        mlp = MLP_Classifier().to(dev)\n",
        "        loaded_L2 = nn.ModuleDict({})\n",
        "        all_polysemes = set([re.search('L2\\.(.*)\\.', polyseme).group(1) for polyseme in L2_state_dict.keys()])\n",
        "        for polyseme in all_polysemes:\n",
        "            num_senses = len([lemma for sense in wn.synsets(polyseme.split('_')[1]) for lemma in sense.lemmas()])\n",
        "            layer = nn.Linear(mlp.H, num_senses)\n",
        "            weights = L2_state_dict['L2.{}.weight'.format(polyseme)]\n",
        "            biases = L2_state_dict['L2.{}.bias'.format(polyseme)]\n",
        "            layer.load_state_dict({\n",
        "                'weight': weights,\n",
        "                'bias': biases\n",
        "            })\n",
        "            loaded_L2.update({polyseme: layer.to(dev)})\n",
        "\n",
        "        mlp.load_state_dict(L1_state_dict)\n",
        "        mlp.L2 = loaded_L2\n",
        "\n",
        "        print(\"Saved model loaded.\")\n",
        "        return mlp\n",
        "    else:\n",
        "        print(\"No saved model found.\")\n",
        "        return False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Saved model found. Loading saved model\n",
            "Saved model loaded.\n"
          ]
        }
      ],
      "source": [
        "do_training = False\n",
        "saved_model_path = \"wsd_model_48.96_f1.pth\"\n",
        "mlp = loadModel(saved_model_path)\n",
        "if not mlp: do_training = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training and Validation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qHLH8ZXWnOhl",
        "outputId": "4b8cf52e-6880-44f9-c80e-eda23ce06359"
      },
      "outputs": [],
      "source": [
        "# Only do training if there is no saved model\n",
        "if do_training:\n",
        "  # Initialize the MLP\n",
        "  mlp = MLP_Classifier().to(dev)\n",
        "\n",
        "  # Define the loss function, optimizer, and scheduler (to reduce learning rate)\n",
        "  loss_function = nn.CrossEntropyLoss().to(dev)\n",
        "  optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-3)\n",
        "  scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=(lambda epoch: 1/(epoch+1)))\n",
        "\n",
        "  min_valid_loss = np.inf\n",
        "  max_f1_score = 0\n",
        "\n",
        "  # Run the training loop\n",
        "  for epoch in range(10):\n",
        "    print(f'Starting epoch {epoch+1}')\n",
        "    \n",
        "    # Each epoch, the model trains on half of the training dataset. \n",
        "    # The first epoch will train on the first half, then the second epoch will train on the second half, and so on, alternating every epoch.\n",
        "    halfway_idx = int(len(training) / 2)\n",
        "    if epoch % 2 == 0:\n",
        "      training_subset = training[:halfway_idx]\n",
        "    else:\n",
        "      training_subset = training[halfway_idx:]\n",
        "\n",
        "    # training_subset = training_subset[:1000]\n",
        "    # Training loop:\n",
        "    train_loss = 0.0\n",
        "    batch_loss = 0.0\n",
        "    batchStart = time.time()\n",
        "    mlp.train().cuda()\n",
        "    for i, data in enumerate(training_subset):\n",
        "      feature_vec = getFeatureVec(data).to(dev)\n",
        "      label = data['label'].to(dev)\n",
        "      optimizer.zero_grad()\n",
        "      outputs = torch.reshape(mlp(feature_vec.to(dev), data), (1, -1))\n",
        "      loss = loss_function(outputs, label)\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "      \n",
        "      train_loss += loss.item()\n",
        "      batch_loss += loss.item()\n",
        "      if i % 500 == 499:\n",
        "          print('Training loss after mini-batch {:5d}/{:d}: {:.3f}'.format(i + 1, len(training_subset), batch_loss / 500))\n",
        "          batchTime = time.time() - batchStart\n",
        "          remainingTime = ((len(training_subset) - (i + 1)) / 500) * batchTime\n",
        "          print(\"Time remaining for current epoch: {}\".format(str(datetime.timedelta(seconds=remainingTime))))\n",
        "          batch_loss = 0.0\n",
        "          batchStart = time.time()\n",
        "\n",
        "    print(\"Training done. Validating...\")\n",
        "    \n",
        "    # Validation loop:\n",
        "    valid_loss, f1 = evaluate(mlp, validation)\n",
        "    \n",
        "    # Decrease learning rate (per epoch)\n",
        "    scheduler.step()\n",
        "    print(\"Learning rate:\", scheduler.get_lr())\n",
        "\n",
        "    # Epoch complete\n",
        "    print(\"Epoch {} complete. Training loss: {:.5f}. Validation loss: {:.5f}. Validation F1 score: {:.5f}\".format(epoch+1, (train_loss / len(training_subset)), (valid_loss / len(validation)), f1) )\n",
        "    if f1 > max_f1_score:\n",
        "      print(f'F1 score increased ({max_f1_score:.6f}--->{f1:.6f}) \\t Saving The Model...')\n",
        "      max_f1_score = f1\n",
        "      # Saving State Dict\n",
        "      torch.save(mlp.state_dict(), 'wsd_model.pth')\n",
        "\n",
        "  # Rename the best model to include F1 score\n",
        "  best_model = 'wsd_model_{:.2f}_f1.pth'.format(max_f1_score)\n",
        "  renameFile('wsd_model.pth', best_model)\n",
        "\n",
        "  # Process is complete. Load best model\n",
        "  print('Training process has finished. Loading best model...')\n",
        "  mlp = loadModel(best_model)\n",
        "  print(\"Best model loaded.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "100/2192 complete.\n",
            "200/2192 complete.\n",
            "300/2192 complete.\n",
            "400/2192 complete.\n",
            "500/2192 complete.\n",
            "600/2192 complete.\n",
            "700/2192 complete.\n",
            "800/2192 complete.\n",
            "900/2192 complete.\n",
            "1000/2192 complete.\n",
            "1100/2192 complete.\n",
            "1200/2192 complete.\n",
            "1300/2192 complete.\n",
            "1400/2192 complete.\n",
            "1500/2192 complete.\n",
            "1600/2192 complete.\n",
            "1700/2192 complete.\n",
            "1800/2192 complete.\n",
            "1900/2192 complete.\n",
            "2000/2192 complete.\n",
            "2100/2192 complete.\n",
            "100/1730 complete.\n",
            "200/1730 complete.\n",
            "300/1730 complete.\n",
            "400/1730 complete.\n",
            "500/1730 complete.\n",
            "600/1730 complete.\n",
            "700/1730 complete.\n",
            "800/1730 complete.\n",
            "900/1730 complete.\n",
            "1000/1730 complete.\n",
            "1100/1730 complete.\n",
            "1200/1730 complete.\n",
            "1300/1730 complete.\n",
            "1400/1730 complete.\n",
            "1500/1730 complete.\n",
            "1600/1730 complete.\n",
            "1700/1730 complete.\n",
            "100/1509 complete.\n",
            "200/1509 complete.\n",
            "300/1509 complete.\n",
            "400/1509 complete.\n",
            "500/1509 complete.\n",
            "600/1509 complete.\n",
            "700/1509 complete.\n",
            "800/1509 complete.\n",
            "900/1509 complete.\n",
            "1000/1509 complete.\n",
            "1100/1509 complete.\n",
            "1200/1509 complete.\n",
            "1300/1509 complete.\n",
            "1400/1509 complete.\n",
            "1500/1509 complete.\n",
            "100/977 complete.\n",
            "200/977 complete.\n",
            "300/977 complete.\n",
            "400/977 complete.\n",
            "500/977 complete.\n",
            "600/977 complete.\n",
            "700/977 complete.\n",
            "800/977 complete.\n",
            "900/977 complete.\n"
          ]
        }
      ],
      "source": [
        "testing_f1_scores = {}\n",
        "for name, dataset in testing.items():\n",
        "    f1 = evaluate(mlp, dataset)\n",
        "    testing_f1_scores[name] = f1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "SE2 F1 score: 52.73236003225463\n",
            "SE3 F1 score: 54.014249905183895\n",
            "SE13 F1 score: 51.13066121281534\n",
            "SE15 F1 score: 51.01071706280993\n"
          ]
        }
      ],
      "source": [
        "for name, score in testing_f1_scores.items():\n",
        "    print(name, \"F1 score:\", score[1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Sense Predictor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Predictor function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {},
      "outputs": [],
      "source": [
        "def predictSense(sentence, polyseme):\n",
        "    idx = 0\n",
        "    polyseme_text = \"\"\n",
        "    split_sentence = list(map(lambda x: re.sub('\\W+', \"\", x), sentence.lower().split()))\n",
        "    \n",
        "    if isinstance(polyseme, str):\n",
        "        try:\n",
        "            idx = split_sentence.index(polyseme.lower())\n",
        "        except ValueError:\n",
        "            print(\"Given string not found in sentence.\")\n",
        "            return\n",
        "        polyseme_text = polyseme.lower()\n",
        "    elif isinstance(polyseme, int):\n",
        "        idx = polyseme\n",
        "        polyseme_text = split_sentence[polyseme].lower()\n",
        "    \n",
        "    data = {\n",
        "        'sentence': sentence,\n",
        "        'polyseme': polyseme_text,\n",
        "        'idx': idx,\n",
        "        'senses': [lemma for sense in wn.synsets(polyseme_text) for lemma in sense.lemmas()]\n",
        "    }\n",
        "    \n",
        "    try:\n",
        "        feature_vec = getFeatureVec(data).to(dev)\n",
        "        output = nn.functional.softmax(torch.reshape(mlp(feature_vec.to(dev), data), (-1,)))\n",
        "        top_preds = output.topk(5)\n",
        "\n",
        "        print(\"Predicted senses for\", polyseme_text + \":\")\n",
        "        for i in range(5):\n",
        "            prob = top_preds[0][i]\n",
        "            sense_idx = top_preds[1][i]\n",
        "            predicted_sense = data['senses'][sense_idx].synset().definition()\n",
        "            print(\"{:3.2%}\".format(prob),\"\\t\", predicted_sense)\n",
        "    except:\n",
        "        print(\"No sense definitions for polyseme found.\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Try it out! "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted senses for lovely:\n",
            "93.78% \t appealing to the emotions as well as the eye\n",
            "3.02% \t lovable especially in a childlike or naive way\n",
            "0.98% \t a very pretty girl who works as a photographer's model\n",
            "0.95% \t lovable especially in a childlike or naive way\n",
            "0.67% \t lovable especially in a childlike or naive way\n"
          ]
        }
      ],
      "source": [
        "predictSense(\"This is such a lovely day!\", \"lovely\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted senses for pay:\n",
            "69.15% \t give money, usually in exchange for goods or services\n",
            "21.36% \t convey, as of a compliment, regards, attention, etc.; bestow\n",
            "5.62% \t something that remunerates\n",
            "0.96% \t cancel or discharge a debt\n",
            "0.54% \t cancel or discharge a debt\n"
          ]
        }
      ],
      "source": [
        "predictSense(\"Please pay with cash.\", \"pay\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted senses for pay:\n",
            "79.48% \t give money, usually in exchange for goods or services\n",
            "11.34% \t convey, as of a compliment, regards, attention, etc.; bestow\n",
            "6.45% \t something that remunerates\n",
            "0.85% \t cancel or discharge a debt\n",
            "0.33% \t do or give something to somebody in return\n"
          ]
        }
      ],
      "source": [
        "predictSense(\"You better pay attention!\", \"pay\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "x5I3y0a98T_H",
        "BJ0wHYFs8bjo",
        "LGHkvHoK8Pqm",
        "lbFaiOfG85dd"
      ],
      "name": "WSD_with_BERT.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}