artcraft/models/tts/NeMo/tutorials/nlp/Entity_Linking_Medical.ipynb

In [25]:

"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.

Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""


## Install dependencies
!pip install wget
!pip install faiss-gpu

## Install NeMo
BRANCH = 'v1.0.2'
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]

In [6]:

import faiss
import torch
import wget
import os
import numpy as np
import pandas as pd

from omegaconf import OmegaConf
from pytorch_lightning import Trainer
from IPython.display import display
from tqdm import tqdm

from nemo.collections import nlp as nemo_nlp
from nemo.utils.exp_manager import exp_manager

Entity Linking

Task Description

Entity linking is the process of connecting concepts mentioned in natural language to their canonical forms stored in a knowledge base. For example, say a knowledge base contained the entity 'ID3452 influenza' and we wanted to process some natural language containing the sentence "The patient has flu like symptoms". An entity linking model would match the word 'flu' to the knowledge base entity 'ID3452 influenza', allowing for disambiguation and normalization of concepts referenced in text. Entity linking applications range from helping automate data ingestion to assisting in real time dialogue concept normalization. We will be focusing on entity linking in the medical domain for this demo, but the entity linking model, dataset, and training code within NVIDIA NeMo can be applied to other domains like finance and retail.

Within NeMo and this tutorial we use the entity linking approach described in Liu et. al's NAACL 2021 "Self-alignment Pre-training for Biomedical Entity Representations". The main idea behind this approach is to reshape an initial concept embedding space such that synonyms of the same concept are pulled closer together and unrelated concepts are pushed further apart. The concept embeddings from this reshaped space can then be used to build a knowledge base embedding index. This index stores concept IDs mapped to their respective concept embeddings in a format conducive to efficient nearest neighbor search. We can link query concepts to their canonical forms in the knowledge base by performing a nearest neighbor search- matching concept query embeddings to the most similar concepts embeddings in the knowledge base index.

In this tutorial we will be using the faiss library to build our concept index.

Self Alignment Pretraining

Self-Alignment pretraining is a second stage pretraining of an existing encoder (called second stage because the encoder model can be further finetuned after this more general pretraining step). The dataset used during training consists of pairs of concept synonyms that map to the same ID. At each training iteration, we only select hard examples present in the mini batch to calculate the loss and update the model weights. In this context, a hard example is an example where a concept is closer to an unrelated concept in the mini batch than it is to the synonym concept it is paired with by some margin. I encourage you to take a look at section 2 of the paper for a more formal and in depth description of how hard examples are selected.

We then use a metric learning loss calculated from the hard examples selected. This loss helps reshape the embedding space. The concept representation space is rearranged to be more suitable for entity matching via embedding cosine similarity.

Now that we have idea of what's going on, let's get started!

Dataset Preprocessing

In [26]:

# Download data into project directory
PROJECT_DIR = "." #Change if you don't want the current directory to be the project dir
DATA_DIR = os.path.join(PROJECT_DIR, "tiny_example_data")

if not os.path.isdir(os.path.join(DATA_DIR)):
    wget.download('https://dldata-public.s3.us-east-2.amazonaws.com/tiny_example_data.zip',
                  os.path.join(PROJECT_DIR, "tiny_example_data.zip"))

    !unzip {PROJECT_DIR}/tiny_example_data.zip -d {PROJECT_DIR}

In this tutorial we will be using a tiny toy dataset to demonstrate how to use NeMo's entity linking model functionality. The dataset includes synonyms for 12 medical concepts. Entity phrases with the same ID are synonyms for the same concept. For example, "chronic kidney failure", "gradual loss of kidney function", and "CKD" are all synonyms of concept ID 5. Here's the dataset before preprocessing:

In [9]:

raw_data = pd.read_csv(os.path.join(DATA_DIR, "tiny_example_dev_data.csv"), names=["ID", "CONCEPT"], index_col=False)
print(raw_data)

    ID                                            CONCEPT
