import torch import json import os import pandas as pd import numpy as np from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer from glob import glob from collections.abc import Iterable from collections import defaultdict pheno_map = {'alcohol.abuse': 0, 'advanced.lung.disease': 1, 'advanced.heart.disease': 2, 'chronic.pain.fibromyalgia': 3, 'other.substance.abuse': 4, 'psychiatric.disorders': 5, 'obesity': 6, 'depression': 7, 'advanced.cancer': 8, 'chronic.neurological.dystrophies': 9, 'none': -1} rev_pheno_map = {v: k for k,v in pheno_map.items()} valid_cats = range(0,9) umls_cats = ['T114', 'T029', 'T073', 'T058', 'T191', 'T200', 'T048', 'T019', 'T046', 'T023', 'T041', 'T059', 'T184', 'T034', 'T116', 'T039', 'T127', 'T201', 'T129', 'T067', 'T109', 'T197', 'T131', 'T130', 'T126', 'T061', 'T203', 'T047', 'T037', 'T074', 'T031', 'T195', 'T168'] umls_map = {s: i for i,s in enumerate(umls_cats)} def gen_splits(args, phenos): np.random.seed(0) if args.task == 'token': files = glob(os.path.join(args.data_dir, 'mimic_decisions/data/**/*')) if args.use_umls: files = ["/".join(x.split('/')[-1:]) for x in files] else: files = ["/".join(x.split('/')[-2:]) for x in files] subjects = np.unique([os.path.basename(x).split('_')[0] for x in files]) elif phenos is not None: subjects = phenos['subject_id'].unique() else: raise ValueError phenos['phenotype_label'] = phenos['phenotype_label'].apply(lambda x: x.lower()) n = len(subjects) train_count = int(0.8*n) val_count = int(0.9*n) - int(0.8*n) test_count = n - int(0.9*n) train, val, test = [], [], [] np.random.shuffle(subjects) subjects = list(subjects) pheno_list = set(pheno_map.keys()) if args.unseen_pheno is not None: test_phenos = {rev_pheno_map[args.unseen_pheno]} unseen_pheno = rev_pheno_map[args.unseen_pheno] train_phenos = pheno_list - test_phenos else: test_phenos = pheno_list train_phenos = pheno_list unseen_pheno = 'null' while len(subjects) > 0: if len(pheno_list) > 0: for pheno in pheno_list: if len(train) < train_count and pheno in train_phenos: el = None for i, subj in enumerate(subjects): row = phenos[phenos.subject_id == subj] if row['phenotype_label'].apply(lambda x: pheno in x and not unseen_pheno in x).any(): el = subjects.pop(i) break if el is not None: train.append(el) elif el is None: pheno_list.remove(pheno) break if len(val) < val_count and (not args.pheno_id or len(val) <= (0.5*val_count)): el = None for i, subj in enumerate(subjects): row = phenos[phenos.subject_id == subj] if row['phenotype_label'].apply(lambda x: pheno in x).any(): el = subjects.pop(i) break if el is not None: val.append(el) elif el is None: pheno_list.remove(pheno) break if len(test) < test_count or (args.unseen_pheno is not None and pheno in test_phenos): el = None for i, subj in enumerate(subjects): row = phenos[phenos.subject_id == subj] if row['phenotype_label'].apply(lambda x: pheno in x).any(): el = subjects.pop(i) break if el is not None: test.append(el) elif el is None: pheno_list.remove(pheno) break else: if len(train) < train_count: el = subjects.pop() if el is not None: train.append(el) if len(val) < val_count: el = subjects.pop() if el is not None: val.append(el) if len(test) < test_count: el = subjects.pop() if el is not None: test.append(el) if args.task == 'token': train = [x for x in files if os.path.basename(x).split('_')[0] in train] val = [x for x in files if os.path.basename(x).split('_')[0] in val] test = [x for x in files if os.path.basename(x).split('_')[0] in test] elif phenos is not None: train = phenos[phenos.subject_id.isin(train)] val = phenos[phenos.subject_id.isin(val)] test = phenos[phenos.subject_id.isin(test)] return train, val, test class MyDataset(Dataset): def __init__(self, args, tokenizer, data_source, phenos, train = False): super().__init__() self.tokenizer = tokenizer self.data = [] self.train = train self.pheno_ids = defaultdict(list) self.dec_ids = {k: [] for k in pheno_map.keys()} if args.task == 'seq': for i, row in data_source.iterrows(): sample = self.load_phenos(args, row, i) self.data.append(sample) else: for i, fn in enumerate(data_source): sample = self.load_decisions(args, fn, i, phenos) self.data.append(sample) def load_phenos(self, args, row, idx): txt_candidates = glob(os.path.join(args.data_dir, f'mimic_decisions/raw_text/{row["subject_id"]}_{row["hadm_id"]}*.txt')) text = open(txt_candidates[0]).read() if args.pheno_n == 500: file_dir = glob(os.path.join(args.data_dir, f'mimic_decisions/data/*/{row["subject_id"]}_{row["hadm_id"]}*.json'))[0] with open(file_dir) as f: data = json.load(f, strict=False) annots = data[0]['annotations'] if args.text_subset: unlabeled_text = np.ones(len(text), dtype=bool) labeled_text = np.zeros(len(text), dtype=bool) for annot in annots: cat = parse_cat(annot['category']) start, end = map(int, (annot['start_offset'], annot['end_offset'])) if cat is not None: unlabeled_text[start:end] = 0 if cat in args.text_subset: labeled_text[start:end] = 1 combined_text = unlabeled_text | labeled_text if args.include_nolabel else labeled_text text = "".join([c for i,c in enumerate(text) if combined_text[i]]) encoding = self.tokenizer.encode_plus(text, truncation=args.truncate_train if self.train else args.truncate_eval) ids = np.zeros((args.num_decs, len(encoding['input_ids']))) for annot in annots: start = int(annot['start_offset']) enc_start = encoding.char_to_token(start) i = 1 while enc_start is None: enc_start = encoding.char_to_token(start+i) i += 1 end = int(annot['end_offset']) enc_end = encoding.char_to_token(end) j = 1 while enc_end is None: enc_end = encoding.char_to_token(end-j) j += 1 if enc_start is None or enc_end is None: raise ValueError cat = parse_cat(annot['category']) if not cat or cat not in valid_cats: continue ids[cat-1, enc_start:enc_end] = 1 else: encoding = self.tokenizer.encode_plus(text, truncation=args.truncate_train if self.train else args.truncate_eval) ids = None labels = np.zeros(args.num_phenos) if args.pheno_n in (500, 800): sample_phenos = row['phenotype_label'] if sample_phenos != 'none': for pheno in sample_phenos.split(','): labels[pheno_map[pheno.lower()]] = 1 elif args.pheno_n == 1500: for k,v in pheno_map.items(): if row[k] == 1: labels[v] = 1 if args.pheno_id is not None: if args.pheno_id == -1: labels = [0.0 if any(labels) else 1.0] else: labels = [labels[args.pheno_id]] return encoding['input_ids'], labels, ids def load_decisions(self, args, fn, idx, phenos): basename = os.path.basename(fn).split("-")[0] if args.use_umls: file_dir = os.path.join(args.data_dir, 'mimic_decisions/umls', basename) else: file_dir = os.path.join(args.data_dir, 'mimic_decisions/data', fn) pheno_id = "_".join(basename.split("_")[:3]) + '.txt' txt_candidates = glob(os.path.join(args.data_dir, f'mimic_decisions/raw_text/{basename}*.txt')) text = open(txt_candidates[0]).read() encoding = self.tokenizer.encode_plus(text, max_length=args.max_len, truncation=args.truncate_train if self.train else args.truncate_eval, padding = 'max_length', ) if pheno_id in phenos.index: sample_phenos = phenos.loc[pheno_id]['phenotype_label'] for pheno in sample_phenos.split(','): self.pheno_ids[pheno].append(idx) with open(file_dir) as f: data = json.load(f, strict=False) if args.use_umls: annots = data else: annots = data[0]['annotations'] if args.label_encoding == 'multiclass': labels = np.full(len(encoding['input_ids']), args.num_labels-1, dtype=int) else: labels = np.zeros((len(encoding['input_ids']), args.num_labels)) for annot in annots: start = int(annot['start_offset']) enc_start = encoding.char_to_token(start) i = 1 while enc_start is None and i < 10: enc_start = encoding.char_to_token(start+i) i += 1 if i == 10: break end = int(annot['end_offset']) enc_end = encoding.char_to_token(end) j = 1 while enc_end is None and j < 10: enc_end = encoding.char_to_token(end-j) j += 1 if j == 10: enc_end = len(encoding.input_ids) if enc_start is None or enc_end is None: raise ValueError if args.label_encoding == 'multiclass' and any([x in [2*y for y in range(args.num_labels//2)] for x in labels[enc_start:enc_end]]): continue if args.use_umls: cat = umls_map.get(annot['category'], None) else: cat = parse_cat(annot['category']) if cat: cat -= 1 if cat is None or (not args.use_umls and cat not in valid_cats): continue if args.label_encoding == 'multiclass': cat1 = cat * 2 cat2 = cat * 2 + 1 labels[enc_start] = cat1 labels[enc_start+1:enc_end] = cat2 elif args.label_encoding == 'bo': cat1 = cat * 2 cat2 = cat * 2 + 1 labels[enc_start, cat1] = 1 labels[enc_start+1:enc_end, cat2] = 1 elif args.label_encoding == 'boe': cat1 = cat * 3 cat2 = cat * 3 + 1 cat3 = cat * 3 + 2 labels[enc_start, cat1] = 1 labels[enc_start+1:enc_end-1, cat2] = 1 labels[enc_end-1, cat3] = 1 else: labels[enc_start:enc_end, cat] = 1 return {'input_ids': encoding['input_ids'], 'labels': labels, 't2c': encoding.token_to_chars} def __getitem__(self, idx): return self.data[idx] def __len__(self): return len(self.data) def parse_cat(cat): for i,c in enumerate(cat): if c.isnumeric(): if cat[i+1].isnumeric(): return int(cat[i:i+2]) return int(c) return None def load_phenos(args): if args.pheno_n == 500: phenos = pd.read_csv(os.path.join(args.data_dir, 'mimic_decisions/phenos500'), sep='\t').rename(lambda x: x.strip(), axis=1) phenos['raw_text'] = phenos['raw_text'].apply(lambda x: os.path.basename(x)) phenos[['SUBJECT_ID', 'HADM_ID', 'ROW_ID']] = \ [os.path.splitext(x)[0].split('_')[:3] for x in phenos['raw_text']] phenos = phenos[phenos['phenotype_label'] != '?'] elif args.pheno_n == 800: phenos = pd.read_csv(os.path.join(args.data_dir, 'mimic_decisions/phenos800.csv')) phenos.rename({'Ham_ID': 'HADM_ID'}, inplace=True, axis=1) phenos = phenos[phenos.phenotype_label != '?'] elif args.pheno_n == 1500: phenos = pd.read_csv(os.path.join(args.data_dir, 'mimic_decisions/phenos1500.csv')) phenos.rename({'Hospital.Admission.ID': 'HADM_ID', 'subject.id': 'SUBJECT_ID'}, inplace=True, axis=1) phenos = phenos[phenos.Unsure != 1] phenos['psychiatric.disorders'] = phenos['Dementia']\ | phenos['Developmental.Delay.Retardation']\ | phenos['Schizophrenia.and.other.Psychiatric.Disorders'] else: raise ValueError phenos.rename(lambda k: k.lower(), inplace=True, axis = 1) return phenos def downsample(dataset): data = dataset.data class0 = [x for x in data if x[1][0] == 0] class1 = [x for x in data if x[1][0] == 1] if len(class0) > len(class1): class0 = resample(class0, replace=False, n_samples=len(class1), random_state=0) else: class1 = resample(class1, replace=False, n_samples=len(class0), random_state=0) dataset.data = class0 + class1 def upsample(dataset): data = dataset.data class0 = [x for x in data if x[1][0] == 0] class1 = [x for x in data if x[1][0] == 1] if len(class0) > len(class1): class1 = resample(class1, replace=True, n_samples=len(class0), random_state=0) else: class0 = resample(class0, replace=True, n_samples=len(class1), random_state=0) dataset.data = class0 + class1 def load_tokenizer(name): return AutoTokenizer.from_pretrained(name) def load_data(args): from sklearn.utils import resample def collate_segment(batch): xs = [] ys = [] t2cs = [] has_ids = 'ids' in batch[0] if has_ids: idss = [] else: ids = None masks = [] for i in range(len(batch)): x = batch[i]['input_ids'] y = batch[i]['labels'] if has_ids: ids = batch[i]['ids'] n = len(x) if n > args.max_len: start = np.random.randint(0, n - args.max_len + 1) x = x[start:start + args.max_len] if args.task == 'token': y = y[start:start + args.max_len] if has_ids: new_ids = [] ids = [x[start:start + args.max_len] for x in ids] for subids in ids: subids = [idx for idx, x in enumerate(subids) if x] new_ids.append(subids) all_ids = set([y for x in new_ids for y in x]) nones = set(range(args.max_len)) - all_ids new_ids.append(list(nones)) mask = [1] * args.max_len elif n < args.max_len: x = np.pad(x, (0, args.max_len - n)) if args.task == 'token': y = np.pad(y, ((0, args.max_len - n), (0, 0))) mask = [1] * n + [0] * (args.max_len - n) else: mask = [1] * n xs.append(x) ys.append(y) t2cs.append(batch[i]['t2c']) if has_ids: idss.append(new_ids) masks.append(mask) xs = torch.tensor(xs) ys = torch.tensor(ys) masks = torch.tensor(masks) return {'input_ids': xs, 'labels': ys, 'ids': ids, 'mask': masks, 't2c': t2cs} def collate_full(batch): lens = [len(x['input_ids']) for x in batch] max_len = max(args.max_len, max(lens)) for i in range(len(batch)): batch[i]['input_ids'] = np.pad(batch[i]['input_ids'], (0, max_len - lens[i])) if args.task == 'token': if args.label_encoding == 'multiclass': batch[i]['labels'] = np.pad(batch[i]['labels'], (0, max_len - lens[i]), constant_values=-100) else: batch[i]['labels'] = np.pad(batch[i]['labels'], ((0, max_len - lens[i]), (0, 0))) mask = [1] * lens[i] + [0] * (max_len - lens[i]) batch[i]['mask'] = mask batch = {k: torch.tensor(np.array([sample[k] for sample in batch])) if isinstance(batch[0][k], Iterable) else [sample[k] for sample in batch] for k in batch[0].keys()} return batch tokenizer = load_tokenizer(args.model_name) args.vocab_size = tokenizer.vocab_size args.max_length = min(tokenizer.model_max_length, 512) if args.mimic_data: from datasets import Dataset df = pd.read_csv('/data/mohamed/data/mimiciii/NOTEEVENTS.csv.gz', usecols=['ROW_ID', 'SUBJECT_ID', 'HADM_ID', 'TEXT']) data = Dataset.from_pandas(df) return data, tokenizer else: phenos = load_phenos(args) train_files, val_files, test_files = gen_splits(args, phenos) phenos.set_index('raw_text', inplace=True) train_dataset = MyDataset(args, tokenizer, train_files, phenos, train=True) if args.resample == 'down': downsample(train_dataset) elif args.resample == 'up': upsample(train_dataset) val_dataset = MyDataset(args, tokenizer, val_files, phenos) test_dataset = MyDataset(args, tokenizer, test_files, phenos) print('Train dataset:', len(train_dataset)) print('Val dataset:', len(val_dataset)) print('Test dataset:', len(test_dataset)) train_ns = DataLoader(train_dataset, 1, False, collate_fn=collate_full, ) train_dataloader = DataLoader(train_dataset, args.batch_size, True, collate_fn=collate_segment, ) val_dataloader = DataLoader(val_dataset, 1, False, collate_fn=collate_full) test_dataloader = DataLoader(test_dataset, 1, False, collate_fn=collate_full) train_files = [os.path.basename(x).split('-')[0] for x in train_files] val_files = [os.path.basename(x).split('-')[0] for x in val_files] test_files = [os.path.basename(x).split('-')[0] for x in test_files] return train_dataloader, val_dataloader, test_dataloader, train_ns, [train_files, val_files, test_files]