data.py

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]