0    1                                          Head ache
1    1                                           Headache
2    1                                           Migraine
3    1                                   Pain in the head
4    1                                          cephalgia
5    1                                        cephalalgia
6    2                                       heart attack
7    2                              Myocardial infraction
8    2                           necrosis of heart muscle
9    2                                                 MI
10   3                                                CAD
11   3                            Coronary artery disease
12   3                      atherosclerotic heart disease
13   3                                      heart disease
14   3                damage of major heart blood vessels
15   4                                myocardial ischemia
16   4                                   cardiac ischemia
17   4                      reduced ability to pump blood
18   5                    gradual loss of kidney function
19   5                        kidneys cannot filter blood
20   5                             chronic kidney disease
21   5                             chronic kidney failure
22   5                                                CKD
23   6                              alchohol intoxication
24   6                        acute alchohol intoxication
25   6                                 alchohol poisoning
26   6                                 severe drunkenness
27   6                        over consumption of alcohol
28   7                                  diabetes mellitus
29   7                                           diabetes
30   7                       inability to process glucose
31   7                            unable to take up sugar
32   7                                    Type 2 diabetes
33   8                                   Hyperinsulinemia
34   8                                   High blood sugar
35   8                  abnormally high levels of insulin
36   9                   Dipeptidyl peptidase-4 inhibitor
37   9                                    dpp-4 inhibitor
38   9                                         alogliptin
39   9                                             Nesina
40   9                                            Vipidia
41  10                                       hypoglycemia
42  10                                    low blood sugar
43  11                                     anticoagulants
44  11                                     blood thinners
45  11                                           Apixaban
46  11                                            Eliquis
47  12                                          Ibuprofen
48  12                                            Aspirin
49  12  over the counter nonsteroidal anti-inflammator...
50  12                                              NSAID

We've already paired off the concepts for this dataset with the format ID concept_synonym1 concept_synonym2. Here are the first ten rows:

In [10]:

training_data = pd.read_table(os.path.join(DATA_DIR, "tiny_example_train_pairs.tsv"), names=["ID", "CONCEPT_SYN1", "CONCEPT_SYN2"], delimiter='\t')
print(training_data.head(10))

   ID      CONCEPT_SYN1      CONCEPT_SYN2
0   1  Pain in the head         cephalgia
1   1  Pain in the head       cephalalgia
2   1          Migraine         cephalgia
3   1         Head ache  Pain in the head
4   1         Head ache          Migraine
5   1         Head ache       cephalalgia
6   1          Headache          Migraine
7   1          Migraine       cephalalgia
8   1         cephalgia       cephalalgia
9   1          Headache  Pain in the head

Use the Unified Medical Language System (UMLS) dataset for full medical domain entity linking training. The data contains over 9 million entities and is a table of medical concepts with their corresponding concept IDs (CUI). After requesting a free license and making a UMLS Terminology Services (UTS) account, the entire UMLS dataset can be downloaded from the NIH's website. If you've cloned the NeMo repo you can run the data processing script located in examples/nlp/entity_linking/data/umls_dataset_processing.py on the full dataset. This script will take in the initial table of UMLS concepts and produce a .tsv file with each row formatted as CUI\tconcept_synonym1\tconcept_synonym2. Once the UMLS dataset .RRF file is downloaded, the script can be run from the examples/nlp/entity_linking directory like so:

python data/umls_dataset_processing.py

Model Training

Second stage pretrain a BERT Base encoder on the self-alignment pretraining task (SAP) for improved entity linking. Using a GPU, the model should take 5 minutes or less to train on this example dataset and training progress will be output below the cell.

In [11]:

# Download config
wget.download("https://raw.githubusercontent.com/vadam5/NeMo/main/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml",
              os.path.join(PROJECT_DIR, "tiny_example_entity_linking_config.yaml"))

# Load in config file
cfg = OmegaConf.load(os.path.join(PROJECT_DIR, "tiny_example_entity_linking_config.yaml"))

# Set config file variables
cfg.project_dir = PROJECT_DIR
cfg.model.nemo_path = os.path.join(PROJECT_DIR, "tiny_example_sap_bert_model.nemo")
cfg.model.train_ds.data_file = os.path.join(DATA_DIR, "tiny_example_train_pairs.tsv")
cfg.model.validation_ds.data_file = os.path.join(DATA_DIR, "tiny_example_validation_pairs.tsv")

In [27]:

# Initialize the trainer and model
trainer = Trainer(**cfg.trainer)
exp_manager(trainer, cfg.get("exp_manager", None))
model = nemo_nlp.models.EntityLinkingModel(cfg=cfg.model, trainer=trainer)

In [28]:

# Train and save the model
trainer.fit(model)
model.save_to(cfg.model.nemo_path)

You can run the script at examples/nlp/entity_linking/self_alignment_pretraining.py to train a model on a larger dataset. Run

python self_alignment_pretraining.py project_dir=.

from the examples/nlp/entity_linking directory.

Model Evaluation

Let's evaluate our freshly trained model and compare its performance with a BERT Base encoder that hasn't undergone self-alignment pretraining. We first need to restore our trained model and load our BERT Base Baseline model.

