Introduce code from Rust monorepo. This is a mess

tts/NeMo/examples/nlp/entity_linking/build_index.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pickle as pkl
+import random
+from argparse import ArgumentParser
+
+import h5py
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from sklearn.decomposition import PCA
+from tqdm import tqdm
+
+from nemo.collections.nlp.models import EntityLinkingModel
+from nemo.utils import logging
+
+try:
+    import faiss
+except ModuleNotFoundError:
+    logging.warning("Faiss is required for building the index. Please install faiss-gpu")
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def build_index(cfg: DictConfig, model: object):
+    """
+    Builds faiss index from index dataset specified in the config.
+        
+    Args:
+        cfg (DictConfig): Config file specifying index parameters
+        model (object): Encoder model
+    """
+
+    # Get index dataset embeddings
+    # PCA model exists and index embeddings have already been PCAed, no need to re-extract/PCA them
+    if cfg.apply_pca and os.path.isfile(cfg.pca.pca_save_name) and os.path.isfile(cfg.pca_embeddings_save_name):
+        logging.info("Loading reduced dimensionality embeddings")
+        embeddings = h5py.File(cfg.pca_embeddings_save_name, "r")
+        embeddings = embeddings[cfg.index_ds.name][:]
+
+    elif os.path.isfile(cfg.embedding_save_name):
+        logging.info("Loading previously extracted index dataset embeddings")
+        embeddings = h5py.File(cfg.embedding_save_name, "r")
+        embeddings = embeddings[cfg.index_ds.name][:]
+
+    else:
+        logging.info("Encoding index dataset, this may take a while")
+        index_dataloader = model.setup_dataloader(cfg.index_ds, is_index_data=True)
+        embeddings, concept_ids = get_index_embeddings(cfg, index_dataloader, model)
+
+    # Create pca model to reduce dimensionality of index dataset and decrease memory footprint
+    if cfg.apply_pca:
+
+        # Need to train PCA model and apply PCA transformation with newly trained model
+        if not os.path.isfile(cfg.pca.pca_save_name):
+            logging.info("Fitting PCA model for embedding dimensionality reduction")
+            pca_train_set = random.sample(list(embeddings), k=int(len(embeddings) * cfg.pca.sample_fraction))
+            pca = PCA(n_components=cfg.pca.output_dim)
+            pca.fit(pca_train_set)
+            pkl.dump(pca, open(cfg.pca.pca_save_name, "wb"))
+            embeddings = reduce_embedding_dim(pca, embeddings, cfg)
+
+        # PCA model already trained, just need to reduce dimensionality of all embeddings
+        elif not os.path.isfile(cfg.pca_embeddings_save_name):
+            pca = pkl.load(open(cfg.pca.pca_save_name, "rb"))
+            embeddings = reduce_embedding_dim(pca, embeddings, cfg)
+
+    # Build faiss index from embeddings
+    logging.info(f"Training index with embedding dim size {cfg.dims} using {faiss.get_num_gpus()} gpus")
+    quantizer = faiss.IndexFlatL2(cfg.dims)
+    index = faiss.IndexIVFFlat(quantizer, cfg.dims, cfg.nlist)
+    index = faiss.index_cpu_to_all_gpus(index)
+    index.train(embeddings)
+
+    logging.info("Adding dataset embeddings to index")
+    for i in tqdm(range(0, embeddings.shape[0], cfg.index_batch_size)):
+        index.add(embeddings[i : i + cfg.index_batch_size])
+
+    logging.info("Saving index")
+    faiss.write_index(faiss.index_gpu_to_cpu(index), cfg.index_save_name)
+    logging.info("Index built and saved")
+
+
+def reduce_embedding_dim(pca, embeddings, cfg):
+    """Apply PCA transformation to index dataset embeddings"""
+
+    logging.info("Applying PCA transformation to entire index dataset")
+    embeddings = np.array(pca.transform(embeddings), dtype=np.float32)
+    emb_file = h5py.File(cfg.pca_embeddings_save_name, "w")
+    emb_file.create_dataset(cfg.index_ds.name, data=embeddings)
+    emb_file.close()
+
+    return embeddings
+
+
+def get_index_embeddings(cfg: DictConfig, dataloader: object, model: object):
+    """Use entity linking encoder to get embeddings for full index dataset"""
+    embeddings = []
+    concept_ids = []
+
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            input_ids, token_type_ids, input_mask, batch_concept_ids = batch
+            input_ids = input_ids.to(device)
+            token_type_ids = token_type_ids.to(device)
+            input_mask = input_mask.to(device)
+            batch_embeddings = model.forward(
+                input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=input_mask
+            )
+
+            embeddings.extend(batch_embeddings.detach().cpu().numpy())
+            concept_ids.extend(batch_concept_ids.numpy())
+
+    emb_file = h5py.File(cfg.embedding_save_name, "w")
+    emb_file.create_dataset(cfg.index_ds.name, data=embeddings)
+    emb_file.close()
+
+    pkl.dump(concept_ids, open(cfg.concept_id_save_name, "wb"))
+
+    return embeddings, concept_ids
+
+
+def load_model(cfg: DictConfig, restore: bool):
+    """
+    Loads encoder model.
+
+    Args:
+        cfg: Config file specifying model parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining. 
+    """
+
+    if restore:
+        model = EntityLinkingModel.restore_from(cfg.nemo_path)
+    else:
+        cfg.train_ds = None
+        cfg.validation_ds = None
+        cfg.test_ds = None
+        model = EntityLinkingModel(cfg)
+
+    model = model.to(device)
+
+    return model
+
+
+def main(cfg: DictConfig, restore: bool):
+    """
+    Builds new index if one hasn't been built yet.
+
+    Args:
+        cfg: Config file specifying index parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining.
+    """
+
+    logging.info("Loading entity linking encoder model")
+    model = load_model(cfg.model, restore)
+
+    if not os.path.isfile(cfg.index.index_save_name) or (
+        cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name)
+    ):
+        logging.info("Building index")
+        build_index(cfg.index, model)
+    else:
+        logging.info("Index and pca model (if required) already exists. Skipping build index step.")
+
+    if not os.path.isfile(cfg.index.idx_to_id):
+        logging.info("Mapping entity index postions to ids")
+        map_idx_to_ids(cfg.index)
+    else:
+        logging.info("Map from concept index to id already exists. Skipping mapping step.")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path"
+    )
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml")
+    args = parser.parse_args()
+
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    main(cfg, args.restore)

