main.py

from torch import optim

from data_util import DataUtil

from models import *
from argument_parsers import *

import time
from time import strftime, localtime
import math
import random

import json

from sklearn.metrics import precision_score, recall_score, f1_score
from sklearn.metrics import confusion_matrix, classification_report

from os.path import exists

import logging


class Experiment:
    def __init__(self, config, sequence_length=20, reload_data=True):
        # Hyper Parameters
        self.sequence_length = sequence_length
        self.hidden_size = 128
        self.num_layers = 1

        self.config = config
        self.data = DataUtil(data_dir=config.data_dir, vocab_dir=config.vocab_dir,
                             split_by_sentence=not config.split_by_section, skip_list=config.skip_list)

        if not self.config.filtered:
            self.data.make_dir(self.config.output_dir + "/models/")

        if reload_data:
            for ds in self.config.textbook_data_sets:
                self.data.load_textbook_train_dev_data(config.data_dir + 'medlit/train/' + ds,
                                                       config.data_dir + 'medlit/dev/' + ds)
            # train
            self.data.load_i2b2_train_data(train_base_dir=config.data_dir + '/i2b2_ehr/')
            # test
            self.data.load_test_data(ref_base_dir=config.data_dir + '/i2b2_ehr/')
            # dev
            self.data.load_test_data(ref_base_dir=config.data_dir + '/i2b2_ehr/', type='dev')

        else:
            self.data.load_split_data()

        self.data.make_dir(self.config.output_dir)

        log_file_name = strftime("log_%Y_%m_%d_%H_%M_%S", localtime())
        self.logger = self.setup_logger(self.config.output_dir + '/%s.txt' % log_file_name)

        if exists(config.vocab_dir + "/NaturalLang.pkl") and not reload_data:
            print("Loading vocab")
            self.data.load_vocab()
        else:
            print("Building vocab")
            self.data.build_vocab(self.data.textbook_train_data, pretrain=False)

        self.model = None
        self.use_cuda = torch.cuda.is_available()

        if not self.config.filtered:
            if self.config.model_type == 'gru_rnn':
                self.model = GRURNN(
                    self.config.embedding_size, self.hidden_size, self.data.input_lang,
                    self.data.pretrained_embeddings, self.num_layers, self.data.input_lang.n_words,
                    self.data.output_lang.n_words, self.config.dropout)
            elif self.config.model_type == 'attn_gru_rnn':
                self.model = AttentionGRURNN(
                    self.config.embedding_size, self.hidden_size, self.data.input_lang,
                    self.data.pretrained_embeddings, self.num_layers, self.data.input_lang.n_words,
                    self.data.output_lang.n_words, self.config.dropout)
            elif self.config.model_type == 'cnn':
                self.model = CNN(
                    self.data.input_lang.n_words, self.data.output_lang.n_words, self.config.embedding_size,
                    self.data.input_lang, self.data.pretrained_embeddings, self.config.dropout)

            self.epoch_start = 1

            if self.use_cuda:
                self.model = self.model.cuda()

    def setup_logger(self, log_file, level=logging.INFO):
        logger = logging.getLogger()
        logger.setLevel(level)
        handler = logging.FileHandler(log_file)
        formatter = logging.Formatter('%(asctime)s %(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        return logger

    def log(self, info):
        print(info)
        if self.logger is not None:
            self.logger.info(info)

    def as_minutes(self, s):
        m = math.floor(s / 60)
        s -= m * 60
        return '%dm %ds' % (m, s)

    def time_since(self, since, percent):
        now = time.time()
        s = now - since
        es = s / percent
        rs = es - s
        return '%s (- %s)' % (self.as_minutes(s), self.as_minutes(rs))

    def train(self, data_setup, save_model_dir,
              print_every=20, plot_every=100, learning_rate=0.001):
        start = time.time()

        plot_losses = []
        print_loss_total = 0
        plot_loss_total = 0

        if self.config.model_type == 'cnn' and self.config.transfer_learning:
            self.model.output_size = self.data.output_lang.n_words
            if self.config.reuse_embedding_layer_only:
                self.model.init_conv1_layer()
                self.model.init_conv2_layer()
                self.model.init_fc_layers()
            if self.config.reuse_embedding_conv1_layers:
                self.model.init_conv2_layer()
                self.model.init_fc_layers()
            if self.use_cuda:
                self.model = self.model.cuda()
        elif self.config.transfer_learning:
            self.model.freeze_layer("fc1")

        if self.config.optimizer == 'sgd':
            optimizer = optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()),
                                  lr=learning_rate, momentum=0.9)
        elif self.config.optimizer == 'adam':
            optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()),
                                   lr=learning_rate)

        self.log('data_setup:' + str(data_setup))

        train_data = []

        for data_set in data_setup:
            data_ratio = data_setup[data_set]
            data = self.data.get_dataset(data_set)
            train_data += self.data.get_data_subset(data, data_ratio)
        print('len train_data:', len(train_data))
        print('training data examples:', train_data[:5])

        if self.config.downsampling:
            train_data = self.data.downsampling(
                train_data, number_samples=self.config.downsampling_size)

        num_train_data = len(train_data)
        print('num_train_data:', num_train_data)
        print('train_data:', train_data[:10])
        num_batches = int(np.ceil(num_train_data / float(self.config.batch_size)))
        self.log('num_batches: ' + str(num_batches))

        if self.config.weighted_loss:
            loss_weight = self.data.get_label_weight(train_data)
            if self.use_cuda:
                loss_weight = loss_weight.cuda()
        else:
            loss_weight = None

        max_dev_acc = 0

        for epoch in range(self.epoch_start, self.config.num_train_epochs + 1):
            batch_start = time.time()
            correct = 0
            total = 0

            random.shuffle(train_data)

            self.model.train()

            for cnt, i in enumerate(random.sample(range(num_batches), num_batches), start=1):
                inputs, seq_lengths, targets, batch = self.data.construct_batch(
                    self.config.batch_size * i, self.config.batch_size * (i + 1),
                    train_data, fixed_length=True if self.config.model_type == 'cnn' else False)

                if self.use_cuda:
                    inputs = inputs.cuda()
                    targets = targets.cuda()

                optimizer.zero_grad()

                if self.config.model_type == 'cnn':
                    outputs = self.model(inputs)  # for CNN
                elif self.config.model_type == 'attn_gru_rnn':
                    outputs = self.model(inputs, self.data.input_lang, seq_lengths)
                else:
                    outputs = self.model(inputs, seq_lengths)

                _, predicted = torch.max(outputs.data, dim=1)

                total += targets.data.size(0)
                correct += (predicted == targets.data).sum()
                batch_train_acc = 100.0 * (predicted == targets.data).sum() / targets.data.size(0)

                loss = F.cross_entropy(outputs, targets, weight=loss_weight)
                loss.backward()
                optimizer.step()
                self.log("Epoch %d, batch %d / %d: train loss = %f, train accuracy = %f %%"
                         % (epoch, cnt, num_batches, loss.data.item(), batch_train_acc))

                print_loss_total += loss.data.item()
                plot_loss_total += loss.data.item()

                if cnt % print_every == 0:
                    print_loss_avg = print_loss_total / print_every
                    print_loss_total = 0
                    self.log('Average batch loss: %s' % str(print_loss_avg))
                    self.log(self.time_since(batch_start, cnt * 1.0 / num_batches))

                if cnt % plot_every == 0:
                    plot_loss_avg = plot_loss_total / plot_every
                    plot_losses.append(plot_loss_avg)
                    plot_loss_total = 0
            self.log('Epoch %d is done' % epoch)
            self.log('Epoch %d Train Accuracy: %f %%' % (epoch, 100.0 * correct / total))
            self.log(self.time_since(start, epoch * 1.0 / self.config.num_train_epochs))

            datasets = []
            print("TUNING SET IS: " + str(self.config.tuning_set))
            if 'ALL' in self.config.tuning_set or 'MedLit' in self.config.tuning_set:
                self.log("Test on MedLit Dev: ")
                datasets.append(self.data.TEXTBOOK_DEV)
            if 'ALL' in self.config.tuning_set or 'i2b2' in self.config.tuning_set:
                self.log("Test on i2b2 EHR Dev: ")
                datasets.append(self.data.i2b2_DEV)

            self.log("Tuning on:")
            self.log(datasets)
            dev_acc = self.test(datasets=datasets, epoch=epoch, calc_confusion_matrix=True)

            # save intermediate training results
            if dev_acc > max_dev_acc:
                save_path = save_model_dir + "/models/best_model.pt"
                torch.save(self.model, save_path)
                self.log('Best Model saved in file: %s' % save_path)
                max_dev_acc = dev_acc

                if 'i2b2' in self.config.test_set:
                    self.log("Test on i2b2 Test:")
                    self.test(datasets=[self.data.i2b2_TEST], epoch=epoch, print_test_results=True)

            save_path = save_model_dir + "/models/epoch_" + str(epoch) + ".pt"
            torch.save(self.model, save_path)
            self.log('Model saved in file: %s' % save_path)

    def test(self, datasets, epoch=-1, calc_confusion_matrix=True,
             generate_reports=True, print_test_results=False, print_examples=False):
        if self.model is None:
            self.log('Restoring model from ' + self.config.reload_model_file)

            if torch.cuda.is_available():
                self.model = torch.load(self.config.reload_model_file)
            else:
                self.model = torch.load(self.config.reload_model_file, map_location='cpu')

            self.log('Model is restored')

        self.model.eval()

        start = time.time()

        data = []

        dataset_name = '_'.join(datasets)

        for dataset in datasets:
            data.extend(self.data.get_dataset(dataset))

        if self.config.downsampling:
            data = []
            for dataset in datasets:
                data.extend(self.data.get_dataset(dataset))
                data = self.data.downsampling(
                    data, number_samples=500)

        num_test_data = len(data)
        self.log("num_test_data: " + str(num_test_data))
        num_batches = int(np.ceil(num_test_data / float(self.config.batch_size)))
        self.log('num_batches: ' + str(num_batches))

        correct = 0
        total = 0
        loss = 0.0
        labels = []
        predictions = []
        examples = []

        for i in range(num_batches):

            inputs, seq_lengths, targets, batch = self.data.construct_batch(
                self.config.batch_size * i, self.config.batch_size * (i + 1),
                data,
                fixed_length=True if self.config.model_type == 'cnn' else False)

            if self.use_cuda:
                inputs = inputs.cuda()
                targets = targets.cuda()

            if self.config.model_type == 'cnn':
                outputs = self.model(inputs)  # for CNN
            elif self.config.model_type == 'attn_gru_rnn':
                outputs = self.model(inputs, self.data.input_lang, seq_lengths)
            else:
                outputs = self.model(inputs, seq_lengths)
            
            _, predicted = torch.max(outputs.data, dim=1)

            ordered = torch.sort(outputs.data)

            total += targets.data.size(0)
            correct += (predicted == targets.data).sum()
            labels.extend(targets.cpu().data.numpy().tolist())
            predictions.extend(predicted.cpu().numpy().tolist())

            loss += F.cross_entropy(outputs, targets).data.item()

            if print_examples or print_test_results:
                for k, d in enumerate(batch):
                    examples.append([d[0], d[1].replace('\r', ' ').replace('\n', ' ').replace('\t', ' '),
                                     d[2], d[3], str(d[4]), str(d[5]),
                                     self.data.output_lang.get_word(predicted[k].cpu().data.item()),
                                     self.data.output_lang.get_word(int(ordered[1][k][outputs.data.shape[1]-2])),
                                     self.data.output_lang.get_word(int(ordered[1][k][outputs.data.shape[1]-3]))])

        if print_examples:
            self.data.make_dir(self.config.output_dir + '/test_saved')
            self.log("Save examples to: " + self.config.output_dir + '/test_saved')
            with open(self.config.output_dir + '/test_saved/' + dataset_name + 'epoch_%d.txt' % epoch, 'w') as f:
                f.write("#\tSentence\tTrue\tHeader String\tLocation\tLine\tPrediction 1\tPrediction 2\tPrediction 3\n")
                for e in examples:
                    f.write('\t'.join(e) + '\n')

        self.log('Epoch %d ' % epoch + 'Time used: ' + str(time.time() - start))
        self.log('Epoch %d ' % epoch + 'Test loss: %f' % loss)
        self.log('Epoch %d ' % epoch + 'Test Accuracy: %f %%' % (100.0 * correct / total))
        self.log('Epoch %d ' % epoch + 'Test Precision: %f %%' %
                 (100.0 * precision_score(labels, predictions, average='micro')))
        self.log('Epoch %d ' % epoch + 'Test Recall: %f %%' %
                 (100.0 * recall_score(labels, predictions, average='micro')))
        self.log('Epoch %d ' % epoch + 'Test F1 Score: %f %%' %
                 (100.0 * f1_score(labels, predictions, average='micro')))

        text_labels = [self.data.output_lang.get_word(l) for l in labels]
        text_preds = [self.data.output_lang.get_word(l) for l in predictions]
        label_set = sorted(list(set(text_labels)))
        if calc_confusion_matrix:
            cm = confusion_matrix(text_labels, text_preds, labels=label_set)
            self.log('confusion_matrix for epoch %d: ' % epoch)
            header = '\t'.join(label_set)
            self.log(header)
            for i, row in enumerate(list(cm)):
                row = [str(num) for num in row]
                self.log('\t'.join([label_set[i]] + row))
            np.savetxt(self.config.output_dir + '/' + dataset_name + '_confusion_matrix_epoch_%d.csv' % epoch,
                       cm, fmt='%d', header=header, delimiter=',')
            self.log('Saved confusion matrix!')

        if generate_reports:
            reports = classification_report(text_labels, text_preds, labels=label_set, target_names=label_set, digits=4)
            self.log(reports)
            with open(self.config.output_dir + '/' + dataset_name + '_report_epoch_%d.txt' % epoch, 'w') as f:
                f.write(reports)
            self.log('Saved report!')

        if print_test_results:
            with open(self.config.output_dir + '/' + dataset_name + '_predictions_epoch_%d.json' % epoch, 'w') as f:
                json.dump(examples, f, indent=4, sort_keys=True)
        return 100.0 * correct / total

    def test_one(self, header, text):
        if self.model is None:
            self.log('Restoring model from ' + self.config.reload_model_file)

            if torch.cuda.is_available():
                self.model = torch.load(self.config.reload_model_file)
            else:
                self.model = torch.load(self.config.reload_model_file, map_location='cpu')
            self.log('Model is restored')
            self.model.eval()
            if self.use_cuda:
                self.model = self.model.cuda()

        inputs, seq_lengths, targets = self.data.construct_one(
            header, text,
            fixed_length=True if self.config.model_type == 'cnn' else False)

        if self.use_cuda:
            inputs = inputs.cuda()
            targets = targets.cuda()

        if self.config.model_type == 'cnn':
            outputs = self.model(inputs)  # for CNN
        elif self.config.model_type == 'attn_gru_rnn':
            outputs = self.model(inputs, self.data.input_lang, seq_lengths)
        else:
            outputs = self.model(inputs, seq_lengths)

        _, predicted = torch.max(outputs.data, dim=1)

        return predicted.cpu().numpy().tolist() == targets.cpu().data.numpy().tolist()