In [29]:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Restore second stage pretrained model
sap_model_cfg = cfg
sap_model_cfg.index.index_save_name = os.path.join(PROJECT_DIR, "tiny_example_entity_linking_index")
sap_model_cfg.index.index_ds.data_file = os.path.join(DATA_DIR, "tiny_example_index_data.tsv")
sap_model = nemo_nlp.models.EntityLinkingModel.restore_from(sap_model_cfg.model.nemo_path).to(device)

# Load original model
base_model_cfg = OmegaConf.load(os.path.join(PROJECT_DIR, "tiny_example_entity_linking_config.yaml"))

# Set train/val datasets to None to avoid loading datasets associated with training
base_model_cfg.model.train_ds = None
base_model_cfg.model.validation_ds = None
base_model_cfg.index.index_save_name = os.path.join(PROJECT_DIR, "base_model_index")
base_model_cfg.index.index_ds.data_file = os.path.join(DATA_DIR, "tiny_example_index_data.tsv")
base_model = nemo_nlp.models.EntityLinkingModel(base_model_cfg.model).to(device)

We are going evaluate our model on a nearest neighbor task using top 1 and top 5 accuracies as our metric. We will be using a tiny example test knowledge base and test queries. For this evaluation we are going to be comparing every test query with every concept vector in our test set knowledge base. We will rank each item in the knowledge base by its cosine similarity with the test query. We'll then compare the IDs of the predicted most similar test knowledge base concepts with our ground truth query IDs to calculate top 1 and top 5 accuracies. For this metric higher is better.

In [15]:

# Helper function to get data embeddings
def get_embeddings(model, dataloader):
    embeddings, cids = [], []

    with torch.no_grad():
        for batch in tqdm(dataloader):
            input_ids, token_type_ids, attention_mask, batch_cids = batch
            batch_embeddings = model.forward(input_ids=input_ids.to(device), 
                                             token_type_ids=token_type_ids.to(device), 
                                             attention_mask=attention_mask.to(device))

            # Accumulate index embeddings and their corresponding IDs
            embeddings.extend(batch_embeddings.cpu().detach().numpy())
            cids.extend(batch_cids)
            
    return embeddings, cids

In [16]:

def evaluate(model, test_kb, test_queries, ks):
    # Initialize knowledge base and query data loaders
    test_kb_dataloader = model.setup_dataloader(test_kb, is_index_data=True)
    test_query_dataloader = model.setup_dataloader(test_queries, is_index_data=True)
    
    # Get knowledge base and query embeddings
    test_kb_embs, test_kb_cids = get_embeddings(model, test_kb_dataloader)
    test_query_embs, test_query_cids = get_embeddings(model, test_query_dataloader)

    # Calculate the cosine distance between each query and knowledge base concept
    score_matrix = np.matmul(np.array(test_query_embs), np.array(test_kb_embs).T)
    accs = {k : 0 for k in ks}
    
    # Compare the knowledge base IDs of the knowledge base entities with 
    # the smallest cosine distance from the query 
    for query_idx in tqdm(range(len(test_query_cids))):
        query_emb = test_query_embs[query_idx]
        query_cid = test_query_cids[query_idx]
        query_scores = score_matrix[query_idx]

        for k in ks:
            topk_idxs = np.argpartition(query_scores, -k)[-k:]
            topk_cids = [test_kb_cids[idx] for idx in topk_idxs]
            
            # If the correct query ID is amoung the top k closest kb IDs
            # the model correctly linked the entity
            match = int(query_cid in topk_cids)
            accs[k] += match

    for k in ks:
        accs[k] /= len(test_query_cids)
                
    return accs

In [30]:

# Create configs for our test data
test_kb = OmegaConf.create({
    "data_file": os.path.join(DATA_DIR, "tiny_example_test_kb.tsv"),
    "max_seq_length": 128,
    "batch_size": 10,
    "shuffle": False,
})

test_queries = OmegaConf.create({
    "data_file": os.path.join(DATA_DIR, "tiny_example_test_queries.tsv"),
    "max_seq_length": 128,
    "batch_size": 10,
    "shuffle": False,
})

ks = [1, 5]

# Evaluate both models on our test data
base_accs = evaluate(base_model, test_kb, test_queries, ks)
base_accs["Model"] = "BERT Base Baseline"

sap_accs = evaluate(sap_model, test_kb, test_queries, ks)
sap_accs["Model"] = "BERT + SAP"