tts/NeMo/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml

@@ -0,0 +1,90 @@
+project_dir: null
+name: SelfAlignmentPretrainingForMedicalEntityLinking
+trainer:
+  gpus: 1
+  num_nodes: 1
+  max_epochs: 2
+  max_steps: null
+  accumulate_grad_batches: 1
+  precision: 16
+  amp_level: O1
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 2
+  checkpoint_callback: false
+  logger: false
+model:
+  nemo_path: ???
+  max_seq_length: 128
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    config_file: null
+    config: null
+    lm_checkpoint: null
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name}
+    vocab_file: null
+    tokenizer_model: null
+    do_lower_case: true
+  loss_params: null
+  train_ds:
+    data_file: ???
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 8
+    shuffle: true
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  validation_ds:
+    data_file: ??? 
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 8
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  optim:
+    name: adam
+    lr: 3.0e-05
+    weight_decay: 0.0
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+index:
+  dims: 768
+  nlist: 2
+  top_n: 3
+  query_num_factor: 20
+  index_save_name: ???
+  index_batch_size: 10
+  index_ds:
+    name: tiny_example
+    data_file: ???
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 100
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  idx_to_id: ${project_dir}/idx_to_id.pkl
+  id_to_string: ${project_dir}/id_to_string.pkl
+  concept_id_save_name: ${project_dir}/tiny_example_concept_ids.pkl
+  embedding_save_name: ${project_dir}/tiny_example_concept_embeddings.hdf5
+  pca_embeddings_save_name: null
+  apply_pca: false
+  pca: null
+exp_manager:
+  exp_dir: .
+  name: ${project_dir}/SelfAlignmentPretrainingTinyExample
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/entity_linking/conf/umls_medical_entity_linking_config.yaml

