Depression_detectoin.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Depression and PTSD Detection using Deep Learning\n",
    "\n",
    "This notebook contains the code to predict *Depression* and *PTSD* based on transcripts of clinical interviews designed to support the diagnosis of psychological distress conditions such as anxiety, depression, and post-traumatic stress disorder.\n",
    "\n",
    "### Datasets used:\n",
    "* **DIAC-WOZ:** This dataset consists of Wizard-of-Oz interviews, conducted\tby an animated virtual interviewer called Ellie, with Veterans.\n",
    "* **ANONYMITY:** This dataset consists of Wizard-of-Oz interviews, conducted\tby an animated virtual interviewer called Ellie, with Students.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 454490,
     "status": "ok",
     "timestamp": 1619065661923,
     "user": {
      "displayName": "Shubham Nagarkar",
      "photoUrl": "",
      "userId": "10621817079012605482"
     },
     "user_tz": 420
    },
    "id": "Lwd4yowDWvOa",
    "outputId": "c6a7930d-3b66-4478-c985-e40e860d3266"
   },
   "outputs": [],
   "source": [
    "# Important Libraries\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from collections import Counter\n",
    "from torch.nn.utils.rnn import pack_padded_sequence\n",
    "from torchsummary import summary\n",
    "import seaborn as sns\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import classification_report\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Paths and Constant Variables:\n",
    "\n",
    "Following containts paths to both the datasets. In order to train the model we make use of *Glove Embeddings*.\n",
    "\n",
    "**GloVe:** <br>\n",
    "GloVe is an unsupervised learning algorithm for obtaining vector representations for words. We are using the embedding file with 840B tokens, 2.2M vocab, consisting of 300d vectors.\n",
    "\n",
    "*Link:* https://nlp.stanford.edu/projects/glove/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cbNRjmMHWWqm"
   },
   "outputs": [],
   "source": [
    "DIAC_WOZ_TRAIN_PATH = 'new-train-split.csv'\n",
    "DIAC_WOZ_TEST_PATH = 'new-test-split.csv'\n",
    "DIAC_WOZ_VAL_PATH = 'new-dev-split.csv'\n",
    "\n",
    "ANONYMITY_DATASET = 'Extra-data.csv'\n",
    "\n",
    "EMBEDDING_FILE = 'glove.42B.300d.txt'\n",
    "\n",
    "VOCAB_TO_INDEX = {\"\":0, \"UNK\":1}\n",
    "\n",
    "# Maximum words in one sentence (participant)\n",
    "SENTENCE_MAX_LEN = 2204\n",
    "\n",
    "PHQ_MAX_VAL = 23\n",
    "PTSD_MAX_VAL = 85\n",
    "\n",
    "# to store model checkpoints and graphical results of the trained model\n",
    "MODEL_CHECKPOINT = 'model_checkpoints/'\n",
    "VISUALIZATION = 'results/'\n",
    "\n",
    "if not os.path.isdir(MODEL_CHECKPOINT):\n",
    "    os.makedirs(MODEL_CHECKPOINT)\n",
    "    print(\"New folder created for checkpoint!\")\n",
    "\n",
    "if not os.path.isdir(VISUALIZATION):\n",
    "    os.makedirs(VISUALIZATION)\n",
    "    print(\"New folder created for results!\")\n",
    "\n",
    "DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Important utility functions along with their descriptions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6XUO1RBNW86u"
   },
   "outputs": [],
   "source": [
    "def get_diacwoz_data(dataset):\n",
    "    \"\"\"\n",
    "    Returns features along with normalized PHQ and PTSD values\n",
    "    :params: dataset - Train/Test/Validation\n",
    "    \"\"\"\n",
    "    features = dataset.Text\n",
    "    PHQ = dataset.PHQ8_Score\n",
    "    PHQ = PHQ/PHQ_MAX_VAL\n",
    "    PTSD = dataset.PTSD\n",
    "    PTSD = PTSD/PTSD_MAX_VAL\n",
    "  \n",
    "    return features, PHQ, PTSD\n",
    "\n",
    "def get_anonymity_data(dataset):\n",
    "    \"\"\"\n",
    "    Returns features along with normalized PHQ values\n",
    "    :params: dataset - Train/Test\n",
    "    \"\"\"\n",
    "    features = dataset.Text\n",
    "    PHQ = dataset.PHQ8_Score\n",
    "    PHQ = PHQ/PHQ_MAX_VAL\n",
    "  \n",
    "    return features, PHQ\n",
    "\n",
    "def get_embeddings():\n",
    "    \"\"\"\n",
    "    Returns a dictionary containing words as keys and glove word vectors as values\n",
    "    \"\"\"\n",
    "    embeddings_index = {}\n",
    "    with open(EMBEDDING_FILE, encoding='utf8') as f:\n",
    "        for line in f:\n",
    "            values = line.rstrip().rsplit(' ')\n",
    "            word = values[0]\n",
    "            coefs = np.asarray(values[1:], dtype='float32')\n",
    "            embeddings_index[word] = coefs\n",
    "    return embeddings_index\n",
    "\n",
    "def tokenize_text(text):\n",
    "    \"\"\"\n",
    "    Splits text into list of words\n",
    "    :params: text - input paragraph\n",
    "    \"\"\"\n",
    "    tokens = []\n",
    "    for txt in text:\n",
    "        tokens.append(txt.split())\n",
    "    return tokens\n",
    "\n",
    "def encode_sentence(text):\n",
    "    \"\"\"\n",
    "    Converts sentences into numerical representation\n",
    "    :params: text - input paragraph\n",
    "    \"\"\"\n",
    "    tokenized = tokenize_text([text])\n",
    "    encoded = np.zeros(SENTENCE_MAX_LEN, dtype=np.float32)\n",
    "    enc1 = np.array([VOCAB_TO_INDEX.get(word, VOCAB_TO_INDEX[\"UNK\"]) for word in tokenized[0]])\n",
    "    length = min(SENTENCE_MAX_LEN, len(enc1))\n",
    "    encoded[:length] = enc1[:length]\n",
    "    return  encoded, length\n",
    "\n",
    "\n",
    "def get_emb_matrix(glove_embed, word_counts, emb_size = 300):\n",
    "    \"\"\"\n",
    "    returns weight matrix for embeddings along with vocabulary\n",
    "    :params: glove_embed - glove dictionary\n",
    "             word_counts - unique words in dataset\n",
    "             emb_size - size of the word embeddings\n",
    "    \"\"\"\n",
    "\n",
    "    vocab_size = len(word_counts) + 2 \n",
    "    vocab = [\"\", \"UNK\"]\n",
    "    W = np.zeros((vocab_size, emb_size), dtype=\"float32\")\n",
    "    W[0] = np.zeros(emb_size, dtype='float32') # adding a vector for padding\n",
    "    W[1] = np.random.uniform(-0.25, 0.25, emb_size) # adding a vector for unknown words \n",
    "    \n",
    "    i = 2\n",
    "    for word in word_counts:\n",
    "        if word in glove_embed:\n",
    "            W[i] = glove_embed[word]\n",
    "        else:\n",
    "            W[i] = np.random.uniform(-0.25,0.25, emb_size)\n",
    "        VOCAB_TO_INDEX[word] = i\n",
    "        vocab.append(word)\n",
    "        i += 1   \n",
    "    return W, np.array(vocab)\n",
    "\n",
    "def write_stats(model, filename, epochs, lr, task):\n",
    "    \"\"\"\n",
    "    Creates a text file with details of the model and hyper parameters.\n",
    "    \"\"\"\n",
    "    f = open(VISUALIZATION + filename, 'w+')\n",
    "    f.write('Model Architecture: \\n'+str(model))\n",
    "    f.write('\\nEpochs: '+str(epochs))\n",
    "    if task == \"Single\":\n",
    "        f.write('\\nlearning rate: '+str(lr[0]))\n",
    "    else:\n",
    "        f.write('\\nLearning rate PHQ head: '+str(lr[1]))\n",
    "        f.write('\\nLearning rate Body: '+str(lr[0]))\n",
    "        \n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance Statistics Functions:\n",
    "\n",
    "These functions basically plots the training and validation loss curves to depict the training process of the model, outputs a confusion matrix and classification report to evaluate the models performance on the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cm(lab, pred, target_names, title, task, TYPE):\n",
    "    \"\"\"\n",
    "    plots normalized confusion matrix\n",
    "    :params: lab - actual labels\n",
    "             pred - predicted labels\n",
    "             target_names - names of classes\n",
    "             title - title of plot\n",
    "             task - name of the task\n",
    "             TYPE - flag for Binary or tertiary classification of PHQ values\n",
    "    \"\"\"\n",
    "    cm = confusion_matrix(lab, pred)\n",
    "    print(cm)\n",
    "    # Normalise\n",
    "    cmn = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]\n",
    "    fig, ax = plt.subplots(figsize=(6,4))\n",
    "    sns.heatmap(cmn, annot=True, fmt='.2f', xticklabels=target_names, yticklabels=target_names, cmap = \"YlGnBu\")\n",
    "    plt.ylabel('Actual')\n",
    "    plt.xlabel('Predicted')\n",
    "    plt.title(title)\n",
    "    if TYPE == \"\":\n",
    "        plt.savefig(VISUALIZATION+\"{}_confusion_matrix.png\".format(task))\n",
    "    else:\n",
    "        plt.savefig(VISUALIZATION+\"{}_{}_confusion_matrix.png\".format(task,TYPE))\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "    \n",
    "def draw_training_curves(train_losses, test_losses, curve_name, title, epoch):\n",
    "    plt.clf()\n",
    "    \n",
    "    plt.xlim([0,epoch])\n",
    "    plt.plot(train_losses, label='Training Loss')\n",
    "    plt.plot(test_losses, label='Testing Loss')\n",
    "    plt.ylabel(curve_name)\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.title(title)\n",
    "    plt.legend(frameon=False)\n",
    "    plt.savefig(VISUALIZATION +\"{}_curve.png\".format(curve_name))\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "    \n",
    "def print_accuracy(true_label, predicted_label):   \n",
    "    \"\"\"\n",
    "    prints accuracy of the task\n",
    "    :params: true_label - actual labels of the dataset\n",
    "             predicted_label - predictions from model\n",
    "    \"\"\"\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for x,y in zip(true_label, predicted_label):\n",
    "        if x == y:\n",
    "            correct+=1\n",
    "        total+=1\n",
    "        \n",
    "    print(\"Accuracy: \", correct/total)\n",
    "    \n",
    "def get_classification_report(actual, predicted, labels, target_names, report_name):\n",
    "    \"\"\"\n",
    "    Creates a csv file consisting of f1-score, recall and precision\n",
    "    \"\"\"\n",
    "    report = pd.DataFrame(classification_report(actual, predicted, labels= labels, target_names=target_names, output_dict = True))\n",
    "    report.to_csv(VISUALIZATION+report_name)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions to classify range:\n",
    "\n",
    "The following cell contains the code to classify the output of the PHQ and PTSD values in the following format.\n",
    "\n",
    "**PHQ**:<br>\n",
    "* If split variable is set to **Two** then we have a **Binary classification**. <br>\n",
    "    * If the values are in the range 0 to 10 (un-normalized) or 0 to 0.434 (normalized) then the person has no depression and the label assigned is **No**<br>\n",
    "    * If the values are in the range 11 to 23 (un-normalized) or 0.478 to 1 (normalized) then the person has depression and the label assigned is **Yes** <br>\n",
    "* If split variable is set to **Three** then we have a **Tertiary classification**. <br>\n",
    "    * If the values are in the range 0 to 10 (un-normalized) or 0 to 0.434 (normalized) then the person has no depression and the label assigned is **No**<br>\n",
    "    * If the values are in the range 11 to 14 (un-normalized) or 0.478 to 0.608 (normalized) then the person has moderate depression and the label assigned is **Moderate** <br> \n",
    "    * If the values are in the range 14 to 23 (un-normalized) or 0.702 to 1 (normalized) then the person has severe depression and the label assigned is **Severe** <br> \n",
    "    \n",
    "\n",
    "**PTSD**:<br>\n",
    "* We have a **Binary classification**. <br>\n",
    "    * If the values are in the range 0 to 33 (un-normalized) or 0 to 0.39 (normalized) then the person has no PTSD and the label assigned is **No**<br>\n",
    "    * If the values are in the range 33 to 85 (un-normalized) or 0.39 to 1 (normalized) then the person has PTSD and the label assigned is **Yes** <br>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_phq(label, split):\n",
    "    \"\"\"\n",
    "    converts probabilites into classes\n",
    "    :params: label - class probabilities (actual/predicted)\n",
    "             split - 2 way split (yes/no) or 3 way split (no/moderate/severe)\n",
    "    \"\"\"\n",
    "    final_label = []\n",
    "    \n",
    "    if split == \"Three\":\n",
    "        for val in label:\n",
    "            if val <=  10:# 0.434:\n",
    "                final_label.append(\"No\")\n",
    "            elif val > 10 and val<15:#>0.4340 and val <= 0.608:\n",
    "                final_label.append(\"Moderate\")\n",
    "            else:\n",
    "                final_label.append(\"Severe\")\n",
    "    else:\n",
    "        for val in label:\n",
    "            if val <= 10:#0.434:\n",
    "                final_label.append(\"No\")\n",
    "            else:\n",
    "                final_label.append(\"Yes\")         \n",
    "    \n",
    "    return final_label\n",
    "\n",
    "def classify_ptsd(label):\n",
    "    \"\"\"\n",
    "    converts probabilites into classes\n",
    "    :params: label - class probabilities (actual/predicted)\n",
    "    \"\"\"\n",
    "    final_label = []\n",
    "    \n",
    "    for val in label:\n",
    "        if val < 34:#0.39:\n",
    "            final_label.append(\"No\")\n",
    "        else:\n",
    "            final_label.append(\"Yes\")\n",
    "    \n",
    "    return final_label\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Early stopping helper class:\n",
    "\n",
    "This class helps to keep a track of the validation loss. Whenever model fails to improve its performance, this class stops the training and returns the best version of the model and avoids overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gbEzH-CqF_xu"
   },
   "outputs": [],
   "source": [
    "class EarlyStopping:\n",
    "    \"\"\"\n",
    "    Early stops the training if validation loss doesn't improve after a given patience.\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self, patience=8, verbose=False, delta=0.001, path= 'early_stopping_model.pth'):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            patience (int): How long to wait after last time validation loss improved.\n",
    "                            Default: 10\n",
    "            verbose (bool): If True, prints a message for each validation loss improvement. \n",
    "                            Default: False\n",
    "            delta (float): Minimum change in the monitored quantity to qualify as an improvement.\n",
    "                            Default: 0\n",
    "            path (str): Path for the checkpoint to be saved to.\n",
    "                            Default: 'early_stopping_vgg16model.pth'   \n",
    "        \"\"\"\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.counter = 0\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.val_loss_min = np.Inf\n",
    "        self.delta = delta\n",
    "        self.path = path\n",
    "    \n",
    "    def __call__(self, val_loss, model):\n",
    "        \n",
    "        score = -val_loss\n",
    "        \n",
    "        if self.best_score is None:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "            \n",
    "        elif score < self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            \n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "                \n",
    "        else:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "            self.counter = 0   \n",
    "    \n",
    "    def save_checkpoint(self, val_loss, model):\n",
    "        \"\"\"\n",
    "        saves the current best version of the model if there is decrease in validation loss\n",
    "        \"\"\"\n",
    "        torch.save(model.state_dict(), self.path)\n",
    "        self.vall_loss_min = val_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DIACWOZ Dataset loader class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bGjKYduu-tpV"
   },
   "outputs": [],
   "source": [
    "class Diacwoz_Dataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    Pytorch implementation of Dataset class to return encoding and their labels\n",
    "    \"\"\"\n",
    "    def __init__(self, encodings, PHQ, PTSD):\n",
    "        self.encodings = encodings\n",
    "        self.PHQ = PHQ\n",
    "        self.PTSD = PTSD\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        encoding = torch.tensor(self.encodings[idx][0])\n",
    "        phq_score = torch.tensor(self.PHQ[idx])\n",
    "        ptsd_score = torch.tensor(self.PTSD[idx])\n",
    "        length = self.encodings[idx][1]\n",
    "        \n",
    "        return encoding, length, phq_score, ptsd_score\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.PHQ)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ANONYMITY Dataset loader class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Anonymity_Dataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    Pytorch implementation of Dataset class to return encoding and their labels\n",
    "    \"\"\"\n",
    "    def __init__(self, encodings, PHQ):\n",
    "        self.encodings = encodings\n",
    "        self.PHQ = PHQ\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        encoding = torch.tensor(self.encodings[idx][0])\n",
    "        phq_score = torch.tensor(self.PHQ[idx])\n",
    "        length = self.encodings[idx][1]\n",
    "        \n",
    "        return encoding, length, phq_score\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.PHQ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bi-LSTM model architecture for Multi-task Learning:\n",
    "\n",
    "This architecture consists of:\n",
    "* 3-layered Bidirectional Long Short Term Memory\n",
    "* Multi-layered Perceptron as a Feature extractor\n",
    "* PHQ head to predict PHQ values\n",
    "* PTSD head to predict PTSD values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3JGwvuDWQPAv"
   },
   "outputs": [],
   "source": [
    "class Multi_Task_Model(torch.nn.Module) :\n",
    "    \"\"\"\n",
    "    Bi-LSTM Multi-task model architecture\n",
    "    \"\"\"\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim, glove_weights) :\n",
    "        super().__init__()\n",
    "\n",
    "        self.dropout = nn.Dropout(0.3)\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.embeddings.weight.data.copy_(torch.from_numpy(glove_weights))\n",
    "        self.embeddings.weight.requires_grad = False\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, num_layers=3, dropout=0.3, bidirectional=True)\n",
    "\n",
    "        \n",
    "        self.MLP = nn.Sequential(\n",
    "            nn.Linear(hidden_dim,128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3)\n",
    "        )\n",
    "\n",
    "        self.ptsd_head = nn.Sequential(\n",
    "            nn.Linear(128,84),\n",
    "            nn.BatchNorm1d(84),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(84, 1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.phq_head = nn.Sequential(\n",
    "            nn.Linear(128,32),\n",
    "            nn.BatchNorm1d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.Linear(32, 1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, sentence, sentence_length):\n",
    "        sentence_embed = self.embeddings(sentence)\n",
    "        sentence_embed = self.dropout(sentence_embed)\n",
    "\n",
    "        sentence_pack = pack_padded_sequence(sentence_embed, sentence_length, batch_first=True, enforce_sorted=False)\n",
    "        out_pack, (ht, ct) = self.lstm(sentence_pack)\n",
    "        features = self.MLP(ht[-1])\n",
    "        ptsd = self.ptsd_head(features)\n",
    "        phq = self.phq_head(features)\n",
    "\n",
    "        return torch.squeeze(phq, 1), torch.squeeze(ptsd, 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bi-LSTM model architecture for Single-task Learning:\n",
    "\n",
    "This architecture consists of:\n",
    "* 2-layered Bidirectional Long Short Term Memory\n",
    "* Multi-layered Perceptron to predict task output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Single_Task_Model(torch.nn.Module) :\n",
    "    \"\"\"\n",
    "    Bi-LSTM architecture for Single task\n",
    "    \"\"\"\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim, glove_weights) :\n",
    "        super().__init__()\n",
    "\n",
    "        self.dropout = nn.Dropout(0.3)\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.embeddings.weight.data.copy_(torch.from_numpy(glove_weights))\n",
    "        self.embeddings.weight.requires_grad = False\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, num_layers=2, dropout=0.4, bidirectional=True)\n",
    "\n",
    "\n",
    "        self.MLP = nn.Sequential(\n",
    "            nn.Linear(hidden_dim,32),\n",
    "            nn.BatchNorm1d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.Linear(32, 1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "    def forward(self, sentence, sentence_length):\n",
    "        sentence_embed = self.embeddings(sentence)\n",
    "        sentence_embed = self.dropout(sentence_embed)\n",
    "\n",
    "        sentence_pack = pack_padded_sequence(sentence_embed, sentence_length, batch_first=True, enforce_sorted=False)\n",
    "        out_pack, (ht, ct) = self.lstm(sentence_pack)\n",
    "        task_op = self.MLP(ht[-1])\n",
    "        return torch.squeeze(task_op, 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Training functions for Multi-task and Single-task for DIAC-WOZ:\n",
    "\n",
    "The following 2 cells contains the code to train the multi-task and single-task models respectively. Each of these functions prints the training and validation/test loss per epoch and checkpoints the best model as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iRi_SHtNS4P7"
   },
   "outputs": [],
   "source": [
    "def multi_task_train(model, optimizer, criterion_phq, criterion_ptsd, train_loader, val_loader, epochs):\n",
    "    \"\"\"\n",
    "    returns trained model along with lossess per epoch for multitask architecture\n",
    "    :params: model - Bi-LSTM model\n",
    "             optimizer - Adam optimizer\n",
    "             criterion_phq, criterion_ptsd - MSE Loss functions\n",
    "             train, val loaders - input data and labels\n",
    "             epochs - number of epochs to train the model\n",
    "    \"\"\"\n",
    "    epoch_train_phq_loss = []\n",
    "    epoch_train_ptsd_loss = []\n",
    "    epoch_train_total_loss = []\n",
    "    epoch_train_ptsd_acc = []\n",
    "\n",
    "    epoch_val_phq_loss = []\n",
    "    epoch_val_ptsd_loss = []\n",
    "    epoch_val_total_loss = []\n",
    "    epoch_val_ptsd_acc = []\n",
    "\n",
    "    epoch_val_loss = []\n",
    "    losses = {'train':{'total':[], 'phq':[], 'ptsd':[]}, 'val':{'total':[], 'phq':[], 'ptsd':[]}}\n",
    "\n",
    "    print(\"Training started...\\n\")\n",
    "    model.to(DEVICE)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.2, last_epoch=-1, verbose=False)\n",
    "    early_stop = EarlyStopping(patience=10, path= MODEL_CHECKPOINT+\"multitask_early_stopping_bilstm_model.pth\")\n",
    "    used_early_stopping = False\n",
    "    epoch_val = 0\n",
    "    for epoch in range(epochs):\n",
    "    \n",
    "        print(\"Epoch : \", epoch+1)\n",
    "        epoch_val = epoch\n",
    "        model.train()\n",
    "        for sentence, length, phq_score, ptsd_score in train_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "            ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            phq_op, ptsd_op = model(sentence, length)\n",
    "\n",
    "            loss_phq = criterion_phq(phq_op, phq_score)\n",
    "            loss_ptsd = criterion_ptsd(ptsd_op, ptsd_score)\n",
    "            \n",
    "            loss = loss_phq + loss_ptsd\n",
    "            loss.backward()\n",
    "\n",
    "\n",
    "            epoch_train_phq_loss.append(loss_phq.item())\n",
    "            epoch_train_ptsd_loss.append(loss_ptsd.item())\n",
    "            epoch_train_total_loss.append(loss.item())\n",
    "            optimizer.step()\n",
    "      \n",
    "        train_total_loss = np.average(epoch_train_total_loss)\n",
    "        train_phq_loss = np.average(epoch_train_phq_loss)\n",
    "        train_ptsd_loss = np.average(epoch_train_ptsd_loss)\n",
    "\n",
    "\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for sentence, length, phq_score, ptsd_score in val_loader:\n",
    "                sentence = sentence.to(DEVICE).long()\n",
    "                phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "                ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "\n",
    "                phq_op, ptsd_op = model(sentence,length)\n",
    "\n",
    "                loss_phq = criterion_phq(phq_op, phq_score)\n",
    "                loss_ptsd = criterion_ptsd(ptsd_op, ptsd_score)\n",
    "                loss = loss_phq + loss_ptsd\n",
    "\n",
    "\n",
    "                epoch_val_phq_loss.append(loss_phq.item())\n",
    "                epoch_val_ptsd_loss.append(loss_ptsd.item())\n",
    "                epoch_val_total_loss.append(loss.item())\n",
    "\n",
    "      \n",
    "        val_total_loss = np.average(epoch_val_total_loss)\n",
    "        val_phq_loss = np.average(epoch_val_phq_loss)\n",
    "        val_ptsd_loss = np.average(epoch_val_ptsd_loss)\n",
    "\n",
    "\n",
    "        \n",
    "        early_stop(val_total_loss, model)\n",
    "        scheduler.step()\n",
    "        \n",
    "        if early_stop.early_stop:\n",
    "            print(\"Early stopping\")\n",
    "            used_early_stopping  = True\n",
    "            break\n",
    "\n",
    "        print(\"Train total loss: {0:.3f}, Train PHQ loss: {1:.3f}, Train PTSD loss: {2:.3f} \".format(train_total_loss, train_phq_loss, train_ptsd_loss))\n",
    "        print(\"Val total loss: {0:.3f}, Val PHQ loss: {1:.3f}, Val PTSD loss: {2:.3f} \".format(val_total_loss, val_phq_loss, val_ptsd_loss))\n",
    "        print('----------------------------------------------------------------------------------------------')\n",
    "\n",
    "        losses['train']['total'].append(train_total_loss) \n",
    "        losses['train']['phq'].append(train_phq_loss) \n",
    "        losses['train']['ptsd'].append(train_ptsd_loss) \n",
    "\n",
    "        losses['val']['total'].append(val_total_loss) \n",
    "        losses['val']['phq'].append(val_phq_loss) \n",
    "        losses['val']['ptsd'].append(val_ptsd_loss) \n",
    "\n",
    "    return model, losses, epoch_val\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_task_train(task, model, optimizer, criterion, train_loader, val_loader, epochs):\n",
    "    \"\"\"\n",
    "    returns trained model along with lossess per epoch for single task\n",
    "    :params: model - Bi-LSTM model\n",
    "             optimizer - Adam optimizer\n",
    "             criterion - MSE Loss function\n",
    "             train, val loaders - input data and labels\n",
    "             epochs - number of epochs to train the model\n",
    "    \"\"\"\n",
    "    epoch_train_loss = []\n",
    "    epoch_val_loss = []\n",
    "    losses = {'train':[], 'val':[]}\n",
    "    used_early_stopping = False\n",
    "    epoch_val = 0\n",
    "    print(\"Training started...\\n\")\n",
    "\n",
    "    model.to(DEVICE)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1, last_epoch=-1, verbose=False)\n",
    "    early_stop = EarlyStopping(patience=15, path=MODEL_CHECKPOINT+'singletask_{}_early_stopping_bilstm_model.pth'.format(task))\n",
    "\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "        print(\"Epoch : \", epoch+1)\n",
    "        epoch_val = epoch\n",
    "        model.train()\n",
    "        for sentence, length, phq_score, ptsd_score in train_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "            ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "            if task == 'PHQ':\n",
    "                label = phq_score\n",
    "            else:\n",
    "                label = ptsd_score\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            task_op = model(sentence, length)\n",
    "\n",
    "            loss = criterion(task_op, label)\n",
    "            loss.backward()\n",
    "            \n",
    "            epoch_train_loss.append(loss.item())\n",
    "            optimizer.step()\n",
    "      \n",
    "        train_loss = np.average(epoch_train_loss)\n",
    "\n",
    "        model.eval()\n",
    "        total = 0\n",
    "        correct =  0\n",
    "        with torch.no_grad():\n",
    "            for sentence, length, phq_score, ptsd_score in val_loader:\n",
    "                sentence = sentence.to(DEVICE).long()\n",
    "                phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "                ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "                if task == 'PHQ':\n",
    "                    label = phq_score\n",
    "                else: \n",
    "                    label = ptsd_score\n",
    "        \n",
    "                task_op = model(sentence, length)\n",
    "\n",
    "                loss = criterion(task_op, label)\n",
    "        \n",
    "                epoch_val_loss.append(loss.item())\n",
    "\n",
    "        val_loss = np.average(epoch_val_loss)\n",
    "\n",
    "        scheduler.step()\n",
    "        early_stop(val_loss, model)\n",
    "\n",
    "\n",
    "        print(\"Train loss: {0:.3f} Val loss: {1:.3f}\".format(train_loss, val_loss))\n",
    "        print('--------------------------------------------------------------------------------------')\n",
    "\n",
    "        if early_stop.early_stop:\n",
    "            print(\"Early stopping\")\n",
    "            used_early_stopping  = True\n",
    "            break\n",
    "\n",
    "        losses['train'].append(train_loss) \n",
    "        losses['val'].append(val_loss) \n",
    "\n",
    "    return model, losses, epoch_val\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Training function for Single-task for ANONYMITY:\n",
    "\n",
    "The following cell contains code to train single-task model on Anonymity dataset. \n",
    "This function prints the training and test loss per epoch and checkpoints the best model as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Anonymity_train(model, optimizer, criterion, train_loader, val_loader, epochs):\n",
    "    \"\"\"\n",
    "    returns trained model along with lossess and accuracies per epoch for single task\n",
    "    :params: model - Bi-LSTM model\n",
    "             optimizer - Adam optimizer\n",
    "             criterion_phq, criterion_ptsd - MSE Loss functions\n",
    "             train, val loaders - input data and labels\n",
    "             epochs - number of epochs to train the model\n",
    "    \"\"\"\n",
    "    epoch_train_loss = []\n",
    "    epoch_val_loss = []\n",
    "    losses = {'train':[], 'val':[]}\n",
    "    used_early_stopping = False\n",
    "    epoch_val = 0\n",
    "    print(\"Training started...\\n\")\n",
    "\n",
    "    model.to(DEVICE)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1, last_epoch=-1, verbose=True)\n",
    "    early_stop = EarlyStopping(patience=16, path='Anonymity_early_stopping_bilstm_model.pth')\n",
    "\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        epoch_val = epoch\n",
    "        print(\"Epoch : \", epoch+1)\n",
    "        model.train()\n",
    "        for sentence, length, phq_score in train_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            label = phq_score.to(DEVICE).type(torch.float32)\n",
    "           \n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            task_op = model(sentence, length)\n",
    "\n",
    "            loss = criterion(task_op, label)\n",
    "            loss.backward()\n",
    "            \n",
    "            epoch_train_loss.append(loss.item())\n",
    "            optimizer.step()\n",
    "      \n",
    "        train_loss = np.average(epoch_train_loss)\n",
    "\n",
    "        model.eval()\n",
    "        total = 0\n",
    "        correct =  0\n",
    "        with torch.no_grad():\n",
    "            for sentence, length, phq_score in val_loader:\n",
    "                sentence = sentence.to(DEVICE).long()\n",
    "                label = phq_score.to(DEVICE).type(torch.float32)\n",
    "               \n",
    "        \n",
    "                task_op = model(sentence, length)\n",
    "\n",
    "                loss = criterion(task_op, label)\n",
    "        \n",
    "                epoch_val_loss.append(loss.item())\n",
    "\n",
    "        val_loss = np.average(epoch_val_loss)\n",
    "\n",
    "        scheduler.step()\n",
    "        early_stop(val_loss, model)\n",
    "\n",
    "\n",
    "        print(\"Train loss: {0:.3f} Val loss: {1:.3f}\".format(train_loss, val_loss))\n",
    "        print('-------------------------------------------------------------------------------')\n",
    "\n",
    "        if early_stop.early_stop:\n",
    "            print(\"Early stopping\")\n",
    "            used_early_stopping  = True\n",
    "            break\n",
    "\n",
    "        losses['train'].append(train_loss) \n",
    "        losses['val'].append(val_loss) \n",
    "\n",
    "    return model, losses, epoch_val\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Evaluation functions:\n",
    "\n",
    "The following cell contains the code to evaluate multi-task and single-task models for DIAC-WOZ and PHQ model for Anynomity. These functions return model predictions on the test dataset along with their true labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_multi_task(model, test_loader):\n",
    "    \"\"\"\n",
    "    returns predictions on test dataset with the best trained model for single task\n",
    "    :params: task - PHQ/PTSD\n",
    "             model - trained model\n",
    "             test_loader - test dataset\n",
    "    \"\"\"\n",
    "    output_phq = []\n",
    "    actual_phq = []\n",
    "    \n",
    "    output_ptsd = []\n",
    "    actual_ptsd = []\n",
    "    \n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for sentence, length, phq_score, ptsd_score in test_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "            ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "            \n",
    "            phqop, ptsdop = model(sentence, length)\n",
    "            \n",
    "            \n",
    "            output_phq.extend(phqop.cpu().tolist())\n",
    "            actual_phq.extend(phq_score.cpu().tolist())\n",
    "            \n",
    "            output_ptsd.extend(ptsdop.cpu().tolist())\n",
    "            actual_ptsd.extend(ptsd_score.cpu().tolist())\n",
    "    \n",
    "\n",
    "    return actual_phq, actual_ptsd, output_phq, output_ptsd\n",
    "\n",
    "\n",
    "def evaluate_single_task(task, model, test_loader):\n",
    "    \"\"\"\n",
    "    returns predictions on test dataset with the best trained model for single task\n",
    "    :params: task - PHQ/PTSD\n",
    "             model - trained model\n",
    "             test_loader - test dataset\n",
    "    \"\"\"\n",
    "    output = []\n",
    "    actual = []\n",
    "    model.eval()\n",
    " \n",
    "    with torch.no_grad():\n",
    "        for sentence, length, phq_score, ptsd_score in test_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            phq_score = phq_score.to(DEVICE).type(torch.float32)\n",
    "            ptsd_score = ptsd_score.to(DEVICE).type(torch.float32)\n",
    "            if task == 'PHQ':\n",
    "                label = phq_score\n",
    "            else: \n",
    "                label = ptsd_score\n",
    "            \n",
    "            task_op = model(sentence, length)\n",
    "\n",
    "            output.extend(task_op.cpu().tolist())\n",
    "            actual.extend(label.cpu().tolist())\n",
    "    \n",
    "    return actual, output\n",
    "\n",
    "def Anonymity_evaluate(model, test_loader):\n",
    "    \"\"\"\n",
    "    returns predictions on test dataset with the best trained model for PHQ task\n",
    "    :params: \n",
    "             model - trained model\n",
    "             test_loader - test dataset\n",
    "    \"\"\"\n",
    "    output = []\n",
    "    actual = []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for sentence, length, phq_score in test_loader:\n",
    "            sentence = sentence.to(DEVICE).long()\n",
    "            label = phq_score.to(DEVICE).type(torch.float32)\n",
    "            \n",
    "            op = model(sentence, length)\n",
    "            \n",
    "            output.extend(op.cpu().tolist())\n",
    "            actual.extend(label.cpu().tolist())\n",
    "\n",
    "    return actual, output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Section:\n",
    "\n",
    "This section of the notebook has ready to execute cells. You will be able to see all the outputs of the above functions here.\n",
    "\n",
    "\n",
    "### 1. Preprocess Datasets and Generate Embeddings:\n",
    "\n",
    "This cell is the main function to read the data by utilizing the above mentioned functions.\n",
    "\n",
    "**Steps:**\n",
    "* Read DIAC-WOZ and ANONYMITY datasets\n",
    "    * Shuffle the ANONYMITY dataset\n",
    "* Split ANONYMITY into Train and Test set\n",
    "    * Split 70% as Train data and 30% as Test data\n",
    "* Generate Features and Labels for each dataset\n",
    "    * Converts the feature and labels into a list\n",
    "* Get Embeddings from GloVe file\n",
    "* Tokenize the datasets\n",
    "* Maintain a Dictionary of Vocabulary in DIAC-WOZ\n",
    "* Update VOCAB_TO_INDEX Dictionary with the vocab generated\n",
    "* Encode the tokenized datasets into their vector representations\n",
    "    * Uses vectors from GloVe file and replaces Unknown words with random numbers\n",
    "* Generate Embeddings Matrix to feed to the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 137108,
     "status": "ok",
     "timestamp": 1619061264697,
     "user": {
      "displayName": "Shubham Nagarkar",
      "photoUrl": "",
      "userId": "10621817079012605482"
     },
     "user_tz": 420
    },
    "id": "w7IaBbgTmBrE",
    "outputId": "d9e00e12-a6ce-4f40-eaaa-f5346b6a4862",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Step 1: Read DIAC_WOZ and ANONYMITY dataset\n",
    "\n",
    "DW_train_dataset = pd.read_csv(DIAC_WOZ_TRAIN_PATH)\n",
    "DW_test_dataset = pd.read_csv(DIAC_WOZ_TEST_PATH)\n",
    "DW_val_dataset = pd.read_csv(DIAC_WOZ_VAL_PATH) \n",
    "\n",
    "AN_dataset = pd.read_csv(ANONYMITY_DATASET)\n",
    "AN_dataset = AN_dataset.sample(frac=1)\n",
    "\n",
    "\n",
    "# Step 2: Split ANONYMITY into train and test\n",
    "\n",
    "split = int(len(AN_dataset)*0.7)\n",
    "AN_train_dataset = AN_dataset.iloc[:split,:]\n",
    "AN_test_dataset = AN_dataset.iloc[split:,:]\n",
    "\n",
    "\n",
    "print(\"****** Anonymity dataset ******\")\n",
    "print(\"Train size      : \",len(AN_train_dataset))\n",
    "print(\"Test size       : \",len(AN_test_dataset))\n",
    "\n",
    "\n",
    "print(\"\\n****** Diac-Woz dataset ******\")\n",
    "print(\"Train size      : \",len(DW_train_dataset))\n",
    "print(\"Validation size : \",len(DW_val_dataset))\n",
    "print(\"Test size       : \",len(DW_test_dataset))\n",
    "\n",
    "\n",
    "# Step 3: Generate Features and their labels\n",
    "DW_X_train, DW_train_PHQ, DW_train_PTSD = get_diacwoz_data(DW_train_dataset)\n",
    "DW_X_val, DW_val_PHQ, DW_val_PTSD = get_diacwoz_data(DW_val_dataset)\n",
    "DW_X_test, DW_test_PHQ, DW_test_PTSD = get_diacwoz_data(DW_test_dataset)\n",
    "\n",
    "DW_X_train = list(DW_X_train.values)\n",
    "DW_X_val = list(DW_X_val.values)\n",
    "DW_X_test = list(DW_X_test.values)\n",
    "\n",
    "AN_X_train, AN_train_PHQ = get_anonymity_data(AN_train_dataset)\n",
    "AN_X_test, AN_test_PHQ = get_anonymity_data(AN_test_dataset)\n",
    "\n",
    "AN_X_train = list(AN_X_train.values)\n",
    "AN_X_test = list(AN_X_test.values)\n",
    "AN_train_PHQ = AN_train_PHQ.values\n",
    "AN_test_PHQ = AN_test_PHQ.values\n",
    "\n",
    "\n",
    "#Step 4: Generate GloVe Embeddings\n",
    "embeddings_dict = get_embeddings()\n",
    "\n",
    "\n",
    "# Step 5: Tokenize the dataset\n",
    "DW_train_tokens = tokenize_text(DW_X_train)\n",
    "DW_val_tokens = tokenize_text(DW_X_val)\n",
    "DW_test_tokens = tokenize_text(DW_X_test)\n",
    "\n",
    "AN_train_tokens = tokenize_text(AN_X_train)\n",
    "AN_test_tokens = tokenize_text(AN_X_test)\n",
    "\n",
    "\n",
    "# Step 6: Maintain the counts of unique words in Diac-Woz\n",
    "counts = Counter()\n",
    "for tok in DW_train_tokens:\n",
    "    counts.update(tok)\n",
    "for tok in DW_val_tokens:  \n",
    "    counts.update(tok)\n",
    "for tok in DW_test_tokens:  \n",
    "    counts.update(tok)\n",
    "\n",
    "    \n",
    "# Step 7: Update the VOCAB_TO_INDEX dictionary with the vocab from Diac-Woz\n",
    "words = [\"\", \"UNK\"]\n",
    "for word in counts:\n",
    "    VOCAB_TO_INDEX[word] = len(words)\n",
    "    words.append(word)\n",
    "\n",
    "\n",
    "# Step 8: Encode the dataset into vector representations obtained from GloVe Embeddings\n",
    "DW_train_encoded = []\n",
    "DW_val_encoded = []\n",
    "DW_test_encoded = []\n",
    "\n",
    "AN_train_encoded = []\n",
    "AN_test_encoded = []\n",
    "\n",
    "for text in DW_X_train:\n",
    "    DW_train_encoded.append(encode_sentence(text))\n",
    "for text in DW_X_val:\n",
    "    DW_val_encoded.append(encode_sentence(text))\n",
    "for text in DW_X_test:\n",
    "    DW_test_encoded.append(encode_sentence(text))\n",
    "\n",
    "for text in AN_X_train:\n",
    "    AN_train_encoded.append(encode_sentence(text))\n",
    "for text in AN_X_test:\n",
    "    AN_test_encoded.append(encode_sentence(text))\n",
    "    \n",
    "    \n",
    "# Step 9: Generate Embedding matrix\n",
    "embed_matrix, vocab= get_emb_matrix(embeddings_dict, counts)\n",
    "\n",
    "print(\"\\n****** Embedding Matrix Details ******\")\n",
    "print(\"Matrix Shape             : \", embed_matrix.shape)\n",
    "print(\"Vocabulary Size          : \", vocab.shape)\n",
    "print(\"Vocabulary to Index Size : \", len(VOCAB_TO_INDEX))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute Multi-task Training For DIAC-WoZ:\n",
    "\n",
    "This following cell contains code required for training the Multi-task model.\n",
    "\n",
    "**Flow:**\n",
    "* Adjust hyper parameters\n",
    "* Define model, optimizer and loss function\n",
    "* Create Dataloaders\n",
    "* Train the model and save it\n",
    "* Plot and save loss curves\n",
    "* Write model and hyper parameter details to a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vkMJXcbpgBdI",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Non-tuneable parameters\n",
    "vocab_size = 8360\n",
    "embedding_dim = 300\n",
    "\n",
    "# Tuneable parameters\n",
    "hidden_dim = 128\n",
    "lr = 1e-3\n",
    "lr_phq = 1e-4\n",
    "learning_rates = [lr, lr_phq]\n",
    "epochs = 500\n",
    "batch_size = 64\n",
    "\n",
    "# define model, optimizer and loss function\n",
    "model = Multi_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix)\n",
    "\n",
    "optimizer = torch.optim.Adam([{'params':model.lstm.parameters()},\n",
    "                            {'params':model.MLP.parameters()},\n",
    "                            {'params':model.ptsd_head.parameters()},\n",
    "                            {'params':model.phq_head.parameters(), 'lr':lr_phq}], lr=lr)\n",
    "\n",
    "criterion_phq = torch.nn.MSELoss().to(DEVICE) \n",
    "criterion_ptsd = torch.nn.MSELoss().to(DEVICE) \n",
    "\n",
    "# create dataloaders\n",
    "train_dataset = Diacwoz_Dataset(DW_train_encoded, DW_train_PHQ, DW_train_PTSD)\n",
    "val_dataset = Diacwoz_Dataset(DW_val_encoded, DW_val_PHQ, DW_val_PTSD)\n",
    "train_loader = torch.utils.data.DataLoader(torch.utils.data.ConcatDataset([train_dataset, val_dataset]), batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(Diacwoz_Dataset(DW_test_encoded, DW_test_PHQ, DW_test_PTSD), batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# train the model\n",
    "model, losses, epoch_val = multi_task_train(model, optimizer, criterion_phq, criterion_ptsd, train_loader, test_loader, epochs)\n",
    "\n",
    "# save the model\n",
    "torch.save(model.state_dict(), MODEL_CHECKPOINT+'multitask_bilstm_model.pth')\n",
    "\n",
    "# plot loss curves\n",
    "curve_name = \"Multitask_combined_loss\" \n",
    "title = \"Combined Multi-task Loss\"\n",
    "draw_training_curves(losses['train']['total'],losses['val']['total'], curve_name, title, epoch_val)\n",
    "\n",
    "curve_name = \"Multitask_PTSD_loss\" \n",
    "title = \"PTSD Multi-task Loss\"\n",
    "draw_training_curves(losses['train']['ptsd'],losses['val']['ptsd'], curve_name , title, epoch_val)\n",
    "\n",
    "curve_name = \"Multitask_PHQ_loss\" \n",
    "title = \"PHQ Multi-task Loss\"\n",
    "draw_training_curves(losses['train']['phq'],losses['val']['phq'], curve_name, title, epoch_val)\n",
    "\n",
    "# write the model stats to a text file\n",
    "write_stats(model,\"Multitask_details.txt\", epoch_val, learning_rates, \"Multi\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute Single-task Training For DIAC-WoZ:\n",
    "\n",
    "This following cell contains code required for training the Single-task model.\n",
    "\n",
    "**Flow:**\n",
    "* Adjust hyper parameters\n",
    "* Set the task for which the model needs to be trained\n",
    "* Define model, optimizer and loss function\n",
    "* Create Dataloaders\n",
    "* Train the model and save it\n",
    "* Plot and save loss curves\n",
    "* Write model and hyper parameter details to a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 41980,
     "status": "ok",
     "timestamp": 1619050145071,
     "user": {
      "displayName": "Shubham Nagarkar",
      "photoUrl": "",
      "userId": "10621817079012605482"
     },
     "user_tz": 420
    },
    "id": "ECbzu6k0MR-j",
    "outputId": "fdb0c7dd-c808-4e44-c9b9-909825259030",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Non-tuneable parameters\n",
    "vocab_size = 8360 \n",
    "embedding_dim = 300\n",
    "\n",
    "# Non-tuneable parameters\n",
    "hidden_dim = 128\n",
    "lr = 1e-4\n",
    "epochs = 250\n",
    "batch_size = 16\n",
    "learning_rates = [lr]\n",
    "\n",
    "# Set the task\n",
    "task = \"PHQ\"\n",
    "\n",
    "# Define model, optimizer and loss function\n",
    "model_task = Single_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix)\n",
    "optimizer = torch.optim.Adam(model_task.parameters(), lr=lr)\n",
    "criterion = torch.nn.MSELoss().to(DEVICE) \n",
    "\n",
    "# Create Dataloaders\n",
    "train_dataset = Diacwoz_Dataset(DW_train_encoded, DW_train_PHQ, DW_train_PTSD)\n",
    "val_dataset = Diacwoz_Dataset(DW_val_encoded, DW_val_PHQ, DW_val_PTSD)\n",
    "train_loader = torch.utils.data.DataLoader(torch.utils.data.ConcatDataset([train_dataset, val_dataset]), batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(Diacwoz_Dataset(DW_test_encoded, DW_test_PHQ, DW_test_PTSD), batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# Train single task model\n",
    "model_task, losses, epoch_val = single_task_train(task, model_task, optimizer, criterion, train_loader, test_loader, epochs)\n",
    "\n",
    "# Save model\n",
    "torch.save(model_task.state_dict(), MODEL_CHECKPOINT + 'singletask_'+task+'_bilstm_model.pth')\n",
    "\n",
    "# plot loss curves\n",
    "curve_name = \"Singletask_{}_loss\".format(task) \n",
    "title = \"{} Single-task Loss\".format(task)\n",
    "\n",
    "draw_training_curves(losses['train'],losses['val'], curve_name, title, epoch_val)\n",
    "\n",
    "# write the model stats to a text file\n",
    "write_stats(model,\"Singletask_details.txt\", epoch_val, learning_rates, \"Single\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Executable Single-task Evaluation for DIAC_WOZ\n",
    "\n",
    "**Flow:**\n",
    "* Select task (PTSD/PHQ)\n",
    "* Load saved best task model\n",
    "* Evaluate task i.e generate predictions on test set\n",
    "* Convert range into classes\n",
    "* Plot confusion matrix, f1-score, recall, accuracy and precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 684,
     "status": "ok",
     "timestamp": 1619045030004,
     "user": {
      "displayName": "Shubham Nagarkar",
      "photoUrl": "",
      "userId": "10621817079012605482"
     },
     "user_tz": 420
    },
    "id": "Bg-MNQ39VSjt",
    "outputId": "2161b1ee-1838-453c-fb6d-539a6ff88802"
   },
   "outputs": [],
   "source": [
    "# Non-tuneable parameters\n",
    "vocab_size = 8360 \n",
    "embedding_dim = 300\n",
    "hidden_dim = 128\n",
    "\n",
    "# Select task\n",
    "task = \"PHQ\"\n",
    "earlystop = False\n",
    "\n",
    "if earlystop:\n",
    "    phq_model  = MODEL_CHECKPOINT + 'singletask_PHQ_early_stopping_bilstm_model.pth'\n",
    "    ptsd_model = MODEL_CHECKPOINT + 'singletask_PTSD_early_stopping_bilstm_model.pth'\n",
    "else:\n",
    "    phq_model  = MODEL_CHECKPOINT + 'singletask_PHQ_bilstm_model.pth'\n",
    "    ptsd_model = MODEL_CHECKPOINT + 'singletask_PTSD_bilstm_model.pth'\n",
    "\n",
    "if task == \"PHQ\":\n",
    "    path =  phq_model\n",
    "else:\n",
    "    path = ptsd_model\n",
    "    \n",
    "# load saved best model\n",
    "loaded_model = Single_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix).to(DEVICE)\n",
    "loaded_model.load_state_dict(torch.load(path))\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(Diacwoz_Dataset(DW_test_encoded, DW_test_PHQ, DW_test_PTSD), batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# evaluate task\n",
    "actual, predicted = evaluate_single_task(task, loaded_model, test_loader)\n",
    "\n",
    "if task == \"PTSD\":\n",
    "    \n",
    "    singletask_ptsd_true_label = classify_ptsd(actual)\n",
    "    singletask_ptsd_pred_label = classify_ptsd(predicted)\n",
    "    \n",
    "    target_names = [\"No\", \"Yes\"]\n",
    "    title = \"PTSD Confusion Matrix\"\n",
    "    TYPE = \"None\" # Always None for PTSD\n",
    "    report_name = \"Singletask_{}_Metrics.txt\".format(task)\n",
    "    \n",
    "    print_accuracy(singletask_ptsd_true_label, singletask_ptsd_pred_label)\n",
    "    plot_cm(singletask_ptsd_true_label, singletask_ptsd_pred_label, target_names, title, \"Singletask_\"+task, TYPE)\n",
    "    get_classification_report(singletask_ptsd_true_label, singletask_ptsd_pred_label, target_names, target_names, report_name)\n",
    "    \n",
    "    \n",
    "elif task == \"PHQ\":\n",
    "    \n",
    "    singletask_binary_phq_true_label = classify_phq(actual, 'Two')\n",
    "    singletask_binary_phq_pred_label = classify_phq(predicted, 'Two')\n",
    "    \n",
    "    singletask_tertiary_phq_true_label = classify_phq(actual, 'Three')    \n",
    "    singletask_tertiary_phq_pred_label = classify_phq(predicted, 'Three')\n",
    "    \n",
    "    target_names = [\"Moderate\", \"No\", \"Severe\"]\n",
    "    title = \"PHQ Tertiary Confusion Matrix\"\n",
    "    TYPE = \"Tertiary\"\n",
    "    report_name = \"Singletask_{}_{}_Metrics.txt\".format(task, TYPE)\n",
    "    \n",
    "    print_accuracy(singletask_tertiary_phq_true_label, singletask_tertiary_phq_pred_label)\n",
    "    plot_cm(singletask_tertiary_phq_true_label, singletask_tertiary_phq_pred_label, target_names, title, \"Singletask_\"+task, 'Tertiary')\n",
    "    get_classification_report(singletask_tertiary_phq_true_label, singletask_tertiary_phq_pred_label, target_names, target_names, report_name)\n",
    "    \n",
    "    target_names = [\"No\", \"Yes\"]\n",
    "    title = \"PHQ Binary Confusion Matrix\"\n",
    "    TYPE = \"Binary\"\n",
    "    report_name = \"Singletask_{}_{}_Metrics.txt\".format(task, TYPE)\n",
    "    \n",
    "    print_accuracy(singletask_binary_phq_true_label, singletask_binary_phq_pred_label)\n",
    "    plot_cm(singletask_binary_phq_true_label, singletask_binary_phq_pred_label, target_names, title, \"Singletask_\"+task, \"Binary\")\n",
    "    get_classification_report(singletask_binary_phq_true_label, singletask_binary_phq_pred_label, target_names, target_names, report_name)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Executable Multi-task Evaluation for DIAC_WOZ\n",
    "\n",
    "**Flow:**\n",
    "* Load saved best task model\n",
    "* Evaluate task i.e generate predictions on test set\n",
    "* Convert range into classes\n",
    "* Plot confusion matrix, f1-score, recall, accuracy and precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "early_stop = True\n",
    "if early_stop:\n",
    "    multitask_model = MODEL_CHECKPOINT + 'multitask_early_stopping_bilstm_model.pth'\n",
    "else:\n",
    "    multitask_model = MODEL_CHECKPOINT + 'multitask_bilstm_model.pth'\n",
    "    \n",
    "# load saved best model\n",
    "load_multi_model = Multi_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix).to(DEVICE)\n",
    "load_multi_model.load_state_dict(torch.load(multitask_model))\n",
    "test_loader = torch.utils.data.DataLoader(Diacwoz_Dataset(DW_test_encoded, DW_test_PHQ, DW_test_PTSD), batch_size=batch_size, shuffle=True)\n",
    "\n",
    "#evaluate the model\n",
    "actual_phq, actual_ptsd, output_phq, output_ptsd = evaluate_multi_task(load_multi_model, test_loader)\n",
    "\n",
    "# convert range to labels\n",
    "multitask_ptsd_true_label = classify_ptsd(actual_ptsd)\n",
    "multitask_ptsd_pred_label = classify_ptsd(output_ptsd)\n",
    "\n",
    "multitask_binary_phq_true_label = classify_phq(actual_phq, 'Two')\n",
    "multitask_tertiary_phq_true_label = classify_phq(actual_phq, 'Three')\n",
    "multitask_binary_phq_pred_label = classify_phq(output_phq, 'Two')\n",
    "multitask_tertiary_phq_pred_label = classify_phq(output_phq, 'Three')\n",
    "\n",
    "# Plot confusion matrix and performance metrics report\n",
    "target_names = [\"No\", \"Yes\"]\n",
    "title = \"Multitask PTSD Confusion Matrix\"\n",
    "report_name = \"Multitask_PTSD_Metrics.txt\"\n",
    "task = \"PTSD\"\n",
    "print_accuracy(multitask_ptsd_true_label, multitask_ptsd_pred_label)\n",
    "plot_cm(multitask_ptsd_true_label, multitask_ptsd_pred_label, target_names, title, \"Multitask_\"+task, \"\")\n",
    "get_classification_report(multitask_ptsd_true_label, multitask_ptsd_pred_label, target_names, target_names, report_name)\n",
    "    \n",
    "\n",
    "target_names = [\"Moderate\", \"No\", \"Severe\"]\n",
    "title = \"Multi-task PHQ Tertiary Confusion Matrix\"\n",
    "report_name = \"Multitask_Tertiary_PHQ_Metrics.txt\"\n",
    "task = \"PHQ\"\n",
    "print_accuracy(multitask_tertiary_phq_true_label, multitask_tertiary_phq_pred_label)\n",
    "plot_cm(multitask_tertiary_phq_true_label, multitask_tertiary_phq_pred_label, target_names, title, \"Multitask_\"+task, 'Tertiary')\n",
    "get_classification_report(multitask_tertiary_phq_true_label, multitask_tertiary_phq_pred_label, target_names, target_names, report_name)\n",
    "\n",
    "target_names = [\"No\", \"Yes\"]\n",
    "title = \"Multi-task PHQ Binary Confusion Matrix\"\n",
    "report_name = \"Multitask_Binary_PHQ_Metrics.txt\"\n",
    "task = \"PHQ\"\n",
    "\n",
    "print_accuracy(multitask_binary_phq_true_label, multitask_binary_phq_pred_label)\n",
    "plot_cm(multitask_binary_phq_true_label, multitask_binary_phq_pred_label, target_names, title, \"Multitask_\"+task, \"Binary\")\n",
    "get_classification_report(multitask_binary_phq_true_label, multitask_binary_phq_pred_label, target_names, target_names, report_name)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer learning Approach:\n",
    "\n",
    "Here, we are using the best model pre-trained on PHQ task of DIAC-WOZ dataset, and retraining it again on ANONYMITY dataset. We are using the same embeddings and vocabulary from DIAC-WOZ and retraining the base and head of the Bi-LSTM model\n",
    "\n",
    "### Execute PHQ task Training For Anonymity:\n",
    "\n",
    "This following cell contains code required for retraining the single task model on Anonymity dataset.\n",
    "\n",
    "**Flow:**\n",
    "* Adjust hyper parameters\n",
    "* Load the pre-trained model trained on DIAC-WOZ\n",
    "* Define model, optimizer and loss function\n",
    "* Create Dataloaders\n",
    "* Train the model and save it\n",
    "* Plot and save loss curves\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Non-tuneable parameters\n",
    "vocab_size = 8360\n",
    "embedding_dim = 300\n",
    "\n",
    "# Tuneable parameters\n",
    "hidden_dim = 128\n",
    "lr = 1e-3\n",
    "epochs = 300\n",
    "batch_size = 32\n",
    "earlystop = True\n",
    "\n",
    "if earlystop:\n",
    "    path = MODEL_CHECKPOINT + 'singletask_PHQ_early_stopping_bilstm_model.pth'\n",
    "else:\n",
    "    path = MODEL_CHECKPOINT + 'singletask_PHQ_bilstm_model.pth'\n",
    "\n",
    "# load saved singletask phq model\n",
    "model_task = Single_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix).to(DEVICE)\n",
    "model_task.load_state_dict(torch.load(path))\n",
    "\n",
    "optimizer = torch.optim.Adam(model_task.parameters(), lr=lr)\n",
    "criterion = torch.nn.MSELoss().to(DEVICE) \n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(Anonymity_Dataset(AN_train_encoded, AN_train_PHQ), batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(Anonymity_Dataset(AN_test_encoded, AN_test_PHQ), batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# train on Anonymity dataset\n",
    "model_task, losses, epoch_val = Anonymity_train(model_task, optimizer, criterion, train_loader, test_loader, epochs)\n",
    "\n",
    "# Save model\n",
    "torch.save(model_task.state_dict(), MODEL_CHECKPOINT + 'Anonymity_bilstm_model.pth')\n",
    "\n",
    "# plot loss curves\n",
    "curve_name = \"Anonymity_PHQ_loss\" \n",
    "title = \"Anonymity PHQ Loss \"\n",
    "\n",
    "draw_training_curves(losses['train'],losses['val'], curve_name, title, epoch_val)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Executable PHQ task Evaluation for ANONYMITY\n",
    "\n",
    "**Flow:**\n",
    "* Load saved best model\n",
    "* Evaluate task i.e generate predictions on test set\n",
    "* Convert range into classes\n",
    "* Plot confusion matrix, f1-score, recall, accuracy and precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\n",
    "saved_model_path  = MODEL_CHECKPOINT + 'Anonymity_bilstm_model.pth'\n",
    "\n",
    "loaded_model = Single_Task_Model(vocab_size, embedding_dim, hidden_dim, embed_matrix).to(DEVICE)\n",
    "loaded_model.load_state_dict(torch.load(saved_model_path))\n",
    "\n",
    "actual, predicted = Anonymity_evaluate(loaded_model, test_loader)\n",
    "\n",
    "binary_phq_true_label = classify_phq(actual, 'Two')\n",
    "binary_phq_pred_label = classify_phq(predicted, 'Two')\n",
    "\n",
    "tertiary_phq_true_label = classify_phq(actual, 'Three')    \n",
    "tertiary_phq_pred_label = classify_phq(predicted, 'Three')\n",
    "\n",
    "target_names = [\"Moderate\", \"No\", \"Severe\"]\n",
    "title = \"PHQ Tertiary Confusion Matrix\"\n",
    "\n",
    "report_name = \"Anonymity_PHQ_Tertiary_Metrics.txt\"\n",
    "\n",
    "print_accuracy(tertiary_phq_true_label, tertiary_phq_pred_label)\n",
    "plot_cm(tertiary_phq_true_label, tertiary_phq_pred_label, target_names, title, \"Anonymity\", 'Tertiary')\n",
    "get_classification_report(tertiary_phq_true_label, tertiary_phq_pred_label, target_names, target_names, report_name)\n",
    "\n",
    "target_names = [\"No\", \"Yes\"]\n",
    "title = \"PHQ Binary Confusion Matrix\"\n",
    "report_name = \"Singletask_PHQ_Binary_Metrics.txt\"\n",
    "\n",
    "print_accuracy(binary_phq_true_label, binary_phq_pred_label)\n",
    "plot_cm(binary_phq_true_label, binary_phq_pred_label, target_names, title, \"Anonymity\", \"Binary\")\n",
    "get_classification_report(binary_phq_true_label, binary_phq_pred_label, target_names, target_names, report_name)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Depression_detectoin.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}