import os
import numpy as np
import pandas as pd
from transformers import AutoTokenizer
from load_umls import UMLS
from torch.utils.data import Dataset, DataLoader
from random import sample
from sampler_util import FixedLengthBatchSampler, my_collate_fn
from torch.utils.data.sampler import RandomSampler
import ipdb
from time import time
import json
def pad(list_ids, pad_length, pad_mark=0):
output = []
for l in list_ids:
if len(l) > pad_length:
output.append(l[0:pad_length])
else:
output.append(l + [pad_mark] * (pad_length - len(l)))
return output
def my_sample(lst, lst_length, start, length):
start = start % lst_length
if start + length < lst_length:
return lst[start:start + length]
return lst[start:] + lst[0:start + length - lst_length]
class UMLSDataset(Dataset):
def __init__(self, umls_folder, model_name_or_path, lang, json_save_path=None, max_lui_per_cui=8, max_length=32):
self.umls = UMLS(umls_folder, lang_range=lang)
self.len = len(self.umls.rel)
self.max_lui_per_cui = max_lui_per_cui
self.max_length = max_length
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
self.json_save_path = json_save_path
self.calculate_class_count()
def calculate_class_count(self):
print("Calculate class count")
self.cui2id = {cui: index for index,
cui in enumerate(self.umls.cui2str.keys())}
self.re_set = set()
self.rel_set = set()
for r in self.umls.rel:
_, _, re, rel = r.split("\t")
self.re_set.update([re])
self.rel_set.update([rel])
self.re_set = list(self.re_set)
self.rel_set = list(self.rel_set)
self.re_set.sort()
self.rel_set.sort()
self.re2id = {re: index for index, re in enumerate(self.re_set)}
self.rel2id = {rel: index for index, rel in enumerate(self.rel_set)}
sty_list = list(set(self.umls.cui2sty.values()))
sty_list.sort()
self.sty2id = {sty: index for index, sty in enumerate(sty_list)}
if self.json_save_path:
with open(os.path.join(self.json_save_path, "re2id.json"), "w") as f:
json.dump(self.re2id, f)
with open(os.path.join(self.json_save_path, "rel2id.json"), "w") as f:
json.dump(self.rel2id, f)
with open(os.path.join(self.json_save_path, "sty2id.json"), "w") as f:
json.dump(self.sty2id, f)
print("CUI:", len(self.cui2id))
print("RE:", len(self.re2id))
print("REL:", len(self.rel2id))
print("STY:", len(self.sty2id))
def tokenize_one(self, string):
return self.tokenizer.encode_plus(string, max_length=self.max_length, truncation=True)['input_ids']
# @profile
def __getitem__(self, index):
cui0, cui1, re, rel = self.umls.rel[index].split("\t")
str0_list = list(self.umls.cui2str[cui0])
str1_list = list(self.umls.cui2str[cui1])
if len(str0_list) > self.max_lui_per_cui:
str0_list = sample(str0_list, self.max_lui_per_cui)
if len(str1_list) > self.max_lui_per_cui:
str1_list = sample(str1_list, self.max_lui_per_cui)
use_len = min(len(str0_list), len(str1_list))
str0_list = str0_list[0:use_len]
str1_list = str1_list[0:use_len]
sty0_index = self.sty2id[self.umls.cui2sty[cui0]]
sty1_index = self.sty2id[self.umls.cui2sty[cui1]]
str2_list = []
cui2_index_list = []
sty2_index_list = []
cui2 = my_sample(self.umls.cui, self.umls.cui_count,
index * self.max_lui_per_cui, use_len * 2)
sample_index = 0
while len(str2_list) < use_len:
if sample_index < len(cui2):
use_cui2 = cui2[sample_index]
else:
sample_index = 0
cui2 = my_sample(self.umls.cui, self.umls.cui_count,
index * self.max_lui_per_cui, use_len * 2)
use_cui2 = cui2[sample_index]
# if not "\t".join([cui0, use_cui2, re, rel]) in self.umls.rel: # TOO SLOW!
if True:
cui2_index_list.append(self.cui2id[use_cui2])
sty2_index_list.append(
self.sty2id[self.umls.cui2sty[use_cui2]])
str2_list.append(sample(self.umls.cui2str[use_cui2], 1)[0])
sample_index += 1
# print(str0_list)
# print(str1_list)
# print(str2_list)
input_ids = [self.tokenize_one(s)
for s in str0_list + str1_list + str2_list]
input_ids = pad(input_ids, self.max_length)
input_ids_0 = input_ids[0:use_len]
input_ids_1 = input_ids[use_len:2 * use_len]
input_ids_2 = input_ids[2 * use_len:]
cui0_index = self.cui2id[cui0]
cui1_index = self.cui2id[cui1]
re_index = self.re2id[re]
rel_index = self.rel2id[rel]
return input_ids_0, input_ids_1, input_ids_2, \
[cui0_index] * use_len, [cui1_index] * use_len, cui2_index_list, \
[sty0_index] * use_len, [sty1_index] * use_len, sty2_index_list, \
[re_index] * use_len, \
[rel_index] * use_len
def __len__(self):
return self.len
def fixed_length_dataloader(umls_dataset, fixed_length=96, num_workers=0):
base_sampler = RandomSampler(umls_dataset)
batch_sampler = FixedLengthBatchSampler(
sampler=base_sampler, fixed_length=fixed_length, drop_last=True)
dataloader = DataLoader(umls_dataset, batch_sampler=batch_sampler,
collate_fn=my_collate_fn, num_workers=num_workers, pin_memory=True)
return dataloader
if __name__ == "__main__":
umls_dataset = UMLSDataset(umls_folder="../umls",
model_name_or_path="../biobert_v1.1",
lang=None)
ipdb.set_trace()
umls_dataloader = fixed_length_dataloader(umls_dataset, num_workers=4)
now_time = time()
for index, batch in enumerate(umls_dataloader):
print(time() - now_time)
now_time = time()
if index < 10:
for item in batch:
print(item.shape)
#print(batch)
else:
import sys
sys.exit()
