#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
import pickle
import pandas as pd
import numpy as np
import sparse
import torch
import model
from tqdm import tqdm
from torch import nn, optim
from torch.utils.data import DataLoader
import torch.nn.functional as F
import sklearn
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, cohen_kappa_score,precision_recall_curve, auc
import matplotlib.pyplot as plt
from baseline import *
from operations import *
from focal_loss.focal_loss import FocalLoss
import os
def eval_metric_stay(eval_set, model, device, encoder = 'normal'):
model.eval()
criterion = torch.nn.BCELoss()
with torch.no_grad():
y_true = np.array([])
y_pred = np.array([])
y_score = np.array([])
for i, batch_data in enumerate(eval_set):
X = batch_data[0]
if encoder == 'HMP':
S = batch_data[1]
elif encoder == 'BERT':
S = batch_data[1]
input_ids = batch_data[2]
attention_mask = batch_data[3]
token_type_ids = batch_data[4]
labels = batch_data[-1]
if encoder == 'normal':
X2 = batch_data[1]
S = batch_data[2]
# print(X.shape)
outputs = model(X, X2, S).squeeze().to(device)
elif encoder == 'HMP':
outputs = model(X,S).squeeze().to(device)
elif encoder == 'BERT':
outputs = model(X,S,input_ids, token_type_ids, attention_mask).squeeze().to(device)
score = outputs
score = score.data.cpu().numpy()
labels = labels.data.cpu().numpy()
pred = np.where(score >= 0.5, 1.0, 0.0)
if labels.shape[0] != 1:
y_true = np.concatenate((y_true, labels))
y_pred = np.concatenate((y_pred, pred))
y_score = np.concatenate((y_score, score))
else:
y_true = np.array(list(y_true) + list(labels))
y_pred = np.array(list(y_pred) + list([pred]))
y_score = np.array(list(y_score) + list([score]))
accuary = accuracy_score(y_true, y_pred)
precision = precision_score(y_true, y_pred)
recall = recall_score(y_true, y_pred)
f1 = f1_score(y_true, y_pred)
roc_auc = roc_auc_score(y_true, y_score)
lr_precision, lr_recall, _ = precision_recall_curve(y_true, y_score)
pr_auc = auc(lr_recall, lr_precision)
kappa = cohen_kappa_score(y_true, y_pred)
loss = criterion(torch.from_numpy(y_true), torch.from_numpy(y_score))
return f1, roc_auc, pr_auc, kappa, loss
def eval_metric_admission(eval_set, model, device, encoder = 'normal'):
model.eval()
criterion = torch.nn.BCELoss()
with torch.no_grad():
y_true = np.array([])
y_pred = np.array([])
y_score = np.array([])
for i, batch_data in enumerate(eval_set):
icd = batch_data[0]
drug = batch_data[1]
X = batch_data[2]
S = batch_data[3]
input_ids = batch_data[4]
attention_mask = batch_data[5]
token_type_ids = batch_data[6]
labels = batch_data[-1]
outputs = model(icd, drug,X,S,input_ids, attention_mask, token_type_ids).squeeze().to(device)
score = outputs
score = score.data.cpu().numpy()
labels = labels.data.cpu().numpy()
# pred = torch.tensor([1 if x > 0.5 else 0 for x in score])
pred = np.where(score >= 0.5, 1.0, 0.0)
y_true = np.concatenate((y_true, labels))
y_pred = np.concatenate((y_pred, pred))
y_score = np.concatenate((y_score, score))
accuary = accuracy_score(y_true, y_pred)
precision = precision_score(y_true, y_pred)
recall = recall_score(y_true, y_pred)
f1 = f1_score(y_true, y_pred)
roc_auc = roc_auc_score(y_true, y_score)
lr_precision, lr_recall, _ = precision_recall_curve(y_true, y_score)
pr_auc = auc(lr_recall, lr_precision)
kappa = cohen_kappa_score(y_true, y_pred)
loss = criterion(torch.from_numpy(y_true), torch.from_numpy(y_score))
return f1, roc_auc, pr_auc, kappa, loss
def icu_trainer(model, train, valid, test, epoch, learn_rate, batch_size, seed, device, encoder = 'normal', patience = 3):
torch.manual_seed(seed)
model.train()
aupr_list = []
criterion = torch.nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(),
momentum=0.9,
lr = learn_rate,
weight_decay = 1e-2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
f1, roc_auc, pr_auc, kappa, valid_loss = eval_metric_stay(valid, model, device, encoder)
best_dev = valid_loss
best_epoc = 0
model.train()
from datetime import datetime
dt = datetime.now()
torch.save(model, "saved_model/hmp_model" + str(dt) +".p")
best_name = "saved_model/hmp_model" + str(dt) +".p"
for epoch in tqdm(range(epoch)):
loss = 0
model.train()
for batch_idx, batch_data in enumerate(train):
X = batch_data[0]
if encoder == 'HMP':
S = batch_data[1]
elif encoder == 'BERT':
S = batch_data[1]
input_ids = batch_data[2]
attention_mask = batch_data[3]
token_type_ids = batch_data[4]
label = batch_data[-1]
optimizer.zero_grad()
if encoder == 'normal':
outputs = model(X).squeeze().to(device)
elif encoder == 'HMP':
outputs = model(X,S).squeeze().to(device)
elif encoder == 'BERT':
outputs = model(X,S,input_ids, attention_mask, token_type_ids).squeeze().to(device)
train_loss = criterion(outputs, label)
train_loss.backward()
optimizer.step()
loss += train_loss.item()
scheduler.step()
print("Training Loss = ", loss)
model.eval()
f1, roc_auc, pr_auc, kappa, valid_loss = eval_metric_stay(valid, model, device, encoder)
print("Dev = ", valid_loss)
dt = datetime.now()
if best_dev > valid_loss:
best_dev = valid_loss
best_epoc = epoch
torch.save(model, "saved_model/icu_model" + str(dt) +".p")
os.remove(best_name)
best_name = "saved_model/icu_model" + str(dt) +".p"
if epoch - best_epoc == patience:
break
model.train()
model = torch.load(best_name)
os.remove(best_name)
return model, aupr_list
def adm_trainer(model, train, valid, test, epoch, learn_rate, batch_size, seed, device, encoder = 'normal', patience = 3):
torch.manual_seed(seed)
model.train()
aupr_list = []
criterion = torch.nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(),
momentum=0.9,
lr = learn_rate,
weight_decay = 1e-2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
f1, roc_auc, pr_auc, kappa, valid_loss = eval_metric_admission(valid, model, device, encoder)
best_dev = valid_loss
best_epoc = 0
model.train()
for epoch in tqdm(range(epoch)):
loss = 0
for batch_idx, batch_data in enumerate(train):
icd = batch_data[0]
drug = batch_data[1]
X = batch_data[2]
S = batch_data[3]
input_ids = batch_data[4]
attention_mask = batch_data[5]
token_type_ids = batch_data[6]
label = batch_data[-1]
optimizer.zero_grad()
outputs = model(icd, drug,X,S,input_ids, attention_mask, token_type_ids).squeeze().to(device)
train_loss = criterion(outputs, label)
train_loss.backward()
optimizer.step()
loss += train_loss.item()
scheduler.step()
print("Training Loss = ", loss)
model.eval()
f1, roc_auc, pr_auc, kappa, valid_loss = eval_metric_admission(valid, model, device, encoder)
print("Dev = ", valid_loss)
if best_dev > valid_loss:
best_dev = valid_loss
best_epoc = epoch
torch.save(model, "saved_model/adm_model.p")
if epoch - best_epoc == patience:
model = torch.load("saved_model/adm_model.p")
break
model.train()
model = torch.load("saved_model/adm_model.p")
return model, aupr_list
Datasets
Models