print("Top 1 and Top 5 Accuracy Comparison:")
results_df = pd.DataFrame([base_accs, sap_accs], columns=["Model", 1, 5])
results_df = results_df.style.set_properties(**{'text-align': 'left', }).set_table_styles([dict(selector='th', props=[('text-align', 'left')])])
display(results_df)

The purpose of this section was to show an example of evaluating your entity linking model. This evaluation set contains very little data, and no serious conclusions should be drawn about model performance. Top 1 accuracy should be between 0.7 and 1.0 for both models and top 5 accuracy should be between 0.8 and 1.0. When evaluating a model trained on a larger dataset, you can use a nearest neighbors index to speed up the evaluation time.

Building an Index

To qualitatively observe the improvement we gain from the second stage pretraining, let's build two indices. One will be built with BERT base embeddings before self-alignment pretraining and one will be built with the model we just trained. Our knowledge base in this tutorial will be in the same domain and have some overlapping concepts as the training set. This data file is formatted as ID\tconcept.

The EntityLinkingDataset class can load the data used for training the entity linking encoder as well as for building the index if the is_index_data flag is set to true.

In [18]:

def build_index(cfg, model):
    # Setup index dataset loader
    index_dataloader = model.setup_dataloader(cfg.index.index_ds, is_index_data=True)
    
    # Get index dataset embeddings
    embeddings, _ = get_embeddings(model, index_dataloader)
    
    # Train IVFFlat index using faiss
    embeddings = np.array(embeddings)
    quantizer = faiss.IndexFlatL2(cfg.index.dims)
    index = faiss.IndexIVFFlat(quantizer, cfg.index.dims, cfg.index.nlist)
    index = faiss.index_cpu_to_all_gpus(index)
    index.train(embeddings)
    
    # Add concept embeddings to index
    for i in tqdm(range(0, embeddings.shape[0], cfg.index.index_batch_size)):
            index.add(embeddings[i:i+cfg.index.index_batch_size])

    # Save index
    faiss.write_index(faiss.index_gpu_to_cpu(index), cfg.index.index_save_name)

In [31]:

build_index(sap_model_cfg, sap_model.to(device))
build_index(base_model_cfg, base_model.to(device))

Entity Linking via Nearest Neighbor Search

Now it's time to query our indices! We are going to query both our index built with embeddings from BERT Base, and our index with embeddings built from the SAP BERT model we trained. Our sample query phrases will be "high blood sugar" and "head pain".

To query our indices, we first need to get the embedding of each query from the corresponding encoder model. We can then pass these query embeddings into the faiss index which will perform a nearest neighbor search, using cosine distance to compare the query embedding with embeddings present in the index. Once we get a list of knowledge base index concept IDs most closely matching our query, all that is left to do is map the IDs to a representative string describing the concept.

In [20]:

def query_index(cfg, model, index, queries, id2string):
    # Get query embeddings from our entity linking encoder model
    query_embs = get_query_embedding(queries, model).cpu().detach().numpy()
    
    # Use query embedding to find closest concept embedding in knowledge base
    distances, neighbors = index.search(query_embs, cfg.index.top_n)
    
    # Get the canonical strings corresponding to the IDs of the query's nearest neighbors in the kb 
    neighbor_concepts = [[id2string[concept_id] for concept_id in query_neighbor] \
                                                for query_neighbor in neighbors]
    
    # Display most similar concepts in the knowledge base. 
    for query_idx in range(len(queries)):
        print(f"\nThe most similar concepts to {queries[query_idx]} are:")
        for cid, concept, dist in zip(neighbors[query_idx], neighbor_concepts[query_idx], distances[query_idx]):
            print(cid, concept, 1 - dist)

    
def get_query_embedding(queries, model):
    # Tokenize our queries
    model_input =  model.tokenizer(queries,
                                   add_special_tokens = True,
                                   padding = True,
                                   truncation = True,
                                   max_length = 512,
                                   return_token_type_ids = True,
                                   return_attention_mask = True)
    
    # Pass tokenized input into model
    query_emb =  model.forward(input_ids=torch.LongTensor(model_input["input_ids"]).to(device),
                               token_type_ids=torch.LongTensor(model_input["token_type_ids"]).to(device),
                               attention_mask=torch.LongTensor(model_input["attention_mask"]).to(device))
    
    return query_emb

In [21]:

# Load indices
sap_index = faiss.read_index(sap_model_cfg.index.index_save_name)
base_index = faiss.read_index(base_model_cfg.index.index_save_name)

In [22]:

