bimodal.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "907d9a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchmetrics import AUROC\n",
    "import pandas as pd\n",
    "\n",
    "from transformers import BertTokenizerFast as BertTokenizer, BertModel, AdamW, RobertaModel, RobertaTokenizer, get_linear_schedule_with_warmup\n",
    "from transformers import AutoConfig, AutoTokenizer\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.metrics.functional import accuracy, f1, auroc\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import classification_report, confusion_matrix, accuracy_score\n",
    "\n",
    "import seaborn as sns\n",
    "from pylab import rcParams\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rc\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'\n",
    "\n",
    "RANDOM_SEED = 42\n",
    "\n",
    "sns.set(style='whitegrid', palette='muted', font_scale=1.2)\n",
    "HAPPY_COLORS_PALETTE = [\"#01BEFE\", \"#FFDD00\", \"#FF7D00\", \"#FF006D\", \"#ADFF02\", \"#8F00FF\"]\n",
    "sns.set_palette(sns.color_palette(HAPPY_COLORS_PALETTE))\n",
    "rcParams['figure.figsize'] = 12, 8\n",
    "pl.seed_everything(RANDOM_SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da900a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e0d48f3",
   "metadata": {},
   "source": [
    "### Text embeddings "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5751efa",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/train.pt\"\n",
    "dev_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/dev.pt\"\n",
    "test_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/test.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acb8e881",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_embeddings = torch.load(train_PATH, map_location=device)\n",
    "dev_embeddings = torch.load(dev_PATH, map_location=device)\n",
    "test_embeddings = torch.load(test_PATH, map_location=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6689a9c",
   "metadata": {},
   "source": [
    "#### Last four layers "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4df6ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/train_lfl.pt\"\n",
    "dev_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/dev_lfl.pt\"\n",
    "test_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/Processed_final/embeddings_v2/test_lfl.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0710e70",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_embeddings = torch.load(train_PATH, map_location=device)\n",
    "dev_embeddings = torch.load(dev_PATH, map_location=device)\n",
    "test_embeddings = torch.load(test_PATH, map_location=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac15c46",
   "metadata": {},
   "source": [
    "### Audio embeddings "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a926649",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_wav_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/wav2vec/embeddings_v2/train_wav.pt\"\n",
    "dev_wav_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/wav2vec/embeddings_v2/dev_wav.pt\"\n",
    "test_wav_PATH = \"C:/Data/Sentiment Analysis/MELD/Processed/wav2vec/embeddings_v2/test_wav.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4efa094",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_wav_embeddings = torch.load(train_wav_PATH, map_location=device)\n",
    "dev_wav_embeddings = torch.load(dev_wav_PATH, map_location=device)\n",
    "test_wav_embeddings = torch.load(test_wav_PATH, map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df0ccf7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Two file_ID is equal => your data are mapped \n",
    "# If this is False, need 2 for loop to find the correct fileID\n",
    "train_embeddings[\"fileID\"] == train_wav_embeddings[\"fileID\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a53dbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The correct shape of embedding should be [768] \n",
    "# [1, 768] is because when creating the embeddings, I used unsqueeze to change the dimension\n",
    "# So better to convert it back to original shape for easier used when training a classifier\n",
    "print(train_embeddings[\"embeddings\"][0].shape)\n",
    "print(train_wav_embeddings[\"embeddings\"][0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dd23ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# So, in the for loop, \n",
    "# I'll add the `[0]` to get the embeddings only (remove the unused dimension)\n",
    "train_embeddings[\"embeddings\"][0][0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a10715",
   "metadata": {},
   "source": [
    "## Train "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6fe4635",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of data\n",
    "train_dataset = []\n",
    "\n",
    "# i is the index of arrays\n",
    "for i in range(len(train_wav_embeddings[\"fileID\"])):\n",
    "    # To check the fileIDs are correct\n",
    "    if train_wav_embeddings[\"fileID\"][i] != train_embeddings[\"fileID\"][i]:\n",
    "        print(i, train_wav_embeddings[\"fileID\"][i], train_embeddings[\"fileID\"][i])\n",
    "    \n",
    "    # To check the labels are correct\n",
    "    if train_wav_embeddings[\"labels\"][i] != train_embeddings[\"labels\"][i]:\n",
    "        print(i, train_wav_embeddings[\"labels\"][i], train_embeddings[\"labels\"][i])\n",
    "    \n",
    "    # Using [i] to get the i-th datapoint\n",
    "    # Use \"dict\" to store data object for easier use when training a classifier\n",
    "    train_dataset.append({\n",
    "        \"fileID\": train_embeddings[\"fileID\"][i], \n",
    "        \"label\": train_embeddings[\"labels\"][i], \n",
    "        \"wav_embeddings\": train_wav_embeddings[\"embeddings\"][i][0],\n",
    "        \"text_embeddings\": train_embeddings[\"embeddings\"][i][0],\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b345625b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43a839b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6d45909",
   "metadata": {},
   "source": [
    "## Dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9968d72a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of data\n",
    "dev_dataset = []\n",
    "\n",
    "# i is the index of arrays\n",
    "for i in range(len(dev_wav_embeddings[\"fileID\"])):\n",
    "    # To check the fileIDs are correct\n",
    "    if dev_wav_embeddings[\"fileID\"][i] != dev_embeddings[\"fileID\"][i]:\n",
    "        print(i, dev_wav_embeddings[\"fileID\"][i], dev_embeddings[\"fileID\"][i])\n",
    "    \n",
    "    # To check the labels are correct\n",
    "    if dev_wav_embeddings[\"labels\"][i] != dev_embeddings[\"labels\"][i]:\n",
    "        print(i, dev_wav_embeddings[\"labels\"][i], dev_embeddings[\"labels\"][i])\n",
    "    \n",
    "    # Using [i] to get the i-th datapoint\n",
    "    # Use \"dict\" to store data object for easier use when training a classifier\n",
    "    dev_dataset.append({\n",
    "        \"fileID\": dev_embeddings[\"fileID\"][i], \n",
    "        \"label\": dev_embeddings[\"labels\"][i], \n",
    "        \"wav_embeddings\": dev_wav_embeddings[\"embeddings\"][i][0],\n",
    "        \"text_embeddings\": dev_embeddings[\"embeddings\"][i][0],\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b398793",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dev_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29db330f",
   "metadata": {},
   "source": [
    "## Test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6e375a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of data\n",
    "test_dataset = []\n",
    "\n",
    "# i is the index of arrays\n",
    "for i in range(len(test_wav_embeddings[\"fileID\"])):\n",
    "    # To check the fileIDs are correct\n",
    "    if test_wav_embeddings[\"fileID\"][i] != test_embeddings[\"fileID\"][i]:\n",
    "        print(i, test_wav_embeddings[\"fileID\"][i], test_embeddings[\"fileID\"][i])\n",
    "    \n",
    "    # To check the labels are correct\n",
    "    if test_wav_embeddings[\"labels\"][i] != test_embeddings[\"labels\"][i]:\n",
    "        print(i, test_wav_embeddings[\"labels\"][i], test_embeddings[\"labels\"][i])\n",
    "    \n",
    "    # Using [i] to get the i-th datapoint\n",
    "    # Use \"dict\" to store data object for easier use when training a classifier\n",
    "    test_dataset.append({\n",
    "        \"fileID\": test_embeddings[\"fileID\"][i], \n",
    "        \"label\": test_embeddings[\"labels\"][i], \n",
    "        \"wav_embeddings\": test_wav_embeddings[\"embeddings\"][i][0],\n",
    "        \"text_embeddings\": test_embeddings[\"embeddings\"][i][0],\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d957da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09d2778c",
   "metadata": {},
   "source": [
    "## Encapsulating all data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc177838",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataModule(pl.LightningDataModule):\n",
    "    \n",
    "    def __init__(self, train_dataset, test_dataset, dev_dataset, batch_size=16):\n",
    "        super().__init__()\n",
    "        self.batch_size = batch_size\n",
    "        self.train = train_dataset\n",
    "        self.dev = dev_dataset\n",
    "        self.test = test_dataset\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(\n",
    "          self.train,\n",
    "          batch_size=self.batch_size,\n",
    "          shuffle=True,\n",
    "          num_workers=0\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(\n",
    "          self.dev,\n",
    "          batch_size=self.batch_size,\n",
    "          num_workers=0\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(\n",
    "          self.test,\n",
    "          batch_size=self.batch_size,\n",
    "          num_workers=0\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ff99015",
   "metadata": {},
   "source": [
    "#### Instance for class DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c63f412c",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCHS = 100\n",
    "BATCH_SIZE = 16\n",
    "\n",
    "data_module = DataModule(train_dataset, test_dataset, dev_dataset, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ed02d74",
   "metadata": {},
   "source": [
    "#### Modelling Prep "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee3f6ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_dataset[0][\"text_embeddings\"].size()\n",
    "train_dataset[0][\"text_embeddings\"].size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0484400",
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_dataset[0][\"wav_embeddings\"].size()\n",
    "train_dataset[0][\"wav_embeddings\"].size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67c2582e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#testtt = torch.cat((train_dataset[0][\"text_embeddings\"], train_dataset[0][\"wav_embeddings\"]))\n",
    "testt = torch.cat((train_dataset[0][\"text_embeddings\"], train_dataset[0][\"wav_embeddings\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "041bc230",
   "metadata": {},
   "outputs": [],
   "source": [
    "testt.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57ed578",
   "metadata": {},
   "source": [
    "## Modelling \n",
    "\n",
    "This is a RoBERTa model wherein multi-modal embeddings are fitted for training into the linear classification layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c87d9733",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tagger(pl.LightningModule):\n",
    "    \n",
    "    def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):\n",
    "        super().__init__()\n",
    "        \n",
    "        # pooler output config\n",
    "#        test 1\n",
    "#        self.linear_1 = nn.Linear(1536, 256)\n",
    "    \n",
    "#        test 2\n",
    "#        self.linear_1 = nn.Linear(1536, 636)\n",
    "#        self.linear_2 = nn.Linear(636, 256)\n",
    "    \n",
    "#        test 3\n",
    "#        self.linear_1 = nn.Linear(1536, 936)\n",
    "#        self.linear_2 = nn.Linear(936, 636)\n",
    "#        self.linear_3 = nn.Linear(636, 256)\n",
    "    \n",
    "#        test 4\n",
    "#        self.linear_1 = nn.Linear(1536, 1236)\n",
    "#        self.linear_2 = nn.Linear(1236, 936)\n",
    "#        self.linear_3 = nn.Linear(936, 636)\n",
    "#        self.linear_4 = nn.Linear(636, 256)\n",
    "\n",
    "#        test 5\n",
    "#        self.linear_1 = nn.Linear(1536, 1236)\n",
    "#        self.linear_2 = nn.Linear(1236, 936)\n",
    "#        self.linear_3 = nn.Linear(936, 636)\n",
    "#        self.linear_4 = nn.Linear(636, 336)\n",
    "#        self.linear_5 = nn.Linear(336, 256)\n",
    "        \n",
    "#        last four hidden layers config\n",
    "#        test 1\n",
    "#        self.linear_1 = nn.Linear(3840, 2928)\n",
    "#        self.linear_2 = nn.Linear(2928, 2016)\n",
    "#        self.linear_3 = nn.Linear(2016, 256)\n",
    "\n",
    "#        test 2\n",
    "#        self.linear_1 = nn.Linear(3840, 2928)\n",
    "#        self.linear_2 = nn.Linear(2928, 2016)\n",
    "#        self.linear_3 = nn.Linear(2016, 1104)\n",
    "#        self.linear_4 = nn.Linear(1104, 256)\n",
    "\n",
    "#        test 3\n",
    "        self.linear_1 = nn.Linear(3840, 2928)\n",
    "        self.linear_2 = nn.Linear(2928, 2016)\n",
    "        self.linear_3 = nn.Linear(2016, 1104)\n",
    "        self.linear_4 = nn.Linear(1104, 526)\n",
    "        self.linear_5 = nn.Linear(526, 256)\n",
    "    \n",
    "        self.classifier = nn.Linear(256, n_classes)\n",
    "        self.n_training_steps = n_training_steps\n",
    "        self.n_warmup_steps = n_warmup_steps\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "        \n",
    "\n",
    "    def forward(self, text_embeddings, wav_embeddings, labels=None):\n",
    "\n",
    "        output = torch.cat((text_embeddings, wav_embeddings), 1)\n",
    "        output = self.linear_1(output)\n",
    "        output = self.linear_2(output)\n",
    "        output = self.linear_3(output)\n",
    "        output = self.linear_4(output)\n",
    "        output = self.linear_5(output)\n",
    "        output = self.classifier(output)\n",
    "        output = torch.softmax(output, dim=1)\n",
    "        loss = 0\n",
    "                \n",
    "        if labels is not None:\n",
    "            labels = labels.flatten() ##\n",
    "            loss = self.criterion(output, labels)\n",
    "        return loss, output\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "\n",
    "        text_embeddings = batch[\"text_embeddings\"]\n",
    "        audio_embeddings = batch[\"wav_embeddings\"]\n",
    "        labels = batch[\"label\"]\n",
    "        loss, outputs = self(text_embeddings, audio_embeddings, labels)\n",
    "        self.log(\"train_loss\", loss, prog_bar=True, logger=True)\n",
    "        return {\"loss\": loss, \"predictions\": outputs, \"labels\": labels}\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        \n",
    "        text_embeddings = batch[\"text_embeddings\"]\n",
    "        audio_embeddings = batch[\"wav_embeddings\"]\n",
    "        labels = batch[\"label\"]\n",
    "        loss, outputs = self(text_embeddings, audio_embeddings, labels)\n",
    "        self.log(\"val_loss\", loss, prog_bar=True, logger=True)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "\n",
    "        text_embeddings = batch[\"text_embeddings\"]\n",
    "        audio_embeddings = batch[\"wav_embeddings\"]\n",
    "        labels = batch[\"label\"]\n",
    "        loss, outputs = self(text_embeddings, audio_embeddings, labels)\n",
    "        self.log(\"test_loss\", loss, prog_bar=True, logger=True)\n",
    "        return loss\n",
    "\n",
    "    def training_epoch_end(self, outputs):\n",
    "\n",
    "        labels = []\n",
    "        predictions = []\n",
    "\n",
    "        for output in outputs:\n",
    "            for out_labels in output[\"labels\"].detach().cpu():\n",
    "                labels.append(out_labels)\n",
    "            for out_predictions in output[\"predictions\"].detach().cpu():\n",
    "                predictions.append(out_predictions)\n",
    "        \n",
    "        labels = torch.stack(labels).int()\n",
    "        predictions = torch.stack(predictions)\n",
    "        pred = torch.argmax(predictions, dim=1)\n",
    "        \n",
    "        train_acc = accuracy(pred, labels, num_classes=3)\n",
    "        print(\"Training Accuracy:\", train_acc)\n",
    "        \n",
    "        label = labels.flatten()\n",
    "        auroc = AUROC(num_classes=3)\n",
    "        auroc = auroc(predictions, label)\n",
    "        print(\"AUROC:\", auroc)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "\n",
    "        optimizer = AdamW(self.parameters(), lr=2e-5)\n",
    "\n",
    "        scheduler = get_linear_schedule_with_warmup(\n",
    "          optimizer,\n",
    "          num_warmup_steps=self.n_warmup_steps,\n",
    "          num_training_steps=self.n_training_steps\n",
    "\n",
    "        )\n",
    "\n",
    "        return dict(\n",
    "            optimizer=optimizer,\n",
    "            lr_scheduler=dict(\n",
    "                scheduler=scheduler,\n",
    "                interval='step'\n",
    "            )\n",
    "\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d817be61",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_epoch=len(train_dataset) // BATCH_SIZE\n",
    "total_training_steps = steps_per_epoch * N_EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f88b9f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "#1/5 of the training steps as warm-up\n",
    "warmup_steps = total_training_steps // 5\n",
    "warmup_steps, total_training_steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "187bc6ee",
   "metadata": {},
   "source": [
    "#### Instance for modelling class Tagger "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199f1681",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Tagger(\n",
    "  n_classes=3,\n",
    "  n_warmup_steps=warmup_steps,\n",
    "  n_training_steps=total_training_steps\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc7f4dfb",
   "metadata": {},
   "source": [
    "## Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4539c5b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Checkpointing that saves the best model (based on validation loss)\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    dirpath=\"C:/Users/id301281/NLP/NLU/MELD/Fusion\",\n",
    "    filename=\"best-checkpoint\",\n",
    "    save_top_k=1,\n",
    "    verbose=True,\n",
    "    monitor=\"val_loss\",\n",
    "    mode=\"min\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06f04ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#And early stopping triggers when the loss hasn’t improved for the last 30 epochs\n",
    "early_stopping_callback = EarlyStopping(monitor='val_loss', patience=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f2db8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#starting training\n",
    "trainer = pl.Trainer(\n",
    "    callbacks=[checkpoint_callback, early_stopping_callback],\n",
    "    max_epochs=N_EPOCHS,\n",
    "    gpus=1,\n",
    "    progress_bar_refresh_rate=30\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e59afb5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e68dc7a7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer.fit(model, data_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b713a752",
   "metadata": {},
   "source": [
    "### Storing preds "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dae5c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "trained_model = Tagger.load_from_checkpoint(\n",
    "  trainer.checkpoint_callback.best_model_path,\n",
    "  n_classes=3\n",
    ")\n",
    "\n",
    "trained_model.eval()\n",
    "trained_model.freeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3ca89c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "trained_model = trained_model.to(device)\n",
    "\n",
    "test_dataset = test_dataset\n",
    "\n",
    "test_predictions = { \"fileID_roberta\": [], \"predicted_roberta\": [], \"actual_roberta\": []}\n",
    "\n",
    "\n",
    "for item in tqdm(test_dataset):\n",
    "    _, prediction = trained_model(\n",
    "        item[\"text_embeddings\"].unsqueeze(dim=0).to(device),\n",
    "        item[\"wav_embeddings\"].unsqueeze(dim=0).to(device)\n",
    "    )\n",
    "    \n",
    "    pred2 = prediction.flatten()\n",
    "    pred3 = torch.argmax(pred2).squeeze().tolist()\n",
    "    test_predictions[\"predicted_roberta\"].append(pred3)\n",
    "    test_predictions[\"actual_roberta\"].append(item[\"label\"])\n",
    "    test_predictions[\"fileID_roberta\"].append(item[\"fileID\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "610cf826",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef5578e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ce0922",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = test_predictions[\"predicted_roberta\"]\n",
    "labels = test_predictions[\"actual_roberta\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a53f153",
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_score(preds, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acb590e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "encode_map = {'negative': 0,'neutral': 1,'positive': 2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f319eaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(classification_report(labels, preds, target_names=encode_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "395b4996",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_confusion_matrix(confusion_matrix):\n",
    "  hmap = sns.heatmap(confusion_matrix, annot=True, fmt=\"d\", cmap= \"YlGnBu\")\n",
    "  hmap.yaxis.set_ticklabels(hmap.yaxis.get_ticklabels(), rotation=0, ha='right', fontsize=30)\n",
    "  hmap.xaxis.set_ticklabels(hmap.xaxis.get_ticklabels(), rotation=30, ha='right', fontsize=30)\n",
    "  hmap.tick_params(axis='both', which='major', labelsize=30) # add this line\n",
    "  plt.ylabel('True sentiment', fontsize=30) # modify fontsize here\n",
    "  plt.xlabel('Predicted sentiment', fontsize=30) # modify fontsize here\n",
    "\n",
    "cm = confusion_matrix(labels, preds)\n",
    "sns.set(font_scale=3.0)\n",
    "df_cm = pd.DataFrame(cm, index=encode_map, columns=encode_map)\n",
    "show_confusion_matrix(df_cm)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-GoogleSTT]",
   "language": "python",
   "name": "conda-env-.conda-GoogleSTT-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}