from transformers import AutoConfig, AutoModel, AutoTokenizer
import torch
from tqdm import tqdm
import numpy as np
import sys
sys.path.append("../../pretrain")
from load_umls import UMLS
from nltk.tokenize import word_tokenize
import os
import ipdb
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
batch_size = 512
max_seq_length = 32
def get_drug_diseases_to_check(concept_filename):
query_to_targets = {}
with open(concept_filename, 'r') as infile:
data = infile.readlines()
for row in data:
drug, diseases = row.strip().split(':')
diseases = diseases.split(',')[:-1]
disease_cui_set = set([])
for disease in diseases:
disease_cui_set.add(disease)
if len(disease_cui_set) > 0:
query_to_targets[drug] = disease_cui_set
cui_list = set()
for query, targets in query_to_targets.items():
cui_list.update([query])
cui_list.update(targets)
return query_to_targets, list(cui_list)
def normalize(tensor):
norm = torch.norm(tensor, p=2, dim=1, keepdim=True).clamp(min=1e-12)
return torch.div(tensor, norm)
def calculate_mrm_ndfrt_origin(term_embedding, cui_list, query_to_targets, k):
return calculate_mrm_ndfrt_delta(term_embedding, cui_list, query_to_targets, None, k)
def calculate_mrm_ndfrt_q2t(term_embedding, cui_list, query_to_targets, k):
delta_list = []
term_embedding = torch.FloatTensor(term_embedding).to(device)
norm_embedding = normalize(term_embedding)
id2cui = {i:cui_list[i] for i in range(len(cui_list))}
cui2id = {cui:index for index, cui in id2cui.items()}
for query, targets in query_to_targets.items():
if query in cui2id:
for target in targets:
if target in cui2id:
delta = term_embedding[cui2id[query]] - term_embedding[cui2id[target]]
delta_list.append(delta)
overall_output = []
for _, delta in tqdm(enumerate(delta_list)):
output = []
for query, targets in query_to_targets.items():
if query in cui2id:
find_embedding = term_embedding[cui2id[query]] - delta
similarity = torch.matmul(norm_embedding, find_embedding)
_, indices = torch.topk(similarity, k=k + 1)
find_cui = [cui_list[index] for index in indices[1:]]
score = 0.
for cui in find_cui:
if cui in targets:
score = 1.
break
output.append(score)
if len(output) > 0:
score = sum(output) / len(output)
else:
score = 0.
overall_output.append(score)
if len(overall_output) > 0:
overall_score = sum(overall_output) / len(overall_output)
overall_max = max(overall_output)
else:
overall_score = 0
overall_max = 0
return overall_score, overall_max
def calculate_mrm_ndfrt_delta(term_embedding, cui_list, query_to_targets, delta=None, k=40):
term_embedding = torch.FloatTensor(term_embedding).to(device)
norm_embedding = normalize(term_embedding)
id2cui = {i:cui_list[i] for i in range(len(cui_list))}
cui2id = {cui:index for index, cui in id2cui.items()}
output = []
check_count = 0
for query, targets in query_to_targets.items():
if query in cui2id:
query_embedding = term_embedding[cui2id[query]]
if delta is None:
find_embedding = query_embedding
else:
find_embedding = query_embedding - torch.FloatTensor(delta).to(device)
similarity = torch.matmul(norm_embedding, find_embedding)
_, indices = torch.topk(similarity, k=k + 1)
find_cui = [cui_list[index] for index in indices[1:]]
score = 0.
for cui in find_cui:
if cui in targets:
score = 1.
break
output.append(score)
check_count += 1
del term_embedding
if len(output) > 0:
score = sum(output) / len(output)
else:
score = 0.
"""
print(f"Check count: {check_count}")
print(score)
"""
return score
def mrm_ndfrt_cui(cui_embedding, umls, cui_list, query_to_targets, k, method):
w, _ = load_embedding(cui_embedding)
new_cui_list = [cui for cui in cui_list if cui in w]
term_embedding = np.array([w[cui] for cui in new_cui_list])
print(f"Cui count:{len(new_cui_list)}")
if method == "origin":
score = calculate_mrm_ndfrt_origin(term_embedding, new_cui_list, query_to_targets, k)
print(f"Origin: {score}")
if method == "all":
score = calculate_mrm_ndfrt_q2t(term_embedding, new_cui_list, query_to_targets, k)
average_score, max_score = score
print(f"Average: {average_score}")
print(f"Max: {max_score}")
return score
def mrm_ndfrt_word(word_embedding, umls, cui_list, query_to_targets, k, method):
w, dim = load_embedding(word_embedding)
print("Tokenize and calculate avg embedding.")
cui_str = [[word for word in word_tokenize(
list(umls.cui2str[cui])[0]) if word in w] for cui in cui_list if cui in umls.cui2str]
new_cui_list = []
check_count = 0
for index, des in enumerate(cui_str):
if len(des) > 0:
tmp_emb = np.zeros((dim))
for word in des:
tmp_emb += w[word]
if check_count == 0:
term_embedding = tmp_emb
else:
term_embedding = np.concatenate(
(term_embedding, tmp_emb), axis=0)
check_count += 1
new_cui_list.append(cui_list[index])
term_embedding = term_embedding.reshape((-1, dim))
print(f"Cui count:{len(new_cui_list)}")
if method == "origin":
score = calculate_mrm_ndfrt_origin(term_embedding, new_cui_list, query_to_targets, k)
print(f"Origin: {score}")
if method == "all":
score = calculate_mrm_ndfrt_q2t(term_embedding, new_cui_list, query_to_targets, k)
average_score, max_score = score
print(f"Average: {average_score}")
print(f"Max: {max_score}")
return score
def mrm_ndfrt_bert(bert_embedding, umls, cui_list, query_to_targets, k, method, summary_method):
print(summary_method)
model, tokenizer = load_bert(bert_embedding)
model.eval()
input_ids = []
new_cui_list = []
for cui in cui_list:
if cui in umls.cui2str:
input_ids.append(tokenizer.encode_plus(
list(umls.cui2str[cui])[
0], max_length=max_seq_length, add_special_tokens=True,
truncation=True, pad_to_max_length=True)['input_ids'])
new_cui_list.append(cui)
count = len(input_ids)
now_count = 0
# with tqdm(total=count) as pbar:
with torch.no_grad():
while now_count < count:
input_gpu_0 = torch.LongTensor(input_ids[now_count:min(
now_count + batch_size, count)]).to(device)
if summary_method == "CLS":
embed = model(input_gpu_0)[1]
if summary_method == "MEAN":
embed = torch.mean(model(input_gpu_0)[0], dim=1)
embed_np = embed.cpu().detach().numpy()
if now_count == 0:
term_embedding = embed_np
else:
term_embedding = np.concatenate((term_embedding, embed_np), axis=0)
update = min(now_count + batch_size, count) - now_count
now_count = now_count + update
# pbar.update(update)
print(f"Cui count:{len(new_cui_list)}")
if method == "origin":
score = calculate_mrm_ndfrt_origin(term_embedding, new_cui_list, query_to_targets, k)
print(f"Origin: {score}")
if method == "all":
score = calculate_mrm_ndfrt_q2t(term_embedding, new_cui_list, query_to_targets, k)
average_score, max_score = score
print(f"Average: {average_score}")
print(f"Max: {max_score}")
if method in ["may_treat", "may_prevent"]:
beta_path = os.path.join(bert_embedding, "run", "1000000", "rel embedding")
with open(os.path.join(beta_path, "metadata.tsv"), "r", encoding="utf-8") as f:
metadata = f.readlines()
metadata = [line.strip() for line in metadata]
with open(os.path.join(beta_path, "tensors.tsv"), "r", encoding="utf-8") as f:
tensor = f.readlines()
tensor = [[float(num) for num in line.split("\t")] for line in tensor]
for index, title in enumerate(metadata):
if title == method:
delta = tensor[index]
score = calculate_mrm_ndfrt_delta(term_embedding, new_cui_list, query_to_targets, delta, k)
print(f"{method}: {score}")
return score
def mrm_ndfrt(embedding_list, embedding_type_list, concept_filename, k=40, check_intersection=True):
if check_intersection:
if not os.path.exists("intersection.txt"):
intersection_cui = get_intersection(
embedding_list, embedding_type_list)
with open("intersection.txt", "w", encoding="utf-8") as f:
for cui in intersection_cui:
f.write(cui.strip() + "\n")
else:
with open("intersection.txt", "r", encoding="utf-8") as f:
lines = f.readlines()
intersection_cui = [line.strip() for line in lines]
query_to_targets, cui_list = get_drug_diseases_to_check(concept_filename)
umls = UMLS("../../umls", only_load_dict=True) # source_range='SNOMEDCT_US')#, only_load_dict=True)
if check_intersection:
cui_list = [cui for cui in cui_list if cui in intersection_cui]
#cui_list = [cui for cui in umls.cui2str if umls.cui2sty[cui] in sty_list]
#cui_list = [cui for cui in cui_list if cui in umls.sty_list]
"""
for cui in cui_list:
if not cui in umls.cui2str:
print(cui)
ipdb.set_trace()
"""
opt = []
"""
# Origin
print("ORIGIN")
for index, embedding in enumerate(embedding_list):
if embedding_type_list[index].lower() == "cui":
opt.append(mrm_ndfrt_cui(embedding, umls, cui_list, query_to_targets, k, "origin"))
if embedding_type_list[index].lower() == "word":
opt.append(mrm_ndfrt_word(embedding, umls, cui_list, query_to_targets, k, "origin"))
if embedding_type_list[index].lower() == "bert":
#opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
# query_to_targets, k, "origin", summary_method="MEAN"))
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, "origin", summary_method="CLS"))
# For UMLSBert
for index, embedding in enumerate(embedding_list):
if embedding_type_list[index].lower() == "bert":
print("BETA")
beta_path = os.path.join(embedding, "run", "1000000", "rel embedding")
if os.path.exists(beta_path):
if concept_filename.find('treat') >= 0:
method = "may_treat"
else:
method = "may_prevent"
#opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
# query_to_targets, k, method, summary_method="MEAN"))
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, method, summary_method="CLS"))
# For average and max
print("ALL")
for index, embedding in enumerate(embedding_list):
if embedding_type_list[index].lower() == "cui":
opt.append(mrm_ndfrt_cui(embedding, umls, cui_list, query_to_targets, k, "all"))
if embedding_type_list[index].lower() == "word":
opt.append(mrm_ndfrt_word(embedding, umls, cui_list, query_to_targets, k, "all"))
if embedding_type_list[index].lower() == "bert":
#opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
# query_to_targets, k, "all", summary_method="MEAN"))
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, "all", summary_method="CLS"))
"""
for index, embedding in enumerate(embedding_list):
if embedding_type_list[index].lower() == "cui":
opt.append(mrm_ndfrt_cui(embedding, umls, cui_list, query_to_targets, k, "origin"))
opt.append(mrm_ndfrt_cui(embedding, umls, cui_list, query_to_targets, k, "all"))
if embedding_type_list[index].lower() == "word":
opt.append(mrm_ndfrt_word(embedding, umls, cui_list, query_to_targets, k, "origin"))
opt.append(mrm_ndfrt_word(embedding, umls, cui_list, query_to_targets, k, "all"))
if embedding_type_list[index].lower() == "bert":
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, "origin", summary_method="CLS"))
beta_path = os.path.join(embedding, "run", "1000000", "rel embedding")
if os.path.exists(beta_path):
if concept_filename.find('treat') >= 0:
method = "may_treat"
else:
method = "may_prevent"
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, method, summary_method="CLS"))
opt.append(mrm_ndfrt_bert(embedding, umls, cui_list,
query_to_targets, k, "all", summary_method="CLS"))
return opt
def load_embedding(filename):
print(filename)
if filename.find('bin') >= 0:
from gensim import models
W = models.KeyedVectors.load_word2vec_format(filename, binary=True)
dim = W.vector_size
return W, dim
if filename.find('pkl') >= 0:
import pickle
with open(filename, 'rb') as f:
W = pickle.load(f)
for key, value in W.items():
W[key] = np.array(list(map(float, value[1:-1].split(","))))
dim = len(list(W.values())[0])
return W, dim
W = {}
with open(filename, 'r') as f:
for i, line in enumerate(f.readlines()):
if i == 0:
continue
toks = line.strip().split()
w = toks[0]
vec = np.array(list(map(float, toks[1:])))
W[w] = vec
dim = len(list(W.values())[0])
return W, dim
def load_bert(model_name_or_path):
print(model_name_or_path)
try:
config = AutoConfig.from_pretrained(model_name_or_path)
model = AutoModel.from_pretrained(
model_name_or_path, config=config).to(device)
except BaseException:
model = torch.load(os.path.join(
model_name_or_path, 'pytorch_model.bin')).to(device)
try:
model.output_hidden_states = False
except BaseException:
pass
try:
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
except BaseException:
tokenizer = AutoTokenizer.from_pretrained(
os.path.join(model_name_or_path, "../"))
return model, tokenizer
def get_intersection(embedding_list, embedding_type_list):
intersection_cui = set()
checker = True
for index, embed in enumerate(embedding_list):
if embedding_type_list[index] == "cui":
w, _ = load_embedding(embed)
if checker:
intersection_cui = set(list(w.keys()))
checker = False
else:
intersection_cui = set(
list(w.keys())).intersection(intersection_cui)
print(f"Intersection count: {len(intersection_cui)}")
return list(intersection_cui)
if __name__ == "__main__":
"""
embedding_list = ["../../embeddings/claims_codes_hs_300.txt",
"../../embeddings/GoogleNews-vectors-negative300.bin",
"../../models/2020_eng"]
embedding_type_list = ["cui", "word", "bert"]
mrm_ndfrt(embedding_list, embedding_type_list, "may_prevent_cui.txt", check_intersection=False)
"""
embedding_list = ["../../embeddings/wikipedia-pubmed-and-PMC-w2v.bin",
"../../embeddings/bio_nlp_vec/PubMed-shuffle-win-2.bin",
"../../embeddings/bio_nlp_vec/PubMed-shuffle-win-30.bin"]
embedding_type_list = ["word", "word", "word"]
embedding_list += ["../../embeddings/DeVine_etal_200.txt",
"/home/yz/pretraining_models/cui2vec.pkl"]
embedding_type_list += ["cui", "cui"]
embedding_list += ["../../models/2020_all",
"/home/yz/pretraining_models/bert-base-cased",
"/home/yz/pretraining_models/biobert_v1.1",
"/home/yz/pretraining_models/BiomedNLP-PubMedBERT-base-uncased-abstract",
"/home/yz/pretraining_models/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
"/home/yz/pretraining_models/kexinghuang_clinical",
"emilyalsentzer/Bio_ClinicalBERT"]
embedding_type_list += ["bert"] * 7
#mrm_ndfrt(embedding_list, embedding_type_list, "may_treat_cui.txt", check_intersection=True)
mrm_ndfrt(embedding_list, embedding_type_list, "may_treat_cui.txt", check_intersection=False)
#mrm_ndfrt(embedding_list[-6:], embedding_type_list[-6:], "may_prevent_cui.txt", check_intersection=True)
#mrm_ndfrt(embedding_list, embedding_type_list, "may_prevent_cui.txt", check_intersection=False)