# Map concept IDs to one canonical string
index_data = open(sap_model_cfg.index.index_ds.data_file, "r", encoding='utf-8-sig')
id2string = {}

for line in index_data:
    cid, concept = line.split("\t")
    id2string[int(cid) - 1] = concept.strip()

In [23]:

id2string

Out [23]:

{0: 'Headache',
 1: 'Myocardial infraction',
 2: 'Coronary artery disease',
 3: 'myocardial ischemia',
 4: 'chronic kidney disease',
 5: 'alchohol intoxication',
 6: 'diabetes',
 7: 'Hyperinsulinemia',
 8: 'Nesina',
 9: 'hypoglycemia',
 10: 'anticoagulants',
 11: 'Ibuprofen'}

In [24]:

# Some sample queries
queries = ["high blood sugar", "head pain"]

# Query BERT Base
print("BERT Base output before Self Alignment Pretraining:")
query_index(base_model_cfg, base_model, base_index, queries, id2string)
print("\n" + "-" * 50 + "\n")

# Query SAP BERT
print("SAP BERT output after Self Alignment Pretraining:")
query_index(sap_model_cfg, sap_model, sap_index, queries, id2string)
print("\n" + "-" * 50 + "\n")

BERT Base output before Self Alignment Pretraining:

The most similar concepts to high blood sugar are:
6 diabetes 0.9095035567879677
0 Headache 0.9046077281236649
8 Nesina 0.8512845635414124

The most similar concepts to head pain are:
1 Myocardial infraction 0.7848672568798065
4 chronic kidney disease 0.7667323648929596
3 myocardial ischemia 0.761662557721138

--------------------------------------------------

SAP BERT output after Self Alignment Pretraining:

The most similar concepts to high blood sugar are:
7 Hyperinsulinemia 0.28035879135131836
3 myocardial ischemia 0.054734230041503906
6 diabetes 0.053402602672576904

The most similar concepts to head pain are:
0 Headache 0.6553303301334381
6 diabetes 0.19157332181930542
2 Coronary artery disease 0.12962108850479126

--------------------------------------------------

Even after only training on this tiny amount of data, the qualitative performance boost from self-alignment pretraining is visible. The baseline model links "high blood sugar" to the entity "6 diabetes" while our SAP BERT model accurately links "high blood sugar" to "Hyperinsulinemia". Similarly, "head pain" and "Myocardial infraction" are not the same concept, but "head pain" and "Headache" are.

For larger knowledge bases keeping the default embedding size might be too large and cause out of memory issues. You can apply PCA or some other dimensionality reduction method to your data to reduce its memory footprint. Code for creating a text file of all the UMLS entities in the correct format needed to build an index and creating a dictionary mapping concept ids to canonical concept strings can be found here examples/nlp/entity_linking/data/umls_dataset_processing.py.

The code for extracting knowledge base concept embeddings, training and applying a PCA transformation to the embeddings, building a faiss index and querying the index from the command line is located at examples/nlp/entity_linking/build_index.py and examples/nlp/entity_linking/query_index.py.

If you've cloned the NeMo repo, both of these steps can be run as follows on the command line from the examples/nlp/entity_linking/ directory.

python data/umls_dataset_processing.py --index
python build_index.py --restore
python query_index.py --restore

By default the project directory will be "." but can be changed by adding the flag --project_dir=<PATH> after each of the above commands. Intermediate steps of the index building process are saved. In the occurrence of an error, previously completed steps do not need to be rerun.

Command Recap

Here is a recap of the commands and steps to repeat this process on the full UMLS dataset.

1. Download the UMLS dataset file MRCONSO.RRF from the NIH website and place it in the examples/nlp/entity_linking/data directory.

2. Run the following commands from the examples/nlp/entity_linking directory

python data/umls_dataset_processing.py
python self_alignment_pretraining.py project_dir=. 
python data/umls_dataset_processing.py --index
python build_index.py --restore
python query_index.py --restore

The model will take ~24hrs to train on two GPUs and ~48hrs to train on one GPU. By default the project directory will be "." but can be changed by adding the flag --project_dir=<PATH> after each of the above commands and changing project_dir=<PATH> in the self_alignment_pretraining.py command. If you change the project directory, you should also move the MRCONOSO.RRF file to a data sub directory within the one you've specified.

As mentioned in the introduction, entity linking within NVIDIA NeMo is not limited to the medical domain. The same data processing and training steps can be applied to a variety of domains and use cases. You can edit the datasets used as well as training and loss function hyperparameters within your config file to better suit your domain.