<jupyter_start><jupyter_text>load snomed<jupyter_code>import networkx as nx
from tqdm import tqdm
from Snomed import Snomed
SNOMED_PATH = '../data/SnomedCT_201907' # you need to download your own SNOMED distribution
snomed = Snomed(SNOMED_PATH)
snomed.load_snomed()
snomed_sf_id_pairs = []

for snomed_id in tqdm(snomed.graph.nodes):
    
    node_descs = snomed.index_definition[snomed_id]
    for d in node_descs:
        snomed_sf_id_pairs.append((d, snomed_id))

print(len(snomed_sf_id_pairs))
snomed_sf_id_pairs[:10]
snomed_sf_id_pairs_100k = snomed_sf_id_pairs[:100000] # for simplicity
all_names = [p[0] for p in snomed_sf_id_pairs_100k]
all_ids = [p[1] for p in snomed_sf_id_pairs_100k]
all_names[:10]
all_ids[:10]<jupyter_output><empty_output><jupyter_text>load sapbert<jupyter_code>from transformers import AutoTokenizer, AutoModel  
tokenizer = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext")  
model = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext") #.cuda(1)
import numpy as np
import torch<jupyter_output><empty_output><jupyter_text>encode snomed labels<jupyter_code>bs = 128
all_reps = []
for i in tqdm(np.arange(0, len(all_names), bs)):
    toks = tokenizer.batch_encode_plus(all_names[i:i+bs], 
                                       padding="max_length", 
                                       max_length=25, 
                                       truncation=True,
                                       return_tensors="pt")
    #toks_cuda = {}
    #for k,v in toks.items():
    #    toks_cuda[k] = v.cuda(1)
    #output = model(**toks_cuda)
    
    output = model(**toks)
    cls_rep = output[0][:,0,:]
    
    all_reps.append(cls_rep.cpu().detach().numpy())
all_reps_emb = np.concatenate(all_reps, axis=0)
print (all_reps_emb.shape)<jupyter_output>(100000, 768)<jupyter_text>encode query<jupyter_code>query = "cardiopathy"
query_toks = tokenizer.batch_encode_plus([query], 
                                       padding="max_length", 
                                       max_length=25, 
                                       truncation=True,
                                       return_tensors="pt")
query_output = model(**query_toks)
query_cls_rep = query_output[0][:,0,:]
query_cls_rep.shape<jupyter_output><empty_output><jupyter_text>find query's nearest neighbour<jupyter_code># for large-scale search, should switch to faiss
from scipy.spatial.distance import cdist
dist = cdist(query_cls_rep.cpu().detach().numpy(), all_reps_emb)
nn_index = np.argmin(dist)
print ("predicted label:", snomed_sf_id_pairs_100k[nn_index])<jupyter_output>predicted label: ('Cardiac complication', '40172005')
{'fraction_non_alphanumeric': 0.10961968680089486, 'fraction_numerical': 0.023862788963460103, 'mean_word_length': 3.398360655737705, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '433601', 'license': 'MIT', 'n_tokens_mistral': 1026, 'n_tokens_neox': 958, 'n_words': 151, 'path': 'inference/inference_on_snomed.ipynb', 'repo_name': 'cambridgeltl/sapbert', 'size': '2682'}