@@ -0,0 +1,95 @@
+project_dir: ???
+name: SelfAlignmentPretrainingForMedicalEntityLinking
+trainer:
+  gpus: 1
+  num_nodes: 1
+  max_epochs: 2
+  max_steps: null
+  accumulate_grad_batches: 1
+  precision: 16
+  amp_level: O1
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1000
+  checkpoint_callback: false
+  logger: false
+model:
+  nemo_path: ${project_dir}/sap_bert_umls.nemo
+  raw_data: ${project_dir}/data/MRCONSO.RRF
+  max_seq_length: 128
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    config_file: null
+    config: null
+    lm_checkpoint: null
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name}
+    vocab_file: null
+    tokenizer_model: null
+    do_lower_case: true
+  train_ds:
+    data_file: ${project_dir}/data/umls_train_pairs.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: true
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  validation_ds:
+    data_file: ${project_dir}/data/umls_validation_pairs.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  optim:
+    name: adam
+    lr: 3.0e-05
+    weight_decay: 0.0
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+index:
+  dims: 256
+  nlist: 300
+  top_n: 5
+  query_num_factor: 20
+  index_save_name: ${project_dir}/medical_entity_linking_index
+  index_batch_size: 1000
+  raw_data: ${model.raw_data}
+  index_ds:
+    name: umls
+    data_file: ${project_dir}/data/umls_index_concepts.tsv
+    max_seq_length: ${model.max_seq_length}
+    batch_size: 128
+    shuffle: false
+    num_workers: 2
+    pin_memory: false
+    drop_last: false
+  idx_to_id: ${project_dir}/data/idx_to_id.pkl
+  id_to_string: ${project_dir}/data/id_to_string.pkl
+  concept_id_save_name: ${project_dir}/data/concept_ids.pkl
+  embedding_save_name: ${project_dir}/data/medical_concept_embeddings.hdf5
+  pca_embeddings_save_name: ${project_dir}/data/medical_concept_reduced_${index.dims}dim_embeddings.hdf5
+  apply_pca: true
+  pca:
+    input_dim: 756
+    output_dim: ${index.dims}
+    sample_fraction: 0.5
+    pca_save_name: ${project_dir}/${index.pca.input_dim}_to_${index.pca.output_dim}_pca_model.pkl
+exp_manager:
+  exp_dir: ${project_dir}/medical_entity_linking_experiments
+  name: sap_bert_umls
+  create_tensorboard_logger: true
+  create_checkpoint_callback: true
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/entity_linking/data/umls_dataset_processing.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import pickle as pkl
+import random
+from argparse import ArgumentParser
+
+import pandas as pd
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+# Info on these headers can be found here on the UMLS website https://www.ncbi.nlm.nih.gov/books/NBK9685/
+# section 3.3.4 Table 1
+HEADERS = [
+    'CUI',
+    'LAT',
+    'TS',
+    'LUI',
+    'STT',
+    'SUI',
+    'ISPREF',
+    'AUI',
+    'SAUI',
+    'SCUI',
+    'SDUI',
+    'SAB',
+    'TTY',
+    'CODE',
+    'STR',
+    'SRL',
+    'SUPPRESS',
+    'CVF',
+]
+
+
+def process_umls_training_dataset(data_path, train_save_name, val_save_name, max_pairs, train_split, headers):
+    """
+    Generates and saves UMLS self alignment pretraining train and validation data. Takes the raw .RRF UMLS 
+    data file and creates different pair combinations for entities with the same CUI. Each row in the output
+    will be formatted as 'CUI EntitySynonym1 EntitySynonym2' with each item in a row separated by tabs.
+    Saves two .tsv output files, one for the train split and one for the validation split.
+    Only data marked as English is added to the train and val splits. 
+
+    Arguments:
+        data_path (str): path to MRCONSO.RRF UMLS data file
+        train_save_name (str): path to where training data will be saved
+        val_save_name (str): path to where validation data will be saved
+        max_pairs (int): max number of pairs for any one CUI added to the train 
+                   or validation splits
+        train_split (float): precentage of raw data to be added to train set split
+        headers (list): column lables within MRCONSO.RRF
+    """
+
+    print("Loading training data file...")
+    df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|')
+    train_file = open(train_save_name, 'w')
+    val_file = open(val_save_name, 'w')
+
+    cui = df["CUI"].iloc[0]
+    names = []
+    random.seed(2021)
+
+    for idx in tqdm(range(len(df))):
+        # Address incorrectly formatted data
+        if type(df["STR"].iloc[idx]) != str or "|" in df["STR"].iloc[idx]:
+            continue
+
+        # Collect all english concept strings matching the current CUI
+        if df["CUI"].iloc[idx] == cui and df["LAT"].iloc[idx] == "ENG":
+            concept_string = df["STR"].iloc[idx]
+            names.append(concept_string)
+
+        else:
+            # Pair off concept synonyms to make training and val sets
+            pairs = list(itertools.combinations(names, 2))
+
+            if len(pairs) == 0:
+                # Not enough concepts gathered to make a pair
+                cui = df["CUI"].iloc[idx]
+                names = [df["STR"].iloc[idx]]
+                continue
+
+            # Removing leading C to convert label string to int
+            cui = int(cui[1:])
+            random.shuffle(pairs)
+
+            # Keep up to max pairs number pairs for any one concept
+            for pair in pairs[:max_pairs]:
+
+                # Want concepts in train and val splits to be randomly selected and mutually exclusive
+                add_to_train = random.random()
+
+                if add_to_train <= train_split:
+                    train_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n')
+                else:
+                    val_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n')
+
+            # Switch to next concept
+            cui = df["CUI"].iloc[idx]
+            names = [df["STR"].iloc[idx]]
+
+    train_file.close()
+    val_file.close()
+    print("Finished making training and validation data")
+
+
+def process_umls_index_dataset(data_path, data_savename, id2string_savename, headers):
+    """
+    Generates data file needed to build a UMLS index and a hash table mapping each
+    CUI to one canonical concept string. Takes the raw .RRF data file and saves 
+    a .tsv indec concept file as well as the a .pkl file of cui to concept string 
+    mappings. Only data marked as English is added to the index data file. 
+
+    Arguments:
+        data_path (str): path to MRCONSO.RRF UMLS data file
+        data_savename (str): path to where .tsv index data will be saved
+        id2string_savename (str): path to where .pkl cui to string mapping will
+                                  be saved
+        headers (list): column lables within MRCONSO.RRF
+    """
+
+    print("Loading index data file...")
+    df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|')
+    id2string = {}
+
+    with open(data_savename, "w") as outfile:
+        for idx, row in tqdm(df.iterrows(), total=df.shape[0]):
+            # Address incorrectly formatted data
+            if type(row["STR"]) != str or "|" in row["STR"]:
+                continue
+
+            cui = row["CUI"]
+            sent = row["STR"]
+
+            # Removing leading C to convert label string to int
+            cui = int(cui[1:])
+
+            # Only keeping english concepts
+            if row["LAT"] == "ENG":
+                outfile.write(f'{cui}\t{sent}\n')
+
+                # Matching each cui to one canonical string represention
+                if cui not in id2string and ":" not in sent:
+                    id2string[cui] = sent
+
+    outfile.close()
+    pkl.dump(id2string, open(id2string_savename, "wb"))
+    print("Finished saving index data and id to concept mapping")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument("--index", action="store_true", help="Whether to process data for building an index")
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="conf/umls_medical_entity_linking_config.yaml")
+    parser.add_argument(
+        "--max_pairs", required=False, type=int, default=50, help="Max number of train pairs for a single concepts"
+    )
+    parser.add_argument(
+        "--train_split", required=False, type=float, default=0.99, help="Precentage of data to add to train set"
+    )
+
+    args = parser.parse_args()
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    if args.index:
+        process_umls_index_dataset(cfg.index.raw_data, cfg.index.index_ds.data_file, cfg.index.id_to_string, HEADERS)
+    else:
+        process_umls_training_dataset(
+            cfg.model.raw_data,
+            cfg.model.train_ds.data_file,
+            cfg.model.validation_ds.data_file,
+            args.max_pairs,
+            args.train_split,
+            HEADERS,
+        )

