import unittest
import random
import sys, os
import re
import nltk
import json
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'allennlp')))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization')))
# print(sys.path)
from ehrkit import ehrkit
from getpass import getpass
try:
from config import USERNAME, PASSWORD
except:
print("Please put your username and password in config.py")
USERNAME = input('DB_username?')
PASSWORD = getpass('DB_password?')
DOC_ID = 1354526 # Temporary!!!
# Number of documents in NOTEEVENTS.
NUM_DOCS = 2083180
# Number of patients in PATIENTS.
NUM_PATIENTS = 46520
# Number of diagnoses in DIAGNOSES_ICD.
NUM_DIAGS = 823933
def select_ehr(ehrdb, requires_long=False, recursing=False):
if recursing:
doc_id = ''
else:
# doc_id = input("MIMIC Document ID [press Enter for random]: ")
doc_id = ''
if doc_id == '':
# Picks random document
ehrdb.cur.execute("SELECT ROW_ID FROM mimic.NOTEEVENTS ORDER BY RAND() LIMIT 1")
doc_id = ehrdb.cur.fetchall()[0][0]
text = ehrdb.get_document(int(doc_id))
if len(text.split()) > 200 or not requires_long:
return doc_id, text
else:
return select_ehr(ehrdb, requires_long, True)
else:
# Get inputted document
try:
text = ehrdb.get_document(int(doc_id))
return doc_id, text
except:
message = 'Error: There is no document with ID \'' + doc_id + '\' in mimic.NOTEEVENTS'
sys.exit(message)
def get_nb_dir(ending, SUMM_DIR):
# Gets path of Naive Bayes model trained on most examples
dir_nums = []
for dir in os.listdir(SUMM_DIR):
if os.path.isdir(os.path.join(SUMM_DIR, dir)) and dir.endswith('_exs_' + ending):
if os.path.exists(os.path.join(SUMM_DIR, dir, 'nb')):
try:
dir_nums.append(int(dir.split('_')[0]))
except:
continue
if len(dir_nums) > 0:
best_dir_name = str(max(dir_nums)) + '_exs_' + ending
return best_dir_name
else:
return None
def show_summary(doc_id, text, summary, model_name):
# x = input('Show full EHR (DOC ID %s)? [DEFAULT=Yes]' % doc_id)
x = ''
if x.lower() in ['y', 'yes', '']:
print('\n\n' + '-'*30 + 'Full EHR' + '-'*30)
print(text + '\n')
print('-'*80 + '\n\n')
print('-'*30 + 'Predicted Summary ' + model_name + '-'*30)
print(summary)
print('-'*80 + '\n\n')
class tests(unittest.TestCase):
def setUp(self):
self.ehrdb = ehrkit.start_session(USERNAME, PASSWORD)
self.ehrdb.get_patients(3)
''' Runs tests 1.1-1.4 '''
class t1(tests):
def test1_1(self):
kit_count = self.ehrdb.count_patients()
print("Patient count: ", kit_count)
self.ehrdb.cur.execute("SELECT COUNT(*) FROM mimic.PATIENTS")
raw = self.ehrdb.cur.fetchall()
test_count = int(raw[0][0])
self.assertEqual(test_count, kit_count)
def test1_2(self):
# Fails! count_docs returns 1573339, but mimic.NOTEEVENTS has 2083180 documents.
# TO DO: Fix whatever is wrong here
kit_count = self.ehrdb.count_docs(['NOTEEVENTS'])
print("Document count: ", kit_count)
self.ehrdb.cur.execute("SELECT COUNT(*) FROM mimic.NOTEEVENTS")
raw = self.ehrdb.cur.fetchall()
test_count = int(raw[0][0])
self.assertEqual(test_count, kit_count)
def test1_3(self):
self.ehrdb.get_note_events()
print('output format: SUBJECT_ID, ROW_ID, NoteEvent length')
lens = [(patient.id, note[0], len(note[1])) for patient in self.ehrdb.patients.values() for note in patient.note_events]
print(lens)
# placeholder, this output cannot be checked easily
self.assertEqual(1, 1)
def test1_4(self):
# Gets longest note among the patient notes queued by get_note_events()
self.ehrdb.get_note_events()
pid, rowid, doclen = self.ehrdb.longest_NE()
print('patient id is:', pid, 'NoteEvent id is:', rowid, 'length: ', doclen)
# placeholder, this output cannot be checked easily
self.assertEqual(1, 1)
class t2(tests):
def test2_1(self):
### There are 2083180 patient records in NOTEEVENTS. ###
record_id = random.randint(1, NUM_DOCS + 1)
kit_rec = self.ehrdb.get_document(record_id)
print("Document with ID %d\n: " % record_id, kit_rec)
self.ehrdb.cur.execute("select TEXT from mimic.NOTEEVENTS where ROW_ID = %d" % record_id)
test_rec = self.ehrdb.cur.fetchall()[0][0]
self.assertEqual(kit_rec, test_rec)
def test2_2(self):
### There are records from 46520 unique patients in MIMIC. ###
patient_id = random.randint(1, NUM_PATIENTS + 1)
kit_ids = self.ehrdb.get_all_patient_document_ids(patient_id)
print('Document IDs related to Patient %d: ' % patient_id, kit_ids)
print("Number of docs related to Patient %d: " % patient_id, len(kit_ids))
self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where SUBJECT_ID = %d" % patient_id)
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
def test2_3(self):
kit_ids = self.ehrdb.list_all_document_ids()
# print(kit_ids)
self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS")
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
def test2_4(self):
kit_ids = self.ehrdb.list_all_patient_ids()
self.ehrdb.cur.execute("select SUBJECT_ID from mimic.PATIENTS")
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
def test2_5(self):
### Select random date from a date in the database.
### Dates are shifted to future but preserve time, weekday, and seasonality.
random_id = random.randint(1, NUM_DOCS + 1)
self.ehrdb.cur.execute("select CHARTDATE from mimic.NOTEEVENTS where ROW_ID = %d" % random_id)
date = self.ehrdb.cur.fetchall()[0][0]
kit_ids = self.ehrdb.get_documents_d(date)
self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where CHARTDATE = \"%s\"" % date)
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
class t3(tests):
def test3_1(self):
# Defines abbreviation as a string of capitalized letters
random_id = random.randint(1, NUM_DOCS + 1)
print("Collecting abbreviations for document %d..." % random_id)
kit_abbs = self.ehrdb.get_abbreviations(random_id)
sents = self.ehrdb.get_document_sents(random_id)
test_abbs = set()
for sent in sents:
for word in ehrkit.word_tokenize(sent):
print(word)
pattern = r'[A-Z]{2}' # Only selects words in ALL CAPS
if re.match(pattern, word):
test_abbs.add(word)
print(kit_abbs)
self.assertEqual(kit_abbs, list(test_abbs))
def test3_2(self):
query = "meningitis"
# print('Printing a list of all document ids including query like ', query)
kit_ids = self.ehrdb.get_documents_q(query)
# print(kit_ids) # Extremely long list of DOC_IDs
query = "%"+query+"%"
self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where TEXT like \'%s\'" % query)
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
def test3_3(self):
### Task 3.3 is the same as task 3.2 with a different query. ###
query = "Service: SURGERY"
# print('Printing a list of all document ids including query like ', query)
kit_ids = self.ehrdb.get_documents_q(query)
# print(kit_ids) # Extremely long list of DOC_IDs
query = "%"+query+"%"
self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where TEXT like \'%s\'" % query)
raw = self.ehrdb.cur.fetchall()
test_ids = ehrkit.flatten(raw)
self.assertEqual(kit_ids, test_ids)
def test3_4(self):
doc_id = random.randint(1, NUM_DOCS + 1)
# print('Kit function printing a numbered list of all sentences in document %d' % doc_id)
# MIMIC EHRs are very messy and sentence tokenizaton often doesn't work
kit_doc = self.ehrdb.get_document_sents(doc_id)
# ehrkit.numbered_print(kit_doc)
self.ehrdb.cur.execute("select TEXT from mimic.NOTEEVENTS where ROW_ID = %d " % doc_id)
raw = self.ehrdb.cur.fetchall()
test_doc = ehrkit.sent_tokenize(raw[0][0])
self.assertEqual(kit_doc, test_doc)
def test3_7(self):
kit_meds = self.ehrdb.count_all_prescriptions()
test_meds = {}
self.ehrdb.cur.execute("select DRUG from mimic.PRESCRIPTIONS")
raw = self.ehrdb.cur.fetchall()
meds_list = ehrkit.flatten(raw)
for med in meds_list:
if med in test_meds:
test_meds[med] += 1
else:
test_meds[med] = 1
self.assertEqual(kit_meds, test_meds)
class t4(tests):
def test4_1(self):
d = self.ehrdb.get_documents_icd9()
print(d)
self.assertIsNotNone(d['code'])
def test4_4(self):
pass
class t5(tests):
def test5_1(self):
doc_id = random.randint(1, NUM_DOCS + 1)
kit_phrases = self.ehrdb.extract_phrases(doc_id)
print("Testing task 5.1\n Check phrases manually: ", kit_phrases)
self.assertIsNotNone(kit_phrases)
def test5_4(self):
gender = random.choice(['M', 'F'])
kit_count = self.ehrdb.count_gender(gender)
self.ehrdb.cur.execute('SELECT COUNT(*) FROM mimic.PATIENTS WHERE GENDER = \'%s\'' % gender)
raw = self.ehrdb.cur.fetchall()
test_count = raw[0][0]
print('Gender:', gender, '\tCount:', str(test_count))
self.assertEqual(kit_count, test_count)
# class t6(tests):
# def test6_1(self):
# import loader
# doc_id, text = select_ehr(self.ehrdb)
# # x = input('GloVe or RoBERTa predictor [g=GloVe, r=RoBERTa]? ')
# x = 'g'
# if x == 'g':
# glove_predictor = loader.load_glove()
# probs = glove_predictor.predict(text)['probs']
# elif x == 'r':
# roberta_predictor = loader.load_roberta()
# try:
# probs = roberta_predictor.predict(text)['probs']
# except:
# print('Document too long for RoBERTa model. Using GLoVe instead.')
# glove_predictor = loader.load_glove()
# probs = glove_predictor.predict(text)['probs']
# else:
# sys.exit('Error: Must input \'g\' or \'r\'')
# classification = 'Positive' if probs[0] >= 0.5 else 'Negative'
# print("Document ID: ", doc_id, "\tPredicted Sentiment: ", classification)
# def test6_2(self):
# import loader
# doc_id, text = select_ehr(self.ehrdb)
# if os.path.exists("../allennlp/elmo-ner/whole_model.pt"):
# predictor = loader.load_ner()
# else:
# predictor = loader.download_ner()
# text = self.ehrdb.get_document(int(doc_id))
# pred = predictor.predict(text)
# # pred = predictor.predict("John likes and Bill hates ice cream")
# # print_results = input("Prediction complete. Print results? (y/n): ")
# print_results='y'
# if print_results == "y":
# print("Document ID: ", doc_id, " Results: ", pred['tags'])
# def test6_3(self):
# import torch
# from transformers import BertTokenizer#, BertModel, BertForMaskedLM
# doc_id, text = select_ehr(self.ehrdb)
# tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# bert_tokenized_text = tokenizer.tokenize(text)
# print('\n' + '-'*20 + 'text' + '-'*20)
# print(text)
# print('\n' + '-'*20 + 'Tokenized text from Huggingface BERT Tokenizer' + '-'*20)
# print(bert_tokenized_text)
# # library function
# ehr_bert_tokenized_text = self.ehrdb.get_bert_tokenize(doc_id)
# self.assertEqual(bert_tokenized_text, ehr_bert_tokenized_text)
class t7(tests):
# Summarization algorithms
def test7_1(self):
from pubmed_naive_bayes import classify_nb
from get_pubmed_nb_data import build_vecs
from sklearn.naive_bayes import GaussianNB
doc_id, text = select_ehr(self.ehrdb)
# body_type = input('Use Naive Bayes model trained from whole body sections or just their body introductions?\n\t'\
# '[w=whole body, j=just intro, DEFAULT=just intro]: ')
body_type = 'j'
if body_type == 'w':
ending = 'body'
elif body_type in ['j', '']:
ending = 'intro'
else:
sys.exit('Error: Must input \'w\' or \'j.\'')
SUMM_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization'))
best_dir_name = get_nb_dir(ending, SUMM_DIR)
if not best_dir_name:
message = 'No Naive Bayes models of this type have been fit. '\
'Would you like to do so now?\n\t[DEFAULT=Yes] '
# response = input(message)
response = ''
if response.lower() in ['y', 'yes', '']:
command = 'python ' + os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization', 'pubmed_naive_bayes.py'))
os.system(command)
best_dir_name = get_nb_dir(ending)
if response.lower() not in ['y', 'yes', ''] or not best_dir_name:
sys.exit('Exiting.')
# Fits model to data
NB_DIR = os.path.join(SUMM_DIR, best_dir_name, 'nb')
with open(os.path.join(NB_DIR, 'feature_vecs.json'), 'r') as f:
data = json.load(f)
xtrain, ytrain = data['train_features'], data['train_outputs']
gnb = GaussianNB()
gnb.fit(xtrain, ytrain)
# Evaluates on model
tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
feature_vecs, _ = build_vecs(text, None, tokenizer)
PCT_SUM = 0.3
preds = classify_nb(feature_vecs, PCT_SUM, gnb)
sents = tokenizer.tokenize(text)
summary = ''
for i in range(len(preds)):
if preds[i] == 1:
summary += sents[i]
show_summary(doc_id, text, summary, 'Naive Bayes')
def test7_2(self):
# Distilbart for summarization. Trained on CNN/ Daily Mail (~4x longer summaries than XSum)
doc_id, text = select_ehr(self.ehrdb, requires_long=True)
model_name = 'sshleifer/distilbart-cnn-12-6'
summary = self.ehrdb.summarize_huggingface(text, model_name)
show_summary(doc_id, text, summary, model_name)
print('Number of Words in Full EHR: %d' % len(text.split()))
print('Number of Words in %s Summary: %d' % (model_name, len(summary.split())))
def test7_3(self):
# T5 for summarization. Trained on CNN/ Daily Mail
doc_id, text = select_ehr(self.ehrdb, requires_long=True)
model_name = 't5-small'
summary = self.ehrdb.summarize_huggingface(text, model_name)
show_summary(doc_id, text, summary, model_name)
print('Number of Words in Full EHR: %d' % len(text.split()))
print('Number of Words in %s Summary: %d' % (model_name, len(summary.split())))
if __name__ == '__main__':
unittest.main()