def main():
    parser = MainArgParser()
    args = parser.parse_args()

    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train:
        raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))

    config = Config(vars(args))

    data_set_names = {
        'DynaMed': 'DM',
        'ElsevierCecil': 'EC',
        'ElsevierText': 'ET',
        'NGC': 'NGC',
        'WikipediaMedical': 'WM',
        'WileyBooks': 'WB',
        'WileyCochrane': 'WC'
    }

    reload_data = config.reload_data

    if config.data_sets == ['All_Textbooks']:
        config.textbook_data_sets = list(data_set_names.keys())
    elif config.data_sets == ['EHR']:
        config.textbook_data_sets = []
    else:
        config.textbook_data_sets = config.data_sets

    experiment = Experiment(config=config, reload_data=reload_data)
    data_setup = {
        experiment.data.TEXTBOOK_TRAIN: config.textbook_data_ratio,
        experiment.data.i2b2_TRAIN: config.i2b2_data_ratio
    }

    if config.do_train:
        if config.reload_model_file:
            print("CONTINUING TO TRAIN")
            if torch.cuda.is_available():
                experiment.model = torch.load(config.reload_model_file)
            else:
                experiment.model = torch.load(config.reload_model_file, map_location='cpu')
            experiment.model.use_cuda = torch.cuda.is_available()
            experiment.log('Model restored.')
            experiment.epoch_start = 1
        experiment.train(learning_rate=config.learning_rate,
                         data_setup=data_setup,
                         save_model_dir=config.output_dir)
    if config.do_eval:
        if "ALL" in config.test_set or "MedLit" in config.test_set:
            print("RUNNING MedLit Dev")
            experiment.test(datasets=[experiment.data.TEXTBOOK_DEV], print_examples=True)
            print("RUNNING DEV")
        if "ALL" in config.test_set or "i2b2" in config.test_set:
            print("RUNNING i2b2 DEV")
            experiment.test(datasets=[experiment.data.i2b2_DEV], print_examples=True, print_test_results=True)
            print("RUNNING TEST")
            print("RUNNING i2b2 TEST")
            experiment.test(datasets=[experiment.data.i2b2_TEST], print_examples=True, print_test_results=True)


if __name__ == "__main__":
    main()