HAN-Multi-Frontier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import re\n",
    "import glob\n",
    "import io\n",
    "import os\n",
    "import time\n",
    "import json\n",
    "import itertools\n",
    "import logging\n",
    "import gensim\n",
    "import csv\n",
    "import warnings\n",
    "import sys\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "from statistics import mean \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import imshow\n",
    "import pandas as pd \n",
    "import mpld3\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "\n",
    "from transformers import BertTokenizer, BertModel\n",
    "\n",
    "from torchtext import data\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence\n",
    "import torch\n",
    "torch.cuda.set_device(1)\n",
    "\n",
    "%matplotlib inline\n",
    "mpld3.enable_notebook()\n",
    "\n",
    "SEED = 77\n",
    "BATCH_SIZE = 16 # batch size\n",
    "\n",
    "OUTPUT_DIM = 5\n",
    "EMBEDDING_DIM = 100\n",
    "N_EPOCHS = 100\n",
    "TRAIN_RATIO = 0.8\n",
    "POS_WEIGHT = torch.tensor([1, 7, 8, 4, 9])\n",
    "\n",
    "MICRO = 'micro'\n",
    "MACRO = 'macro'\n",
    "FINE_TUNE_EMBEDDING = True  # fine-tune word embeddings?\n",
    "\n",
    "SENTENCE_LIMIT = 100\n",
    "WORD_LIMIT = 1000\n",
    "MIN_WORD_COUNT = 5\n",
    "DATA_FOLDER ='.\\\\HAN'\n",
    "WORKER = 0\n",
    "MAX_VOCAB_SIZE = 25000\n",
    "\n",
    "WORD_RNN_SIZE = 250 \n",
    "SENTENCE_RNN_SIZE = 200\n",
    "WORD_RNN_LAYERS = 2 \n",
    "SENTENCE_RNN_LAYERS = 2\n",
    "WORD_ATTENTION_SIZE = 250\n",
    "SENTENCE_ATTENTION_SIZE = 200\n",
    "DROPOUT = 0.5 \n",
    "\n",
    "ORIGINAL_TRAIN = r'..\\Datasets\\NTUH\\train_preprocessing.txt'\n",
    "TEST = r'..\\Datasets\\NTUH\\test_preprocessing.txt'\n",
    "TRAIN = r'..\\Datasets\\NTUH\\train_preprocessing.txt'\n",
    "VALID = r'..\\Datasets\\NTUH\\validate_preprocessing.txt'\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Required Functions and Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tokenizers\n",
    "def sentence_segment(sentence):\n",
    "    sentences = re.split(r'\\s*<sep>(?:\\s*<sep>)*\\s*', sentence)\n",
    "    filtered_sentence = list(filter(lambda sent: '<unk>' not in sent, sentences))\n",
    "    return [sent for sent in filtered_sentence]\n",
    "    return tokens\n",
    "\n",
    "sent_tokenizer = sentence_segment\n",
    "word_tokenizer = str.split\n",
    "\n",
    "def preprocess(text):\n",
    "    if isinstance(text, float):\n",
    "        return ''\n",
    "\n",
    "    return text.lower()\n",
    "\n",
    "def read_all_text_from_corpus(tsv_file, sentence_limit, word_limit):\n",
    "    docs = []\n",
    "    labels = []\n",
    "    word_counter = Counter()\n",
    "    with io.open(tsv_file, 'r', encoding=\"utf-8\") as file:\n",
    "        for i, line in enumerate(tqdm(file)):\n",
    "            sentences = list()\n",
    "            pid, bh_text, ep_text, major_d, sc, bp, minor_d, de = preprocess(line).strip().split('\\t')\n",
    "            all_text = \"%s <sep> %s\" % (bh_text, ep_text)\n",
    "            sentences.extend([s for s in sent_tokenizer(all_text)])\n",
    "            \n",
    "            words = list()\n",
    "            for s in sentences[:sentence_limit]:\n",
    "                w = word_tokenizer(s)[:word_limit]\n",
    "                # If sentence is empty (due to ?)\n",
    "                if len(w) == 0:\n",
    "                    continue\n",
    "                words.append(w)\n",
    "                word_counter.update(w)\n",
    "            # If all sentences were empty\n",
    "            if len(words) == 0:\n",
    "                continue\n",
    "\n",
    "            labels.append([float(major_d), float(sc), float(bp), float(minor_d), float(de)]) \n",
    "            docs.append(words)\n",
    "\n",
    "    return docs, labels, word_counter\n",
    "\n",
    "def create_input_files(train, test, output_folder, sentence_limit, word_limit, read_tsv, \n",
    "                       min_word_count=5, save_word2vec_data = False, valid = None):\n",
    "    # Read training data\n",
    "    print(f'\\nReading and preprocessing training data {train}...\\n')\n",
    "    train_docs, train_labels, word_counter = read_tsv(train, sentence_limit, word_limit)\n",
    "    # Save text data for word2vec\n",
    "    if save_word2vec_data:\n",
    "        torch.save(train_docs, os.path.join(output_folder, 'word2vec_data.pth.tar'))\n",
    "        print('\\nText data for word2vec saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    # Create word map\n",
    "    word_map = dict()\n",
    "    word_map['<pad>'] = 0\n",
    "    for word, count in word_counter.items():\n",
    "        if count >= min_word_count:\n",
    "            word_map[word] = len(word_map)\n",
    "    word_map['<unk>'] = len(word_map)\n",
    "    print('\\nDiscarding words with counts less than %d, the size of the vocabulary is %d.\\n' % (\n",
    "        min_word_count, len(word_map)))\n",
    "\n",
    "    with open(os.path.join(output_folder, 'word_map.json'), 'w') as j:\n",
    "        json.dump(word_map, j)\n",
    "    print('Word map saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    # Encode and pad\n",
    "    print('Encoding and padding training data...\\n')\n",
    "    encoded_train_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: list(map(lambda w: word_map.get(w, word_map['<unk>']), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), train_docs))\n",
    "    sentences_per_train_document = list(map(lambda doc: len(doc), train_docs))\n",
    "    words_per_train_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), train_docs))\n",
    "\n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    output_content = {'docs': encoded_train_docs,\n",
    "                'labels': train_labels,\n",
    "                'sentences_per_document': sentences_per_train_document,\n",
    "                'words_per_sentence': words_per_train_sentence}\n",
    "    torch.save(output_content,\n",
    "               os.path.join(output_folder, 'TRAIN_data.pth.tar'))\n",
    "    print('Encoded, padded training data saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del train_docs, encoded_train_docs, train_labels, sentences_per_train_document, words_per_train_sentence\n",
    "\n",
    "    print(f'Reading and preprocessing test data {test}...\\n')\n",
    "    test_docs, test_labels, _ = read_tsv(test, sentence_limit, word_limit)\n",
    "\n",
    "    # Encode and pad\n",
    "    print('\\nEncoding and padding test data...\\n')\n",
    "    encoded_test_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: list(map(lambda w: word_map.get(w, word_map['<unk>']), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), test_docs))\n",
    "    sentences_per_test_document = list(map(lambda doc: len(doc), test_docs))\n",
    "    words_per_test_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), test_docs))\n",
    "\n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    torch.save({'docs': encoded_test_docs,\n",
    "                'labels': test_labels,\n",
    "                'sentences_per_document': sentences_per_test_document,\n",
    "                'words_per_sentence': words_per_test_sentence},\n",
    "               os.path.join(output_folder, 'TEST_data.pth.tar'))\n",
    "    print('Encoded, padded test data saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del test_docs, encoded_test_docs, test_labels, sentences_per_test_document, words_per_test_sentence\n",
    "    if valid:\n",
    "        print(f'Reading and preprocessing validation data {valid}...\\n')\n",
    "        valid_docs, valid_labels, _ = read_tsv(valid, sentence_limit, word_limit)\n",
    "        # Encode and pad\n",
    "        print('\\nEncoding and padding validation data...\\n')\n",
    "        encoded_valid_docs = list(map(lambda doc: list(\n",
    "            map(lambda s: list(map(lambda w: word_map.get(w, word_map['<unk>']), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), valid_docs))\n",
    "        sentences_per_valid_document = list(map(lambda doc: len(doc), valid_docs))\n",
    "        words_per_valid_sentence = list(\n",
    "            map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), valid_docs))\n",
    "        # Save\n",
    "        print('Saving...\\n')\n",
    "        torch.save({'docs': encoded_valid_docs,\n",
    "                'labels': valid_labels,\n",
    "                'sentences_per_document': sentences_per_valid_document,\n",
    "                'words_per_sentence': words_per_valid_sentence},\n",
    "                   os.path.join(output_folder, 'VALID_data.pth.tar'))\n",
    "        print('Encoded, padded valid data saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "    print('All done!\\n')\n",
    "    \n",
    "    \n",
    "class NTUH_HANDataset(Dataset):\n",
    "    diagnosis_types = ['major_depressive', 'schizophrenia', 'biploar', 'minor_depressive', 'dementia']\n",
    "\n",
    "    def __init__(self, data_folder, split):\n",
    "        split = split.upper()\n",
    "        self.split = split\n",
    "\n",
    "        # Load data\n",
    "        self.data = torch.load(os.path.join(data_folder, split + '_data.pth.tar'))\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return torch.LongTensor(self.data['docs'][i]), \\\n",
    "               torch.LongTensor([self.data['sentences_per_document'][i]]), \\\n",
    "               torch.LongTensor(self.data['words_per_sentence'][i]), \\\n",
    "               torch.FloatTensor(self.data['labels'][i])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data['labels'])    \n",
    "    \n",
    "class NTUHDataset(data.Dataset):\n",
    "    #urls = ['Datasets\\\\NTUH\\\\corpus.txt']\n",
    "    name = 'ntuh'\n",
    "    dirname = 'ntuh'\n",
    "    diagnosis_types = ['major_depressive', 'schizophrenia', 'biploar', 'minor_depressive', 'dementia']\n",
    "    \n",
    "    @staticmethod\n",
    "    def sort_key(ex):\n",
    "        return len(ex.all_text) # TODO add ep_text?\n",
    "\n",
    "    def __init__(self, path, id_field, bh_text_field, ep_text_field, all_text_field,\n",
    "                 major_label_field, sch_label_field, bipolar_label_field, minor_label_field, dementia_label_field,\n",
    "                 **kwargs):\n",
    "        fields = [('patient_id', id_field), \n",
    "                  ('bh_text', bh_text_field),\n",
    "                  ('ep_text', ep_text_field),\n",
    "                  ('all_text', all_text_field),\n",
    "                  ('major_depressive', major_label_field),\n",
    "                  ('schizophrenia', sch_label_field),\n",
    "                  ('biploar', bipolar_label_field),\n",
    "                  ('minor_depressive', minor_label_field),\n",
    "                  ('dementia', dementia_label_field)]\n",
    "        examples = []\n",
    "        \n",
    "        for fname in glob.iglob(path + '.txt'):\n",
    "            with io.open(fname, 'r', encoding=\"utf-8\") as f:\n",
    "                for line in f:\n",
    "                    pid, bh_text, ep_text, major_d, sc, bp, minor_d, de = line.strip().split('\\t')\n",
    "                    all_text = \"%s <sep> %s\" % (bh_text, ep_text)\n",
    "                    examples.append(data.Example.fromlist([pid, bh_text, ep_text, all_text, major_d, sc, bp, minor_d, de], \n",
    "                                                          fields))\n",
    "        super(NTUHDataset, self).__init__(examples, fields, **kwargs)\n",
    "\n",
    "    @classmethod\n",
    "    def splits(cls, id_field,\n",
    "               bh_text_field, ep_text_field, all_text_field,\n",
    "               major_label_field, sch_label_field, bipolar_label_field, minor_label_field, dementia_label_field,\n",
    "               root='..\\\\Datasets\\\\NTUH',\n",
    "               train='train_preprocessing', test='test_preprocessing', **kwargs):\n",
    "        return super(NTUHDataset, cls).splits(\n",
    "            path = root, root=root, id_field=id_field,\n",
    "            bh_text_field = bh_text_field, ep_text_field = ep_text_field, all_text_field = all_text_field, \n",
    "            major_label_field = major_label_field, sch_label_field = sch_label_field, \n",
    "            bipolar_label_field = bipolar_label_field, minor_label_field = minor_label_field, \n",
    "            dementia_label_field = dementia_label_field,\n",
    "            train=train, validation=None, test=test, **kwargs)\n",
    "    \n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "def create_input_files_for_glove(glove_stoi, train, test, output_folder, sentence_limit, word_limit, read_tsv, \n",
    "                       min_word_count=5, valid = None):\n",
    "    # Read training data\n",
    "    print(f'\\nReading and preprocessing training data {train}...\\n')\n",
    "    train_docs, train_labels, word_counter = read_tsv(train, sentence_limit, word_limit)\n",
    "    \n",
    "    # Create word map\n",
    "    word_map = dict(glove_stoi)\n",
    "    print('\\nThe size of the vocabulary is %d.\\n' % (len(word_map)))\n",
    "\n",
    "    with open(os.path.join(output_folder, 'glove_word_map.json'), 'w') as j:\n",
    "        json.dump(word_map, j)\n",
    "    print('Word map saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    # Encode and pad\n",
    "    print('Encoding and padding training data...\\n')\n",
    "    encoded_train_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: list(map(lambda w: word_map.get(w, UNK_IDX), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), train_docs))\n",
    "    sentences_per_train_document = list(map(lambda doc: len(doc), train_docs))\n",
    "    words_per_train_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), train_docs))\n",
    "\n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    # Because of the large data, saving as a JSON can be very slow\n",
    "    output_content = {'docs': encoded_train_docs,\n",
    "                'labels': train_labels,\n",
    "                'sentences_per_document': sentences_per_train_document,\n",
    "                'words_per_sentence': words_per_train_sentence}\n",
    "    torch.save(output_content,\n",
    "               os.path.join(output_folder, 'GLOVE_TRAIN_data.pth.tar'))\n",
    "    print('Encoded, padded training data (GLOVE_TRAIN_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del train_docs, encoded_train_docs, train_labels, sentences_per_train_document, words_per_train_sentence\n",
    "\n",
    "    # Read test data\n",
    "    print(f'Reading and preprocessing test data {test}...\\n')\n",
    "    test_docs, test_labels, _ = read_tsv(test, sentence_limit, word_limit)\n",
    "\n",
    "    # Encode and pad\n",
    "    print('\\nEncoding and padding test data...\\n')\n",
    "    encoded_test_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: list(map(lambda w: word_map.get(w, UNK_IDX), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), test_docs))\n",
    "    sentences_per_test_document = list(map(lambda doc: len(doc), test_docs))\n",
    "    words_per_test_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), test_docs))\n",
    "\n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    torch.save({'docs': encoded_test_docs,\n",
    "                'labels': test_labels,\n",
    "                'sentences_per_document': sentences_per_test_document,\n",
    "                'words_per_sentence': words_per_test_sentence},\n",
    "               os.path.join(output_folder, 'GLOVE_TEST_data.pth.tar'))\n",
    "    print('Encoded, padded test data (GLOVE_TEST_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del test_docs, encoded_test_docs, test_labels, sentences_per_test_document, words_per_test_sentence\n",
    "    if valid:\n",
    "        print(f'Reading and preprocessing validation data {valid}...\\n')\n",
    "        valid_docs, valid_labels, _ = read_tsv(valid, sentence_limit, word_limit)\n",
    "        # Encode and pad\n",
    "        print('\\nEncoding and padding validation data...\\n')\n",
    "        encoded_valid_docs = list(map(lambda doc: list(\n",
    "            map(lambda s: list(map(lambda w: word_map.get(w, UNK_IDX), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc)), valid_docs))\n",
    "        sentences_per_valid_document = list(map(lambda doc: len(doc), valid_docs))\n",
    "        words_per_valid_sentence = list(\n",
    "            map(lambda doc: list(map(lambda s: len(s), doc)) + [0] * (sentence_limit - len(doc)), valid_docs))\n",
    "        # Save\n",
    "        print('Saving...\\n')\n",
    "        torch.save({'docs': encoded_valid_docs,\n",
    "                'labels': valid_labels,\n",
    "                'sentences_per_document': sentences_per_valid_document,\n",
    "                'words_per_sentence': words_per_valid_sentence},\n",
    "                   os.path.join(output_folder, 'GLOVE_VALID_data.pth.tar'))\n",
    "        print('Encoded, padded valid data (GLOVE_VALID_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "    print('All done!\\n')\n",
    "    \n",
    "def load_glove_embeddings():\n",
    "    random.seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    torch.manual_seed(SEED)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    \n",
    "    ID_TEXT = data.Field(batch_first = True, lower = True)\n",
    "    BH_TEXT = data.Field(batch_first = True, lower = True)\n",
    "    EP_TEXT = data.Field(batch_first = True, lower = True)\n",
    "    ALL_TEXT = data.Field(batch_first = True, lower = True)\n",
    "\n",
    "    MAJ_LABEL = data.LabelField(dtype = torch.float)\n",
    "    SCH_LABEL = data.LabelField(dtype = torch.float)\n",
    "    BIP_LABEL = data.LabelField(dtype = torch.float)\n",
    "    MIN_LABEL = data.LabelField(dtype = torch.float)\n",
    "    DEM_LABEL = data.LabelField(dtype = torch.float)\n",
    "    full_train_data, test_data = NTUHDataset.splits(ID_TEXT, BH_TEXT, EP_TEXT, ALL_TEXT, \n",
    "                                               MAJ_LABEL, SCH_LABEL, BIP_LABEL, MIN_LABEL, DEM_LABEL)\n",
    "    ALL_TEXT.build_vocab(full_train_data, max_size = MAX_VOCAB_SIZE, vectors = \"glove.6B.300d\", \n",
    "                         unk_init = torch.Tensor.normal_)\n",
    "\n",
    "    print(\"\\nEmbedding length is %d.\\n\" % len(ALL_TEXT.vocab.vectors))\n",
    "\n",
    "    # Create tensor to hold embeddings for words that are in-corpus\n",
    "    embeddings = ALL_TEXT.vocab.vectors\n",
    "\n",
    "    return embeddings, ALL_TEXT.vocab.vectors.shape[1], ALL_TEXT.vocab, ALL_TEXT.pad_token, ALL_TEXT.unk_token\n",
    "\n",
    "def analysis_plotter(fig, ax, train, valid, title, param_dict1, param_dict2):\n",
    "    out = ax.plot(train, **param_dict1)\n",
    "    out = ax.plot(valid, **param_dict2)\n",
    "    ax.title.set_text(title)\n",
    "    ax.legend()\n",
    "    pv = float('inf')\n",
    "    x = []\n",
    "    y = []\n",
    "    for k, v in enumerate(valid):\n",
    "        if v > pv:\n",
    "            x.append(k)\n",
    "            y.append(v)\n",
    "        pv = v\n",
    "    scatter = ax.scatter(x, y)\n",
    "    labels = []\n",
    "    for x, y in zip(x,y):\n",
    "        labels.append(f'{x}: {y}')\n",
    "    tooltip = mpld3.plugins.PointLabelTooltip(scatter, labels=labels)\n",
    "    mpld3.plugins.connect(fig, tooltip)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "\n",
    "def init_embedding(input_embedding):\n",
    "    \"\"\"\n",
    "    Initialize embedding tensor with values from the uniform distribution.\n",
    "    :param input_embedding: embedding tensor\n",
    "    \"\"\"\n",
    "    random.seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    torch.manual_seed(SEED)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "    bias = np.sqrt(3.0 / input_embedding.size(1))\n",
    "    nn.init.uniform_(input_embedding, -bias, bias)\n",
    "\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names =['ID', 'BH Text', 'EP Text', 'Major Depressive', 'Schizophrenia', 'Biploar', 'Minor Depressive', 'Dementia']\n",
    "\n",
    "label_map = {k: v for v, k in enumerate(names[3:])}\n",
    "rev_label_map = {v: k for k, v in label_map.items()}\n",
    "n_classes = len(label_map)\n",
    "\n",
    "print(label_map)\n",
    "print(rev_label_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HierarchialAttentionNetwork(nn.Module):\n",
    "    def __init__(self, n_classes, vocab_size, emb_size, word_rnn_size, sentence_rnn_size, word_rnn_layers,\n",
    "                 sentence_rnn_layers, word_att_size, sentence_att_size, dropout=0.5):\n",
    "        super(HierarchialAttentionNetwork, self).__init__()\n",
    "\n",
    "        self.sentence_attention = SentenceAttention(vocab_size, emb_size, word_rnn_size, sentence_rnn_size,\n",
    "                                                    word_rnn_layers, sentence_rnn_layers, word_att_size,\n",
    "                                                    sentence_att_size, dropout)\n",
    "        self.fc = nn.Linear(2 * sentence_rnn_size, n_classes)\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, documents, sentences_per_document, words_per_sentence):\n",
    "        document_embeddings, word_alphas, sentence_alphas = self.sentence_attention(documents, sentences_per_document,\n",
    "                                                                                    words_per_sentence)  \n",
    "        scores = self.fc(self.dropout(document_embeddings))  # (n_documents, n_classes)\n",
    "\n",
    "        return scores, word_alphas, sentence_alphas\n",
    "\n",
    "\n",
    "class SentenceAttention(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_size, word_rnn_size, sentence_rnn_size, word_rnn_layers, sentence_rnn_layers,\n",
    "                 word_att_size, sentence_att_size, dropout):\n",
    "        super(SentenceAttention, self).__init__()\n",
    "\n",
    "        self.word_attention = WordAttention(vocab_size, emb_size, word_rnn_size, word_rnn_layers, word_att_size,\n",
    "                                            dropout)\n",
    "\n",
    "        self.sentence_rnn = nn.GRU(2 * word_rnn_size, sentence_rnn_size, num_layers=sentence_rnn_layers,\n",
    "                                   bidirectional=True, dropout=dropout, batch_first=True)\n",
    "\n",
    "        self.sentence_attention = nn.Linear(2 * sentence_rnn_size, sentence_att_size)\n",
    "\n",
    "        self.sentence_context_vector = nn.Linear(sentence_att_size, 1,\n",
    "                                                 bias=False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, documents, sentences_per_document, words_per_sentence):\n",
    "        packed_sentences = pack_padded_sequence(documents,\n",
    "                                                lengths=sentences_per_document.tolist(), \n",
    "                                                batch_first=True,\n",
    "                                                enforce_sorted=False)  \n",
    "        packed_words_per_sentence = pack_padded_sequence(words_per_sentence,\n",
    "                                                         lengths=sentences_per_document.tolist(),\n",
    "                                                         batch_first=True,\n",
    "                                                         enforce_sorted=False)  \n",
    "        sentences, word_alphas = self.word_attention(packed_sentences.data,\n",
    "                                                     packed_words_per_sentence.data)\n",
    "        sentences = self.dropout(sentences)\n",
    "        \n",
    "        packed_sentences, _ = self.sentence_rnn(PackedSequence(data=sentences,\n",
    "                                                               batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                               sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                               unsorted_indices=packed_sentences.unsorted_indices)) \n",
    "\n",
    "        att_s = self.sentence_attention(packed_sentences.data)\n",
    "        att_s = torch.tanh(att_s)\n",
    "        att_s = self.sentence_context_vector(att_s).squeeze(1) \n",
    "\n",
    "        max_value = att_s.max() \n",
    "        att_s = torch.exp(att_s - max_value) \n",
    "\n",
    "        att_s, _ = pad_packed_sequence(PackedSequence(data=att_s,\n",
    "                                                      batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                      sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                      unsorted_indices=packed_sentences.unsorted_indices),\n",
    "                                       batch_first=True)\n",
    "\n",
    "        sentence_alphas = att_s / torch.sum(att_s, dim=1, keepdim=True) \n",
    "\n",
    "        documents, _ = pad_packed_sequence(packed_sentences,\n",
    "                                           batch_first=True) \n",
    "\n",
    "        documents = documents * sentence_alphas.unsqueeze(2)\n",
    "        documents = documents.sum(dim=1) \n",
    "\n",
    "        word_alphas, _ = pad_packed_sequence(PackedSequence(data=word_alphas,\n",
    "                                                            batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                            sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                            unsorted_indices=packed_sentences.unsorted_indices),\n",
    "                                             batch_first=True)\n",
    "\n",
    "        return documents, word_alphas, sentence_alphas\n",
    "\n",
    "class WordAttention(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_size, word_rnn_size, word_rnn_layers, word_att_size, dropout):\n",
    "        super(WordAttention, self).__init__()\n",
    "\n",
    "        self.embeddings = nn.Embedding(vocab_size, emb_size)\n",
    "\n",
    "        self.word_rnn = nn.GRU(emb_size, word_rnn_size, num_layers=word_rnn_layers, bidirectional=True,\n",
    "                               dropout=dropout, batch_first=True)\n",
    "\n",
    "        self.word_attention = nn.Linear(2 * word_rnn_size, word_att_size)\n",
    "\n",
    "        self.word_context_vector = nn.Linear(word_att_size, 1, bias=False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def init_embeddings(self, embeddings):\n",
    "        self.embeddings.weight = nn.Parameter(embeddings)\n",
    "\n",
    "    def fine_tune_embeddings(self, fine_tune=False):\n",
    "        for p in self.embeddings.parameters():\n",
    "            p.requires_grad = fine_tune\n",
    "\n",
    "    def forward(self, sentences, words_per_sentence):\n",
    "        sentences = self.dropout(self.embeddings(sentences)) \n",
    "        \n",
    "        packed_words = pack_padded_sequence(sentences,\n",
    "                                            lengths=words_per_sentence.tolist(),\n",
    "                                            batch_first=True,\n",
    "                                            enforce_sorted=False) \n",
    "        \n",
    "        packed_words, _ = self.word_rnn(\n",
    "            packed_words)\n",
    "        \n",
    "        att_w = self.word_attention(packed_words.data) \n",
    "        att_w = torch.tanh(att_w)\n",
    "        att_w = self.word_context_vector(att_w).squeeze(1)  # (n_words)\n",
    "\n",
    "        max_value = att_w.max()\n",
    "        att_w = torch.exp(att_w - max_value)\n",
    "        \n",
    "        att_w, _ = pad_packed_sequence(PackedSequence(data=att_w,\n",
    "                                                      batch_sizes=packed_words.batch_sizes,\n",
    "                                                      sorted_indices=packed_words.sorted_indices,\n",
    "                                                      unsorted_indices=packed_words.unsorted_indices),\n",
    "                                       batch_first=True) \n",
    "        \n",
    "        word_alphas = att_w / torch.sum(att_w, dim=1, keepdim=True)\n",
    "        \n",
    "        \n",
    "        sentences, _ = pad_packed_sequence(packed_words,\n",
    "                                           batch_first=True)\n",
    "\n",
    "        sentences = sentences * word_alphas.unsqueeze(2) \n",
    "        sentences = sentences.sum(dim=1)\n",
    "        \n",
    "        return sentences, word_alphas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_measure(predictions, labels):\n",
    "    diagnoses = {}\n",
    "    MICRO = 'micro'\n",
    "    MACRO = 'macro'\n",
    "    predicts = []\n",
    "    diagnoses[MICRO] = {}\n",
    "    \n",
    "    rounded_preds = torch.round(torch.sigmoid(predictions))\n",
    "    predicts.extend(rounded_preds.data.tolist())\n",
    "    \n",
    "    for index, value in enumerate(rounded_preds):\n",
    "        for did, dvalue in enumerate(rounded_preds[index]):\n",
    "            v = dvalue.item()                    \n",
    "            if v == 1:\n",
    "                if dvalue == labels[index, did]:\n",
    "                    if did not in diagnoses:\n",
    "                        diagnoses[did] = {}                                \n",
    "                    diagnoses[did]['tp'] = diagnoses[did].get('tp', 0) + 1\n",
    "                    diagnoses[MICRO]['tp'] = diagnoses[MICRO].get('tp', 0) + 1\n",
    "                else:\n",
    "                    if did not in diagnoses:\n",
    "                        diagnoses[did] = {}\n",
    "                    diagnoses[did]['fp'] = diagnoses[did].get('fp', 0) + 1\n",
    "                    diagnoses[MICRO]['fp'] = diagnoses[MICRO].get('fp', 0) + 1\n",
    "            elif v == 0:\n",
    "                if 1 == labels[index, did].item():\n",
    "                    if did not in diagnoses:\n",
    "                        diagnoses[did] = {}\n",
    "                    diagnoses[did]['fn'] = diagnoses[did].get('fn', 0) + 1\n",
    "                    diagnoses[MICRO]['fn'] = diagnoses[MICRO].get('fn', 0) + 1\n",
    "    diagnoses[MACRO] = {}\n",
    "    for d in diagnoses:        \n",
    "        if d is MACRO:\n",
    "            continue\n",
    "        try:\n",
    "            diagnoses[d]['p']=diagnoses[d].get('tp', 0)/(diagnoses[d].get('tp', 0)+diagnoses[d].get('fp', 0))            \n",
    "        except:            \n",
    "            diagnoses[d]['p']=0.0\n",
    "        if d is not MICRO:\n",
    "                diagnoses[MACRO]['p']=diagnoses[MACRO].get('p', 0.0)+diagnoses[d]['p']                \n",
    "            \n",
    "        try:\n",
    "            diagnoses[d]['r']=diagnoses[d].get('tp', 0)/(diagnoses[d].get('tp', 0)+diagnoses[d].get('fn', 0))            \n",
    "        except:\n",
    "            diagnoses[d]['r']=0.0\n",
    "        if d is not MICRO:\n",
    "            diagnoses[MACRO]['r']=diagnoses[MACRO].get('r', 0.0)+diagnoses[d]['r']\n",
    "        \n",
    "        try:\n",
    "            diagnoses[d]['f']=2/(1/diagnoses[d]['p']+1/diagnoses[d]['r'])            \n",
    "        except:\n",
    "            diagnoses[d]['f']=0.0\n",
    "        if d is not MICRO:\n",
    "                diagnoses[MACRO]['f']=diagnoses[MACRO].get('f', 0.0)+diagnoses[d]['f']\n",
    "\n",
    "    diagnoses[MACRO]['f']=diagnoses[MACRO]['f']/float(len(diagnoses)-2)\n",
    "    diagnoses[MACRO]['p']=diagnoses[MACRO]['p']/float(len(diagnoses)-2)\n",
    "    diagnoses[MACRO]['r']=diagnoses[MACRO]['r']/float(len(diagnoses)-2)\n",
    "    return diagnoses, predicts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "        \n",
    "def clip_gradient(optimizer, grad_clip):\n",
    "    for group in optimizer.param_groups:\n",
    "        for param in group['params']:\n",
    "            if param.grad is not None:\n",
    "                param.grad.data.clamp_(-grad_clip, grad_clip)\n",
    "\n",
    "def update_fscores(new, overall):\n",
    "    MICRO = 'micro'\n",
    "    MACRO = 'macro'\n",
    "    \n",
    "    for k in new:\n",
    "        if k not in overall:\n",
    "            overall[k] = {}\n",
    "        overall[k]['tp'] = overall[k].get('tp', 0) + new[k].get('tp', 0)\n",
    "        overall[k]['fp'] = overall[k].get('fp', 0) + new[k].get('fp', 0)\n",
    "        overall[k]['fn'] = overall[k].get('fn', 0) + new[k].get('fn', 0)\n",
    "        overall[MICRO]['tp'] = overall[MICRO].get('tp', 0) + new[k].get('tp', 0)\n",
    "        overall[MICRO]['fp'] = overall[MICRO].get('fp', 0) + new[k].get('fp', 0)\n",
    "        overall[MICRO]['fn'] = overall[MICRO].get('fn', 0) + new[k].get('fn', 0)\n",
    "        \n",
    "    overall[MACRO] = {}\n",
    "    for d in overall:        \n",
    "        if d is MACRO:\n",
    "            continue\n",
    "        try:\n",
    "            overall[d]['p']=overall[d].get('tp', 0)/(overall[d].get('tp', 0)+overall[d].get('fp', 0))            \n",
    "        except:            \n",
    "            overall[d]['p']=0.0\n",
    "        if d is not MICRO:\n",
    "            overall[MACRO]['p']=overall[MACRO].get('p', 0.0)+overall[d]['p']                \n",
    "            \n",
    "        try:\n",
    "            overall[d]['r']=overall[d].get('tp', 0)/(overall[d].get('tp', 0)+overall[d].get('fn', 0))            \n",
    "        except:\n",
    "            overall[d]['r']=0.0\n",
    "        if d is not MICRO:\n",
    "            overall[MACRO]['r']=overall[MACRO].get('r', 0.0)+overall[d]['r']\n",
    "        \n",
    "        try:\n",
    "            overall[d]['f']=2/(1/overall[d]['p']+1/overall[d]['r'])            \n",
    "        except:\n",
    "            overall[d]['f']=0.0\n",
    "        if d is not MICRO:\n",
    "                overall[MACRO]['f']=overall[MACRO].get('f', 0.0)+overall[d]['f']\n",
    "\n",
    "    overall[MACRO]['f']=overall[MACRO]['f']/float(len(overall)-2)\n",
    "    overall[MACRO]['p']=overall[MACRO]['p']/float(len(overall)-2)\n",
    "    overall[MACRO]['r']=overall[MACRO]['r']/float(len(overall)-2)\n",
    "    return overall\n",
    "    \n",
    "def train(train_loader, model, criterion, optimizer, grad_clip): #, epoch, total_epoch, print_freq = 100):\n",
    "    model.train()\n",
    "\n",
    "    losses = AverageMeter() \n",
    "    accs = {}\n",
    "\n",
    "    # Batches\n",
    "    for i, (documents, sentences_per_document, words_per_sentence, labels) in enumerate(train_loader):\n",
    "        # Back prop.\n",
    "        optimizer.zero_grad()        \n",
    "\n",
    "        documents = documents.to(device) \n",
    "        \n",
    "        sentences_per_document = sentences_per_document.squeeze(1).to(device) \n",
    "        \n",
    "        words_per_sentence = words_per_sentence.to(device)  \n",
    "        \n",
    "        labels = labels.to(device)  \n",
    "        \n",
    "        scores, word_alphas, sentence_alphas = model(documents, sentences_per_document,\n",
    "                                                     words_per_sentence)  \n",
    "        loss = criterion(scores, labels) \n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        if grad_clip is not None:\n",
    "            clip_gradient(optimizer, grad_clip)\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        fscores, _ = f_measure(scores, labels)            \n",
    "        losses.update(loss.item(), labels.size(0))\n",
    "        accs = update_fscores(fscores, accs)\n",
    "        \n",
    "    return losses.avg, accs['micro'][\"f\"]        \n",
    "\n",
    "def save_checkpoint(epoch, model, optimizer, word_map, data_folder, model_name):\n",
    "    state = {'epoch': epoch,\n",
    "             'model': model,\n",
    "             'optimizer': optimizer,\n",
    "             'word_map': word_map}\n",
    "    filename = os.path.join(data_folder, model_name)#'checkpoint_han.pth.tar'\n",
    "    torch.save(state, filename)\n",
    "    \n",
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs\n",
    "\n",
    "def train_epoch(start_epoch, epochs, data_loader, model, criterion, optimizer, word_map, model_name, grad_clip, valid_iterator = None, \n",
    "                interval = 50, early_stop = False, period = 20, gap = 0.005, threshold = 0.5):\n",
    "    best_valid_loss = float('inf')\n",
    "    best_valid_fscore = 0\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    train_accs = []\n",
    "    valid_accs = []\n",
    "    observed_time = 0\n",
    "    \n",
    "    # Epochs\n",
    "    for epoch in range(start_epoch, epochs):\n",
    "        start_time = time.time()\n",
    "        # One epoch's training\n",
    "        #try:\n",
    "        train_loss, train_acc = train(train_loader=data_loader, model=model, \n",
    "                                      criterion=criterion, optimizer=optimizer, grad_clip=grad_clip)#, epoch=epoch, \n",
    "              #total_epoch = epochs)\n",
    "        train_losses.append(train_loss)\n",
    "        train_accs.append(train_acc)\n",
    "\n",
    "        end_time = time.time()\n",
    "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "\n",
    "        if valid_iterator:\n",
    "            valid_loss, valid_acc = evaluate(model, valid_iterator, criterion)#, model_type)\n",
    "            valid_losses.append(valid_loss)\n",
    "            valid_accs.append(valid_acc)\n",
    "        else:\n",
    "            valid_loss = 0 \n",
    "\n",
    "        if (epoch + 1) % interval == 0:\n",
    "            print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "            print(f'\\tTrain Loss: {train_loss:.3f} | Train micro-F-score: {train_acc*100:.2f}%')\n",
    "            if valid_iterator:\n",
    "                print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. micro-F-score: {valid_acc*100:.2f}%')\n",
    "        elif epoch == epochs - 1:\n",
    "            print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "            print(f'\\tTrain Loss: {train_loss:.3f} | Train micro-F-score: {train_acc*100:.2f}%')\n",
    "            if valid_iterator:\n",
    "                print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. micro-F-score: {valid_acc*100:.2f}%')\n",
    "\n",
    "        # Save checkpoint\n",
    "        if valid_loss < best_valid_loss:\n",
    "            best_valid_loss = valid_loss\n",
    "            save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_loss.pt')\n",
    "\n",
    "        if early_stop and best_valid_fscore > threshold and best_valid_fscore - valid_acc > gap:\n",
    "            observed_time += 1\n",
    "            print(f'\\rBest validation F-measure: {best_valid_fscore:.3f}/Current F-measure: {valid_acc:.3f} [Times: {observed_time}/{period}]')  \n",
    "            if observed_time >= period:\n",
    "                print(f'Early stop at epoch {epoch+1:02}.')\n",
    "                break                        \n",
    "        if valid_acc > best_valid_fscore:\n",
    "            best_valid_fscore = valid_acc\n",
    "            save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_fscore.pt')\n",
    "            observed_time = 0\n",
    "        #except RuntimeError as e:\n",
    "        #    print(f'Runtime error at epoch {epoch+1:02}: {str(e)}')\n",
    "        #    save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_except.pt')\n",
    "        #    return train_losses, valid_losses, train_accs, valid_accs\n",
    "            \n",
    "    # Decay learning rate every epoch\n",
    "    #adjust_learning_rate(optimizer, 0.1)\n",
    "\n",
    "    #return model\n",
    "    return train_losses, valid_losses, train_accs, valid_accs   \n",
    "\n",
    "\n",
    "def evaluate(model, iterator, criterion):#, model_type):    \n",
    "    epoch_loss = 0\n",
    "    epoch_fscore = 0\n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():    \n",
    "        for i, (documents, sentences_per_document, words_per_sentence, labels) in enumerate(iterator):\n",
    "            documents = documents.to(device)\n",
    "            sentences_per_document = sentences_per_document.squeeze(1).to(device) \n",
    "            words_per_sentence = words_per_sentence.to(device)\n",
    "            labels = labels.to(device)  # (batch_size)\n",
    "            \n",
    "            # Forward prop.\n",
    "            scores, word_alphas, sentence_alphas = model(documents, sentences_per_document,\n",
    "                    words_per_sentence) \n",
    "            \n",
    "            loss = criterion(scores, labels) \n",
    "            \n",
    "            fscores, _ = f_measure(scores, labels)                    \n",
    "        \n",
    "            epoch_loss += loss.item()\n",
    "            epoch_fscore += fscores['micro'][\"f\"]\n",
    "        \n",
    "    return epoch_loss / len(iterator), epoch_fscore / len(iterator)\n",
    "\n",
    "def test(model, iterator, model_type, model_path = None):    \n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    if model_path:\n",
    "        model.load_state_dict(torch.load(model_path))\n",
    "\n",
    "    model.eval()\n",
    "    \n",
    "    diagnoses = {}\n",
    "    predicts = []\n",
    "    MICRO = 'micro'\n",
    "    MACRO = 'macro'\n",
    "    diagnoses[MICRO] = {}\n",
    "    with torch.no_grad():     \n",
    "        for i, (documents, sentences_per_document, words_per_sentence, labels) in enumerate(\n",
    "            tqdm(iterator, desc='Evaluating')):\n",
    "            \n",
    "            documents = documents.to(device)  # (batch_size, sentence_limit, word_limit)\n",
    "            sentences_per_document = sentences_per_document.squeeze(1).to(device)  # (batch_size)\n",
    "            words_per_sentence = words_per_sentence.to(device)  # (batch_size, sentence_limit)\n",
    "            labels = labels.to(device)  # (batch_size)\n",
    "            \n",
    "            # Forward prop.\n",
    "            scores, word_alphas, sentence_alphas = model(documents, sentences_per_document,\n",
    "                    words_per_sentence)  # (n_documents, n_classes), (n_documents, max_doc_len_in_batch, max_sent_len_in_batch), (n_documents, max_doc_len_in_batch)\n",
    "            \n",
    "            rounded_preds = torch.round(torch.sigmoid(scores))\n",
    "            predicts.extend(rounded_preds.data.tolist())\n",
    "            \n",
    "            for index, value in enumerate(rounded_preds):\n",
    "                for did, dvalue in enumerate(rounded_preds[index]):\n",
    "                    v = dvalue.item()                    \n",
    "                    if v == 1:\n",
    "                        if dvalue == labels[index, did]:\n",
    "                            if did not in diagnoses:\n",
    "                                diagnoses[did] = {}                                \n",
    "                            diagnoses[did]['tp'] = diagnoses[did].get('tp', 0) + 1\n",
    "                            diagnoses[MICRO]['tp'] = diagnoses[MICRO].get('tp', 0) + 1 \n",
    "                        else:\n",
    "                            if did not in diagnoses:\n",
    "                                diagnoses[did] = {}\n",
    "                            diagnoses[did]['fp'] = diagnoses[did].get('fp', 0) + 1\n",
    "                            diagnoses[MICRO]['fp'] = diagnoses[MICRO].get('fp', 0) + 1\n",
    "                    elif v == 0:\n",
    "                        if 1 == labels[index, did].item():\n",
    "                            if did not in diagnoses:\n",
    "                                diagnoses[did] = {}\n",
    "                            diagnoses[did]['fn'] = diagnoses[did].get('fn', 0) + 1\n",
    "                            diagnoses[MICRO]['fn'] = diagnoses[MICRO].get('fn', 0) + 1\n",
    "                        else:\n",
    "                            if did not in diagnoses:\n",
    "                                diagnoses[did] = {}\n",
    "                            diagnoses[did]['tn'] = diagnoses[did].get('tn', 0) + 1\n",
    "                            diagnoses[MICRO]['tn'] = diagnoses[MICRO].get('tn', 0) + 1\n",
    "    diagnoses[MACRO] = {}\n",
    "    for d in diagnoses:        \n",
    "        if d is MACRO:\n",
    "            continue\n",
    "        try:\n",
    "            diagnoses[d]['p']=diagnoses[d].get('tp', 0)/(diagnoses[d].get('tp', 0)+diagnoses[d].get('fp', 0))\n",
    "            if d is not MICRO:\n",
    "                diagnoses[MACRO]['p']=diagnoses[MACRO].get('p', 0.0)+diagnoses[d]['p']                \n",
    "        except:            \n",
    "            diagnoses[d]['p']=0.0\n",
    "            \n",
    "        try:\n",
    "            diagnoses[d]['r']=diagnoses[d].get('tp', 0)/(diagnoses[d].get('tp', 0)+diagnoses[d].get('fn', 0))\n",
    "            if d is not MICRO:\n",
    "                diagnoses[MACRO]['r']=diagnoses[MACRO].get('r', 0.0)+diagnoses[d]['r']\n",
    "        except:\n",
    "            diagnoses[d]['r']=0.0\n",
    "        \n",
    "        try:\n",
    "            diagnoses[d]['f']=2/(1/diagnoses[d]['p']+1/diagnoses[d]['r'])\n",
    "            if d is not MICRO:\n",
    "                diagnoses[MACRO]['f']=diagnoses[MACRO].get('f', 0.0)+diagnoses[d]['f']\n",
    "        except:\n",
    "            diagnoses[d]['f']=0.0\n",
    "    diagnoses[MACRO]['f']=diagnoses[MACRO]['f']/float(len(diagnoses)-2)\n",
    "    diagnoses[MACRO]['p']=diagnoses[MACRO]['p']/float(len(diagnoses)-2)\n",
    "    diagnoses[MACRO]['r']=diagnoses[MACRO]['r']/float(len(diagnoses)-2)\n",
    "    return diagnoses, predicts\n",
    "\n",
    "def initialize_embeddings(embedding_dim, word_map):\n",
    "    print(\"\\nEmbedding length is %d.\\n\" % embedding_dim)\n",
    "\n",
    "    # Create tensor to hold embeddings\n",
    "    embeddings = torch.FloatTensor(len(word_map), embedding_dim)\n",
    "    init_embedding(embeddings)\n",
    "\n",
    "    print(\"Done.\\n Embedding vocabulary: %d.\\n\" % len(word_map))\n",
    "\n",
    "    return embeddings, embedding_dim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset Analysis\n",
    "## Training Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntuhdataset = pd.read_csv('../Datasets/NTUH/train_preprocessing.txt', sep ='\\t', names = names)\n",
    "bh_text = ntuhdataset['BH Text'].apply(lambda x: len(re.sub(r'\\s+', ' ', re.sub(r'<[^>]+>', '', str(x))).split()))\n",
    "plt.figure(figsize=(5,5))\n",
    "avg_bh_text = mean(bh_text)\n",
    "plt.title(f'BH Token Length Distribution: Average legnth: {avg_bh_text}')\n",
    "plt.hist(bh_text, bins = 50)\n",
    "plt.show()\n",
    "ep_text = ntuhdataset['EP Text'].apply(lambda x: len(re.sub(r'\\s+', ' ', re.sub(r'<[^>]+>', '', str(x))).split()))\n",
    "avg_ep_text = mean(ep_text)\n",
    "plt.figure(figsize=(5,5))\n",
    "plt.title(f'EP Token Length Distribution: Average length: {avg_ep_text}')\n",
    "plt.hist(ep_text, bins = 50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "bh_text = ntuhdataset['BH Text'].apply(lambda x: len(re.sub(r'\\s+', ' ', str(x)).split('<sep>')))\n",
    "plt.figure(figsize=(5,5))\n",
    "avg_bh_text = mean(bh_text)\n",
    "plt.title(f'BH Sentence Number Distribution: Average legnth: {avg_bh_text}')\n",
    "plt.hist(bh_text, bins = 50)\n",
    "plt.show()\n",
    "ep_text = ntuhdataset['EP Text'].apply(lambda x: len(re.sub(r'\\s+', ' ', str(x)).split('<sep>')))\n",
    "avg_ep_text = mean(ep_text)\n",
    "plt.figure(figsize=(5,5))\n",
    "plt.title(f'EP Sentence Number Distribution: Average length: {avg_ep_text}')\n",
    "plt.hist(ep_text, bins = 50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocessing\n",
    "\n",
    "Create validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NTUHDataset(data.Dataset):\n",
    "    #urls = ['Datasets\\\\NTUH\\\\corpus.txt']\n",
    "    name = 'ntuh'\n",
    "    dirname = 'ntuh'\n",
    "    diagnosis_types = ['major_depressive', 'schizophrenia', 'biploar', 'minor_depressive', 'dementia']\n",
    "    \n",
    "    @staticmethod\n",
    "    def sort_key(ex):\n",
    "        return len(ex.all_text) # TODO add ep_text?\n",
    "\n",
    "    def __init__(self, path, id_field, bh_text_field, ep_text_field, all_text_field,\n",
    "                 major_label_field, sch_label_field, bipolar_label_field, minor_label_field, dementia_label_field,\n",
    "                 **kwargs):\n",
    "        fields = [('patient_id', id_field), \n",
    "                  ('bh_text', bh_text_field),\n",
    "                  ('ep_text', ep_text_field),\n",
    "                  ('all_text', all_text_field),\n",
    "                  ('major_depressive', major_label_field),\n",
    "                  ('schizophrenia', sch_label_field),\n",
    "                  ('biploar', bipolar_label_field),\n",
    "                  ('minor_depressive', minor_label_field),\n",
    "                  ('dementia', dementia_label_field)]\n",
    "        examples = []\n",
    "        \n",
    "        for fname in glob.iglob(path + '.txt'):\n",
    "            with io.open(fname, 'r', encoding=\"utf-8\") as f:\n",
    "                for line in f:\n",
    "                    pid, bh_text, ep_text, major_d, sc, bp, minor_d, de = line.strip().split('\\t')\n",
    "                    all_text = \"%s <sep> %s\" % (bh_text, ep_text)\n",
    "                    examples.append(data.Example.fromlist([pid, bh_text, ep_text, all_text, major_d, sc, bp, minor_d, de], \n",
    "                                                          fields))\n",
    "        super(NTUHDataset, self).__init__(examples, fields, **kwargs)\n",
    "\n",
    "    @classmethod\n",
    "    def splits(cls, id_field,\n",
    "               bh_text_field, ep_text_field, all_text_field,\n",
    "               major_label_field, sch_label_field, bipolar_label_field, minor_label_field, dementia_label_field,\n",
    "               root='..\\\\Datasets\\\\NTUH',\n",
    "               train='train_preprocessing', test='test_preprocessing', **kwargs):\n",
    "        \"\"\"Create dataset objects for splits of the NTUH dataset.\n",
    "        Arguments:\n",
    "            text_field: The field that will be used for the sentence.\n",
    "            label_field: The field that will be used for label data.\n",
    "            root: Root dataset storage directory. Default is '.data'.\n",
    "            train: The directory that contains the training examples\n",
    "            test: The directory that contains the test examples\n",
    "            Remaining keyword arguments: Passed to the splits method of\n",
    "                Dataset.\n",
    "        \"\"\"\n",
    "        return super(NTUHDataset, cls).splits(\n",
    "            path = root, root=root, id_field=id_field,\n",
    "            bh_text_field = bh_text_field, ep_text_field = ep_text_field, all_text_field = all_text_field, \n",
    "            major_label_field = major_label_field, sch_label_field = sch_label_field, \n",
    "            bipolar_label_field = bipolar_label_field, minor_label_field = minor_label_field, \n",
    "            dementia_label_field = dementia_label_field,\n",
    "            train=train, validation=None, test=test, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ID_TEXT = data.Field(batch_first = True)\n",
    "BH_TEXT = data.Field(batch_first = True)\n",
    "EP_TEXT = data.Field(batch_first = True)\n",
    "ALL_TEXT = data.Field(batch_first = True)\n",
    "\n",
    "MAJ_LABEL = data.LabelField(dtype = torch.float)\n",
    "SCH_LABEL = data.LabelField(dtype = torch.float)\n",
    "BIP_LABEL = data.LabelField(dtype = torch.float)\n",
    "MIN_LABEL = data.LabelField(dtype = torch.float)\n",
    "DEM_LABEL = data.LabelField(dtype = torch.float)\n",
    "full_train_data, test_data = NTUHDataset.splits(ID_TEXT, BH_TEXT, EP_TEXT, ALL_TEXT, \n",
    "                                           MAJ_LABEL, SCH_LABEL, BIP_LABEL, MIN_LABEL, DEM_LABEL)\n",
    "train_data, valid_data = full_train_data.split(random_state = random.seed(SEED), split_ratio = TRAIN_RATIO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dump Split Training and Validation Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(TRAIN, 'wt', encoding=\"utf-8\") as out_file:\n",
    "    for example in train_data.examples:\n",
    "        out_file.write('{example.patient_id[0]}\\t{0}\\t{1}\\t{example.major_depressive}\\t{example.schizophrenia}\\t{example.biploar}\\t{example.minor_depressive}\\t{example.dementia}\\n'\n",
    "                            .format(' '.join(example.bh_text), ' '.join(example.ep_text), example=example))\n",
    "        \n",
    "with open(VALID, 'wt', encoding=\"utf-8\") as out_file:\n",
    "    for example in valid_data.examples:\n",
    "        out_file.write('{example.patient_id[0]}\\t{0}\\t{1}\\t{example.major_depressive}\\t{example.schizophrenia}\\t{example.biploar}\\t{example.minor_depressive}\\t{example.dementia}\\n'\n",
    "                            .format(' '.join(example.bh_text), ' '.join(example.ep_text), example=example)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntuhdataset = pd.read_csv(TRAIN, sep ='\\t', names = names)\n",
    "ntuhdataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntuhdataset = pd.read_csv(VALID, sep ='\\t', names = names)\n",
    "ntuhdataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Input Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_input_files(TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, read_all_text_from_corpus,\n",
    "                       MIN_WORD_COUNT, True, VALID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Randmly Initialized Word Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(embeddings.shape)\n",
    "print(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'<pad>: {word_map[\"<pad>\"]}\\n<unk>: {word_map[\"<unk>\"]}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset and DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "len(train_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "start_epoch = 0  # start at this epoch\n",
    "grad_clip = None\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = HierarchialAttentionNetwork(n_classes=n_classes,\n",
    "                                            vocab_size=len(word_map),\n",
    "                                            emb_size=emb_size,\n",
    "                                            word_rnn_size=WORD_RNN_SIZE,\n",
    "                                            sentence_rnn_size=SENTENCE_RNN_SIZE,\n",
    "                                            word_rnn_layers=WORD_RNN_LAYERS,\n",
    "                                            sentence_rnn_layers=SENTENCE_RNN_LAYERS,\n",
    "                                            word_att_size=WORD_ATTENTION_SIZE,\n",
    "                                            sentence_att_size=SENTENCE_ATTENTION_SIZE,\n",
    "                                            dropout=DROPOUT)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "cudnn.benchmark = True\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map)\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'han_rand', grad_clip, valid_loader, early_stop = True, period = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, figsize=(15,10))\n",
    "analysis_plotter(fig, ax1, train_losses, valid_losses, 'Training/Validation Loss', {'label': 'Training Loss'}, {'label': 'Validation Loss'})\n",
    "analysis_plotter(fig, ax2, train_accs, valid_accs, 'Training/Validation Micro-F-Measure', {'label': 'Training F-Measure'}, {'label': 'Validation F-Measure'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, criterion, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'han_rand_fscore.pt')\n",
    "\n",
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\n",
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deal with Imbalance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "model = HierarchialAttentionNetwork(n_classes=n_classes,\n",
    "                                            vocab_size=len(word_map),\n",
    "                                            emb_size=emb_size,\n",
    "                                            word_rnn_size=WORD_RNN_SIZE,\n",
    "                                            sentence_rnn_size=SENTENCE_RNN_SIZE,\n",
    "                                            word_rnn_layers=WORD_RNN_LAYERS,\n",
    "                                            sentence_rnn_layers=SENTENCE_RNN_LAYERS,\n",
    "                                            word_att_size=WORD_ATTENTION_SIZE,\n",
    "                                            sentence_att_size=SENTENCE_ATTENTION_SIZE,\n",
    "                                            dropout=DROPOUT)\n",
    "\n",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT)\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'han_rand', grad_clip, valid_loader, early_stop = True, period = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, figsize=(15,10))\n",
    "analysis_plotter(fig, ax1, train_losses, valid_losses, 'Training/Validation Loss', {'label': 'Training Loss'}, {'label': 'Validation Loss'})\n",
    "analysis_plotter(fig, ax2, train_accs, valid_accs, 'Training/Validation Micro-F-Measure', {'label': 'Training F-Measure'}, {'label': 'Validation F-Measure'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'han_rand_fscore.pt')\n",
    "\n",
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\n",
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word2vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_word2vec_model(data_folder, embedding_dim, algorithm='skipgram'):\n",
    "    sg = 1 if algorithm is 'skipgram' else 0\n",
    "\n",
    "    sentences = torch.load(os.path.join(data_folder, 'word2vec_data.pth.tar'))\n",
    "    sentences = list(itertools.chain.from_iterable(sentences))\n",
    "\n",
    "    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)\n",
    "\n",
    "    model = gensim.models.word2vec.Word2Vec(sentences=sentences, size=embedding_dim, workers=8, window=10, min_count=5,\n",
    "                                            sg=sg, iter = 50, seed = SEED)\n",
    "    model.init_sims(True)\n",
    "    model.wv.save(os.path.join(data_folder, f'word2vec_{algorithm}_model')) \n",
    "    model.wv.save_word2vec_format(os.path.join(data_folder, f'word2vec_{algorithm}_model.bin'))\n",
    "\n",
    "train_word2vec_model(data_folder=DATA_FOLDER, embedding_dim= EMBEDDING_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2vec_file = os.path.join(DATA_FOLDER, 'word2vec_skipgram_model')  # path to pre-trained word2vec embeddings\n",
    "\n",
    "def load_word2vec_embeddings(word2vec_file, word_map):\n",
    "    w2v = gensim.models.KeyedVectors.load(word2vec_file, mmap='r')\n",
    "\n",
    "    print(\"\\nEmbedding length is %d.\\n\" % w2v.vector_size)\n",
    "    embeddings = torch.FloatTensor(len(word_map), w2v.vector_size)\n",
    "    init_embedding(embeddings)\n",
    "\n",
    "    print(\"Loading embeddings...\")\n",
    "    for word in word_map:\n",
    "        if word in w2v.vocab:\n",
    "            embeddings[word_map[word]] = torch.FloatTensor(w2v[word])\n",
    "\n",
    "    print(\"Done.\\n Embedding vocabulary: %d.\\n\" % len(word_map))\n",
    "\n",
    "    return embeddings, w2v.vector_size\n",
    "\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "embeddings, emb_size = load_word2vec_embeddings(word2vec_file, word_map)\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize and Train Our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model = HierarchialAttentionNetwork(n_classes=n_classes,\n",
    "                                            vocab_size=len(word_map),\n",
    "                                            emb_size=emb_size,\n",
    "                                            word_rnn_size=WORD_RNN_SIZE,\n",
    "                                            sentence_rnn_size=SENTENCE_RNN_SIZE,\n",
    "                                            word_rnn_layers=WORD_RNN_LAYERS,\n",
    "                                            sentence_rnn_layers=SENTENCE_RNN_LAYERS,\n",
    "                                            word_att_size=WORD_ATTENTION_SIZE,\n",
    "                                            sentence_att_size=SENTENCE_ATTENTION_SIZE,\n",
    "                                            dropout=DROPOUT)\n",
    "\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings) \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT)\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'han_w2v', grad_clip, valid_loader, early_stop = True, period = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, figsize=(15,10))\n",
    "analysis_plotter(fig, ax1, train_losses, valid_losses, 'Training/Validation Loss', {'label': 'Training Loss'}, {'label': 'Validation Loss'})\n",
    "analysis_plotter(fig, ax2, train_accs, valid_accs, 'Training/Validation Micro-F-Measure', {'label': 'Training F-Measure'}, {'label': 'Validation F-Measure'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'han_w2v_fscore.pt')\n",
    "\n",
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\n",
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_f_scores, predicts = test(model, test_loader, criterion, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Glove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings, emb_size, vocab, PAD, GLOVE_UNK_TOKEN = load_glove_embeddings()\n",
    "\n",
    "GLOVE_INPUT_DIM = len(vocab)\n",
    "GLOVE_EMBEDDING_DIM = emb_size\n",
    "GLOVE_PAD_IDX = vocab[PAD]\n",
    "GLOVE_UNK_IDX = vocab[GLOVE_UNK_TOKEN]\n",
    "\n",
    "print(\"Input dimension: %s\\nUnknown word (%s) index: %s\\nPadding index: %s\\nEmbedding dimension: %s\" % \n",
    "      (GLOVE_INPUT_DIM, GLOVE_UNK_TOKEN, GLOVE_UNK_IDX, GLOVE_PAD_IDX, GLOVE_EMBEDDING_DIM))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_input_files_for_glove(vocab.stoi, TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, \n",
    "                             read_all_text_from_corpus, MIN_WORD_COUNT, VALID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(DATA_FOLDER, 'glove_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "print(embeddings)\n",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "print(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 8 # batch size\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_TRAIN'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model = HierarchialAttentionNetwork(n_classes=n_classes,\n",
    "                                            vocab_size=len(word_map),\n",
    "                                            emb_size=emb_size,\n",
    "                                            word_rnn_size=WORD_RNN_SIZE,\n",
    "                                            sentence_rnn_size=SENTENCE_RNN_SIZE,\n",
    "                                            word_rnn_layers=WORD_RNN_LAYERS,\n",
    "                                            sentence_rnn_layers=SENTENCE_RNN_LAYERS,\n",
    "                                            word_att_size=WORD_ATTENTION_SIZE,\n",
    "                                            sentence_att_size=SENTENCE_ATTENTION_SIZE,\n",
    "                                            dropout=DROPOUT)\n",
    "\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.sentence_attention.word_attention.embeddings.weight)\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT)\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "start_epoch = 0\n",
    "grad_clip = None\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'glove', grad_clip, valid_loader, early_stop=True, period = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, figsize=(15,10))\n",
    "analysis_plotter(fig, ax1, train_losses, valid_losses, 'Training/Validation Loss', {'label': 'Training Loss'}, {'label': 'Validation Loss'})\n",
    "analysis_plotter(fig, ax2, train_accs, valid_accs, 'Training/Validation Micro-F-Measure', {'label': 'Training F-Measure'}, {'label': 'Validation F-Measure'})                \n",
    "\n",
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_fscore.pt')\n",
    "\n",
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "DROPOUT = 0  # dropout\n",
    "\n",
    "BATCH_SIZE = 2\n",
    "\n",
    "BERT_MODEL = 'bert-base-uncased'\n",
    "\n",
    "bert_tokenizer = BertTokenizer.from_pretrained(BERT_MODEL, do_lower_case = True)\n",
    "bert = BertModel.from_pretrained(BERT_MODEL, output_hidden_states = True)\n",
    "\n",
    "BERT_BOS_TOKEN = bert_tokenizer.cls_token\n",
    "BERT_EOS_TOKEN = bert_tokenizer.sep_token\n",
    "BERT_PAD = bert_tokenizer.pad_token\n",
    "BERT_UNK = bert_tokenizer.unk_token\n",
    "\n",
    "BERT_BOS_IDX = bert_tokenizer.convert_tokens_to_ids(BERT_BOS_TOKEN)\n",
    "BERT_EOS_IDX = bert_tokenizer.convert_tokens_to_ids(BERT_EOS_TOKEN)\n",
    "BERT_PAD_IDX = bert_tokenizer.convert_tokens_to_ids(BERT_PAD)\n",
    "BERT_UNK_IDX = bert_tokenizer.convert_tokens_to_ids(BERT_UNK)\n",
    "\n",
    "print(f'{BERT_BOS_TOKEN}: {BERT_BOS_IDX}, {BERT_EOS_TOKEN}:{BERT_EOS_IDX}, {BERT_PAD}:{BERT_PAD_IDX}, \\\n",
    "{BERT_UNK}:{BERT_UNK_IDX}') #, bos_token, eos_token')\n",
    "\n",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]\n",
    "print(BERT_WORD_LIMIT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_input_files_for_bert(bert_tokenizer, train, test, output_folder, sentence_limit, word_limit, read_tsv, \n",
    "                       valid = None):\n",
    "    print(f'\\nReading and preprocessing training data {train}...\\n')\n",
    "    train_docs, train_labels, word_counter = read_tsv(bert_tokenizer, train, sentence_limit, word_limit)\n",
    "    print(len(train_docs))\n",
    "    \n",
    "    word_map = dict()\n",
    "    for word, count in word_counter.items():\n",
    "        word_map[word] = bert_tokenizer.convert_tokens_to_ids(word)    \n",
    "    \n",
    "    print('Encoding and padding training data...\\n')\n",
    "    encoded_train_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: [BERT_BOS_IDX] + list(map(lambda w: word_map.get(w, BERT_UNK_IDX), s)) + [BERT_EOS_IDX]\n",
    "            + [BERT_PAD_IDX] * (word_limit - len(s) - 2), doc)) + [[BERT_PAD_IDX] * word_limit] * (sentence_limit - len(doc)), \n",
    "                                  train_docs))\n",
    "    \n",
    "    sentences_per_train_document = list(map(lambda doc: len(doc), train_docs))\n",
    "    words_per_train_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s)+2, doc)) + [0] * (sentence_limit - len(doc)), train_docs))\n",
    "    \n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    output_content = {'docs': encoded_train_docs,\n",
    "                'labels': train_labels,\n",
    "                'sentences_per_document': sentences_per_train_document,\n",
    "                'words_per_sentence': words_per_train_sentence}\n",
    "    torch.save(output_content,\n",
    "               os.path.join(output_folder, 'BERT_TRAIN_data.pth.tar'))\n",
    "    print('Encoded, padded training data (BERT_TRAIN_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del train_docs, encoded_train_docs, train_labels, sentences_per_train_document, words_per_train_sentence\n",
    "\n",
    "    # Read test data\n",
    "    print(f'Reading and preprocessing test data {test}...\\n')\n",
    "    test_docs, test_labels, word_counter = read_tsv(bert_tokenizer, test, sentence_limit, word_limit)\n",
    "    print(f'Updating word map based on test data {test}...\\n')\n",
    "    for word, count in word_counter.items():\n",
    "        word_map[word] = bert_tokenizer.convert_tokens_to_ids(word)    \n",
    "    \n",
    "    # Encode and pad\n",
    "    print('\\nEncoding and padding test data...\\n')\n",
    "    encoded_test_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: [BERT_BOS_IDX] + list(map(lambda w: word_map.get(w, BERT_UNK_IDX), s)) + [BERT_EOS_IDX]\n",
    "            + [BERT_PAD_IDX] * (word_limit - len(s) - 2), doc)) + [[BERT_PAD_IDX] * word_limit] * (sentence_limit - len(doc)), \n",
    "                                 test_docs))\n",
    "    sentences_per_test_document = list(map(lambda doc: len(doc), test_docs))\n",
    "    \n",
    "    words_per_test_sentence = list(\n",
    "        map(lambda doc: list(map(lambda s: len(s)+2, doc)) + [0] * (sentence_limit - len(doc)), test_docs))\n",
    "\n",
    "    # Save\n",
    "    print('Saving...\\n')\n",
    "    torch.save({'docs': encoded_test_docs,\n",
    "                'labels': test_labels,\n",
    "                'sentences_per_document': sentences_per_test_document,\n",
    "                'words_per_sentence': words_per_test_sentence},\n",
    "               os.path.join(output_folder, 'BERT_TEST_data.pth.tar'))\n",
    "    print('Encoded, padded test data (BERT_TEST_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    del test_docs, encoded_test_docs, test_labels, sentences_per_test_document, words_per_test_sentence\n",
    "    if valid:\n",
    "        print(f'Reading and preprocessing validation data {valid}...\\n')\n",
    "        valid_docs, valid_labels, word_counter = read_tsv(bert_tokenizer, valid, sentence_limit, word_limit)\n",
    "        print(f'Updating word map based on validation data {valid}...\\n')\n",
    "        for word, count in word_counter.items():\n",
    "            word_map[word] = bert_tokenizer.convert_tokens_to_ids(word)    \n",
    "    \n",
    "        # Encode and pad\n",
    "        print('\\nEncoding and padding validation data...\\n')\n",
    "        encoded_valid_docs = list(map(lambda doc: list(\n",
    "        map(lambda s: [BERT_BOS_IDX] + list(map(lambda w: word_map.get(w, BERT_UNK_IDX), s)) + [BERT_EOS_IDX]\n",
    "            + [BERT_PAD_IDX] * (word_limit - len(s) - 2), doc)) + [[BERT_PAD_IDX] * word_limit] * (sentence_limit - len(doc)), \n",
    "                                 valid_docs))\n",
    "        sentences_per_valid_document = list(map(lambda doc: len(doc), valid_docs))\n",
    "        words_per_valid_sentence = list(\n",
    "            map(lambda doc: list(map(lambda s: len(s)+2, doc)) + [0] * (sentence_limit - len(doc)), valid_docs))\n",
    "        # Save\n",
    "        print('Saving...\\n')\n",
    "        torch.save({'docs': encoded_valid_docs,\n",
    "                'labels': valid_labels,\n",
    "                'sentences_per_document': sentences_per_valid_document,\n",
    "                'words_per_sentence': words_per_valid_sentence},\n",
    "                   os.path.join(output_folder, 'BERT_VALID_data.pth.tar'))\n",
    "        print('Encoded, padded valid data (BERT_VALID_data.pth.tar) saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "    \n",
    "    print('\\nThe size of the vocabulary is %d.\\n' % (len(word_map)))\n",
    "    with open(os.path.join(output_folder, 'BERT_word_map.json'), 'w') as j:\n",
    "        json.dump(word_map, j)\n",
    "    print('Word map saved to %s.\\n' % os.path.abspath(output_folder))\n",
    "\n",
    "    print('All done!\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_all_text_from_corpus_with_bert(word_tokenizer, tsv_file, sentence_limit, word_limit):\n",
    "    docs = []\n",
    "    labels = []\n",
    "    word_counter = Counter()\n",
    "    with io.open(tsv_file, 'r', encoding=\"utf-8\") as file:\n",
    "        for i, line in enumerate(tqdm(file)):\n",
    "            sentences = list()\n",
    "            pid, bh_text, ep_text, major_d, sc, bp, minor_d, de = preprocess(line).strip().split('\\t')\n",
    "            all_text = \"%s <sep> %s\" % (bh_text, ep_text)\n",
    "            sentences.extend([s for s in sent_tokenizer(all_text)])\n",
    "            \n",
    "            words = list()\n",
    "            for s in sentences[:sentence_limit]:\n",
    "                w = word_tokenizer.tokenize(s)[:word_limit-2]\n",
    "                # If sentence is empty (due to ?)\n",
    "                if len(w) == 0:\n",
    "                    continue\n",
    "                words.append(w)\n",
    "                word_counter.update(w)\n",
    "            # If all sentences were empty\n",
    "            if len(words) == 0:\n",
    "                continue\n",
    "            \n",
    "            labels.append([float(major_d), float(sc), float(bp), float(minor_d), float(de)]) \n",
    "            docs.append(words)\n",
    "    return docs, labels, word_counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None \n",
    "\n",
    "with open(os.path.join(DATA_FOLDER, 'bert_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "han_bert_cache = torch.load('han_cache_train.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTHierarchialAttentionNetwork(nn.Module):\n",
    "    def __init__(self, bert, n_classes, vocab_size, word_rnn_size, sentence_rnn_size, word_rnn_layers,\n",
    "                 sentence_rnn_layers, word_att_size, sentence_att_size, han_bert_cache, dropout=0.5):\n",
    "        super(BERTHierarchialAttentionNetwork, self).__init__()\n",
    "\n",
    "        self.sentence_attention = BERTSentenceAttention(bert, vocab_size, word_rnn_size, sentence_rnn_size,\n",
    "                                                    word_rnn_layers, sentence_rnn_layers, word_att_size,\n",
    "                                                    sentence_att_size, dropout, han_bert_cache)\n",
    "\n",
    "        self.fc = nn.Linear(2 * sentence_rnn_size, n_classes)\n",
    "        self.han_bert_cache = han_bert_cache\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, documents, sentences_per_document, words_per_sentence):\n",
    "        document_embeddings, word_alphas, sentence_alphas = self.sentence_attention(documents, sentences_per_document,\n",
    "                                                                                    words_per_sentence)  \n",
    "        scores = self.fc(self.dropout(document_embeddings))  # (n_documents, n_classes)\n",
    "\n",
    "        return scores, word_alphas, sentence_alphas\n",
    "\n",
    "\n",
    "class BERTSentenceAttention(nn.Module):\n",
    "    def __init__(self, bert, vocab_size, word_rnn_size, sentence_rnn_size, word_rnn_layers, sentence_rnn_layers,\n",
    "                 word_att_size, sentence_att_size, dropout, han_bert_cache):\n",
    "        super(BERTSentenceAttention, self).__init__()\n",
    "\n",
    "        self.word_attention = BERTWordAttention(bert, vocab_size, word_rnn_size, word_rnn_layers, word_att_size,\n",
    "                                            dropout, han_bert_cache)\n",
    "\n",
    "        self.sentence_rnn = nn.GRU(2 * word_rnn_size, sentence_rnn_size, num_layers=sentence_rnn_layers,\n",
    "                                   bidirectional=True, dropout=dropout, batch_first=True)\n",
    "\n",
    "        self.sentence_attention = nn.Linear(2 * sentence_rnn_size, sentence_att_size)\n",
    "\n",
    "        self.sentence_context_vector = nn.Linear(sentence_att_size, 1,\n",
    "                                                 bias=False) \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, documents, sentences_per_document, words_per_sentence):\n",
    "        packed_sentences = pack_padded_sequence(documents,\n",
    "                                                lengths=sentences_per_document.tolist(), \n",
    "                                                batch_first=True,\n",
    "                                                enforce_sorted=False)  \n",
    "        packed_words_per_sentence = pack_padded_sequence(words_per_sentence,\n",
    "                                                         lengths=sentences_per_document.tolist(),\n",
    "                                                         batch_first=True,\n",
    "                                                         enforce_sorted=False)  \n",
    "        sentences, word_alphas = self.word_attention(packed_sentences.data,\n",
    "                                                     packed_words_per_sentence.data)  \n",
    "        sentences = self.dropout(sentences)\n",
    "        \n",
    "        packed_sentences, _ = self.sentence_rnn(PackedSequence(data=sentences,\n",
    "                                                               batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                               sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                               unsorted_indices=packed_sentences.unsorted_indices))\n",
    "\n",
    "        att_s = self.sentence_attention(packed_sentences.data) \n",
    "        att_s = torch.tanh(att_s) \n",
    "        att_s = self.sentence_context_vector(att_s).squeeze(1)\n",
    "\n",
    "        max_value = att_s.max() \n",
    "        att_s = torch.exp(att_s - max_value)\n",
    "\n",
    "        att_s, _ = pad_packed_sequence(PackedSequence(data=att_s,\n",
    "                                                      batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                      sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                      unsorted_indices=packed_sentences.unsorted_indices),\n",
    "                                       batch_first=True) \n",
    "\n",
    "        sentence_alphas = att_s / torch.sum(att_s, dim=1, keepdim=True)\n",
    "\n",
    "        documents, _ = pad_packed_sequence(packed_sentences, batch_first=True)\n",
    "\n",
    "        documents = documents * sentence_alphas.unsqueeze(2)\n",
    "        documents = documents.sum(dim=1)\n",
    "\n",
    "        word_alphas, _ = pad_packed_sequence(PackedSequence(data=word_alphas,\n",
    "                                                            batch_sizes=packed_sentences.batch_sizes,\n",
    "                                                            sorted_indices=packed_sentences.sorted_indices,\n",
    "                                                            unsorted_indices=packed_sentences.unsorted_indices),\n",
    "                                             batch_first=True)\n",
    "\n",
    "        return documents, word_alphas, sentence_alphas\n",
    "\n",
    "\n",
    "class BERTWordAttention(nn.Module):\n",
    "    def __init__(self, bert, vocab_size, word_rnn_size, word_rnn_layers, word_att_size, dropout, han_bert_cache):\n",
    "        super(BERTWordAttention, self).__init__()\n",
    "        self.bert = bert\n",
    "        self.han_bert_cache = han_bert_cache\n",
    "        self.bert.eval()\n",
    "        self.embedding_dim = self.bert.config.to_dict()['hidden_size']\n",
    "\n",
    "        self.word_rnn = nn.GRU(self.embedding_dim, word_rnn_size, num_layers=word_rnn_layers, bidirectional=True,\n",
    "                               dropout=dropout, batch_first=True)\n",
    "\n",
    "        self.word_attention = nn.Linear(2 * word_rnn_size, word_att_size)\n",
    "\n",
    "        self.word_context_vector = nn.Linear(word_att_size, 1, bias=False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def create_attention_masks(self, ids):\n",
    "        attention_masks = []\n",
    "        for id in ids:\n",
    "            id_mask = [float(i>0) for i in id]            \n",
    "            attention_masks.append(id_mask)\n",
    "        return torch.tensor(attention_masks).to(device)\n",
    "            \n",
    "    def embeddings(self, batch):\n",
    "        # ID:102 is used to separate sentences\n",
    "        # [batch size, sent len]\n",
    "        batch_embeddings = []\n",
    "        for sents in batch:\n",
    "            key = ' '.join(str(x) for x in sents.data.tolist())\n",
    "            if key in self.han_bert_cache:\n",
    "                sent_embeddings = self.han_bert_cache[key].to(device)\n",
    "            else:            \n",
    "                print('Not Found')\n",
    "                sep_idxes = (sents == BERT_EOS_IDX).nonzero().squeeze(1).data.tolist()\n",
    "                seq_lengths = []\n",
    "                sents_ids = []\n",
    "                pv = -1\n",
    "                for k, v in enumerate(sep_idxes):                \n",
    "                    sent_embedding = [PAD_IDX]*WORD_LIMIT\n",
    "                    if k == 0:\n",
    "                        seq_lengths.append(v+1)\n",
    "                        sent_embedding[:v+1] = sents[:v+1].data.tolist()\n",
    "                    else:\n",
    "                        seq_lengths.append(v-pv)\n",
    "                        sent_embedding[:v-pv] = sents[pv+1:v+1].data.tolist()\n",
    "                    sents_ids.append(sent_embedding)\n",
    "                    pv = v\n",
    "                attention_masks = self.create_attention_masks(sents_ids)\n",
    "                sents_ids = torch.tensor(sents_ids).to(device)\n",
    "                sent_embeddings = []\n",
    "                with torch.no_grad():\n",
    "                    last_hidden_state, _, hidden_states = self.bert(sents_ids, attention_masks)\n",
    "                    token_embeddings = torch.stack(hidden_states[:-1], dim=0)\n",
    "                    token_embeddings = token_embeddings.permute(1, 2, 0, 3)\n",
    "                    for id, tks in enumerate(token_embeddings):\n",
    "                        token_vecs = []\n",
    "                        for i in range(seq_lengths[id]):\n",
    "                            #cat_vec = torch.cat((tks[i][-1], tks[i][-2], tks[i][-3], tks[i][-4]), dim =0)\n",
    "                            cat_vec = tks[i][-1] + tks[i][-2] + tks[i][-3] + tks[i][-4]\n",
    "                            token_vecs.append(cat_vec)\n",
    "                        token_vecs=torch.stack(token_vecs, 0)\n",
    "                        sent_embeddings.append(token_vecs)\n",
    "                    sent_embeddings = torch.cat(sent_embeddings, 0)                \n",
    "\n",
    "                    if sent_embeddings.shape[0] != WORD_LIMIT:\n",
    "                        sent_embeddings = torch.cat((sent_embeddings, \\\n",
    "                                torch.zeros(WORD_LIMIT - sent_embeddings.shape[0], self.embedding_dim).to(device)), 0)\n",
    "                    # # sentences, # words, # layers, # features\n",
    "                self.han_bert_cache[key] = sent_embeddings\n",
    "            batch_embeddings.append(sent_embeddings)\n",
    "        batch_embeddings = torch.stack(batch_embeddings, 0)\n",
    "        return batch_embeddings\n",
    "\n",
    "    def forward(self, sentences, words_per_sentence):\n",
    "        sentences = self.dropout(self.embeddings(sentences))  \n",
    "        packed_words = pack_padded_sequence(sentences,\n",
    "                                            lengths=words_per_sentence.tolist(),\n",
    "                                            batch_first=True,\n",
    "                                            enforce_sorted=False)\n",
    "        \n",
    "        packed_words, _ = self.word_rnn(\n",
    "            packed_words)\n",
    "        \n",
    "        att_w = self.word_attention(packed_words.data) \n",
    "        att_w = torch.tanh(att_w) \n",
    "        att_w = self.word_context_vector(att_w).squeeze(1) \n",
    "\n",
    "        max_value = att_w.max() \n",
    "        att_w = torch.exp(att_w - max_value)\n",
    "        \n",
    "        att_w, _ = pad_packed_sequence(PackedSequence(data=att_w,\n",
    "                                                      batch_sizes=packed_words.batch_sizes,\n",
    "                                                      sorted_indices=packed_words.sorted_indices,\n",
    "                                                      unsorted_indices=packed_words.unsorted_indices),\n",
    "                                       batch_first=True)\n",
    "        \n",
    "        word_alphas = att_w / torch.sum(att_w, dim=1, keepdim=True)  \n",
    "        sentences, _ = pad_packed_sequence(packed_words,\n",
    "                                           batch_first=True) \n",
    "\n",
    "        sentences = sentences * word_alphas.unsqueeze(2)  \n",
    "        sentences = sentences.sum(dim=1) \n",
    "        \n",
    "        return sentences, word_alphas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bert_train_epoch(start_epoch, epochs, data_loader, model, criterion, optimizer, word_map, model_name, grad_clip, valid_iterator = None, \n",
    "                interval = 1, early_stop = False, period = 20, gap = 0.005, threshold = 0.5, best_valid_fscore = 0):\n",
    "    best_valid_loss = float('inf')\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    train_accs = []\n",
    "    valid_accs = []\n",
    "    observed_time = 0\n",
    "    \n",
    "    # Epochs\n",
    "    for epoch in range(start_epoch, epochs):\n",
    "        start_time = time.time()\n",
    "        # One epoch's training\n",
    "        try:\n",
    "            train_loss, train_acc = train(train_loader=data_loader, model=model, \n",
    "                                          criterion=criterion, optimizer=optimizer, grad_clip=grad_clip)#, epoch=epoch, \n",
    "                  #total_epoch = epochs)\n",
    "            train_losses.append(train_loss)\n",
    "            train_accs.append(train_acc)\n",
    "\n",
    "            end_time = time.time()\n",
    "            epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "\n",
    "            if valid_iterator:\n",
    "                valid_loss, valid_acc = evaluate(model, valid_iterator, criterion)#, model_type)\n",
    "                valid_losses.append(valid_loss)\n",
    "                valid_accs.append(valid_acc)\n",
    "            else:\n",
    "                valid_loss = 0 \n",
    "\n",
    "            if (epoch + 1) % interval == 0:\n",
    "                print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "                print(f'\\tTrain Loss: {train_loss:.3f} | Train micro-F-score: {train_acc*100:.2f}%')\n",
    "                if valid_iterator:\n",
    "                    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. micro-F-score: {valid_acc*100:.2f}%')\n",
    "            elif epoch == epochs - 1:\n",
    "                print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "                print(f'\\tTrain Loss: {train_loss:.3f} | Train micro-F-score: {train_acc*100:.2f}%')\n",
    "                if valid_iterator:\n",
    "                    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. micro-F-score: {valid_acc*100:.2f}%')\n",
    "\n",
    "            # Save checkpoint\n",
    "            if valid_loss < best_valid_loss:\n",
    "                best_valid_loss = valid_loss\n",
    "                cache = model.han_bert_cache\n",
    "                model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = None\n",
    "                save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_loss.pt')\n",
    "                model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = cache\n",
    "            if early_stop and best_valid_fscore > threshold and best_valid_fscore - valid_acc > gap:\n",
    "                observed_time += 1\n",
    "                print(f'\\rBest validation F-measure: {best_valid_fscore:.3f}/Current F-measure: {valid_acc:.3f} [Times: {observed_time}/{period}]')  \n",
    "                if observed_time >= period:\n",
    "                    print(f'Early stop at epoch {epoch+1:02}.')\n",
    "                    break                        \n",
    "            if valid_acc > best_valid_fscore:\n",
    "                best_valid_fscore = valid_acc\n",
    "                cache = model.han_bert_cache\n",
    "                model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = None\n",
    "                save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_fscore.pt')            \n",
    "                model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = cache\n",
    "                observed_time = 0            \n",
    "        except RuntimeError as e:\n",
    "            print(f'Runtime error at epoch {epoch+1:02}: {str(e)}')\n",
    "            model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = None\n",
    "            save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_except.pt')\n",
    "            return train_losses, valid_losses, train_accs, valid_accs\n",
    "        \n",
    "        model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = None\n",
    "        save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_current.pt')\n",
    "        model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = cache\n",
    "    # Decay learning rate every epoch\n",
    "    #adjust_learning_rate(optimizer, 0.1)\n",
    "\n",
    "    #return model\n",
    "    return train_losses, valid_losses, train_accs, valid_accs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_TRAIN'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model = BERTHierarchialAttentionNetwork(bert=bert, n_classes=n_classes,\n",
    "                                    vocab_size=len(word_map),\n",
    "                                    word_rnn_size=WORD_RNN_SIZE,\n",
    "                                    sentence_rnn_size=SENTENCE_RNN_SIZE,\n",
    "                                    word_rnn_layers=WORD_RNN_LAYERS,\n",
    "                                    sentence_rnn_layers=SENTENCE_RNN_LAYERS,\n",
    "                                    word_att_size=WORD_ATTENTION_SIZE,\n",
    "                                    sentence_att_size=SENTENCE_ATTENTION_SIZE,\n",
    "                                    han_bert_cache = han_bert_cache, dropout=DROPOUT)\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_TRAIN'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_current.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCHS = 500\n",
    "BATCH_SIZE = 12\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT)\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "start_epoch = 113 #89\n",
    "best_valid_fscore = 0.5258\n",
    "period = 6\n",
    "grad_clip = None\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    bert_train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'bert', grad_clip, valid_loader, early_stop = True, period = period, best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "han_bert_cache_test = torch.load('han_cache_test.pt')\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_fscore.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache_test\n",
    "test_f_scores, predicts = test(model, test_loader, 0)\n",
    "\n",
    "for f in test_f_scores:\n",
    "    if f is MICRO or f is MACRO:\n",
    "        print(f'{f}-average:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')\n",
    "    else:\n",
    "        print(f'{NTUH_HANDataset.diagnosis_types[f]}:\\n\\tprecision: {test_f_scores[f][\"p\"]:0.3f}\\n\\trecall: {test_f_scores[f][\"r\"]:0.3f}\\n\\tf-score: {test_f_scores[f][\"f\"]:0.3f}\\n')                \n",
    "                                            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify(document):\n",
    "    doc = list()\n",
    "\n",
    "    sentences = list()    \n",
    "    pid, bh_text, ep_text, major_d, sc, bp, minor_d, de = preprocess(document).strip().split('\\t')\n",
    "    all_text = \"%s <sep> %s\" % (bh_text, ep_text)\n",
    "    sentences.extend([s for s in sent_tokenizer(all_text)])\n",
    "    \n",
    "    for s in sentences[:sentence_limit]:\n",
    "        w = word_tokenizer(s)[:word_limit]\n",
    "        if len(w) == 0:\n",
    "            continue\n",
    "        doc.append(w)\n",
    "\n",
    "    sentences_in_doc = len(doc)\n",
    "    sentences_in_doc = torch.LongTensor([sentences_in_doc]).to(device)  # (1)\n",
    "\n",
    "    words_in_each_sentence = list(map(lambda s: len(s), doc))\n",
    "    words_in_each_sentence = torch.LongTensor(words_in_each_sentence).unsqueeze(0).to(device)  # (1, n_sentences)\n",
    "\n",
    "    encoded_doc = list(\n",
    "        map(lambda s: list(map(lambda w: word_map.get(w, word_map['<unk>']), s)) + [0] * (word_limit - len(s)),\n",
    "            doc)) + [[0] * word_limit] * (sentence_limit - len(doc))\n",
    "    encoded_doc = torch.LongTensor(encoded_doc).unsqueeze(0).to(device)\n",
    "\n",
    "    scores, word_alphas, sentence_alphas = model(encoded_doc, sentences_in_doc,\n",
    "                                                 words_in_each_sentence)  # (1, n_classes), (1, n_sentences, max_sent_len_in_document), (1, n_sentences)\n",
    "        \n",
    "    word_alphas = word_alphas.squeeze(0)  # (n_sentences, max_sent_len_in_document)\n",
    "    sentence_alphas = sentence_alphas.squeeze(0)  # (n_sentences)\n",
    "    words_in_each_sentence = words_in_each_sentence.squeeze(0)  # (n_sentences)\n",
    "\n",
    "    return doc, scores.squeeze(0), word_alphas, sentence_alphas, words_in_each_sentence\n",
    "\n",
    "\n",
    "def visualize_attention(doc, scores, word_alphas, sentence_alphas, words_in_each_sentence):\n",
    "    # Find best prediction\n",
    "    rounded_preds = torch.round(torch.sigmoid(scores)).squeeze(0) # (n_classes)\n",
    "    prediction = ''\n",
    "    for i, score in enumerate(rounded_preds):\n",
    "        if score == 1:\n",
    "            prediction += '{category} ({score:.2f}\\n) '.format(category=rev_label_map[i], score=scores.tolist()[i])\n",
    "    \n",
    "    alphas = (sentence_alphas.unsqueeze(1) * word_alphas * words_in_each_sentence.unsqueeze(\n",
    "        1).float() / words_in_each_sentence.max().float())\n",
    "    alphas = alphas.to('cpu')\n",
    "\n",
    "    min_font_size = 15 \n",
    "    max_font_size = 55 \n",
    "    space_size = ImageFont.truetype(\"./calibril.ttf\", max_font_size).getsize(' ') \n",
    "    line_spacing = 15 \n",
    "    left_buffer = 300 \n",
    "    top_buffer = 2 * min_font_size + 3 * line_spacing  \n",
    "    image_width = left_buffer\n",
    "    image_height = top_buffer + line_spacing  \n",
    "    word_loc = [image_width, image_height] \n",
    "    rectangle_height = 0.75 * max_font_size \n",
    "    max_rectangle_width = 0.8 * left_buffer \n",
    "    rectangle_loc = [0.9 * left_buffer,\n",
    "                     image_height + rectangle_height]\n",
    "    word_viz_properties = list()\n",
    "    sentence_viz_properties = list()\n",
    "    for s, sentence in enumerate(doc):\n",
    "        sentence_factor = sentence_alphas[s].item() / sentence_alphas.max().item()\n",
    "        rectangle_saturation = str(int(sentence_factor * 100))\n",
    "        rectangle_lightness = str(25 + 50 - int(sentence_factor * 50))\n",
    "        rectangle_color = 'hsl(0,' + rectangle_saturation + '%,' + rectangle_lightness + '%)'\n",
    "        rectangle_bounds = [rectangle_loc[0] - sentence_factor * max_rectangle_width,\n",
    "                            rectangle_loc[1] - rectangle_height] + rectangle_loc\n",
    "\n",
    "        sentence_viz_properties.append({'bounds': rectangle_bounds.copy(),\n",
    "                                        'color': rectangle_color})\n",
    "\n",
    "        for w, word in enumerate(sentence):\n",
    "            word_factor = alphas[s, w].item() / alphas.max().item()\n",
    "\n",
    "            word_saturation = str(int(word_factor * 100))\n",
    "            word_lightness = str(25 + 50 - int(word_factor * 50))\n",
    "            word_color = 'hsl(0,' + word_saturation + '%,' + word_lightness + '%)'\n",
    "\n",
    "            word_font_size = int(min_font_size + word_factor * (max_font_size - min_font_size))\n",
    "            word_font = ImageFont.truetype(\"./calibril.ttf\", word_font_size)\n",
    "\n",
    "            word_viz_properties.append({'loc': word_loc.copy(),\n",
    "                                        'word': word,\n",
    "                                        'font': word_font,\n",
    "                                        'color': word_color})\n",
    "\n",
    "            word_size = word_font.getsize(word)\n",
    "            word_loc[0] += word_size[0] + space_size[0]\n",
    "            image_width = max(image_width, word_loc[0])\n",
    "        word_loc[0] = left_buffer\n",
    "        word_loc[1] += max_font_size + line_spacing\n",
    "        image_height = max(image_height, word_loc[1])\n",
    "        rectangle_loc[1] += max_font_size + line_spacing\n",
    "\n",
    "    img = Image.new('RGB', (image_width, image_height), (255, 255, 255))\n",
    "    \n",
    "    draw = ImageDraw.Draw(img)\n",
    "    for viz in word_viz_properties:\n",
    "        draw.text(xy=viz['loc'], text=viz['word'], fill=viz['color'], font=viz['font'])\n",
    "    for viz in sentence_viz_properties:\n",
    "        draw.rectangle(xy=viz['bounds'], fill=viz['color'])\n",
    "    category_font = ImageFont.truetype(\"./calibril.ttf\", min_font_size)\n",
    "    draw.text(xy=[line_spacing, line_spacing], text='Detected Category:', fill='grey', font=category_font)\n",
    "    draw.text(xy=[line_spacing, line_spacing + category_font.getsize('Detected Category:')[1] + line_spacing],\n",
    "              text=prediction.upper(), fill='black',\n",
    "              font=category_font)\n",
    "    del draw\n",
    "    \n",
    "    fig = plt.figure(figsize = (15, 15)) # create a 5 x 5 figure \n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.imshow(np.asarray(img), interpolation='none')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_lines = []\n",
    "with io.open(TEST, 'r', encoding=\"utf-8\") as file:\n",
    "    for i, line in enumerate(tqdm(file)):\n",
    "        test_lines.append(line)\n",
    "test_lines[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classify(test_lines[0])[1].max(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mpld3.enable_notebook()\n",
    "visualize_attention(*classify(test_lines[4]))"
   ]
  }
 ],
 "metadata": {
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}