HAN-Binary_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 = 1\n",
    "EMBEDDING_DIM = 100\n",
    "N_EPOCHS = 500\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-Bin'\n",
    "WORKER = 0  # number of workers for loading data in the DataLoader\n",
    "MAX_VOCAB_SIZE = 25000\n",
    "\n",
    "WORD_RNN_SIZE = 250  # word RNN size\n",
    "SENTENCE_RNN_SIZE = 200  # character RNN size\n",
    "WORD_RNN_LAYERS = 2  # number of layers in character RNN\n",
    "SENTENCE_RNN_LAYERS = 2  # number of layers in word RNN\n",
    "WORD_ATTENTION_SIZE = 250  # size of the word attention layer\n",
    "SENTENCE_ATTENTION_SIZE = 200  # size of the sentence attention layer\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",
    "\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, label):\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 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",
    "            if label == 0:\n",
    "                labels.append([float(major_d)])\n",
    "            elif label == 1:\n",
    "                labels.append([float(sc)])\n",
    "            elif label == 2:\n",
    "                labels.append([float(bp)])\n",
    "            elif label == 3:\n",
    "                labels.append([float(minor_d)])\n",
    "            elif label == 4:\n",
    "                labels.append([float(de)])\n",
    "            else:\n",
    "                assert False;\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, label,\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, label)\n",
    "    \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",
    "    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, f'TRAIN_{label}_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",
    "    # 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, label)\n",
    "\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",
    "    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, f'TEST_{label}_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",
    "    \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, label)\n",
    "    \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",
    "        \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, f'VALID_{label}_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",
    "def init_embedding(input_embedding):\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",
    "    \n",
    "class NTUH_HANDataset(Dataset):\n",
    "    diagnosis_types = ['major_depressive', 'schizophrenia', 'biploar', 'minor_depressive', 'dementia']\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 getlabel(self, i):\n",
    "        return self.data['labels'][i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data['labels'])\n",
    "    \n",
    "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 = 1 #len(label_map)\n",
    "\n",
    "print(label_map)\n",
    "print(rev_label_map)    \n",
    "\n",
    "def load_word2vec_embeddings(word2vec_file, word_map):\n",
    "    # Load word2vec model\n",
    "    w2v = gensim.models.KeyedVectors.load(word2vec_file, mmap='r')\n",
    "\n",
    "    print(\"\\nEmbedding length is %d.\\n\" % w2v.vector_size)\n",
    "\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"
   ]
  },
  {
   "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",
    "\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",
    "    \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, 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",
    "\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",
    "        att_s = self.sentence_attention(packed_sentences.data)  # (n_sentences, att_size)\n",
    "        att_s = torch.tanh(att_s)  \n",
    "        att_s = self.sentence_context_vector(att_s).squeeze(1)  # (n_sentences)\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), 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_documents, max(sentences_per_document), max(words_per_sentence))\n",
    "        return documents, word_alphas, sentence_alphas\n",
    "\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_sentences, word_pad_len, emb_size)\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",
    "        \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), batch_first=True)\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\n",
    "    \n",
    "def create_input_files_for_bert(bert_tokenizer, train, test, output_folder, sentence_limit, word_limit, read_tsv, label,\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, label)\n",
    "    \n",
    "    # Create word map\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",
    "    # Encode and pad\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",
    "    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, f'BERT_TRAIN_{label}_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",
    "    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, label)\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",
    "    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, f'BERT_TEST_{label}_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, label)\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, f'BERT_VALID_{label}_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')\n",
    "    \n",
    "def read_all_text_from_corpus_with_bert(word_tokenizer, tsv_file, sentence_limit, word_limit, label):\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",
    "            if label == 0:\n",
    "                labels.append([float(major_d)])\n",
    "            elif label == 1:\n",
    "                labels.append([float(sc)])\n",
    "            elif label == 2:\n",
    "                labels.append([float(bp)])\n",
    "            elif label == 3:\n",
    "                labels.append([float(minor_d)])\n",
    "            elif label == 4:\n",
    "                labels.append([float(de)])\n",
    "            else:\n",
    "                assert False;\n",
    "            docs.append(words)            \n",
    "    return docs, labels, word_counter  \n",
    "\n",
    "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",
    "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, 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, 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",
    "        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 = 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",
    "                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",
    "        \n",
    "        packed_words = pack_padded_sequence(sentences,\n",
    "                                            lengths=words_per_sentence.tolist(),\n",
    "                                            batch_first=True,\n",
    "                                            enforce_sorted=False)\n",
    "        packed_words, _ = self.word_rnn(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",
    "        word_alphas = att_w / torch.sum(att_w, dim=1, keepdim=True) \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\n",
    "    \n",
    "def bert_train_epoch(start_epoch, epochs, data_loader, model, criterion, optimizer, word_map, model_name, grad_clip, \n",
    "                     valid_iterator = None, \n",
    "                interval = 1, early_stop = False, period = 20, gap = 0.005, threshold = 0.5, best_valid_fscore = 0,\n",
    "                    zero_time = 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",
    "        # The first epoch 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",
    "                print(f'Update best validation fscore {valid_acc}')\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",
    "            if valid_loss == 0:\n",
    "                zero_time += 1\n",
    "                if zero_time >= 30:\n",
    "                    print(f'Early stop at epoch {epoch+1:02} (loss is zero for 10 times).')\n",
    "                    break\n",
    "            elif valid_loss > 0:\n",
    "                zero_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",
    "\n",
    "    return train_losses, valid_losses, train_accs, valid_accs        \n",
    "\n",
    "\n",
    "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",
    "    if len(diagnoses)>2:\n",
    "        diagnoses[MACRO]['f']=diagnoses[MACRO].get('f', 0)/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",
    "    else:\n",
    "        diagnoses[MACRO]['f']=\"n/a\"\n",
    "        diagnoses[MACRO]['p']=\"n/a\"\n",
    "        diagnoses[MACRO]['r']=\"n/a\"\n",
    "    return diagnoses, predicts\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",
    "    if len(overall) > 2:\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",
    "    else:\n",
    "        overall[MACRO]['f']=\"n/a\"\n",
    "        overall[MACRO]['p']=\"n/a\"\n",
    "        overall[MACRO]['r']=\"n/a\"\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",
    "    for i, (documents, sentences_per_document, words_per_sentence, labels) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()        \n",
    "\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",
    "        \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",
    "        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",
    "    return losses.avg, accs['micro'][\"f\"]\n",
    "\n",
    "class AverageMeter(object):\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 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 = 5, early_stop = False, period = 20, gap = 0.005, threshold = 0.5, best_valid_fscore = 0):\n",
    "    \n",
    "    best_valid_loss = float('inf')\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    train_accs = []\n",
    "    valid_accs = []\n",
    "    observed_time = 0\n",
    "    loss_zero = 0\n",
    "    \n",
    "    for epoch in range(start_epoch, epochs):\n",
    "        start_time = time.time()\n",
    "        # The first 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",
    "        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",
    "            print(f'Update best validation score to {valid_acc}')\n",
    "            observed_time = 0\n",
    "        if train_loss == 0:\n",
    "            if loss_zero == 5:            \n",
    "                if observed_time >= period:\n",
    "                    print(f'Loss is zero for {loss_zero} times.\\nEarly stop at epoch {epoch+1:02}.')\n",
    "                break\n",
    "            else:\n",
    "                loss_zero += 1\n",
    "        \n",
    "        save_checkpoint(epoch, model, optimizer, word_map, DATA_FOLDER, model_name + '_current.pt')        \n",
    "    return train_losses, valid_losses, train_accs, valid_accs\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) \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) \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)  \n",
    "            \n",
    "            scores, word_alphas, sentence_alphas = model(documents, sentences_per_document, words_per_sentence)\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].get('f', 0)/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",
    "    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    \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)    \n",
    "   \n",
    "    \n",
    "class NTUHDataset(data.Dataset):\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', test='test', **kwargs):\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",
    "    \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, label,\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, label)\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",
    "    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, GLOVE_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",
    "    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, f'GLOVE_TRAIN_{label}_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",
    "    print(f'Reading and preprocessing test data {test}...\\n')\n",
    "    test_docs, test_labels, _ = read_tsv(test, sentence_limit, word_limit, label)\n",
    "\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, GLOVE_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",
    "    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, f'GLOVE_TEST_{label}_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, label)\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, GLOVE_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",
    "        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, f'GLOVE_VALID_{label}_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",
    "    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    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocessing\n",
    "\n",
    "Create validation set"
   ]
  },
  {
   "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",
    "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": "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, 0,\n",
    "                       MIN_WORD_COUNT, True, VALID)\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if none\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": [
    "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": [
    "# 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>\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "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 #1 #None  # clip gradients at this value\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train_0'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_0'), 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)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune"
   ]
  },
  {
   "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",
    "\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "criterion = nn.BCEWithLogitsLoss()#pos_weight=POS_WEIGHT)\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",
    "                'rand_0', 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": [
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_0_fscore.pt')\n",
    "\n",
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "model.eval()\n",
    "\n",
    "# Load test data\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_0'), 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": [
    "rev_word_map = {v: k for k, v in word_map.items()}\n",
    "rev_word_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gensim.models import KeyedVectors\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)   \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",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_0_fscore.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "trained_embedding = model.sentence_attention.word_attention.embeddings.weight.to('cpu').data.numpy()\n",
    "trained_embedding_vec = KeyedVectors(len(trained_embedding.tolist()[0]))\n",
    "trained_embedding_vec.add(list(word_map.keys()), trained_embedding.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for w, s in trained_embedding_vec.most_similar('depressed'):\n",
    "    print(f'{w:15s} {s:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = ['depressed', 'manic', 'hallucination', 'zyprexa']\n",
    "\n",
    "embedding_clusters = []\n",
    "word_clusters = []\n",
    "\n",
    "for word in keys:\n",
    "    embeddings = []\n",
    "    words = []\n",
    "    for similar_word, _ in trained_embedding_vec.most_similar(word, topn=30):\n",
    "        words.append(similar_word)\n",
    "        embeddings.append(trained_embedding_vec[similar_word])\n",
    "    embedding_clusters.append(embeddings)\n",
    "    word_clusters.append(words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.manifold import TSNE\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "\n",
    "embedding_clusters = np.array(embedding_clusters)\n",
    "n, m, k = embedding_clusters.shape\n",
    "tsne_model_aicup = TSNE(perplexity=15, n_components=2, init='pca', n_iter=3500, random_state=32)\n",
    "\n",
    "embeddings_aicup = np.array(tsne_model_aicup.fit_transform(embedding_clusters.reshape(n * m, k))).reshape(n, m, 2)\n",
    "\n",
    "def tsne_plot_similar_words(title, labels, embedding_clusters, word_clusters, a, filename=None):\n",
    "    plt.figure(figsize=(20, 20))\n",
    "    colors = cm.rainbow(np.linspace(0, 1, len(labels)))\n",
    "    for label, embeddings, words, color in zip(labels, embedding_clusters, word_clusters, colors):\n",
    "        x = embeddings[:, 0]\n",
    "        y = embeddings[:, 1]\n",
    "        plt.scatter(x, y, color=color, alpha=a, label=label)\n",
    "        for i, word in enumerate(words):\n",
    "            plt.annotate(word, alpha=0.5, xy=(x[i], y[i]), xytext=(5, 2),\n",
    "                         textcoords='offset points', ha='right', va='bottom', size=8)\n",
    "    plt.legend(loc=4)\n",
    "    plt.title(title)\n",
    "    plt.grid(True)\n",
    "    if filename:\n",
    "        plt.savefig(filename, format='png', dpi=150, bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "tsne_plot_similar_words('Similar Words', keys, embeddings_aicup, word_clusters, 0.7,\n",
    "                        'similar_words_rand.png')"
   ]
  },
  {
   "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",
    "with open(os.path.join(DATA_FOLDER, 'glove_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_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(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_0_fscore.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "trained_embedding = model.sentence_attention.word_attention.embeddings.weight.to('cpu').data.numpy()\n",
    "trained_embedding_vec = KeyedVectors(len(trained_embedding.tolist()[0]))\n",
    "trained_embedding_vec.add(list(word_map.keys()), trained_embedding.tolist())\n",
    "for w, s in trained_embedding_vec.most_similar('hallucination', topn=30):\n",
    "    print(f'{w:15s} {s:.4f}')\n",
    "    \n",
    "embedding_clusters = []\n",
    "word_clusters = []\n",
    "\n",
    "for word in keys:\n",
    "    embeddings = []\n",
    "    words = []\n",
    "    for similar_word, _ in trained_embedding_vec.most_similar(word, topn=30):\n",
    "        words.append(similar_word)\n",
    "        embeddings.append(trained_embedding_vec[similar_word])\n",
    "    embedding_clusters.append(embeddings)\n",
    "    word_clusters.append(words)\n",
    "    \n",
    "embedding_clusters = np.array(embedding_clusters)\n",
    "n, m, k = embedding_clusters.shape\n",
    "tsne_model_aicup = TSNE(perplexity=15, n_components=2, init='pca', n_iter=3500, random_state=32)\n",
    "\n",
    "embeddings_aicup = np.array(tsne_model_aicup.fit_transform(embedding_clusters.reshape(n * m, k))).reshape(n, m, 2)\n",
    "\n",
    "def tsne_plot_similar_words(title, labels, embedding_clusters, word_clusters, a, filename=None):\n",
    "    plt.figure(figsize=(20, 20))\n",
    "    colors = cm.rainbow(np.linspace(0, 1, len(labels)))\n",
    "    for label, embeddings, words, color in zip(labels, embedding_clusters, word_clusters, colors):\n",
    "        x = embeddings[:, 0]\n",
    "        y = embeddings[:, 1]\n",
    "        plt.scatter(x, y, color=color, alpha=a, label=label)\n",
    "        for i, word in enumerate(words):\n",
    "            plt.annotate(word, alpha=0.5, xy=(x[i], y[i]), xytext=(5, 2),\n",
    "                         textcoords='offset points', ha='right', va='bottom', size=8)\n",
    "    plt.legend(loc=4)#, prop=fontprop)\n",
    "    plt.title(title)\n",
    "    plt.grid(True)\n",
    "    if filename:\n",
    "        plt.savefig(filename, format='png', dpi=150, bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "tsne_plot_similar_words('Similar Words', keys, embeddings_aicup, word_clusters, 0.7,\n",
    "                        'similar_words_glove.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other Disease"
   ]
  },
  {
   "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, 1,\n",
    "                       MIN_WORD_COUNT, True, VALID)\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)   \n",
    "\n",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "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 #1 #None  # clip gradients at this value\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train_1'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_1'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = 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",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "cudnn.benchmark = True  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[1])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'rand_1', grad_clip, valid_loader, early_stop = True, period = 30, threshold=0.4)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_1'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[1]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_1_loss.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[1]}:\\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",
    "\n",
    "              \n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_1'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_1_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[1]}:\\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": [
    "create_input_files(TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, read_all_text_from_corpus, 2,\n",
    "                       MIN_WORD_COUNT, True, VALID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)   \n",
    "    \n",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "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 #1 #None  # clip gradients at this value\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train_2'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_2'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = 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",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "cudnn.benchmark = True  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[2])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'rand_2', grad_clip, valid_loader, early_stop = True, period = 30, threshold=0.3)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_2'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[2]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_2_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[2]}:\\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": [
    "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_2'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[2]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_2_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[2]}:\\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": [
    "create_input_files(TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, read_all_text_from_corpus, 3,\n",
    "                       MIN_WORD_COUNT, True, VALID)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)   \n",
    "\n",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "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 #1 #None  # clip gradients at this value\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train_3'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_3'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\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",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "cudnn.benchmark = True  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[3])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'rand_3', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.3)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_3'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[3]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_3_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[3]}:\\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": [
    "create_input_files(TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, read_all_text_from_corpus, 4,\n",
    "                       MIN_WORD_COUNT, True, VALID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "with open(os.path.join(DATA_FOLDER, 'word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)   \n",
    "\n",
    "embeddings, emb_size = initialize_embeddings(EMBEDDING_DIM, word_map)\n",
    "PAD_IDX = word_map[\"<pad>\"]\n",
    "UNK_IDX = word_map[\"<unk>\"]\n",
    "embeddings[PAD_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[UNK_IDX] = torch.zeros(EMBEDDING_DIM)\n",
    "\n",
    "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 #1 #None  # clip gradients at this value\n",
    "\n",
    "# DataLoaders\n",
    "train_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'train_4'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_4'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "\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",
    "model.sentence_attention.word_attention.init_embeddings(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "cudnn.benchmark = True  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[4])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'rand_4', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.5)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_4'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[4]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'rand_4_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[4]}:\\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, 'rand_0_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[4]}:\\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",
    "    # Read data\n",
    "    sentences = torch.load(os.path.join(data_folder, 'word2vec_data.pth.tar'))\n",
    "    sentences = list(itertools.chain.from_iterable(sentences))\n",
    "\n",
    "    # Activate logging for verbose training\n",
    "    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)\n",
    "\n",
    "    # Initialize and train the model (this will take some time)\n",
    "    #model = gensim.models.word2vec.Word2Vec(sentences=sentences, size=EMBEDDING_DIM, workers=8, window=10, min_count=5,\n",
    "    #                                        sg=sg)\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",
    "    # Normalize vectors and save model\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",
    "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)  # load pre-trained word2vec embeddings\n",
    "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": [
    "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, 'train_0'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_0'), 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(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\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",
    "\n",
    "#start_epoch = 0  # start at this epoch\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",
    "                'w2v_0', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.5)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_0'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[0]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'w2v_0_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[0]}:\\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": [
    "word2vec_file = os.path.join(DATA_FOLDER, 'word2vec_skipgram_model')  # path to pre-trained word2vec embeddings\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)  # load pre-trained word2vec embeddings\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "\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, 'train_1'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_1'), 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(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[1])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "start_epoch = 0  # start at this epoch\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",
    "                'w2v_1', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.35)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_1'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[1]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'w2v_1_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[1]}:\\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": [
    "word2vec_file = os.path.join(DATA_FOLDER, 'word2vec_skipgram_model')  # path to pre-trained word2vec embeddings\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)  # load pre-trained word2vec embeddings\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "\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, 'train_2'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_2'), 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(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[2])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "start_epoch = 0  # start at this epoch\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",
    "                'w2v_2', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.35)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_2'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[2]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'w2v_2_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[2]}:\\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": [
    "word2vec_file = os.path.join(DATA_FOLDER, 'word2vec_skipgram_model')  # path to pre-trained word2vec embeddings\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)  # load pre-trained word2vec embeddings\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "\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, 'train_3'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_3'), 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(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[3])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "start_epoch = 0  # start at this epoch\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",
    "                'w2v_3', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.35)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_3'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[3]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'w2v_3_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[3]}:\\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": [
    "word2vec_file = os.path.join(DATA_FOLDER, 'word2vec_skipgram_model')  # path to pre-trained word2vec embeddings\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)  # load pre-trained word2vec embeddings\n",
    "embeddings[word_map[\"<pad>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "embeddings[word_map[\"<unk>\"]] = torch.zeros(EMBEDDING_DIM)\n",
    "\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, 'train_4'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'valid_4'), 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(\n",
    "            embeddings)  # initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[4])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "start_epoch = 0  # start at this epoch\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",
    "                'w2v_4', grad_clip, valid_loader, early_stop = True, period = 30, threshold = 0.5)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'test_4'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[4]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'w2v_4_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[4]}:\\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, 0, 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",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "\n",
    "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_0'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID_0'), 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",
    "# initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_0_current.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses, valid_losses, train_accs, valid_accs = \\\n",
    "    train_epoch(start_epoch, N_EPOCHS, train_loader, model, criterion, optimizer, word_map, \n",
    "                'glove_0', grad_clip, valid_loader, early_stop=True, period = period, \n",
    "                best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test_0'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[0]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_0_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[0]}:\\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": [
    "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]"
   ]
  },
  {
   "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, 1, 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",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)"
   ]
  },
  {
   "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, 'GLOVE_TRAIN_1'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID_1'), 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",
    "# initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[1])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_1_current.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\n",
    "start_epoch = 65\n",
    "best_valid_fscore = 0.793\n",
    "period = 10\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_1', grad_clip, valid_loader, early_stop=True, period = period, \n",
    "                best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test_1'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[1]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_1_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[1]}:\\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": [
    "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",
    "with open(os.path.join(DATA_FOLDER, 'glove_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)"
   ]
  },
  {
   "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, 'GLOVE_TRAIN_2'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID_2'), 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",
    "# initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[2])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_2_current.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\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_2', grad_clip, valid_loader, early_stop=True, period = period, \n",
    "                best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test_2'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[2]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_2_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[2]}:\\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",
    "              \n",
    "              "
   ]
  },
  {
   "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",
    "create_input_files_for_glove(vocab.stoi, TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, \n",
    "                             read_all_text_from_corpus, 3, MIN_WORD_COUNT, VALID)\n",
    "\n",
    "with open(os.path.join(DATA_FOLDER, 'glove_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "\n",
    "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, 'GLOVE_TRAIN_3'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID_3'), 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",
    "# initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[3])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "#checkpoint = os.path.join(DATA_FOLDER, 'glove_1_current.pt')\n",
    "#checkpoint = torch.load(checkpoint)\n",
    "#model = checkpoint['model']\n",
    "\n",
    "start_epoch = 0#65\n",
    "best_valid_fscore = 0#0.793\n",
    "period = 30\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_3', grad_clip, valid_loader, early_stop=True, period = period, \n",
    "                best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test_3'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[3]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_3_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[3]}:\\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": [
    "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",
    "\n",
    "#create_input_files_for_glove(vocab.stoi, TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, WORD_LIMIT, \n",
    "#                             read_all_text_from_corpus, 4, MIN_WORD_COUNT, VALID)\n",
    "\n",
    "with open(os.path.join(DATA_FOLDER, 'glove_word_map.json'), 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "\n",
    "embeddings[GLOVE_PAD_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "embeddings[GLOVE_UNK_IDX] = torch.zeros(GLOVE_EMBEDDING_DIM)\n",
    "\n",
    "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, 'GLOVE_TRAIN_4'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'GLOVE_VALID_4'), 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",
    "# initialize embedding layer with pre-trained embeddings\n",
    "model.sentence_attention.word_attention.init_embeddings(embeddings)  \n",
    "model.sentence_attention.word_attention.fine_tune_embeddings(FINE_TUNE_EMBEDDING)  # fine-tune\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[4])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_4_current.pt')\n",
    "checkpoint = torch.load(checkpoint)\n",
    "model = checkpoint['model']\n",
    "\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_4', grad_clip, valid_loader, early_stop=True, period = period, \n",
    "                best_valid_fscore = best_valid_fscore)"
   ]
  },
  {
   "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",
    "model.eval()\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'glove_test_4'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\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[4]}:\\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",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'glove_4_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[4]}:\\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",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_input_files_for_bert(bert_tokenizer, TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, BERT_WORD_LIMIT, \n",
    "                            read_all_text_from_corpus_with_bert, 0, VALID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if 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": [
    "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_0'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID_0'), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_0_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 = 52\n",
    "best_valid_fscore = 0.583\n",
    "period = 1\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_0', grad_clip, valid_loader, \n",
    "                     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_0'), 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[0]}:\\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_0_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[0]}:\\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": [
    "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",
    "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",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]\n",
    "\n",
    "#create_input_files_for_bert(bert_tokenizer, TRAIN, TEST, DATA_FOLDER, SENTENCE_LIMIT, BERT_WORD_LIMIT, \n",
    "#                            read_all_text_from_corpus_with_bert, 1, VALID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = None  # path to model checkpoint, None if 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')    \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, 'BERT_TRAIN_1'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID_1'), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_1_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\n",
    "\n",
    "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[1])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\n",
    "start_epoch = 44\n",
    "best_valid_fscore = 0.5133333333333333\n",
    "period = 5\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_1', grad_clip, valid_loader, \n",
    "                     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\n",
    "\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test_1'), 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[1]}:\\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",
    "\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_1_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[1]}:\\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": [
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test_1'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_1_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[1]}:\\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": [
    "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",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if 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')    \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, 'BERT_TRAIN_2'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID_2'), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_2_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\n",
    "\n",
    "N_EPOCHS = 200\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[2])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\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_2', grad_clip, valid_loader, \n",
    "                     early_stop = True, period = period, best_valid_fscore = best_valid_fscore, threshold = 0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test_2'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache_test\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_2_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[2]}:\\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": [
    "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",
    "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",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if 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')    \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, 'BERT_TRAIN_3'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID_3'), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_3_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\n",
    "\n",
    "optimizer = optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()))\n",
    "criterion = nn.BCEWithLogitsLoss(pos_weight=POS_WEIGHT[3])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\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_3', grad_clip, valid_loader, \n",
    "                     early_stop = True, period = period, best_valid_fscore = best_valid_fscore, threshold = 0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "han_bert_cache_test = torch.load('han_cache_test.pt')\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test_3'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache_test\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_3_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[3]}:\\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": [
    "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",
    "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",
    "BERT_WORD_LIMIT = bert_tokenizer.max_model_input_sizes[BERT_MODEL]\n",
    "\n",
    "checkpoint = None  # path to model checkpoint, None if 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')    \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, 'BERT_TRAIN_4'), batch_size=BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=WORKER, pin_memory=True)\n",
    "valid_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'BERT_VALID_4'), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_4_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\n",
    "\n",
    "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[4])\n",
    "model = model.to(device)\n",
    "criterion = criterion.to(device)\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_4', grad_clip, valid_loader, \n",
    "                     early_stop = True, period = period, best_valid_fscore = best_valid_fscore, threshold = 0.5, \n",
    "                     zero_time = zero_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "han_bert_cache_test = torch.load('han_cache_test.pt')\n",
    "test_loader = DataLoader(NTUH_HANDataset(DATA_FOLDER, 'bert_test_4'), batch_size=BATCH_SIZE, shuffle=False,\n",
    "                                          num_workers=WORKER, pin_memory=True)\n",
    "\n",
    "model.han_bert_cache = model.sentence_attention.word_attention.han_bert_cache = han_bert_cache_test\n",
    "\n",
    "checkpoint = os.path.join(DATA_FOLDER, 'bert_4_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[3]}:\\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",
    "                                            "
   ]
  }
 ],
 "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
}