tts/NeMo/examples/nlp/entity_linking/query_index.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pickle as pkl
+from argparse import ArgumentParser
+from collections import OrderedDict
+from typing import Dict
+
+import numpy as np
+import torch
+from build_index import load_model
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.utils import logging
+
+try:
+    import faiss
+except ModuleNotFoundError:
+    logging.warning("Faiss is required for building the index. Please install faiss-gpu")
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def get_query_embedding(query, model):
+    """Use entity linking encoder to get embedding for index query"""
+    model_input = model.tokenizer(
+        query,
+        add_special_tokens=True,
+        padding=True,
+        truncation=True,
+        max_length=512,
+        return_token_type_ids=True,
+        return_attention_mask=True,
+    )
+
+    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
+
+
+def query_index(
+    query: str, cfg: DictConfig, model: object, index: object, pca: object, idx2id: dict, id2string: dict,
+) -> Dict:
+
+    """
+    Query the nearest neighbor index of entities to find the 
+    concepts in the index dataset that are most similar to the 
+    query.
+
+    Args:
+        query (str): entity to look up in the index
+        cfg (DictConfig): config object to specifiy query parameters
+        model (EntityLinkingModel): entity linking encoder model
+        index (object): faiss index
+        pca (object): sklearn pca transformation to be applied to queries 
+        idx2id (dict): dictionary mapping unique concept dataset index to 
+                       its CUI
+        id2string (dict): dictionary mapping each unqiue CUI to a 
+                          representative english description of
+                          the concept
+    Returns:
+        A dictionary with the concept ids of the index's most similar 
+        entities as the keys and a tuple containing the string 
+        representation of that concept and its cosine similarity to 
+        the query as the values. 
+    """
+    query_emb = get_query_embedding(query, model).detach().cpu().numpy()
+
+    if cfg.apply_pca:
+        query_emb = pca.transform(query_emb)
+
+    dist, neighbors = index.search(query_emb.astype(np.float32), cfg.query_num_factor * cfg.top_n)
+    dist, neighbors = dist[0], neighbors[0]
+    unique_ids = OrderedDict()
+    neighbor_idx = 0
+
+    # Many of nearest neighbors could map to the same concept id, their idx is their unique identifier
+    while len(unique_ids) < cfg.top_n and neighbor_idx < len(neighbors):
+        concept_id_idx = neighbors[neighbor_idx]
+        concept_id = idx2id[concept_id_idx]
+
+        # Only want one instance of each unique concept
+        if concept_id not in unique_ids:
+            concept = id2string[concept_id]
+            unique_ids[concept_id] = (concept, 1 - dist[neighbor_idx])
+
+        neighbor_idx += 1
+
+    unique_ids = dict(unique_ids)
+
+    return unique_ids
+
+
+def main(cfg: DictConfig, restore: bool):
+    """
+    Loads faiss index and allows commandline queries 
+    to the index. Builds new index if one hasn't been built yet.
+
+    Args:
+        cfg: Config file specifying index parameters
+        restore: Whether to restore model weights trained
+                 by the user. Otherwise will load weights
+                 used before self alignment pretraining.
+    """
+
+    if not os.path.isfile(cfg.index.index_save_name) or (
+        cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name) or not os.path.isfile(cfg.index.idx_to_id)
+    ):
+        logging.warning("Either no index and/or no mapping from entity idx to ids exists. Please run `build_index.py`")
+        return
+
+    logging.info("Loading entity linking encoder model")
+    model = load_model(cfg.model, restore)
+
+    logging.info("Loading index and associated files")
+    index = faiss.read_index(cfg.index.index_save_name)
+    idx2id = pkl.load(open(cfg.index.idx_to_id, "rb"))
+    id2string = pkl.load(open(cfg.index.id_to_string, "rb"))  # Should be created during dataset prep
+
+    if cfg.index.apply_pca:
+        pca = pkl.load(open(cfg.index.pca.pca_save_name, "rb"))
+
+    while True:
+        query = input("enter index query: ")
+        output = query_index(query, cfg.top_n, cfg.index, model, index, pca, idx2id, id2string)
+
+        if query == "exit":
+            break
+
+        for concept_id in output:
+            concept_details = output[concept_id]
+            concept_id = "C" + str(concept_id).zfill(7)
+            print(concept_id, concept_details)
+
+        print("----------------\n")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path"
+    )
+    parser.add_argument("--project_dir", required=False, type=str, default=".")
+    parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml")
+    args = parser.parse_args()
+
+    cfg = OmegaConf.load(args.cfg)
+    cfg.project_dir = args.project_dir
+
+    main(cfg, args.restore)

tts/NeMo/examples/nlp/entity_linking/self_alignment_pretraining.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# Please see tutorial at Nemo/tutorials/nlp/Entity_Linking_Medical.ipynb for
+# more information on entity linking and self alignment pretraining.
+
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.nlp.models import EntityLinkingModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="umls_medical_entity_linking_config.yaml")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"\nConfig Params:\n{OmegaConf.to_yaml(cfg)}")
+    trainer = Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    logging.info(f"Loading weights from pretrained model {cfg.model.language_model.pretrained_model_name}")
+    model = EntityLinkingModel(cfg=cfg.model, trainer=trainer)
+    logging.info("===========================================================================================")
+    logging.info('Starting training...')
+    trainer.fit(model)
+    logging.info('Training finished!')
+    logging.info("===========================================================================================")
+
+    if cfg.model.nemo_path:
+        # '.nemo' file contains the last checkpoint and the params to initialize the model
+        model.save_to(cfg.model.nemo_path)
+        logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}')
+
+
+if __name__ == '__main__':
+    main()