Introduce code from Rust monorepo. This is a mess

tts/NeMo/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml

@@ -0,0 +1,115 @@
+# Copyright (c) 2020, 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.
+
+
+pretrained_model: null # pretrained QAModel model from list_available_models()
+do_training: true # training mode, for testing change to false
+trainer:
+  gpus: 8 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1]
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 1.0
+  amp_level: O1 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2 to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 5  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-cased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  encoder:
+    dropout: 0.1
+
+  dataset:
+    max_seq_length: 128 # the maximum length BERT supports is 512
+    max_num_cat_slot: 6 # maximum number of different categorical slots per service in dataset 
+    max_num_noncat_slot: 12 # maximum number of different non-categorical slots per service in dataset 
+    max_value_per_cat_slot: 12 # maximum number of different categorical slot values per service in dataset 
+    max_num_intent: 4 # maximum number of different intents per service in dataset 
+    task_name: sgd_single_domain # or from [sgd_all, sgd_all_single, sgd_multi_domain, debug_sample]
+    state_tracker: nemotracker # or baseline
+    data_dir: ???
+    dialogues_example_dir: ??? # cache location to store or load preprocessed data from. To use cache, specify `use_cache`
+    subsample: false # balances negative and positive training examples for improved performance
+    use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up
+    use_fuzzy_match: true # Whether to use fuzzy string matching when comparing non-categorical slot values. Should be set to False when conducting multiwoz style evaluation.
+    joint_acc_across_turn: false # Whether to compute joint goal accuracy across turn instead of across service. Should be set to True when conducting multiwoz style evaluation.
+
+  train_ds:
+    ds_item: "train"
+    batch_size: 128
+    shuffle: true
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    ds_item: ["dev"] # ["dev", "test"]
+    batch_size: 8
+    shuffle: false
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    ds_item: ["test"]
+    batch_size: 8
+    shuffle: false
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 1e-4
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: PolynomialDecayAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.02
+      last_epoch: -1
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "SGDQA"  # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback
+  resume_if_exists: false
+  resume_ignore_no_checkpoint: false

tts/NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+# Copyright 2019 The Google Research Authors.
+#
+# 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.
+
+"""
+This script contains an example of how to train and test the NeMo SGD-QA Model (https://arxiv.org/abs/2105.08049).
+The SGD-QA model is a fast multi-pass schema-guided state-tracking model, that is trained on the Google schema-guided state tracking dataset (https://arxiv.org/abs/1909.05855).
+The model takes dialogue as input and outputs the dialogue state, which includes slot-value pairs. 
+The model consists of two components: a neural natural language understanding model (NLU), and a rule-based state tracker.
+The NLU takes in a dialogue turn and different schema (entity) information options and outputs their match score. The state tracker takes the highest rated entities and composes
+the dialogue state across turns.
+
+***Data format***
+The SGD-QA requires a JSON schema and dialogue files for each dataset split. 
+In the following we will show an example for a service entry in the schema file.
+* service_name
+* description
+* slots
+    * name
+    * description
+    * is_categorical
+    * possible values
+* intents
+    * name
+    * description
+    * required_slots (not used)
+    * is_transactional (not used)
+    * optional_slots (not used)
+    * result_slots (not used)
+
+
+In the following we will show an example for a dialogue. More information at https://arxiv.org/abs/1909.05855
+* dialogue_id
+* services
+* turns
+    * frames
+        * actions
+            * act
+            * slot
+            * values
+        * service
+        * slots
+            * exclusive_end
+            * slot
+            * start
+        * state
+            * active_intent
+            * requeste_slots
+            * slot_values 
+    * speaker - [USER, SYSTEM]
+    * utterance
+
+
+***Downloading the dataset***
+#   git clone https://github.com/google-research-datasets/dstc8-schema-guided-dialogue.git
+
+***Setting the configs***
+The model and the PT trainer are defined in a config file that declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, SGD-QA encoder and decoder, loss, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+
+This script uses the `/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml` config file
+by default. You may update the config file from the file directly. The other option is to set another config file via command-line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+
+***Model Training***
+# python sgd_qa.py
+    do_training=True
+    model.dataset.data_dir=<DATA_DIR_WITH_JSON_DATA>
+    model.dataset.dialogues_example_dir=<DATA_DIR_WITH_PREPROCESSED_DATA>
+    model.validation_ds.ds_item=<LIST_OF_SPLITS>
+    trainer.max_epochs=<NUM_EPOCHS>
+    trainer.gpus=[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+
+***Model Evaluation***
+#   python sgd_qa.py
+    do_training=False
+    model.test_ds.ds_item=<LIST_OF_SPLITS>
+
+To load a pretrained checkpoint from the cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=<MODEL_NAME>. You can find all pretrained model names by using 
+SGDQAModel.list_available_models(). To load a local checkpoint use model.restore_from(<PATH_TO_CHECKPOINT>)
+
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.dialogue_state_tracking.sgdqa_model import SGDQAModel
+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="sgdqa_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if cfg.pretrained_model or (cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path)):
+        if cfg.pretrained_model:
+            logging.info(f'Loading pretrained model {cfg.pretrained_model}')
+            model = SGDQAModel.from_pretrained(cfg.pretrained_model)
+        else:
+            logging.info(f'Restoring model from {cfg.model.nemo_path}')
+            model = SGDQAModel.restore_from(cfg.model.nemo_path)
+        if cfg.do_training:
+            model.setup_training_data(train_data_config=cfg.model.train_ds)
+            model.setup_multiple_validation_data(val_data_config=cfg.model.validation_ds)
+    else:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = SGDQAModel(cfg.model, trainer=trainer)
+
+    if cfg.do_training:
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.ds_item is not None:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+        trainer = pl.Trainer(gpus=gpu)
+        model.setup_multiple_test_data(test_data_config=cfg.model.test_ds)
+        if model.prepare_test(trainer):
+            trainer.test(model)
+
+
+if __name__ == '__main__':
+    main()

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()

tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2020, 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.
+
+"""
+## Tasks
+This script works with all GLUE Benchmark tasks, more details about the GLUE Benchmark could be found at
+https://gluebenchmark.com/
+
+More details on how to use this script could be found in tutorials/nlp/GLUE_Benchmark.ipynb
+
+## Model Training
+
+To train GLUEModel with the default config file, run:
+    python glue_benchmark.py \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+    model.task_name=TASK_NAME \
+    trainer.max_epochs=<NUM_EPOCHS> \
+    trainer.gpus="[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+Supported task names:
+["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"]
+Note, MNLI task includes both matched and mismatched dev sets
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import GLUEModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_name="glue_benchmark_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager_cfg = cfg.get("exp_manager", None)
+
+    if exp_manager_cfg:
+        exp_manager_cfg.name = cfg.model.task_name
+        logging.info(f'Setting task_name to {exp_manager_cfg.name} in exp_manager')
+    exp_manager(trainer, exp_manager_cfg)
+
+    if cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path):
+        model = GLUEModel.restore_from(cfg.model.nemo_path)
+        logging.info(f'Restoring model from {cfg.model.nemo_path}')
+        model.update_data_dir(data_dir=cfg.model.dataset.data_dir)
+        model.setup_training_data()
+        model.setup_multiple_validation_data()
+        trainer.fit(model)
+    else:
+        model = GLUEModel(cfg.model, trainer=trainer)
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark_config.yaml

@@ -0,0 +1,82 @@
+# GLUE Benchmark with pre-trained BERT models
+supported_tasks: &supported_tasks ['cola', 'sst-2', 'mrpc', 'sts-b', 'qqp', 'mnli', 'qnli', 'rte', 'wnli']
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16
+  accelerator: ddp
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  task_name: &task_name mrpc # choose from: ["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"] GLUE task name, MNLI includes both matched and mismatched dev sets
+  supported_tasks: *supported_tasks
+  output_dir: null # dir to write write predictions
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  dataset:
+    data_dir: ??? # /path/to/data
+    max_seq_length: 128
+    use_cache: true
+
+    # shared across dataloaders:
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    ds_item: 'train.tsv'
+    shuffle: true
+    num_samples: -1
+    batch_size: 32
+
+  validation_ds:
+    ds_item: 'dev.tsv' # for MNLI 'dev_matched.tsv' and 'dev_mismatched.tsv' will de used
+    shuffle: false
+    num_samples: -1
+    batch_size: 32
+
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  optim:
+    name: adam
+    lr: 5e-5
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./NeMo_experiments"
+  name: *task_name # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertDPRModel
+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="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_dpr_model = BertDPRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_dpr_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertJointIRModel
+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="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_joint_ir_model = BertJointIRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_joint_ir_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml

@@ -0,0 +1,99 @@
+# Fine-tuning BERT model for information retrieval
+name: &name BertIR
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 0.05  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  checkpoint_callback: false  # provided by exp_manager
+  logger: false  # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    sim_score_dropout: 0.1
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+
+  train_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with training queries and their indices
+    query_to_passages: null
+    # path to file with training examples which have the form of
+    # (query_id, relevant_passage_id, irrelevant_passage_1_id, ..., irrelevant_passage_n_id)
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: npz
+    shuffle: true
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with validation queries and their indices
+    query_to_passages: null # path to file with passages to re-rank for each validation query
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: pkl
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adam
+    lr: 1e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0
+
+    sched:
+      name: WarmupAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.05
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # where to store logs and checkpoints
+  name: *name  # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2020, 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 argparse
+
+import numpy as np
+import pandas as pd
+
+
+def construct_negatives(input_file, output_file, num_passages, num_negatives):
+    qrels = pd.read_csv(input_file, delimiter="\t", header=None)
+    with open(output_file, "w") as f:
+        for i in range(len(qrels)):
+            query_id, rel_passage_id = qrels[0][i], qrels[2][i]
+            negatives = np.random.randint(num_passages, size=num_negatives)
+            output_ids = [query_id, rel_passage_id] + negatives.tolist()
+            output_str = [str(id_) for id_ in output_ids]
+            print("\t".join(output_str), file=f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Negative passages construction")
+    parser.add_argument("--data", type=str, default="msmarco_dataset", help="path to folder with data")
+    parser.add_argument("--num_passages", type=int, default=8841823, help="total number of passages")
+    parser.add_argument("--num_negatives", type=int, default=10, help="number of negatives per positive")
+    args = parser.parse_args()
+
+    for mode in ["train", "dev"]:
+        construct_negatives(
+            input_file=f"{args.data}/qrels.{mode}.tsv",
+            output_file=f"{args.data}/query2passages.{mode}.tsv",
+            num_passages=args.num_passages,
+            num_negatives=args.num_negatives,
+        )
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh

@@ -0,0 +1,22 @@
+echo "=== Acquiring MSMARCO dataset ==="
+echo "---"
+
+mkdir -p msmarco_dataset
+cd msmarco_dataset
+
+echo "- Downloading passages"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz
+tar -xzvf collection.tar.gz
+rm collection.tar.gz
+
+echo "- Downloading queries"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz
+tar -xzvf queries.tar.gz
+rm queries.tar.gz
+rm queries.eval.tsv
+
+echo "- Downloading relevance labels"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.train.tsv
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.dev.tsv
+
+echo "---"

tts/NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml

@@ -0,0 +1,111 @@
+# Intent and Slot classification with pretrained BERT models
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 50
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 amp_level to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+
+model:
+  nemo_path: null # filename to save the model and associated artifacts to .nemo file
+  data_dir: ??? # /path/to/data
+  class_labels:
+    intent_labels_file: intent_labels.csv
+    slot_labels_file: slot_labels.csv
+  class_balancing: null # or weighted_loss
+  intent_loss_weight: 0.6 # relation of intent to slot loss in total loss (between 0 to 1)
+  pad_label: -1 # if -1 not slot token will be used
+  ignore_extra_tokens: false
+  ignore_start_end: true # do not use first and last token for slot training
+
+  train_ds:
+    prefix: train
+    batch_size: 32
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    prefix: test
+    batch_size: 32
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    prefix: test
+    batch_size: 32
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    max_seq_length: 50
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  head:
+    num_output_layers: 2
+    fc_dropout: 0.1
+
+  optim:
+    name: adam
+    lr: 2e-5
+    args:
+      name: auto
+      params:
+        weight_decay: 0.01
+
+    sched:
+      name: WarmupAnnealing
+      iters_per_batch: null # computed at runtime
+      max_steps: null # computed at runtime or explicitly set here
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+      # scheduler config override
+      args:
+        name: auto
+        params:
+          warmup_steps: null
+          warmup_ratio: 0.1
+          last_epoch: -1
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "IntentSlot"  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a modelcheckpoint callback
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/intent_slot_classification/data/assistant_utils.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2020, 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 re
+import shutil
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import DATABASE_EXISTS_TMP, if_exist, write_files
+from nemo.utils import logging
+
+
+def copy_input_files(infold):
+    """ Put training files in convenient place for conversion to our format. """
+    our_infold = infold + "/dataset"
+
+    if os.path.exists(our_infold + "/trainset") and os.path.exists(our_infold + "/testset"):
+        logging.info("Input folders exists")
+        return
+
+    logging.info(f"Copying files to input folder: {our_infold}")
+    os.makedirs(infold, exist_ok=True)
+
+    old_infold = (
+        infold + '/CrossValidation/autoGeneFromRealAnno/autoGene_2018_03_22-13_01_25_169/CrossValidation/KFold_1'
+    )
+    if not os.path.exists(our_infold + "/trainset"):
+        shutil.copytree(old_infold + '/trainset', our_infold + '/trainset')
+
+    if not os.path.exists(our_infold + "/testset"):
+        shutil.copytree(old_infold + '/testset/csv', our_infold + '/testset')
+
+
+def get_intents(infold):
+    """ Get list of intents from file names. """
+    intents = [f[:-4] for f in os.listdir(infold)]
+    intents.sort()
+    print(f'Found {len(intents)} intents')
+    return intents
+
+
+def get_intent_queries(infold, intent_names, mode):
+    """ Get list of queries with their corresponding intent number. """
+    intent_queries = ['sentence\tlabel\n']
+
+    for index, intent in enumerate(intent_names):
+        queries = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines()
+        for query in queries[1:]:
+            phrases = query.split(";")
+            intent_query = phrases[4][1:-1] + "\t" + str(index)
+            intent_queries.append(intent_query)
+
+    return intent_queries
+
+
+def get_slots(infold, modes):
+    """
+    Find a lost of unique slot types in training and testing data.
+    We use a single slot type name both for starting and continuation tokes (not using B-, I- notation).
+    """
+    slots = set()
+
+    for mode in modes:
+        path = f'{infold}/{mode}set'
+        for filename in os.listdir(path):
+            lines = open(f'{path}/{filename}', 'r').readlines()
+            for line in lines[1:]:
+                query = line.split(";")[3]
+                slot_phrases = re.findall('\[.*?\]', query)
+                for slot_phrase in slot_phrases:
+                    slot = slot_phrase.split(" : ")[0][1:]
+                    slots.add(slot)
+
+    slots = sorted(slots)
+    slots.append("O")
+    print(f'Found {len(slots)} slot types')
+    return slots
+
+
+def get_slot_queries(infold, slot_dict, mode, intent_names):
+    """ Convert each word in a query to corresponding slot number. """
+    slot_queries = []
+    outside_slot = len(slot_dict) - 1
+
+    # keep the same order of files/queries as for intents
+    for intent in intent_names:
+        lines = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines()
+        for line in lines[1:]:
+            slot_query = ""
+            query = line.split(";")[3]
+            words = query.split(" ")
+            current_slot = outside_slot
+            for word in words:
+                if word[0] == "[":
+                    current_slot = slot_dict[word[1:]]
+                elif word[0] == ":":
+                    continue
+                else:
+                    slot_query += str(current_slot) + " "
+                    if word[-1] == ']':
+                        current_slot = outside_slot
+
+            slot_queries.append(slot_query.strip())
+
+    return slot_queries
+
+
+def process_assistant(infold, outfold, modes=['train', 'test']):
+    """
+    https://github.com/xliuhw/NLU-Evaluation-Data - this dataset includes
+    about 25 thousand examples with 66 various multi-domain intents and 57 entity types.
+    """
+    if if_exist(outfold, [f'{mode}_slots.tsv' for mode in modes]):
+        logging.info(DATABASE_EXISTS_TMP.format('robot', outfold))
+        return outfold
+
+    logging.info(f'Processing assistant commands dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    # copy train/test files to the convenient directory to work with
+    copy_input_files(infold)
+    infold += "/dataset"
+
+    # get list of intents from train folder (test folder supposed to be the same)
+    intent_names = get_intents(infold + "/trainset")
+    write_files(intent_names, f'{outfold}/dict.intents.csv')
+
+    # get all train and test queries with their intent
+    for mode in modes:
+        intent_queries = get_intent_queries(infold, intent_names, mode)
+        write_files(intent_queries, f'{outfold}/{mode}.tsv')
+
+    # get list of all unique slots in training and testing files
+    slot_types = get_slots(infold, modes)
+    write_files(slot_types, f'{outfold}/dict.slots.csv')
+
+    # create files of slot queries
+    slot_dict = {k: v for v, k in enumerate(slot_types)}
+    for mode in modes:
+        slot_queries = get_slot_queries(infold, slot_dict, mode, intent_names)
+        write_files(slot_queries, f'{outfold}/{mode}_slots.tsv')

tts/NeMo/examples/nlp/intent_slot_classification/data/import_datasets.py

@@ -0,0 +1,265 @@
+# Copyright (c) 2020, 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 argparse
+import os
+import shutil
+from os.path import exists
+
+from assistant_utils import process_assistant
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import (
+    DATABASE_EXISTS_TMP,
+    MODE_EXISTS_TMP,
+    create_dataset,
+    get_dataset,
+    get_vocab,
+    if_exist,
+)
+from nemo.utils import logging
+
+
+def ids2text(ids, vocab):
+    return ' '.join([vocab[int(id_)] for id_ in ids])
+
+
+def process_atis(infold, outfold, modes=['train', 'test'], do_lower_case=False):
+    """ MSFT's dataset, processed by Kaggle
+    https://www.kaggle.com/siddhadev/atis-dataset-from-ms-cntk
+    """
+    vocab = get_vocab(f'{infold}/atis.dict.vocab.csv')
+
+    if if_exist(outfold, [f'{mode}.tsv' for mode in modes]):
+        logging.info(DATABASE_EXISTS_TMP.format('ATIS', outfold))
+        return outfold
+    logging.info(f'Processing ATIS dataset and storing at {outfold}.')
+
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        outfiles[mode].write('sentence\tlabel\n')
+        outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w')
+
+        queries = open(f'{infold}/atis.{mode}.query.csv', 'r').readlines()
+        intents = open(f'{infold}/atis.{mode}.intent.csv', 'r').readlines()
+        slots = open(f'{infold}/atis.{mode}.slots.csv', 'r').readlines()
+
+        for i, query in enumerate(queries):
+            sentence = ids2text(query.strip().split()[1:-1], vocab)
+            if do_lower_case:
+                sentence = sentence.lower()
+            outfiles[mode].write(f'{sentence}\t{intents[i].strip()}\n')
+            slot = ' '.join(slots[i].strip().split()[1:-1])
+            outfiles[mode + '_slots'].write(slot + '\n')
+
+    shutil.copyfile(f'{infold}/atis.dict.intent.csv', f'{outfold}/dict.intents.csv')
+    shutil.copyfile(f'{infold}/atis.dict.slots.csv', f'{outfold}/dict.slots.csv')
+    for mode in modes:
+        outfiles[mode].close()
+
+
+def process_snips(infold, outfold, do_lower_case, modes=['train', 'test'], dev_split=0.1):
+    if not os.path.exists(infold):
+        link = 'https://github.com/snipsco/spoken-language-understanding-research-datasets'
+        raise ValueError(f'Data not found at {infold}. ' f'You may request to download the SNIPS dataset from {link}.')
+
+    exist = True
+    for dataset in ['light', 'speak', 'all']:
+        if if_exist(f'{outfold}/{dataset}', [f'{mode}.tsv' for mode in modes]):
+            logging.info(DATABASE_EXISTS_TMP.format('SNIPS-' + dataset, outfold))
+        else:
+            exist = False
+    if exist:
+        return outfold
+
+    logging.info(f'Processing SNIPS dataset and storing at folders "speak", "light" and "all" under {outfold}.')
+    logging.info(
+        f'Processing and importing "smart-speaker-en-close-field" -> "speak" and "smart-speaker-en-close-field" -> "light".'
+    )
+
+    os.makedirs(outfold, exist_ok=True)
+
+    speak_dir = 'smart-speaker-en-close-field'
+    light_dir = 'smart-lights-en-close-field'
+
+    light_files = [f'{infold}/{light_dir}/dataset.json']
+    speak_files = [f'{infold}/{speak_dir}/training_dataset.json']
+    speak_files.append(f'{infold}/{speak_dir}/test_dataset.json')
+
+    light_train, light_dev, light_slots, light_intents = get_dataset(light_files, dev_split)
+    speak_train, speak_dev, speak_slots, speak_intents = get_dataset(speak_files)
+
+    create_dataset(light_train, light_dev, light_slots, light_intents, do_lower_case, f'{outfold}/light')
+    create_dataset(speak_train, speak_dev, speak_slots, speak_intents, do_lower_case, f'{outfold}/speak')
+    create_dataset(
+        light_train + speak_train,
+        light_dev + speak_dev,
+        light_slots | speak_slots,
+        light_intents | speak_intents,
+        do_lower_case,
+        f'{outfold}/all',
+    )
+
+
+def process_jarvis_datasets(
+    infold, outfold, modes=['train', 'test', 'dev'], do_lower_case=False, ignore_prev_intent=False
+):
+    """ process and convert Jarvis datasets into NeMo's BIO format
+    """
+    dataset_name = "jarvis"
+    if if_exist(outfold, ['dict.intents.csv', 'dict.slots.csv']):
+        logging.info(DATABASE_EXISTS_TMP.format(dataset_name, outfold))
+        return outfold
+
+    logging.info(f'Processing {dataset_name} dataset and storing at {outfold}')
+
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    intents_list = {}
+    slots_list = {}
+    slots_list_all = {}
+
+    outfiles['dict_intents'] = open(f'{outfold}/dict.intents.csv', 'w')
+    outfiles['dict_slots'] = open(f'{outfold}/dict.slots.csv', 'w')
+
+    outfiles['dict_slots'].write('O\n')
+    slots_list["O"] = 0
+    slots_list_all["O"] = 0
+
+    for mode in modes:
+        if if_exist(outfold, [f'{mode}.tsv']):
+            logging.info(MODE_EXISTS_TMP.format(mode, dataset_name, outfold, mode))
+            continue
+
+        if not if_exist(infold, [f'{mode}.tsv']):
+            logging.info(f'{mode} mode of {dataset_name}' f' is skipped as it was not found.')
+            continue
+
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        outfiles[mode].write('sentence\tlabel\n')
+        outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w')
+
+        queries = open(f'{infold}/{mode}.tsv', 'r').readlines()
+
+        for i, query in enumerate(queries):
+            line_splits = query.strip().split("\t")
+            if len(line_splits) == 3:
+                intent_str, slot_tags_str, sentence = line_splits
+            else:
+                intent_str, sentence = line_splits
+                slot_tags_str = ""
+
+            if intent_str not in intents_list:
+                intents_list[intent_str] = len(intents_list)
+                outfiles['dict_intents'].write(f'{intent_str}\n')
+
+            if ignore_prev_intent:
+                start_token = 2
+            else:
+                start_token = 1
+
+            if do_lower_case:
+                sentence = sentence.lower()
+            sentence_cld = " ".join(sentence.strip().split()[start_token:-1])
+            outfiles[mode].write(f'{sentence_cld}\t' f'{str(intents_list[intent_str])}\n')
+
+            slot_tags_list = []
+            if slot_tags_str.strip():
+                slot_tags = slot_tags_str.strip().split(",")
+                for st in slot_tags:
+                    if not st.strip():
+                        continue
+                    [start_i, end_i, slot_name] = st.strip().split(":")
+                    slot_tags_list.append([int(start_i), int(end_i), slot_name])
+                    if slot_name not in slots_list:
+                        slots_list[slot_name] = len(slots_list)
+                        slots_list_all[f'B-{slot_name}'] = len(slots_list_all)
+                        slots_list_all[f'I-{slot_name}'] = len(slots_list_all)
+                        outfiles['dict_slots'].write(f'B-{slot_name}\n')
+                        outfiles['dict_slots'].write(f'I-{slot_name}\n')
+
+            slot_tags_list.sort(key=lambda x: x[0])
+            slots = []
+            processed_index = 0
+            for tag_start, tag_end, tag_str in slot_tags_list:
+                if tag_start > processed_index:
+                    words_list = sentence[processed_index:tag_start].strip().split()
+                    slots.extend([str(slots_list_all['O'])] * len(words_list))
+                words_list = sentence[tag_start:tag_end].strip().split()
+                slots.append(str(slots_list_all[f'B-{tag_str}']))
+                slots.extend([str(slots_list_all[f'I-{tag_str}'])] * (len(words_list) - 1))
+                processed_index = tag_end
+
+            if processed_index < len(sentence):
+                words_list = sentence[processed_index:].strip().split()
+                slots.extend([str(slots_list_all['O'])] * len(words_list))
+
+            slots = slots[1:-1]
+            slot = ' '.join(slots)
+            outfiles[mode + '_slots'].write(slot + '\n')
+
+        outfiles[mode + '_slots'].close()
+        outfiles[mode].close()
+
+    outfiles['dict_slots'].close()
+    outfiles['dict_intents'].close()
+
+    return outfold
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.")
+    parser.add_argument(
+        "--dataset_name", required=True, type=str, choices=['atis', 'snips', 'jarvis', 'assistant'],
+    )
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help='path to the folder containing the dataset files'
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str, help='path to save the processed dataset')
+    parser.add_argument("--do_lower_case", action='store_true')
+    parser.add_argument(
+        "--ignore_prev_intent",
+        action='store_true',
+        help='ignores previous intent while importing datasets in jarvis\'s format',
+    )
+
+    args = parser.parse_args()
+
+    dataset_name = args.dataset_name
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+
+    if not exists(source_dir):
+        raise FileNotFoundError(f"{source_dir} does not exist.")
+
+    if dataset_name == 'atis':
+        process_atis(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case)
+    elif dataset_name == 'snips':
+        process_snips(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case)
+    elif dataset_name == 'jarvis':
+        process_jarvis_datasets(
+            infold=source_dir,
+            outfold=target_dir,
+            modes=["train", "test", "dev"],
+            do_lower_case=args.do_lower_case,
+            ignore_prev_intent=args.ignore_prev_intent,
+        )
+    elif dataset_name == 'assistant':
+        process_assistant(infold=source_dir, outfold=target_dir)
+    else:
+        raise ValueError(f'Dataset {dataset_name} is not supported.')

tts/NeMo/examples/nlp/intent_slot_classification/intent_slot_classification.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2020, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import IntentSlotClassificationModel
+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="intent_slot_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config Params:\n {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    # initialize the model using the config file
+    model = IntentSlotClassificationModel(cfg.model, trainer=trainer)
+
+    # training
+    logging.info("================================================================================================")
+    logging.info('Starting training...')
+    trainer.fit(model)
+    logging.info('Training finished!')
+
+    # Stop further testing as fast_dev_run does not save checkpoints
+    if trainer.fast_dev_run:
+        return
+
+    # after model training is done, you can load the model from the saved checkpoint
+    # and evaluate it on a data file or on given queries.
+    logging.info("================================================================================================")
+    logging.info("Starting the testing of the trained model on test set...")
+    logging.info("We will load the latest model saved checkpoint from the training...")
+
+    # for evaluation and inference you can load the previously trained model saved in .nemo file
+    # like this in your code, but we will just reuse the trained model here
+    # eval_model = IntentSlotClassificationModel.restore_from(restore_path=checkpoint_path)
+    eval_model = model
+
+    # we will setup testing data reusing the same config (test section)
+    eval_model.update_data_dir_for_testing(data_dir=cfg.model.data_dir)
+    eval_model.setup_test_data(test_data_config=cfg.model.test_ds)
+
+    trainer.test(model=eval_model, ckpt_path=None, verbose=False)
+    logging.info("Testing finished!")
+
+    # run an inference on a few examples
+    logging.info("======================================================================================")
+    logging.info("Evaluate the model on the given queries...")
+
+    # this will work well if you train the model on Assistant dataset
+    # for your own dataset change the examples appropriately
+    queries = [
+        'set alarm for seven thirty am',
+        'lower volume by fifty percent',
+        'what is my schedule for tomorrow',
+    ]
+
+    pred_intents, pred_slots = eval_model.predict_from_examples(queries, cfg.model.test_ds)
+
+    logging.info('The prediction results of some sample queries with the trained model:')
+    for query, intent, slots in zip(queries, pred_intents, pred_slots):
+        logging.info(f'Query : {query}')
+        logging.info(f'Predicted Intent: {intent}')
+        logging.info(f'Predicted Slots: {slots}')
+
+    logging.info("Inference finished!")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/language_modeling/bert_pretraining.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2020, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning.plugins import DDPPlugin
+
+from nemo.collections.nlp.models.language_modeling import BERTLMModel
+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="bert_pretraining_from_text_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config:\n {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(plugins=[DDPPlugin(find_unused_parameters=True)], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_model = BERTLMModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_model)
+    if cfg.model.nemo_path:
+        bert_model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml

@@ -0,0 +1,80 @@
+# BERT Pretraining from Preprocessed (tokenized) data
+name: &name PretrainingBERTFromPreprocessed
+trainer:
+  gpus: 8 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_steps: 2285714 # precedence over max_epochs
+  num_sanity_val_steps: 0 # needed for bert pretraining from preproc
+  replace_sampler_ddp: false # needed for bert pretraining from preproc
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 1.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  only_mlm_loss: true # only use masked language model without next sentence prediction
+  num_tok_classification_layers: 1 # number of token classification head output layers
+  num_seq_classification_layers: 2 # number of sequence classification head output layers
+
+
+  language_model:
+    pretrained_model_name: bert-base-uncased # huggingface model name
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer: null
+
+  train_ds:
+    data_file: null # path to hdf5 file (or directory)
+    max_predictions_per_seq: 80
+    batch_size: 16
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 0.4375e-4
+    weight_decay: 0.01
+
+    sched:
+      name: SquareRootAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.01
+      min_lr: 0.0
+      last_epoch: -1
+
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml

@@ -0,0 +1,109 @@
+# BERT Pretraining from Text
+name: &name PretrainingBERTFromText
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  only_mlm_loss: false # only use masked language model without next sentence prediction
+  num_tok_classification_layers: 1 # number of token classification head output layers
+  num_seq_classification_layers: 2 # number of sequence classification head output layers
+  max_seq_length: 128
+  # The maximum total input sequence length after tokenization. Sequences longer than this
+  # will be truncated, and sequences shorter than this will be padded.
+  mask_prob: 0.15
+  # Probability of masking a token in the input text during data processing.
+  short_seq_prob: 0.1
+  # Probability of having a sequence shorter than the maximum sequence length `max_seq_length` in data processing.",
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+        unk_token: '[UNK]'
+        sep_token: '[SEP]'
+        pad_token: '[PAD]'
+        bos_token: '[CLS]'
+        mask_token: '[MASK]'
+        eos_token: '[SEP]'
+        cls_token: '[CLS]'
+
+  train_ds:
+    data_file: null # path to data file
+    max_seq_length: ${model.max_seq_length}
+    mask_prob: ${model.mask_prob}
+    short_seq_prob: ${model.short_seq_prob}
+    batch_size: 16 # per GPU
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    data_file: null # path to data file
+    max_seq_length: ${model.max_seq_length}
+    mask_prob: ${model.mask_prob}
+    short_seq_prob: ${model.short_seq_prob}
+    batch_size: 16 # per GPU
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 3e-5
+    weight_decay: 0.0
+
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.1
+      min_lr: 0.0
+      last_epoch: -1
+
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/language_modeling/conf/transformer_lm_config.yaml

@@ -0,0 +1,96 @@
+name: TransformerLanguageModel
+do_training: True # set to False if only preprocessing data
+
+model:
+  label_smoothing: 0.0
+  preproc_out_dir: null # path to store data preprocessing outputs
+  
+  train_ds:
+    use_tarred_dataset: true
+    tar_files: ??? # path to tarred files
+    metadata_file: ??? # metadata for tarred dataset
+    shard_strategy: scatter
+    tokens_in_batch: 4096
+    tar_shuffle_n: 256
+    
+  validation_ds:
+    file_name: ??? # path to file with validation data
+    tokens_in_batch: 512
+    clean: false
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    
+  test_ds:
+    file_name: ??? # path to file with test data
+    tokens_in_batch: 512
+    clean: false
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  tokenizer:
+    tokenizer_name: yttm
+    tokenizer_model: ???
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  encoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    num_layers: 6
+    inner_size: 2048
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    mask_future: true
+    pre_ln: false
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: TransformerLM
+  files_to_copy: []

tts/NeMo/examples/nlp/language_modeling/convert_weights_to_nemo1.0.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2020, 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.
+
+"""
+Converts BERT NeMo0.* checkpoints to NeMo1.0 format.
+"""
+
+from argparse import ArgumentParser
+
+import torch
+
+parser = ArgumentParser()
+parser.add_argument("--bert_encoder", required=True, help="path to BERT encoder, e.g. /../BERT-STEP-2285714.pt")
+parser.add_argument(
+    "--bert_token_classifier",
+    required=True,
+    help="path to BERT token classifier, e.g. /../BertTokenClassifier-STEP-2285714.pt",
+)
+parser.add_argument(
+    "--bert_sequence_classifier",
+    required=False,
+    default=None,
+    help="path to BERT sequence classifier, e.g /../SequenceClassifier-STEP-2285714.pt",
+)
+parser.add_argument(
+    "--output_path", required=False, default="converted_model.pt", help="output path to newly converted model"
+)
+args = parser.parse_args()
+
+bert_in = torch.load(args.bert_encoder)
+tok_in = torch.load(args.bert_token_classifier)
+if args.bert_sequence_classifier:
+    seq_in = torch.load(args.bert_sequence_classifier)
+
+new_dict = {}
+new_model = {"state_dict": new_dict}
+for k in bert_in:
+    new_name = k.replace("bert.", "bert_model.")
+    new_dict[new_name] = bert_in[k]
+
+for k in tok_in:
+    new_name = "mlm_classifier." + k
+    new_dict[new_name] = tok_in[k]
+
+if args.bert_sequence_classifier:
+    for k in seq_in:
+        new_name = "nsp_classifier." + k
+        new_dict[new_name] = seq_in[k]
+
+torch.save(new_model, args.output_path)

tts/NeMo/examples/nlp/language_modeling/get_wkt2.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+"""
+This file is adapted from
+https://github.com/salesforce/awd-lstm-lm/blob/master/getdata.sh
+Copyright by the AWD LSTM authors.
+"""
+DATA_DIR=$1
+echo "- Downloading WikiText-2"
+
+wget --continue -P $DATA_DIR https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip
+unzip -q $DATA_DIR/wikitext-2-v1.zip -d $DATA_DIR
+cd $DATA_DIR/wikitext-2
+mv wiki.train.tokens train.txt
+sed -i -e "s/<unk>/[UNK]/g" train.txt
+mv wiki.valid.tokens valid.txt
+sed -i -e "s/<unk>/[UNK]/g" valid.txt
+mv wiki.test.tokens test.txt
+sed -i -e "s/<unk>/[UNK]/g" test.txt
+cd ..
+rm wikitext-2-v1.zip
+
+echo "- WikiText-2 saved at $DATA_DIR/wikitext-2"

tts/NeMo/examples/nlp/language_modeling/transformer_lm.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.language_modeling import TransformerLMModel
+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="transformer_lm_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    transformer_lm = TransformerLMModel(cfg.model, trainer=trainer)
+    trainer.fit(transformer_lm)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml

@@ -0,0 +1,156 @@
+name: AttentionIsAllYouNeed
+do_training: True # set to False if only preprocessing data
+do_testing: False # set to True to run evaluation on test data after training
+
+model:
+  beam_size: 4
+  len_pen: 0.6
+  multilingual: False
+  max_generation_delta: 5
+  label_smoothing: 0.1
+  shared_tokenizer: True # train tokenizer model across src and tgt train data
+  preproc_out_dir: null # path to store data preprocessing outputs
+  src_language: 'en'
+  tgt_language: 'de'
+
+  train_ds:
+    src_file_name: null
+    tgt_file_name: null
+    use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) 
+    # config for preprocessing training data and creating a tarred datset automatically
+    tar_file_prefix: parallel # prefix for tar file names
+    tar_files: null # if data has already been preprocessed (rest of config ignored)
+    metadata_file: null # metadata for tarred dataset
+    lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
+    num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
+    tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
+    shard_strategy: scatter # tarred dataset shard distribution strategy
+    n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
+    tokens_in_batch: 512
+    clean: true
+    max_seq_length: 512
+    shuffle: true
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+    concat_sampling_technique: temperature # only used with ConcatTranslationDataset 
+    concat_sampling_temperature: 5 # only used with ConcatTranslationDataset 
+    concat_sampling_probabilities: null # only used with ConcatTranslationDataset 
+
+  validation_ds:
+    src_file_name: ???
+    tgt_file_name: ???
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  test_ds:
+    src_file_name: ???
+    tgt_file_name: ???
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  encoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_size: null # vocab size for training bpe
+    bpe_dropout: null
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  decoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_size: null # vocab size for training bpe
+    bpe_dropout: null
+    vocab_file: null
+    special_tokens: null
+    training_sample_size: null # valid for sentencepiece tokenizer
+
+  encoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    num_layers: 6
+    inner_size: 2048
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    mask_future: false
+    pre_ln: false
+    pre_ln_final_layer_norm: true
+
+  decoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 0
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    inner_size: 2048
+    num_layers: 6
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    pre_ln: false
+    pre_ln_final_layer_norm: true
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: AAYNBase
+  files_to_copy: []

tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml

@@ -0,0 +1,132 @@
+name: HuggingFaceEncoder
+do_training: True # set to False if only preprocessing data
+do_testing: False # set to True to run evaluation on test data after training
+
+model:
+  beam_size: 4
+  len_pen: 0.6
+  max_generation_delta: 5
+  label_smoothing: 0.1
+  shared_tokenizer: false
+  preproc_out_dir: null
+  src_language: 'en'
+  tgt_language: 'de'
+
+  train_ds:
+    src_file_name: null
+    tgt_file_name: null
+    use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) 
+    # config for preprocessing training data and creating a tarred datset automatically
+    tar_file_prefix: parallel # prefix for tar file names
+    tar_files: null # if data has already been preprocessed (rest of config ignored)
+    metadata_file: null # metadata for tarred dataset
+    lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
+    num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
+    tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
+    shard_strategy: scatter # tarred dataset shard distribution strategy
+    n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
+    tokens_in_batch: 512
+    clean: true
+    max_seq_length: 512
+    shuffle: true
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  validation_ds:
+    src_file_name: null
+    tgt_file_name: null
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  test_ds:
+    src_file_name: null
+    tgt_file_name: null
+    tokens_in_batch: 512
+    clean: false
+    max_seq_length: 512
+    shuffle: false
+    num_samples: -1
+    drop_last: false
+    pin_memory: false
+    num_workers: 8
+
+  optim:
+    name: adam
+    lr: 0.001
+    betas:
+      - 0.9
+      - 0.98
+    weight_decay: 0.0
+    sched:
+      name: InverseSquareRootAnnealing
+      min_lr: 0.0
+      last_epoch: -1
+      warmup_ratio: 0.1
+
+  encoder_tokenizer:
+    library: huggingface
+    tokenizer_model: null
+    vocab_file: null
+    special_tokens: null
+    vocab_size: null
+
+  decoder_tokenizer:
+    library: yttm
+    tokenizer_model: null
+    vocab_file: null
+    special_tokens: null
+    vocab_size: null
+
+  encoder:
+    library: huggingface
+    model_name: bert-base-uncased
+    pretrained: false
+
+  decoder:
+    library: nemo
+    model_name: null
+    pretrained: false
+    max_sequence_length: 512
+    num_token_types: 2
+    embedding_dropout: 0.1
+    learn_positional_encodings: false
+    hidden_size: 512
+    inner_size: 2048
+    num_layers: 6
+    num_attention_heads: 8
+    ffn_dropout: 0.1
+    attn_score_dropout: 0.1
+    attn_layer_dropout: 0.1
+    hidden_act: relu
+    pre_ln: false
+
+  head:
+    num_layers: 1
+    activation: relu
+    log_softmax: true
+    dropout: 0.0
+    use_transformer_init: true
+
+trainer:
+  gpus: 4
+  num_nodes: 1
+  max_epochs: 200
+  amp_level: O2 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False
+  logger: False
+  log_every_n_steps: 50  # Interval of logging.
+  check_val_every_n_epoch: 1
+
+exp_manager:
+  name: HuggingFaceEncoder
+  files_to_copy: []

tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, 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 argparse
+import os
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
+    parser.add_argument('--tokenizer_model', type=str, required=True, help='Path to tokenizer model')
+    parser.add_argument('--tokenizer_name', type=str, default='yttm', help='BPE Tokenizer Name, Options: [yttm]')
+    parser.add_argument('--bpe_droput', type=float, default=0.0, help='BPE dropout to use')
+    parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
+    parser.add_argument('--pkl_file_prefix', type=str, default='parallel', help='Prefix for tar and pickle files')
+    parser.add_argument('--fname', type=str, required=True, help='Path to monolingual data file')
+    parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
+    parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
+    parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
+    parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
+    parser.add_argument(
+        '--lines_per_dataset_fragment',
+        type=int,
+        default=1000000,
+        help='Number of lines to consider for bucketing and padding',
+    )
+    parser.add_argument(
+        '--num_batches_per_tarfile',
+        type=int,
+        default=1000,
+        help='Number of batches (pickle files) within each tarfile',
+    )
+
+    args = parser.parse_args()
+    if not os.path.exists(args.out_dir):
+        os.mkdir(args.out_dir)
+    if not os.path.exists(args.tokenizer_model):
+        assert FileNotFoundError("Could not find tokenizer model %s" % (args.tokenizer))
+
+    tokenizer_model = MTDataPreproc.get_monolingual_tokenizer(
+        tokenizer_name=args.tokenizer_name, tokenizer_model=args.tokenizer_model, bpe_dropout=args.bpe_droput
+    )
+
+    MTDataPreproc.preprocess_monolingual_dataset(
+        clean=args.clean,
+        fname=args.fname,
+        out_dir=args.out_dir,
+        tokenizer=tokenizer_model,
+        max_seq_length=args.max_seq_length,
+        min_seq_length=args.min_seq_length,
+        tokens_in_batch=args.tokens_in_batch,
+        lines_per_dataset_fragment=args.lines_per_dataset_fragment,
+        num_batches_per_tarfile=args.num_batches_per_tarfile,
+        pkl_file_prefix=args.pkl_file_prefix,
+        global_rank=0,
+        world_size=1,
+    )

tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2020, 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 argparse
+import os
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
+    parser.add_argument('--shared_tokenizer', action="store_true", help='Whether to share encoder/decoder tokenizers')
+    parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
+    parser.add_argument('--tar_file_prefix', type=str, default='parallel', help='Prefix for tar files')
+    parser.add_argument('--bpe_dropout', type=float, default=0.1, help='Whether to share encoder/decoder tokenizers')
+    parser.add_argument('--src_fname', type=str, required=True, help='Path to the source file')
+    parser.add_argument('--tgt_fname', type=str, required=True, help='Path to the target file')
+    parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
+    parser.add_argument('--encoder_model_name', type=str, default=None, help='For use with pretrained encoders')
+    parser.add_argument(
+        '--decoder_model_name', type=str, default=None, help='For use with pretrained decoders (not yet supported)'
+    )
+    parser.add_argument(
+        '--encoder_tokenizer_model', type=str, default='None', help='Path to pre-trained encoder tokenizer model'
+    )
+    parser.add_argument(
+        '--encoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
+    )
+    parser.add_argument('--encoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
+    parser.add_argument(
+        '--encoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
+    )
+    parser.add_argument('--encoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
+    parser.add_argument(
+        '--decoder_tokenizer_model', type=str, default='None', help='Path to pre-trained decoder tokenizer model'
+    )
+    parser.add_argument(
+        '--decoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
+    )
+    parser.add_argument('--decoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
+    parser.add_argument(
+        '--decoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
+    )
+    parser.add_argument('--decoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
+    parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
+    parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
+    parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
+    parser.add_argument('--coverage', type=float, default=0.999, help='BPE character coverage [0-1]')
+    parser.add_argument(
+        '--lines_per_dataset_fragment',
+        type=int,
+        default=1000000,
+        help='Number of lines to consider for bucketing and padding',
+    )
+    parser.add_argument(
+        '--num_batches_per_tarfile',
+        type=int,
+        default=1000,
+        help='Number of batches (pickle files) within each tarfile',
+    )
+    parser.add_argument(
+        '--n_preproc_jobs', type=int, default=-2, help='Number of processes to use for creating the tarred dataset.',
+    )
+
+    args = parser.parse_args()
+    if not os.path.exists(args.out_dir):
+        os.mkdir(args.out_dir)
+
+    if (
+        args.encoder_tokenizer_model != 'None'
+        and args.decoder_tokenizer_model == 'None'
+        or args.decoder_tokenizer_model != 'None'
+        and args.encoder_tokenizer_model == 'None'
+    ):
+        if args.shared_tokenizer:
+            raise ValueError(
+                '''
+                If using a pre-trained shared tokenizer,
+                both encoder and decoder tokenizers must be the same
+                '''
+            )
+        else:
+            raise ValueError('Both encoder and decoder pre-trained tokenizer models must be specified')
+
+    if args.encoder_tokenizer_model == 'None' and args.decoder_tokenizer_model == 'None':
+        encoder_tokenizer_model, decoder_tokenizer_model = MTDataPreproc.train_tokenizers(
+            out_dir=args.out_dir,
+            src_fname=args.src_fname,
+            tgt_fname=args.tgt_fname,
+            shared_tokenizer=args.shared_tokenizer,
+            encoder_tokenizer_name=args.encoder_tokenizer_name,
+            encoder_tokenizer_vocab_size=args.encoder_tokenizer_vocab_size,
+            encoder_tokenizer_coverage=args.encoder_tokenizer_coverage,
+            decoder_tokenizer_name=args.decoder_tokenizer_name,
+            decoder_tokenizer_vocab_size=args.decoder_tokenizer_vocab_size,
+            decoder_tokenizer_coverage=args.decoder_tokenizer_coverage,
+            global_rank=0,
+        )
+    else:
+        encoder_tokenizer_model, decoder_tokenizer_model = args.encoder_tokenizer_model, args.decoder_tokenizer_model
+
+    encoder_tokenizer, decoder_tokenizer = MTDataPreproc.get_enc_dec_tokenizers(
+        encoder_tokenizer_name=args.encoder_tokenizer_name,
+        encoder_tokenizer_model=encoder_tokenizer_model,
+        encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
+        decoder_tokenizer_name=args.decoder_tokenizer_name,
+        decoder_tokenizer_model=decoder_tokenizer_model,
+        decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
+    )
+
+    _, _ = MTDataPreproc.preprocess_parallel_dataset(
+        clean=args.clean,
+        src_fname=args.src_fname,
+        tgt_fname=args.tgt_fname,
+        out_dir=args.out_dir,
+        encoder_tokenizer_name=args.encoder_tokenizer_name,
+        encoder_model_name=args.encoder_model_name,
+        encoder_tokenizer_model=encoder_tokenizer_model,
+        encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
+        decoder_tokenizer_name=args.decoder_tokenizer_name,
+        decoder_model_name=args.decoder_model_name,
+        decoder_tokenizer_model=decoder_tokenizer_model,
+        decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
+        max_seq_length=args.max_seq_length,
+        min_seq_length=args.min_seq_length,
+        tokens_in_batch=args.tokens_in_batch,
+        lines_per_dataset_fragment=args.lines_per_dataset_fragment,
+        num_batches_per_tarfile=args.num_batches_per_tarfile,
+        tar_file_prefix=args.tar_file_prefix,
+        global_rank=0,
+        world_size=1,
+        n_jobs=args.n_preproc_jobs,
+    )

tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2020, 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.
+
+from dataclasses import dataclass
+from typing import Optional
+
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTEncDecModelConfig
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
+from nemo.core.config import hydra_runner
+from nemo.core.config.modelPT import NemoConfig
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.utils import logging
+from nemo.utils.config_utils import update_model_config
+from nemo.utils.exp_manager import ExpManagerConfig, exp_manager
+
+
+"""
+Usage:
+ 1. If you need to start docker and install NeMo, otherwise skip this step:
+ 
+    a. ```docker run --gpus all -it --rm -v /home/okuchaiev/repos/NeMo/:/NeMo -p 6006:6006  -v /mnt:/mnt --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:20.11-py3```
+    b. ```cd /NeMo```
+    c. ```./reinstall.sh```
+ 
+ 2. Train a new tokenizer (or use pre-trained one):
+    ```yttm bpe --data /mnt/D1/Data/NMT/wmt16_de_en/train.clean.en-de.shuffled.common --model tokenizer.BPE.8192.model --vocab_size 8192```
+
+(To use WANDB, optionally, do login first)
+``wandb login [YOUR WANDB login]``
+    
+ 3. Start training:
+ 
+
+ (This example for "base" model on 2 GPUs for 150000 steps with batch size of 12500 tokens per GPU)
+ 
+ python enc_dec_nmt.py \
+      --config-path=conf \
+      --config-name=aayn_base \
+      trainer.gpus=[0,1] \
+      ~trainer.max_epochs \
+      +trainer.max_steps=150000 \
+      model.beam_size=4 \
+      model.max_generation_delta=5 \
+      model.label_smoothing=0.1 \
+      model.encoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model  \
+      model.decoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model  \
+      model.encoder.num_layers=6 \
+      model.encoder.hidden_size=512 \
+      model.encoder.inner_size=2048 \
+      model.encoder.num_attention_heads=8 \
+      model.encoder.ffn_dropout=0.1 \
+      model.decoder.num_layers=6 \
+      model.decoder.hidden_size=512 \
+      model.decoder.inner_size=2048 \
+      model.decoder.num_attention_heads=8 \
+      model.decoder.ffn_dropout=0.1 \
+      model.train_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.de.shuffled \
+      model.train_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.en.shuffled \
+      model.train_ds.tokens_in_batch=12500 \
+      model.validation_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
+      model.validation_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
+      model.validation_ds.tokens_in_batch=8192 \
+      model.test_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
+      model.test_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
+      model.optim.lr=0.001  \
+      model.optim.sched.warmup_ratio=0.05 \
+      +exp_manager.create_wandb_logger=True \
+      +exp_manager.wandb_logger_kwargs.name=TEST-nmt-base \
+      +exp_manager.wandb_logger_kwargs.project=nmt-de-en \
+      +exp_manager.create_checkpoint_callback=True \
+      +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \
+      +exp_manager.exp_dir=nmt_base \
+      +exp_manager.checkpoint_callback_params.mode=max 
+"""
+
+
+@dataclass
+class MTEncDecConfig(NemoConfig):
+    name: Optional[str] = 'MTEncDec'
+    do_training: bool = True
+    do_testing: bool = False
+    model: MTEncDecModelConfig = MTEncDecModelConfig()
+    trainer: Optional[TrainerConfig] = TrainerConfig()
+    exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTEncDec', files_to_copy=[])
+
+
+@hydra_runner(config_path="conf", config_name="aayn_base")
+def main(cfg: MTEncDecConfig) -> None:
+    # merge default config with user specified config
+    default_cfg = MTEncDecConfig()
+    cfg = update_model_config(default_cfg, cfg)
+    logging.info("\n\n************** Experiment configuration ***********")
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+
+    # training is managed by PyTorch Lightning
+    trainer = Trainer(**cfg.trainer)
+
+    # tokenizers will be trained and and tarred training data will be created if needed
+    # model config is then updated
+    if cfg.model.preproc_out_dir is not None:
+        MTDataPreproc(cfg=cfg.model, trainer=trainer)
+
+    # experiment logs, checkpoints, and auto-resume are managed by exp_manager and PyTorch Lightning
+    exp_manager(trainer, cfg.exp_manager)
+
+    # everything needed to train translation models is encapsulated in the NeMo MTEncdDecModel
+    mt_model = MTEncDecModel(cfg.model, trainer=trainer)
+
+    logging.info("\n\n************** Model parameters and their sizes ***********")
+    for name, param in mt_model.named_parameters():
+        print(name, param.size())
+    logging.info("***********************************************************\n\n")
+
+    if cfg.do_training:
+        trainer.fit(mt_model)
+
+    if cfg.do_testing:
+        trainer.test(mt_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2020, 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.
+
+"""
+Given NMT model's .nemo file, this script can be used to translate text.
+USAGE Example:
+1. Obtain text file in src language. You can use sacrebleu to obtain standard test sets like so:
+    sacrebleu -t wmt14 -l de-en --echo src > wmt14-de-en.src
+2. Translate:
+    python nmt_transformer_infer.py --model=[Path to .nemo file] --srctext=wmt14-de-en.src --tgtout=wmt14-de-en.pre
+"""
+
+
+from argparse import ArgumentParser
+
+import torch
+
+import nemo.collections.nlp as nemo_nlp
+from nemo.collections.nlp.modules.common.transformer import (
+    BeamSearchSequenceGenerator,
+    BeamSearchSequenceGeneratorWithLanguageModel,
+)
+from nemo.utils import logging
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--model", type=str, required=True, help="")
+    parser.add_argument("--srctext", type=str, required=True, help="")
+    parser.add_argument("--tgtout", type=str, required=True, help="")
+    parser.add_argument("--batch_size", type=int, default=256, help="")
+    parser.add_argument("--beam_size", type=int, default=4, help="")
+    parser.add_argument("--len_pen", type=float, default=0.6, help="")
+    parser.add_argument("--max_delta_length", type=int, default=5, help="")
+    parser.add_argument("--target_lang", type=str, default=None, help="")
+    parser.add_argument("--source_lang", type=str, default=None, help="")
+    # shallow fusion specific parameters
+    parser.add_argument("--lm_model", type=str, default=None, help="")
+    parser.add_argument("--fusion_coef", type=float, default=0.0, help="")
+
+    args = parser.parse_args()
+    torch.set_grad_enabled(False)
+    if args.model.endswith(".nemo"):
+        logging.info("Attempting to initialize from .nemo file")
+        model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=args.model)
+        src_text = []
+        tgt_text = []
+    else:
+        raise NotImplemented(f"Only support .nemo files, but got: {args.model}")
+
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    if args.lm_model is not None:
+        lm_model = nemo_nlp.models.language_modeling.TransformerLMModel.restore_from(restore_path=args.lm_model).eval()
+        model.beam_search = BeamSearchSequenceGeneratorWithLanguageModel(
+            embedding=model.decoder.embedding,
+            decoder=model.decoder.decoder,
+            log_softmax=model.log_softmax,
+            bos=model.decoder_tokenizer.bos_id,
+            pad=model.decoder_tokenizer.pad_id,
+            eos=model.decoder_tokenizer.eos_id,
+            language_model=lm_model,
+            fusion_coef=args.fusion_coef,
+            max_sequence_length=model.decoder.max_sequence_length,
+            beam_size=args.beam_size,
+            len_pen=args.len_pen,
+            max_delta_length=args.max_delta_length,
+        )
+    else:
+        model.beam_search = BeamSearchSequenceGenerator(
+            embedding=model.decoder.embedding,
+            decoder=model.decoder.decoder,
+            log_softmax=model.log_softmax,
+            bos=model.decoder_tokenizer.bos_id,
+            pad=model.decoder_tokenizer.pad_id,
+            eos=model.decoder_tokenizer.eos_id,
+            max_sequence_length=model.decoder.max_sequence_length,
+            beam_size=args.beam_size,
+            len_pen=args.len_pen,
+            max_delta_length=args.max_delta_length,
+        )
+
+    logging.info(f"Translating: {args.srctext}")
+
+    count = 0
+    with open(args.srctext, 'r') as src_f:
+        for line in src_f:
+            src_text.append(line.strip())
+            if len(src_text) == args.batch_size:
+                res = model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
+                if len(res) != len(src_text):
+                    print(len(res))
+                    print(len(src_text))
+                    print(res)
+                    print(src_text)
+                tgt_text += res
+                src_text = []
+            count += 1
+            # if count % 300 == 0:
+            #    print(f"Translated {count} sentences")
+        if len(src_text) > 0:
+            tgt_text += model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
+
+    with open(args.tgtout, 'w') as tgt_f:
+        for line in tgt_text:
+            tgt_f.write(line + "\n")
+
+
+if __name__ == '__main__':
+    main()  # noqa pylint: disable=no-value-for-parameter

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst

@@ -0,0 +1,7 @@
+**STEPS 2 RUN**
+===============
+0. (Make sure you have Flask installed - ``pip install -r requirements.txt``)
+1. Edit "config.json" file to only contain models you need. If model's location starts with "NGC/" - it will load this model from NVIDIA's NGC. Otherwise, specify full path to .nemo file.
+4. To run: ``python nmt_service.py``
+5. To translate: ``http://127.0.0.1:5000/translate?text=Frohe%20Weihnachten`` (here %20 means space)
+6. Run web UI: ``python -m http.server``

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json

@@ -0,0 +1,12 @@
+{
+  "en-de": "NGC/nmt_en_de_transformer12x2",
+  "de-en": "NGC/nmt_de_en_transformer12x2",
+  "en-es": "NGC/nmt_en_es_transformer12x2",
+  "es-en": "NGC/nmt_es_en_transformer12x2",
+  "en-ru": "NGC/nmt_en_ru_transformer6x6",
+  "ru-en": "NGC/nmt_ru_en_transformer6x6",
+  "en-fr": "NGC/nmt_en_fr_transformer12x2",
+  "fr-en": "NGC/nmt_fr_en_transformer12x2",
+  "en-zh": "NGC/nmt_en_zh_transformer6x6",
+  "zh-en": "NGC/nmt_zh_en_transformer6x6"
+}

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html

@@ -0,0 +1,140 @@
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>Neural Machine Translation Demo App</title>
+  <link rel="stylesheet" href="//code.jquery.com/ui/1.12.1/themes/base/jquery-ui.css">  
+  <link rel="stylesheet" href="/resources/demos/style.css">
+  <style>
+    .ui-resizable-se {
+      bottom: 17px;
+    }
+  </style>
+  <link rel="stylesheet" href="style.css">
+  <script src="https://code.jquery.com/jquery-1.12.4.js"></script>
+  <script src="https://code.jquery.com/ui/1.12.1/jquery-ui.js"></script>
+  <script>
+  $( function() {
+    $( "#tgt_text" ).resizable({
+      handles: "se"
+    });
+  } );
+  $( function() {
+    $( "#src_text" ).resizable({
+      handles: "se"
+    });
+  } );
+  $( function() {
+    $( "#translate_button" ).button();
+    $( "#translate_button" ).click( function( event ) {
+        $.ajax({
+          url: "http://10.110.40.187:5000/translate",
+          type: "GET",                    
+          data: {
+            text: $( "#src_text" ).val(),
+            langpair: $('input[type=radio][name=langpair]:checked').attr('id'),
+            do_moses: true
+          },
+          success: function (result) {            
+            $( "#tgt_text" ).val(result.translation);
+          },
+          error: function (err) {
+            alert('Sorry, can not reach translation service');            
+          },
+        })
+    } );
+  } );
+  </script>
+</head>
+<body>
+<div class="contact-us">
+  <h1>Neural Machine Translation Demo</h1>
+  <p>
+  <br/>
+  <h4>Running NeMo commit id: v1.0.2  </h4>
+  <br/>
+  </p>
+  <p>
+  <fieldset>
+    <legend>Select translation direction: </legend>
+    <table>
+      <tr>
+        <td>
+          <label for="en-de">English to German</label>
+          <input type="radio" name="langpair" id="en-de">
+        </td>
+        <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+        <td>
+          <label for="de-en">German to English</label>
+          <input type="radio" name="langpair" id="de-en">
+        </td>
+        <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+        <td>
+          <label for="es-en">Spanish to English</label>
+          <input type="radio" name="langpair" id="es-en">
+        </td>
+	<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+	<td>
+	  <label for="en-es" selected="selected">English to Spanish</label>
+	  <input type="radio" name="langpair" id="en-es" checked>
+  </td> 
+      </tr>
+      <tr>
+         <td>
+           <label for="en-ru">English to Russian</label>
+           <input type="radio" name="langpair" id="en-ru">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+           <label for="ru-en">Russian to English</label>
+           <input type="radio" name="langpair" id="ru-en">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+           <label for="en-fr">English to French</label>
+	   <input type="radio" name="langpair" id="en-fr">
+	 </td>
+	 <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+	 <td>
+	  <label for="fr-en">French to English</label>
+          <input type="radio" name="langpair" id="fr-en">
+         </td>
+         <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+          <label for="en-zh">English to Chinese</label>
+          <input type="radio" name="langpair" id="en-zh">
+         </td>
+	 <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
+         <td>
+          <label for="zh-en">Chinese to English</label>
+          <input type="radio" name="langpair" id="zh-en">
+         </td>
+      </tr>
+    </table>
+  </fieldset>
+  <div>
+    <table>
+      <tr>
+        <td>
+          <textarea id="src_text" rows="12" cols="60">Type something here, chose language pair and click "Translate" button.</textarea>
+        </td>
+      </tr>
+      <tr>
+        <td>
+          <button id="translate_button">Translate !</button>
+        </td>
+      </tr>
+      <tr>
+        <td>
+          <textarea id="tgt_text" rows="12" cols="60" readonly></textarea>
+        </td>
+      </tr>
+    </table>
+  </div>
+  </p>
+</div>
+ 
+ 
+</body>
+</html>

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2020, 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 json
+import time
+
+import flask
+import torch
+from flask import Flask, json, request
+from flask_cors import CORS
+
+import nemo.collections.nlp as nemo_nlp
+from nemo.utils import logging
+
+MODELS_DICT = {}
+
+model = None
+api = Flask(__name__)
+CORS(api)
+
+
+def initialize(config_file_path: str):
+    """
+    Loads 'language-pair to NMT model mapping'
+    """
+    __MODELS_DICT = None
+
+    logging.info("Starting NMT service")
+    logging.info(f"I will attempt to load all the models listed in {config_file_path}.")
+    logging.info(f"Edit {config_file_path} to disable models you don't need.")
+    if torch.cuda.is_available():
+        logging.info("CUDA is available. Running on GPU")
+    else:
+        logging.info("CUDA is not available. Defaulting to CPUs")
+
+    # read config
+    with open(config_file_path) as f:
+        __MODELS_DICT = json.load(f)
+
+    if __MODELS_DICT is not None:
+        for key, value in __MODELS_DICT.items():
+            logging.info(f"Loading model for {key} from file: {value}")
+            if value.startswith("NGC/"):
+                model = nemo_nlp.models.machine_translation.MTEncDecModel.from_pretrained(model_name=value[4:])
+            else:
+                model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=value)
+            if torch.cuda.is_available():
+                model = model.cuda()
+            MODELS_DICT[key] = model
+    else:
+        raise ValueError("Did not find the config.json or it was empty")
+    logging.info("NMT service started")
+
+
+@api.route('/translate', methods=['GET', 'POST', 'OPTIONS'])
+def get_translation():
+    try:
+        time_s = time.time()
+        langpair = request.args["langpair"]
+        src = request.args["text"]
+        do_moses = request.args.get('do_moses', False)
+        if langpair in MODELS_DICT:
+            if do_moses:
+                result = MODELS_DICT[langpair].translate(
+                    [src], source_lang=langpair.split('-')[0], target_lang=langpair.split('-')[1]
+                )
+            else:
+                result = MODELS_DICT[langpair].translate([src])
+
+            duration = time.time() - time_s
+            logging.info(
+                f"Translated in {duration}. Input was: {request.args['text']} <############> Translation was: {result[0]}"
+            )
+            res = {'translation': result[0]}
+            response = flask.jsonify(res)
+            response.headers.add('Access-Control-Allow-Origin', '*')
+            return response
+
+        else:
+            logging.error(f"Got the following langpair: {langpair} which was not found")
+    except Exception as ex:
+        res = {'translation': str(ex)}
+        response = flask.jsonify(res)
+        response.headers.add('Access-Control-Allow-Origin', '*')
+        return res
+
+
+if __name__ == '__main__':
+    initialize('config.json')
+    api.run(host='0.0.0.0')

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt

@@ -0,0 +1,3 @@
+nemo_toolkit[all]>=1.0.0rc1
+flask
+flask_cors

tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css

@@ -0,0 +1,111 @@
+@import url("https://fonts.googleapis.com/css?family=Spartan&display=swap");
+* {
+  margin: 0;
+  padding: 0;
+}
+
+body {
+  /*background: url("https://img.wallpapersafari.com/desktop/1680/1050/65/18/8AYkSt.jpg") center center;*/
+    background: url("https://images5.alphacoders.com/550/550948.jpg") center center;
+  background-size: cover;
+  width: 100vw;
+  height: 100vh;
+  display: grid;
+  align-items: center;
+  justify-items: center;
+  font-size: 12pt;
+  font-family: "Spartan";
+  color: #ffffff;
+}
+
+.contact-us {
+  /*background: #1f1d1d;*/
+  padding: 50px 100px;
+  border-top: 10px solid #229955;
+}
+
+label, input, textarea {
+  display: block;
+  width: 100%;
+  font-size: 12pt;
+  line-height: 24pt;
+  font-family: "Spartan";
+}
+
+input {
+  margin-bottom: 24pt;
+}
+
+h1 {
+  color: #ffffff;
+  font-weight: bold;
+  font-size: 22pt;
+  line-height: 24pt;
+  font-style: italic;
+  margin: 0 0 0.5em 0;
+}
+
+span {
+  font-size: 8pt;
+}
+
+em {
+  color: #229955;
+  font-weight: bold;
+}
+
+input, textarea {
+  border: none;
+  border: 1px solid rgba(0, 0, 0, 0.1);
+  border-radius: 2px;
+  background: #f8f4e5;
+  padding-left: 5px;
+  outline: none;
+}
+
+input:focus, textarea:focus {
+  border: 1px solid #6bd4b1;
+}
+
+textarea {
+  resize: none;
+}
+
+fieldset {
+      overflow: hidden
+    }
+
+button {
+  display: block;
+  float: left;
+  line-height: 24pt;
+  padding: 0 20px;
+  border: none;
+  background: #f45702;
+  color: white;
+  letter-spacing: 2px;
+  transition: 0.2s all ease-in-out;
+  border-bottom: 2px solid transparent;
+  outline: none;
+}
+button:hover {
+  background: inherit;
+  color: #ea6153;
+  border-bottom: 2px solid #f45702;
+}
+
+::selection {
+  background: #ffc7b8;
+}
+
+input:-webkit-autofill,
+input:-webkit-autofill:hover,
+input:-webkit-autofill:focus,
+textarea:-webkit-autofill,
+textarea:-webkit-autofill:hover,
+textarea:-webkit-autofill:focus {
+  border: 1px solid #6bd4b1;
+  -webkit-text-fill-color: #2A293E;
+  -webkit-box-shadow: 0 0 0px 1000px #f8f4e5 inset;
+  transition: background-color 5000s ease-in-out 0s;
+}

tts/NeMo/examples/nlp/machine_translation/translate_ddp.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2020, 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
+from argparse import ArgumentParser
+
+import torch
+import torch.multiprocessing as mp
+from torch.utils.data import DataLoader
+
+from nemo.collections.nlp.data.language_modeling import TarredSentenceDataset
+from nemo.collections.nlp.data.machine_translation import TarredTranslationDataset
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
+from nemo.utils import logging
+
+
+def get_args():
+    parser = ArgumentParser(description='Batch translation of sentences from a pre-trained model on multiple GPUs')
+    parser.add_argument("--model", type=str, required=True, help="Path to the .nemo translation model file")
+    parser.add_argument(
+        "--text2translate", type=str, required=True, help="Path to the pre-processed tarfiles for translation"
+    )
+    parser.add_argument("--result_dir", type=str, required=True, help="Folder to write translation results")
+    parser.add_argument(
+        "--twoside", action="store_true", help="Set flag when translating the source side of a parallel dataset"
+    )
+    parser.add_argument(
+        '--metadata_path', type=str, required=True, help="Path to the JSON file that contains dataset info"
+    )
+    parser.add_argument('--topk', type=int, default=500, help="Value of k for topk sampling")
+    parser.add_argument('--src_language', type=str, required=True, help="Source lang ID for detokenization")
+    parser.add_argument('--tgt_language', type=str, required=True, help="Target lang ID for detokenization")
+    parser.add_argument(
+        '--reverse_lang_direction',
+        action="store_true",
+        help="Reverse source and target language direction for parallel dataset",
+    )
+    parser.add_argument('--n_gpus', type=int, default=-1, help="Number of GPUs to use")
+    args = parser.parse_args()
+    return args
+
+
+def translate(rank, world_size, args):
+    if args.model.endswith(".nemo"):
+        logging.info("Attempting to initialize from .nemo file")
+        model = MTEncDecModel.restore_from(restore_path=args.model, map_location=f"cuda:{rank}")
+    elif args.model.endswith(".ckpt"):
+        logging.info("Attempting to initialize from .ckpt file")
+        model = MTEncDecModel.load_from_checkpoint(checkpoint_path=args.model, map_location=f"cuda:{rank}")
+    model.replace_beam_with_sampling(topk=args.topk)
+    model.eval()
+    if args.twoside:
+        dataset = TarredTranslationDataset(
+            text_tar_filepaths=args.text2translate,
+            metadata_path=args.metadata_path,
+            encoder_tokenizer=model.encoder_tokenizer,
+            decoder_tokenizer=model.decoder_tokenizer,
+            shuffle_n=100,
+            shard_strategy="scatter",
+            world_size=world_size,
+            global_rank=rank,
+            reverse_lang_direction=args.reverse_lang_direction,
+        )
+    else:
+        dataset = TarredSentenceDataset(
+            text_tar_filepaths=args.text2translate,
+            metadata_path=args.metadata_path,
+            tokenizer=model.encoder_tokenizer,
+            shuffle_n=100,
+            shard_strategy="scatter",
+            world_size=world_size,
+            global_rank=rank,
+        )
+    loader = DataLoader(dataset, batch_size=1)
+    result_dir = os.path.join(args.result_dir, f'rank{rank}')
+    os.makedirs(result_dir, exist_ok=True)
+    originals_file_name = os.path.join(result_dir, 'originals.txt')
+    translations_file_name = os.path.join(result_dir, 'translations.txt')
+    num_translated_sentences = 0
+
+    with open(originals_file_name, 'w') as of, open(translations_file_name, 'w') as tf:
+        for batch_idx, batch in enumerate(loader):
+            for i in range(len(batch)):
+                if batch[i].ndim == 3:
+                    batch[i] = batch[i].squeeze(dim=0)
+                batch[i] = batch[i].to(rank)
+            if args.twoside:
+                src_ids, src_mask, _, _, _ = batch
+            else:
+                src_ids, src_mask = batch
+            if batch_idx % 100 == 0:
+                logging.info(
+                    f"{batch_idx} batches ({num_translated_sentences} sentences) were translated by process with "
+                    f"rank {rank}"
+                )
+            num_translated_sentences += len(src_ids)
+            inputs, translations = model.batch_translate(src=src_ids, src_mask=src_mask)
+            for src, translation in zip(inputs, translations):
+                of.write(src + '\n')
+                tf.write(translation + '\n')
+
+
+def main() -> None:
+    args = get_args()
+    world_size = torch.cuda.device_count() if args.n_gpus == -1 else args.n_gpus
+    mp.spawn(translate, args=(world_size, args), nprocs=world_size, join=True)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml

@@ -0,0 +1,135 @@
+# Question Answering with SQUAD
+name: &name QA
+
+pretrained_model: null # pretrained QAModel model from list_available_models()
+do_training: true # training mode, for testing change to false
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+  num_sanity_val_steps: 0
+  log_every_n_steps: 1  # Interval of logging.
+
+model:
+  nemo_path: null # exported .nemo path
+  dataset:
+    version_2_with_negative: false
+    # If true, the examples contain some that do not have an answer.
+    doc_stride: 128
+    # When splitting up a long document into chunks,
+    # how much stride to take between chunks.
+    max_query_length: 64
+    # The maximum number of tokens for the question.
+    # Questions longer than this will be truncated to
+    # this length.
+    max_seq_length: 384
+    # The maximum total input sequence length after
+    # WordPiece tokenization. Sequences longer than this
+    # will be truncated, and sequences shorter than this
+    # will be padded.
+    max_answer_length: 30
+    # The maximum length of an answer that can be
+    # generated. This is needed because the start
+    # and end predictions are not conditioned
+    # on one another.
+    null_score_diff_threshold: 0.0
+    # If null_score - best_non_null is greater than the threshold predict null.
+    n_best_size: 20
+    # The total number of n-best predictions to generate at testing.
+    use_cache: true
+    do_lower_case: true
+    # if true does lower case
+
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: true
+    num_samples: -1
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+
+  validation_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: false
+    num_samples: -1    
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+
+  test_ds:
+    file: null # .json file
+    batch_size: 24 # per GPU
+    shuffle: false
+    num_samples: -1
+    # Default values for the following params are retrieved from dataset config section, but you may override them
+    num_workers: ${model.dataset.num_workers}
+    drop_last: ${model.dataset.drop_last}
+    pin_memory: ${model.dataset.pin_memory}
+    
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null # expand the following to a dictionary if special tokens need to be added.
+    #  only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently.
+
+  language_model:
+    pretrained_model_name: bert-base-uncased # BERT-like model name
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null # if specified initializes model from scratch
+
+  token_classifier:
+    num_layers: 1
+    dropout: 0.0
+    num_classes: 2
+    activation: relu
+    log_softmax: false
+    use_transformer_init: true
+
+
+  optim:
+    name: adamw
+    lr: 3e-5
+    weight_decay: 0.0
+    sched:
+      name: SquareRootAnnealing
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+      # scheduler config override
+      warmup_steps: null
+      warmup_ratio: 0.0
+      last_epoch: -1
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/question_answering/get_squad.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, 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 argparse
+import os
+import urllib.request
+
+from nemo.utils import logging
+
+
+class SquadDownloader:
+    def __init__(self, save_path):
+        self.save_path = save_path + '/squad'
+
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+
+        if not os.path.exists(self.save_path + '/v1.1'):
+            os.makedirs(self.save_path + '/v1.1')
+
+        if not os.path.exists(self.save_path + '/v2.0'):
+            os.makedirs(self.save_path + '/v2.0')
+
+        self.download_urls = {
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v1.1.json': 'v1.1/train-v1.1.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v1.1.json': 'v1.1/dev-v1.1.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v2.0.json': 'v2.0/train-v2.0.json',
+            'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v2.0.json': 'v2.0/dev-v2.0.json',
+        }
+
+    def download(self):
+        for item in self.download_urls:
+            url = item
+            file = self.download_urls[item]
+
+            logging.info('Downloading: %s', url)
+            if os.path.isfile(self.save_path + '/' + file):
+                logging.info('** Download file already exists, skipping download')
+            else:
+                response = urllib.request.urlopen(url)
+                with open(self.save_path + '/' + file, "wb") as handle:
+                    handle.write(response.read())
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Download Squad')
+    parser.add_argument(
+        '--destDir',
+        type=str,
+        required=False,
+        help='directory to store data',
+        default=os.path.split(os.path.abspath(__file__))[0],
+    )
+    args = parser.parse_args()
+    logging.info(args.destDir)
+    squad_dl = SquadDownloader(args.destDir)
+    squad_dl.download()

tts/NeMo/examples/nlp/question_answering/question_answering_squad.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020, 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.
+
+
+"""
+This script contains an example on how to train, evaluate and perform inference with the question answering model.
+The QAModel in NeMo supports extractive question answering problems for data in the SQuAD (https://rajpurkar.github.io/SQuAD-explorer/) format.
+
+***Data format***
+The QAModel requires a JSON file for each dataset split. 
+In the following we will show example for a training file. Each title has one or multiple paragraph entries, each consisting of the text - "context", and question-answer entries. Each question-answer entry has:
+* a question
+* a globally unique id
+* a boolean flag "is_impossible" which shows if the question is answerable or not
+* in case the question is answerable one answer entry, which contains the text span and its starting character index in the context. If not answerable, the "answers" list is empty
+
+The evaluation file follows the above format except for it can provide more than one answers to the same question. 
+The inference file follows the above format except for it does not require the "answers" and "is_impossible" keywords.
+
+
+***Downloading the dataset***
+Run ./NeMo/examples/nlp/question_answering/get_squad.py to download the SQuAD dataset:
+
+#   python get_squad.py --destDir=<PATH_TO_DATA>
+
+***Setting the configs***
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+
+This script uses the `/examples/nlp/question_answering/conf/question_answering_squad_config.yaml` config file
+by default. You may update the config file from the file directly. The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+
+***Model Training***
+# python question_answering_squad.py
+    model.train_ds.file=<TRAIN_JSON_FILE>
+    model.validation_ds=<VAL_JSON_FILE>
+    trainer.max_epochs=<NUM_EPOCHS>
+    trainer.gpus=[<CHANGE_TO_GPU_YOU_WANT_TO_USE>]
+
+
+***Model Evaluation***
+Set `do_training=False` in the script and run:
+
+#   python question_answering_squad.py
+    model.test_file=<TEST_JSON_FILE>
+
+To load a pretrained checkpoint from cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=<MODEL_NAME>,
+e.g. MODEL_NAME='BERTBaseUncasedSQuADv1.1'. You can find all pretrained model names by using 
+QAModel.list_available_models(). To load a local checkpoint use qa_model.restore_from(<PATH_TO_CHECKPOINT>)
+
+
+***Model Inference***
+For inference use 
+    qa_model.inference(
+        file=<INFERENCE_JSON_FILE>,
+        batch_size=<BATCH_SIZE>,
+        output_nbest_file=<OPTIONAL_OUTPUT_FILE_FOR_NBEST_LIST>,
+        output_prediction_file=<OPTIONAL_OUTPUT_FILE_FOR_PREDICTION>
+    )
+
+More details on how to use this script can be found in
+./NeMo/tutorials/nlp/Question_Answering_Squad.ipynb
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.question_answering.qa_model import QAModel
+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="question_answering_squad_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.pretrained_model:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = QAModel(cfg.model, trainer=trainer)
+    else:
+        logging.info(f'Loading pretrained model {cfg.pretrained_model}')
+        model = QAModel.from_pretrained(cfg.pretrained_model)
+        if cfg.do_training:
+            model.setup_training_data(train_data_config=cfg.model.train_ds)
+            model.setup_validation_data(val_data_config=cfg.model.validation_ds)
+
+    if cfg.do_training:
+        trainer.fit(model)
+        if cfg.model.nemo_path:
+            model.save_to(cfg.model.nemo_path)
+
+    if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.file is not None:
+        trainer.test(model)
+
+    # change to path if you want results to be written to file e.g. os.path.join(exp_dir, "output_nbest_file.txt")
+    output_nbest_file = None
+    # change to path if you want results to be written to file e.g.  os.path.join(exp_dir, "output_prediction_file.txt")
+    output_prediction_file = None
+    inference_samples = 5  # for test purposes. To use entire inference dataset set to -1
+    all_preds, all_nbests = model.inference(
+        file=cfg.model.validation_ds.file,
+        batch_size=1,
+        num_samples=inference_samples,
+        output_nbest_file=output_nbest_file,
+        output_prediction_file=output_prediction_file,
+    )
+
+    for _, item in all_preds.items():
+        print(f"question: {item[0]} answer: {item[1]}")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/text2sparql/conf/text2sparql_config.yaml

@@ -0,0 +1,105 @@
+# Text2Sparql with BART
+
+name: &name Text2Sparql
+
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  accelerator: ddp
+  gradient_clip_val: 0.0
+  log_every_n_steps: 1
+  val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch
+  checkpoint_callback: false # provided by exp_manager
+  logger: false # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+  max_seq_length: 150
+  batch_size: 16
+  convert_labels: true # true if Bart, false otherwise (converts pad_id to -100 for masked loss)
+  data_dir: null
+
+  language_model:
+    pretrained_model_name: facebook/bart-base # huggingface end-to-end model name
+    pretrained_encoder_model_name: null # huggingface encoder model name
+    pretrained_decoder_model_name: null # huggingface decoder model name
+    lm_checkpoint: null
+    config: null
+    config_file: null # json file, precedence over config
+
+  encoder_tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+    add_special_tokens: true
+
+  decoder_tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+    add_special_tokens: true
+
+  train_ds:
+    filepath: ${model.data_dir}/train.tsv # path to data file
+    shuffle: true
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    filepath: ${model.data_dir}/test_easy.tsv # path to data file
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    filepath: ${model.data_dir}/test_hard.tsv # path to data file
+    shuffle: false
+    num_samples: -1
+    num_workers: 2
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adamw
+    lr: 4e-5
+    weight_decay: 0.0
+
+    sched:
+      name: CosineAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.06
+      min_lr: 0.0
+      last_epoch: -1
+
+  generate:
+    max_length: ${model.max_seq_length}
+    num_beams: 1
+    length_penalty: 2.0
+    early_stopping: true
+    repetition_penalty: 1.0
+    do_sample: false
+    top_k: null
+    top_p: null
+    num_return_sequences: 1
+
+exp_manager:
+  exp_dir: null # where to store logs and checkpoints
+  name: *name # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/text2sparql/data/import_datasets.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, 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.
+
+"""
+This script downloads Text2Sparql data and processes it into NeMo's neural machine translation dataset format.
+
+Text2Sparql data consists of 3 files which are saved to source_data_dir:
+    - train_queries_v3.tsv
+    - test_easy_queries_v3.tsv
+    - test_hard_queries_v3.tsv
+
+After processing, the script saves them to the target_data_dir as:
+    - train.tsv
+    - test_easy.tsv
+    - test_hard.tsv
+
+
+You may run it with:
+
+python import_datasets \
+    --source_data_dir ./text2sparql_src \
+    --target_data_dir ./text2sparql_tgt
+"""
+
+import argparse
+import csv
+import os
+from urllib.request import Request, urlopen
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import MODE_EXISTS_TMP, if_exist
+from nemo.utils import logging
+
+base_url = "https://m.meetkai.com/public_datasets/knowledge/"
+prefix_map = {
+    "train_queries_v3.tsv": "train.tsv",
+    "test_easy_queries_v3.tsv": "test_easy.tsv",
+    "test_hard_queries_v3.tsv": "test_hard.tsv",
+}
+
+
+def download_text2sparql(infold: str):
+    """Downloads text2sparql train, test_easy, and test_hard data
+
+    Args:
+        infold: save directory path
+    """
+    os.makedirs(infold, exist_ok=True)
+
+    for prefix in prefix_map:
+        url = base_url + prefix
+
+        logging.info(f"Downloading: {url}")
+        if if_exist(infold, [prefix]):
+            logging.info("** Download file already exists, skipping download")
+        else:
+            req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
+            with open(os.path.join(infold, prefix), "wb") as handle:
+                handle.write(urlopen(req, timeout=20).read())
+
+
+def process_text2sparql(infold: str, outfold: str, do_lower_case: bool):
+    """ Process and convert MeetKai's text2sparql datasets to NeMo's neural machine translation format.
+
+    Args:
+        infold: directory path to raw text2sparql data containing
+            train.tsv, test_easy.tsv, test_hard.tsv
+        outfold: output directory path to save formatted data for NeuralMachineTranslationDataset
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        do_lower_case: if true, convert all sentences and labels to lower
+    """
+    logging.info(f"Processing Text2Sparql dataset and storing at: {outfold}")
+
+    os.makedirs(outfold, exist_ok=True)
+
+    dataset_name = "Text2Sparql"
+    for prefix in prefix_map:
+        input_file = os.path.join(infold, prefix)
+        output_file = os.path.join(outfold, prefix_map[prefix])
+
+        if if_exist(outfold, [prefix_map[prefix]]):
+            logging.info(f"** {MODE_EXISTS_TMP.format(prefix_map[prefix], dataset_name, output_file)}")
+            continue
+
+        if not if_exist(infold, [prefix]):
+            logging.info(f"** {prefix} of {dataset_name}" f" is skipped as it was not found")
+            continue
+
+        assert input_file != output_file, "input file cannot equal output file"
+        with open(input_file, "r") as in_file:
+            with open(output_file, "w") as out_file:
+                reader = csv.reader(in_file, delimiter="\t")
+
+                # replace headers
+                out_file.write("sentence\tlabel\n")
+                next(reader)
+
+                for line in reader:
+                    sentence = line[0]
+                    label = line[1]
+                    if do_lower_case:
+                        sentence = sentence.lower()
+                        label = label.lower()
+                    out_file.write(f"{sentence}\t{label}\n")
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo's format")
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help="Path to the folder containing the dataset files"
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str, help="Path to save the processed dataset")
+    parser.add_argument("--do_lower_case", action="store_true")
+    args = parser.parse_args()
+
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+    do_lower_case = args.do_lower_case
+
+    download_text2sparql(infold=source_dir)
+    process_text2sparql(infold=source_dir, outfold=target_dir, do_lower_case=do_lower_case)

tts/NeMo/examples/nlp/text2sparql/evaluate_text2sparql.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, 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.
+
+"""
+This script contains an example on how to evaluate a NeuralMachineTranslationModel.
+To load the example Text2Sparql dataset, please refer to ./data/import_datasets.py.
+To train a model, please refer to text2sparql.py.
+
+
+***Setting the configs***
+This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default.
+You may update the config file from the file directly or by using the command line arguments.
+Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+Please refer to text2sparql.py for detailed instructions on setting the configuration.
+
+
+***How to run the script?***
+- To reload and evaluate the model, run:
+
+python evaluate_text2sparql.py \
+    model.test_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \
+    model.batch_size=16 \
+    model.nemo_path=./NeMo_logs/bart.nemo \
+    exp_manager.exp_dir=./NeMo_logs
+"""
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@hydra_runner(config_path="conf", config_name="text2sparql_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
+    trainer = pl.Trainer(gpus=cfg.trainer.gpus)
+    nmt_model = NeuralMachineTranslationModel.restore_from(restore_path=cfg.model.nemo_path)
+    nmt_model.setup_test_data(cfg.model.test_ds)
+    results = trainer.test(nmt_model)
+
+    with open(cfg.model.test_ds.filepath, "r") as f:
+        lines = f.readlines()
+
+    lines[0] = lines[0].strip() + f"\tpredictions\n"
+    for i, res in enumerate(results[0]["texts"]):
+        lines[i + 1] = lines[i + 1].strip() + f"\t{res}\n"
+
+    savepath = os.path.join(cfg.exp_manager.exp_dir, os.path.basename(cfg.model.test_ds.filepath))
+    with open(savepath, "w") as f:
+        f.writelines(lines)
+        logging.info(f"Predictions saved to {savepath}")
+
+
+if __name__ == "__main__":
+    main()

tts/NeMo/examples/nlp/text2sparql/text2sparql.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, 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.
+
+"""
+This script contains an example on how to train and save a NeuralMachineTranslationModel.
+NeuralMachineTranslationModel in NeMo supports sequence to sequence problems such as language translation
+and text summarization, provided the data follows the format specified below.
+
+
+***Data format***
+NeuralMachineTranslationModel requires the data to be stored in TAB separated files (.tsv) with two columns of
+sentence and label, where the first line is a header of format:
+    sentence[TAB]label
+And each line is of the format:
+    [SENTENCE][TAB][LABEL]
+
+If your dataset is stored in another format, you need to convert it to this format to use a
+NeuralMachineTranslationModel.
+
+
+***Setting the configs***
+This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default.
+You may update the config file from the file directly or by using the command line arguments.
+Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+A NeuralMachineTranslationModel's config file declares multiple import sections. They are:
+    - trainer: Arguments to be passed to PyTorch Lightning.
+    - model: All arguments that relate to the Model - language_model, tokenizers, datasets, optimizer, generate.
+    - exp_manager: Arguments to be passed to NeMo's experiment manager.
+    - hydra: Arguments to be passed to Hydra.
+
+If using text2sparql_config.yaml, you must first update the following fields in the config:
+    - model.nemo_path: Model save path. Eg. [PATH]/bart.nemo
+    - model.data_dir: Path to data directory. Alternatively, you can adjust the file paths directly:
+        - model.train_ds.filepath
+        - model.validation_ds.filepath
+        - model.test_ds.filepath
+    - exp_manager.exp_dir: Directory to log results from the experiment.
+
+It is highly recommended to also adjust these parameters as necessary:
+    - trainer.gpus: Set to 0 to use CPU. Otherwise the number denotes the number of GPUs.
+    - trainer.max_epochs: Maximum number of epochs to train for.
+    - model.batch_size: 8 is sufficient to train a decent Bart model for Text2Sparql.
+    - model.max_seq_length: Maximum (tokenized) sequence length. 150 is sufficient for Text2Sparql.
+    - model.language_model.pretrained_model_name: End2end pretrained model name from huggingface.
+        - model.encoder_tokenizer.tokenizer_name: Pretrained tokenizer name from huggingface.
+        - model.decoder_tokenizer.tokenizer_name: The same as above, as the tokenizer will handle encoding and decoding.
+    - model.optim.lr: Learning rate.
+
+You can also specify an encoder and decoder rather than using an end2end model like Bart by defining these parameters:
+    - model.language_model.pretrained_encoder_model_name: Pretrained huggingface encoder model name.
+        - model.encoder_tokenizer.tokenizer_name: Pretrained huggingface encoder tokenizer name.
+    - model.language_model.pretrained_decoder_model_name: Pretrained huggingface decoder model name.
+        - model.decoder_tokenizer.tokenizer_name: Pretrained huggingface decoder tokenizer name.
+    - model.language_model.pretrained_model_name: Set this to null.
+
+
+***How to run the script?***
+- First, download the data to TGT_DATA_DIR (see ./data/import_datasets.py):
+
+SRC_DATA_DIR=./data/text2sparql_src
+TGT_DATA_DIR=./data/text2sparql_tgt
+
+python ./data/import_datasets.py \
+    --source_data_dir $SRC_DATA_DIR \
+    --target_data_dir $TGT_DATA_DIR
+
+- And run the following to train and save the model:
+
+python text2sparql.py \
+    model.train_ds.filepath="$TGT_DATA_DIR"/train.tsv \
+    model.validation_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \
+    model.test_ds.filepath="$TGT_DATA_DIR"/test_hard.tsv \
+    model.batch_size=16 \
+    model.nemo_path=./NeMo_logs/bart.nemo \
+    exp_manager.exp_dir=./NeMo_logs
+"""
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel
+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="text2sparql_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    nmt_model = NeuralMachineTranslationModel(cfg.model, trainer=trainer)
+    trainer.fit(nmt_model)
+    if cfg.model.nemo_path:
+        nmt_model.save_to(cfg.model.nemo_path)
+
+
+if __name__ == "__main__":
+    main()

tts/NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml

@@ -0,0 +1,119 @@
+# Copyright (c) 2020, 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.
+
+# Config file for text classification with pre-trained BERT models
+
+trainer:
+  gpus: 1 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1]
+  num_nodes: 1
+  max_epochs: 100
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
+
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+
+model:
+  nemo_path: text_classification_model.nemo # filename to save the model and associated artifacts to .nemo file
+  tokenizer:
+      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+      vocab_file: null # path to vocab file
+      tokenizer_model: null # only used if tokenizer is sentencepiece
+      special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  classifier_head:
+    num_output_layers: 2
+    fc_dropout: 0.1
+
+  class_labels:
+    class_labels_file : null # optional to specify a file containing the list of the labels
+
+  dataset:
+    num_classes: ??? # The number of classes. 0 < Label <num_classes.
+    do_lower_case: false # true for uncased models, false for cased models, will be set automatically if pre-trained tokenizer model is used
+    max_seq_length: 256 # the maximum length BERT supports is 512
+    class_balancing: null # null or 'weighted_loss'. 'weighted_loss' enables the weighted class balancing of the loss, may be used for handling unbalanced classes
+    use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up
+
+  train_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: true
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  test_ds:
+    file_path: null
+    batch_size: 64
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 3
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adam
+    lr: 2e-5
+    # optimizer arguments
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+
+    # scheduler setup
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+  # List of some sample queries for inference after training is done
+  infer_samples: [
+    'by the end of no such thing the audience , like beatrice , has a watchful affection for the monster .',
+    'director rob marshall went out gunning to make a great one .',
+    'uneasy mishmash of styles and genres .',
+  ]
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: "TextClassification"  # The name of your model
+  create_tensorboard_logger: True  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: True  # Whether you want exp_manager to create a modelcheckpoint callback

tts/NeMo/examples/nlp/text_classification/data/import_datasets.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2020, 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.
+
+"""
+This script can be used to process and import IMDB, ChemProt, SST-2, and THUCnews datasets into NeMo's format.
+You may run it as the following:
+
+python import_datasets.py \
+        --dataset_name DATASET_NAME \
+        --target_data_dir TARGET_PATH \
+        --source_data_dir SOURCE_PATH
+
+The dataset should be specified by "DATASET_NAME" which can be from ["sst-2", "chemprot", "imdb", "thucnews"].
+It reads the data from "SOURCE_PATH" folder, processes and converts the data into NeMo's format.
+Then writes the results into "TARGET_PATH" folder.
+
+"""
+
+import argparse
+import csv
+import glob
+import os
+from os.path import exists
+
+import tqdm
+
+from nemo.utils import logging
+
+
+def process_imdb(infold, outfold, uncased, modes=['train', 'test']):
+    if not os.path.exists(infold):
+        link = 'https://ai.stanford.edu/~amaas/data/sentiment/'
+        raise ValueError(
+            f'Data not found at {infold}. '
+            f'Please download IMDB reviews dataset from {link} and '
+            f'extract it into the folder specified by source_data_dir argument.'
+        )
+
+    logging.info(f'Processing IMDB dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w')
+        for sent in ['neg', 'pos']:
+            if sent == 'neg':
+                label = 0
+            else:
+                label = 1
+            files = glob.glob(f'{infold}/{mode}/{sent}/*.txt')
+            for file in files:
+                with open(file, 'r') as f:
+                    review = f.read().strip()
+                if uncased:
+                    review = review.lower()
+                review = review.replace("<br />", "")
+                outfiles[mode].write(f'{review}\t{label}\n')
+
+    for mode in modes:
+        outfiles[mode].close()
+
+    class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
+    class_labels_file.write('negative\npositive\n')
+    class_labels_file.close()
+
+
+def process_sst2(infold, outfold, uncased, splits=['train', 'dev']):
+    """Process sst2 dataset."""
+    # "test" split doesn't have labels, so it is skipped
+    if not os.path.exists(infold):
+        link = 'https://dl.fbaipublicfiles.com/glue/data/SST-2.zip'
+        raise ValueError(
+            f'Data not found at {infold}. Please download SST-2 dataset from `{link}` and '
+            f'extract it into the folder specified by `source_data_dir` argument.'
+        )
+
+    logging.info(f'Processing SST-2 dataset')
+    os.makedirs(outfold, exist_ok=True)
+
+    def _read_tsv(input_file, quotechar=None):
+        """Read a tab separated value file."""
+        with open(input_file, "r") as f:
+            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
+            lines = []
+            for line in reader:
+                lines.append(line)
+            return lines
+
+    for split in splits:
+        # Load input file.
+        input_file = os.path.join(infold, split + '.tsv')
+        lines = _read_tsv(input_file)
+        # Create output.
+        outfile = open(os.path.join(outfold, split + '.tsv'), 'w')
+
+        # Copy lines, skip the header (line 0).
+        for line in lines[1:]:
+            text = line[0]
+            label = line[1]
+            # Lowercase when required.
+            if uncased:
+                text = text.lower()
+            # Write output.
+            outfile.write(f'{text}\t{label}\n')
+        # Close file.
+        outfile.close()
+
+    class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w')
+    class_labels_file.write('negative\npositive\n')
+    class_labels_file.close()
+
+    logging.info(f'Result stored at {outfold}')
+
+
+def process_chemprot(source_dir, target_dir, uncased, modes=['train', 'test', 'dev']):
+    if not os.path.exists(source_dir):
+        link = 'https://github.com/arwhirang/recursive_chemprot/tree/master/Demo/tree_LSTM/data'
+        raise ValueError(f'Data not found at {source_dir}. ' f'Please download ChemProt from {link}.')
+
+    logging.info(f'Processing Chemprot dataset and store at {target_dir}')
+    os.makedirs(target_dir, exist_ok=True)
+
+    naming_map = {'train': 'trainingPosit_chem', 'test': 'testPosit_chem', 'dev': 'developPosit_chem'}
+
+    def _read_tsv(input_file, quotechar=None):
+        """Reads a tab separated value file."""
+        with open(input_file, "r") as f:
+            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
+            lines = []
+            for line in reader:
+                lines.append(line)
+            return lines
+
+    outfiles = {}
+    label_mapping = {}
+    out_label_mapping = open(os.path.join(target_dir, 'label_mapping.tsv'), 'w')
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(target_dir, mode + '.tsv'), 'w')
+        input_file = os.path.join(source_dir, naming_map[mode])
+        lines = _read_tsv(input_file)
+        for line in lines:
+            text = line[1]
+            label = line[2]
+            if label == "True":
+                label = line[3]
+            if uncased:
+                text = text.lower()
+            if label not in label_mapping:
+                out_label_mapping.write(f'{label}\t{len(label_mapping)}\n')
+                label_mapping[label] = len(label_mapping)
+            label = label_mapping[label]
+            outfiles[mode].write(f'{text}\t{label}\n')
+    for mode in modes:
+        outfiles[mode].close()
+    out_label_mapping.close()
+
+
+def process_thucnews(infold, outfold):
+    modes = ['train', 'test']
+    train_size = 0.8
+    if not os.path.exists(infold):
+        link = 'thuctc.thunlp.org/'
+        raise ValueError(f'Data not found at {infold}. ' f'Please download THUCNews from {link}.')
+
+    logging.info(f'Processing THUCNews dataset and store at {outfold}')
+    os.makedirs(outfold, exist_ok=True)
+
+    outfiles = {}
+    for mode in modes:
+        outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'a+', encoding='utf-8')
+    categories = ['体育', '娱乐', '家居', '彩票', '房产', '教育', '时尚', '时政', '星座', '游戏', '社会', '科技', '股票', '财经']
+    for category in categories:
+        label = categories.index(category)
+        category_files = glob.glob(f'{infold}/{category}/*.txt')
+        test_num = int(len(category_files) * (1 - train_size))
+        test_files = category_files[:test_num]
+        train_files = category_files[test_num:]
+
+        for mode in modes:
+            logging.info(f'Processing {mode} data of the category {category}')
+            if mode == 'test':
+                files = test_files
+            else:
+                files = train_files
+
+            if len(files) == 0:
+                logging.info(f'Skipping category {category} for {mode} mode')
+                continue
+
+            for file in tqdm.tqdm(files):
+                with open(file, 'r', encoding='utf-8') as f:
+                    news = f.read().strip().replace('\r', '')
+                    news = news.replace('\n', '').replace('\t', ' ')
+                    outfiles[mode].write(f'{news}\t{label}\n')
+    for mode in modes:
+        outfiles[mode].close()
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.")
+    parser.add_argument("--dataset_name", required=True, type=str, choices=['imdb', 'thucnews', 'chemprot'])
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help='The path to the folder containing the dataset files.'
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str)
+    parser.add_argument("--do_lower_case", action='store_true')
+    args = parser.parse_args()
+
+    dataset_name = args.dataset_name
+    do_lower_case = args.do_lower_case
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+
+    if not exists(source_dir):
+        raise FileNotFoundError(f"{source_dir} does not exist.")
+
+    if dataset_name == 'imdb':
+        process_imdb(source_dir, target_dir, do_lower_case)
+    elif dataset_name == 'thucnews':
+        process_thucnews(source_dir, target_dir)
+    elif dataset_name == "chemprot":
+        process_chemprot(source_dir, target_dir, do_lower_case)
+    elif dataset_name == "sst-2":
+        process_sst2(source_dir, target_dir, do_lower_case)
+    else:
+        raise ValueError(
+            f'Dataset {dataset_name} is not supported.'
+            + "Please make sure that you build the preprocessing process for it. "
+            + "NeMo's format assumes that a data file has a header and each line of the file follows "
+            + "the format: text [TAB] label. Label is assumed to be an integer."
+        )

tts/NeMo/examples/nlp/text_classification/model_parallel_text_classification_evaluation.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""
+This script runs model parallel text classification evaluation.
+"""
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.text_classification import TextClassificationModel
+from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
+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="text_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    # TODO: can we drop strict=False
+    model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer, strict=False)
+    model.setup_test_data(test_data_config=cfg.model.test_ds)
+
+    trainer.test(model=model, ckpt_path=None)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/text_classification/text_classification_with_bert.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2020, 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.
+
+"""
+This script contains an example on how to train, evaluate and perform inference with the TextClassificationModel.
+TextClassificationModel in NeMo supports text classification problems such as sentiment analysis or
+domain/intent detection for dialogue systems, as long as the data follows the format specified below.
+
+***Data format***
+TextClassificationModel requires the data to be stored in TAB separated files (.tsv) with two columns of sentence and
+label. Each line of the data file contains text sequences, where words are separated with spaces and label separated
+with [TAB], i.e.:
+
+[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]
+
+For example:
+
+hide new secretions from the parental units[TAB]0
+that loves its characters and communicates something rather beautiful about human nature[TAB]1
+...
+
+If your dataset is stored in another format, you need to convert it to this format to use the TextClassificationModel.
+
+
+***Setting the configs***
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, head classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+
+This script uses the `/examples/nlp/text_classification/conf/text_classification_config.yaml` default config file
+by default. You may update the config file from the file directly or by using the command line arguments.
+Other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+You first need to set the num_classes in the config file which specifies the number of classes in the dataset.
+Notice that some config lines, including `model.dataset.classes_num`, have `???` as their value, this means that values
+for these fields are required to be specified by the user. We need to specify and set the `model.train_ds.file_name`,
+`model.validation_ds.file_name`, and `model.test_ds.file_name` in the config file to the paths of the train, validation,
+ and test files if they exist. We may do it by updating the config file or by setting them from the command line.
+
+
+***How to run the script?***
+For example the following would train a model for 50 epochs in 2 GPUs on a classification task with 2 classes:
+
+# python text_classification_with_bert.py
+        model.dataset.num_classes=2
+        model.train_ds=PATH_TO_TRAIN_FILE
+        model.validation_ds=PATH_TO_VAL_FILE
+        trainer.max_epochs=50
+        trainer.gpus=2
+
+This script would also reload the last checkpoint after the training is done and does evaluation on the dev set,
+then performs inference on some sample queries.
+
+By default, this script uses examples/nlp/text_classification/conf/text_classifciation_config.py config file, and
+you may update all the params in the config file from the command line. You may also use another config file like this:
+
+# python text_classification_with_bert.py --config-name==PATH_TO_CONFIG_FILE
+        model.dataset.num_classes=2
+        model.train_ds=PATH_TO_TRAIN_FILE
+        model.validation_ds=PATH_TO_VAL_FILE
+        trainer.max_epochs=50
+        trainer.gpus=2
+
+***Load a saved model***
+This script would save the model after training into '.nemo' checkpoint file specified by nemo_path of the model config.
+You may restore the saved model like this:
+    model = TextClassificationModel.restore_from(restore_path=NEMO_FILE_PATH)
+
+***Evaluation a saved model on another dataset***
+# If you wanted to evaluate the saved model on another dataset, you may restore the model and create a new data loader:
+    eval_model = TextClassificationModel.restore_from(restore_path=checkpoint_path)
+
+# Then, you may create a dataloader config for evaluation:
+    eval_config = OmegaConf.create(
+        {'file_path': cfg.model.test_ds.file_path, 'batch_size': 64, 'shuffle': False, 'num_workers': 3}
+    )
+    eval_model.setup_test_data(test_data_config=eval_config)
+
+# You need to create a new trainer:
+    eval_trainer = pl.Trainer(gpus=1)
+    eval_model.set_trainer(eval_trainer)
+    eval_trainer.test(model=eval_model, verbose=False)
+"""
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models.text_classification import TextClassificationModel
+from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
+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="text_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.model.train_ds.file_path:
+        raise ValueError("'train_ds.file_path' need to be set for the training!")
+
+    model = TextClassificationModel(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}')
+
+    # We evaluate the trained model on the test set if test_ds is set in the config file
+    if cfg.model.test_ds.file_path:
+        logging.info("===========================================================================================")
+        logging.info("Starting the testing of the trained model on test set...")
+        trainer.test(model=model, ckpt_path=None, verbose=False)
+        logging.info("Testing finished!")
+        logging.info("===========================================================================================")
+
+    # perform inference on a list of queries.
+    if "infer_samples" in cfg.model and cfg.model.infer_samples:
+        logging.info("===========================================================================================")
+        logging.info("Starting the inference on some sample queries...")
+
+        # max_seq_length=512 is the maximum length BERT supports.
+        results = model.classifytext(queries=cfg.model.infer_samples, batch_size=16, max_seq_length=512)
+        logging.info('The prediction results of some sample queries with the trained model:')
+        for query, result in zip(cfg.model.infer_samples, results):
+            logging.info(f'Query : {query}')
+            logging.info(f'Predicted label: {result}')
+
+        logging.info("Inference finished!")
+        logging.info("===========================================================================================")
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml

@@ -0,0 +1,133 @@
+# 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.
+
+# Punctuation and capitalization model with pretrained BERT-like models
+
+pretrained_model: null # pretrained Punctuation and Capitalization model from list_available_models(), for example:
+# punctuation_en_bert or punctuation_en_distilbert
+# or your_model.nemo
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: Punctuation_and_Capitalization  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a model checkpoint callback
+
+model:
+  punct_label_ids: null # will be filled during training
+  capit_label_ids: null # will be filled during training
+  class_labels:
+    punct_labels_file: punct_label_ids.csv
+    capit_labels_file: capit_label_ids.csv
+  dataset:
+    data_dir: ??? # /path/to/training/data
+    # parameters for dataset preprocessing
+    max_seq_length: 128
+    pad_label: 'O'
+    ignore_extra_tokens: false
+    ignore_start_end: false
+    use_cache: true
+    # shared among dataloaders
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    text_file: text_train.txt
+    labels_file: labels_train.txt
+    shuffle: true
+    num_samples: -1
+    batch_size: 64
+
+  validation_ds:
+    # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for
+    # evaluation is needed, specify ds_item, otherwise by default data_dir is used
+    ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets)
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  test_ds:
+    # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for
+    # evaluation is needed, specify ds_item, otherwise by default data_dir is used
+    ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets)
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  punct_head:
+    punct_num_fc_layers: 1
+    fc_dropout: 0.1
+    activation: 'relu'
+    use_transformer_init: True
+
+  capit_head:
+    capit_num_fc_layers: 1
+    fc_dropout: 0.1
+    activation: 'relu'
+    use_transformer_init: true
+
+  optim:
+    name: adam
+    lr: 1e-4
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml

@@ -0,0 +1,119 @@
+# 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.
+
+# Token Classification tasks (for example, Named Entity Recognition) with pretrained BERT-like models
+
+pretrained_model: null # pretrained TokenClassification model from list_available_models() or path to a .nemo file,
+# for example: ner_en_bert or your_model.nemo
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 5
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: token_classification_model  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manager to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a model checkpoint callback
+
+model:
+  label_ids: null # will be filled during training
+  class_labels:
+    class_labels_file: label_ids.csv # will be generated during training and saved in .nemo file
+  dataset:
+    data_dir: ??? # /path/to/data
+    class_balancing: null # choose from [null, weighted_loss]. Weighted_loss enables the weighted class balancing of the loss, may be used for handling unbalanced classes
+    max_seq_length: 128
+    pad_label: 'O'
+    ignore_extra_tokens: false
+    ignore_start_end: false
+    use_cache: true
+    # shared among dataloaders
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    text_file: text_train.txt
+    labels_file: labels_train.txt
+    shuffle: true
+    num_samples: -1
+    batch_size: 64
+
+  validation_ds:
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  test_ds:
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+
+  head:
+    num_fc_layers: 2
+    fc_dropout: 0.5
+    activation: 'relu'
+    use_transformer_init: True
+
+  optim:
+    name: adam
+    lr: 5e-5
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/token_classification/data/get_tatoeba_data.py

@@ -0,0 +1,242 @@
+# 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 argparse
+import logging
+import os
+import random
+import re
+import string
+import subprocess
+
+from nemo.utils import logging
+
+URL = {'tatoeba': 'https://downloads.tatoeba.org/exports/sentences.csv'}
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if not exists.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+    """
+    source = URL[source]
+    if not os.path.exists(destination):
+        logging.info(f'Downloading {source} to {destination}')
+        subprocess.run(['wget', '-O', destination, source])
+    else:
+        logging.info(f'{destination} found. Skipping download')
+
+
+def __process_english_sentences(
+    in_file: str, out_file: str, percent_to_cut: float = 0, num_to_combine: int = 1, num_samples: int = -1
+):
+    """
+    Extract English sentences from the Tatoeba dataset.
+
+    Expected in_file format
+    that
+    contrain letters and punctuation marks (,.?).
+    Chop and combine sentences.
+    Args:
+        in_file: local filepath to the tatoeba dataset.
+    Format: id [TAB] region_name [TAB] sentence,
+    for example: "1276\teng\tLet's try something.\n"
+        out_file: local filepath to the clean dataset
+        percent_to_cut: Percent of sentences to cut in the middle
+            to get examples of incomplete sentences.
+            This could be useful since ASR output not always
+            represents a complete sentence
+        num_to_combine: Number of sentences to combine into
+            a single example
+        num_samples: Number of samples in the final dataset
+    """
+    if not os.path.exists(in_file):
+        raise FileNotFoundError(f'{in_file} not found.')
+
+    in_file = open(in_file, 'r')
+    out_file = open(out_file, 'w')
+    lines_to_combine = []
+    samples_count = 0
+
+    for line in in_file:
+        line = line.split('\t')
+        # use only English sentences
+        if line[1] == 'eng':
+            line = line[2].strip()
+            if re.match("^[A-Z][A-Za-z.,'?\s]+$", line):  # nopep8
+                # chop some sentences in the middle
+                if percent_to_cut > 0:
+                    line = line.split()
+                    if random.random() < percent_to_cut:
+                        line = line[: len(line) // 2]
+                    line = ' '.join(line)
+
+                # combine multiple sentences into a single example
+                # to make it harder for the model to learn eos punctuation
+                if len(lines_to_combine) >= num_to_combine:
+                    if samples_count == num_samples:
+                        return
+                    out_file.write(' '.join(lines_to_combine) + '\n')
+                    lines_to_combine = []
+                    samples_count += 1
+                lines_to_combine.append(line)
+
+    if len(lines_to_combine) > 0 and (samples_count < num_samples or num_samples < 0):
+        out_file.write(' '.join(lines_to_combine) + '\n')
+
+
+def __split_into_train_dev(in_file: str, train_file: str, dev_file: str, percent_dev: float):
+    """
+    Create train and dev split of the dataset.
+    Args:
+        in_file: local filepath to the dataset
+        train_file: local filepath to the train dataset
+        dev_file: local filepath to the dev dataset
+        percent_dev: Percent of the sentences in the dev set
+    """
+    if not os.path.exists(in_file):
+        raise FileNotFoundError(f'{in_file} not found.')
+
+    lines = open(in_file, 'r').readlines()
+    train_file = open(train_file, 'w')
+    dev_file = open(dev_file, 'w')
+
+    dev_size = int(len(lines) * percent_dev)
+    train_file.write(' '.join(lines[:-dev_size]))
+    dev_file.write(' '.join(lines[-dev_size:]))
+
+
+def remove_punctuation(word: str):
+    """
+    Removes all punctuation marks from a word except for '
+    that is often a part of word: don't, it's, and so on
+    """
+    all_punct_marks = string.punctuation.replace("'", '')
+    return re.sub('[' + all_punct_marks + ']', '', word)
+
+
+def create_text_and_labels(output_dir: str, file_path: str, punct_marks: str = ',.?'):
+    """
+    Create datasets for training and evaluation.
+
+    Args:
+      output_dir: path to the output data directory
+      file_path: path to file name
+      punct_marks: supported punctuation marks
+
+    The data will be split into 2 files: text.txt and labels.txt. \
+    Each line of the text.txt file contains text sequences, where words\
+    are separated with spaces. The labels.txt file contains \
+    corresponding labels for each word in text.txt, the labels are \
+    separated with spaces. Each line of the files should follow the \
+    format:  \
+    [WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and \
+    [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).'
+    """
+    if not os.path.exists(file_path):
+        raise ValueError(f'{file_path} not found')
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    base_name = os.path.basename(file_path)
+    labels_file = os.path.join(output_dir, 'labels_' + base_name)
+    text_file = os.path.join(output_dir, 'text_' + base_name)
+
+    with open(file_path, 'r') as f:
+        with open(text_file, 'w') as text_f:
+            with open(labels_file, 'w') as labels_f:
+                for line in f:
+                    line = line.split()
+                    text = ''
+                    labels = ''
+                    for word in line:
+                        label = word[-1] if word[-1] in punct_marks else 'O'
+                        word = remove_punctuation(word)
+                        if len(word) > 0:
+                            if word[0].isupper():
+                                label += 'U'
+                            else:
+                                label += 'O'
+
+                            word = word.lower()
+                            text += word + ' '
+                            labels += label + ' '
+
+                    text_f.write(text.strip() + '\n')
+                    labels_f.write(labels.strip() + '\n')
+
+    print(f'{text_file} and {labels_file} created from {file_path}.')
+
+
+def __delete_file(file_to_del: str):
+    """
+    Deletes the file
+    Args:
+        file_to_del: local filepath to the file to delete
+    """
+    if os.path.exists(file_to_del):
+        os.remove(file_to_del)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Prepare tatoeba dataset')
+    parser.add_argument("--data_dir", required=True, type=str)
+    parser.add_argument("--dataset", default='tatoeba', type=str)
+    parser.add_argument("--num_samples", default=-1, type=int, help='-1 to use the whole dataset')
+    parser.add_argument("--percent_to_cut", default=0, type=float, help='Percent of sentences to cut in the middle')
+    parser.add_argument(
+        "--num_lines_to_combine", default=1, type=int, help='Number of lines to combine into single example'
+    )
+    parser.add_argument("--percent_dev", default=0.2, type=float, help='Size of the dev set, float')
+    parser.add_argument("--clean_dir", action='store_true')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.data_dir):
+        os.makedirs(args.data_dir)
+
+    if args.dataset != 'tatoeba':
+        raise ValueError("Unsupported dataset.")
+
+    logging.info(f'Downloading tatoeba dataset')
+    tatoeba_dataset = os.path.join(args.data_dir, 'sentences.csv')
+    __maybe_download_file(tatoeba_dataset, args.dataset)
+
+    logging.info(f'Processing English sentences...')
+    clean_eng_sentences = os.path.join(args.data_dir, 'clean_eng_sentences.txt')
+    __process_english_sentences(
+        tatoeba_dataset, clean_eng_sentences, args.percent_to_cut, args.num_lines_to_combine, args.num_samples
+    )
+
+    train_file = os.path.join(args.data_dir, 'train.txt')
+    dev_file = os.path.join(args.data_dir, 'dev.txt')
+
+    logging.info(
+        f'Splitting the {args.dataset} dataset into train and dev sets' + ' and creating labels and text files'
+    )
+    __split_into_train_dev(clean_eng_sentences, train_file, dev_file, args.percent_dev)
+
+    logging.info(f'Creating text and label files for training')
+    create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'train.txt'))
+    create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'dev.txt'))
+
+    if args.clean_dir:
+        logging.info(f'Cleaning up {args.data_dir}')
+        __delete_file(clean_eng_sentences)
+        __delete_file(tatoeba_dataset)
+        __delete_file(train_file)
+        __delete_file(dev_file)
+    logging.info(f'Processing of the {args.dataset} is complete')

tts/NeMo/examples/nlp/token_classification/data/import_from_iob_format.py

@@ -0,0 +1,124 @@
+# 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 argparse
+import os
+
+from nemo.utils import logging
+
+
+def __convert_data(in_file: str, out_text_f: str, out_labels_f: str, max_length: int):
+    """
+    Convert data from the IOB format to NeMo accepted format described below.
+    in_file should be in the IOB format, see example here:
+    https://www.clips.uantwerpen.be/conll2003/ner/.
+
+    Args:
+        in_file: input file name
+        out_text_f: output file with text
+        out_labels_f: output file with labels
+        max_length: use -1 to leave the examples' length as is, otherwise long examples will be split into multiple
+            examples
+    After the conversion, the dataset is split into 2 files: text.txt
+    and labels.txt.
+    Each line of the text.txt file contains text sequences, where words
+    are separated with spaces. The labels.txt file contains corresponding
+    labels for each word in text.txt, the labels are separated with spaces.
+    Each line of the files should follow the format:
+    [WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and
+    [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).
+
+    """
+    in_file = open(in_file, 'r')
+
+    if max_length == -1:
+        with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
+            for line in in_file:
+                if line == '\n':
+                    out_text.write(line)
+                    out_labels.write(line)
+                else:
+                    line = line.split()
+                    out_text.write(line[0] + ' ')
+                    out_labels.write(line[-1] + ' ')
+
+    else:
+        words = []
+        labels = []
+        with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
+            lines = in_file.readlines()
+            for line_id, line in enumerate(lines):
+                logging.info(f"{line_id} {len(lines)}")
+                contends = line.strip()
+                if len(contends) == 0:
+                    assert len(words) == len(labels)
+                    if len(words) > max_length:
+                        # split if the sentence is longer than max_length
+                        while len(words) > max_length:
+                            tmplabel = labels[:max_length]
+                            for iidx in range(len(tmplabel)):
+                                if tmplabel.pop() == 'O':
+                                    break
+                            l = ' '.join([label for label in labels[: len(tmplabel) + 1] if len(label) > 0])
+                            w = ' '.join([word for word in words[: len(tmplabel) + 1] if len(word) > 0])
+
+                            out_text.write(w + "\n")
+                            out_labels.write(l + "\n")
+                            words = words[len(tmplabel) + 1 :]
+                            labels = labels[len(tmplabel) + 1 :]
+
+                    if len(words) == 0:
+                        continue
+                    l = ' '.join([label for label in labels if len(label) > 0])
+                    w = ' '.join([word for word in words if len(word) > 0])
+
+                    out_text.write(w + "\n")
+                    out_labels.write(l + "\n")
+                    words = []
+                    labels = []
+                    continue
+
+                word = line.strip().split()[0]
+                label = line.strip().split()[-1]
+                words.append(word)
+                labels.append(label)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Convert data from IOB format to the format compatible with \
+        nlp/examples/token_classification/scripts/token_classification_train.py and \
+        token_classification_evaluate.py'
+    )
+    parser.add_argument("--data_file", required=True, type=str, help='path to a file in IOB format')
+    parser.add_argument(
+        "--max_length",
+        default=-1,
+        type=int,
+        help='use -1 to leave the examples\'s length as is, '
+        'otherwise long examples will be split into multiple examples',
+    )
+    args = parser.parse_args()
+
+    data_dir, basename = os.path.split(args.data_file)
+    prefix = os.path.splitext(basename)[0]
+    if not os.path.exists(args.data_file):
+        raise FileNotFoundError(f"{args.data_file} not found")
+
+    logging.info(f'Processing {args.data_file}')
+    out_text = os.path.join(data_dir, 'text_' + prefix + '.txt')
+    out_labels = os.path.join(data_dir, 'labels_' + prefix + '.txt')
+
+    __convert_data(args.data_file, out_text, out_labels, args.max_length)
+    logging.info(f'Processing of the {args.data_file} is complete')

tts/NeMo/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py

@@ -0,0 +1,99 @@
+# 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.
+
+"""
+The script converts raw text to the NeMo format for punctuation and capitalization task.
+
+Raw Data Format
+---------------
+
+The Punctuation and Capitalization model can work with any text dataset, although it is recommended to balance the data, especially for the punctuation task.
+Before pre-processing the data to the format expected by the model, the data should be split into train.txt and dev.txt (and optionally test.txt).
+Each line in the **train.txt/dev.txt/test.txt** should represent one or more full and/or truncated sentences.
+
+Example of the train.txt/dev.txt file:
+    When is the next flight to New York?
+    The next flight is ...
+    ....
+
+
+The `source_data_dir` structure should look like this:
+   .
+   |--sourced_data_dir
+     |-- dev.txt
+     |-- train.txt
+
+
+
+NeMo Data Format for training the model
+---------------------------------------
+
+The punctuation and capitalization model expects the data in the following format:
+
+The training and evaluation data is divided into 2 files: text.txt and labels.txt. \
+Each line of the **text.txt** file contains text sequences, where words are separated with spaces, i.e.
+
+[WORD] [SPACE] [WORD] [SPACE] [WORD], for example:
+        when is the next flight to new york
+        the next flight is ...
+        ...
+
+The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. \
+Each label in labels.txt file consists of 2 symbols:
+
+* the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed);
+* the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.)
+
+By default the following punctuation marks are considered: commas, periods, and question marks; the rest punctuation marks were removed from the data.
+This can be changed by introducing new labels in the labels.txt files
+
+Each line of the labels.txt should follow the format: [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). \
+For example, labels for the above text.txt file should be:
+
+        OU OO OO OO OO OO OU ?U
+        OU OO OO OO ...
+        ...
+
+The complete list of all possible labels for this task used in this tutorial is: OO, ,O, .O, ?O, OU, ,U, .U, ?U.
+
+Converting Raw data to NeMo format
+----------------------------------
+
+To pre-process the raw text data, stored under :code:`sourced_data_dir` (see the :ref:`raw_data_format_punct`
+section), run the following command:
+
+    python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \
+           -s <PATH_TO_THE_SOURCE_FILE>
+           -o <PATH_TO_THE_OUTPUT_DIRECTORY>
+
+"""
+
+import argparse
+import os
+
+from get_tatoeba_data import create_text_and_labels
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Prepare data for punctuation and capitalization tasks')
+    parser.add_argument("-s", "--source_file", required=True, type=str, help="Path to the source file")
+    parser.add_argument("-o", "--output_dir", required=True, type=str, help="Path to the output directory")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.source_file):
+        raise ValueError(f'{args.source_file} was not found')
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    create_text_and_labels(args.output_dir, args.source_file)
+
+    print(f'Processing of the {args.source_file} is complete')

tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py

@@ -0,0 +1,127 @@
+# 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 pytorch_lightning as pl
+from omegaconf import DictConfig
+
+from nemo.collections.nlp.models import PunctuationCapitalizationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+This script shows how to perform evaluation and runs inference of a few examples.
+
+More details on the task and data format could be found in tutorials/nlp/Punctuation_and_Capitalization.ipynb
+
+*** Setting the configs ***
+
+This script uses the `/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml` config file
+by default. You may update the config file from the file directly. 
+The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb
+
+
+*** Model Evaluation ***
+
+    python punctuation_capitalization_evaluate.py \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+    pretrained_model=punctuation_en_distilbert 
+
+<PATH_TO_DATA_DIR> - a directory that contains test_ds.text_file and test_ds.labels_file (see the config)
+pretrained_model   - pretrained PunctuationCapitalizationModel model from list_available_models() or 
+                     path to a .nemo file, for example: punctuation_en_bert or your_model.nemo
+
+"""
+
+
+@hydra_runner(config_path="conf", config_name="punctuation_capitalization_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(
+        'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU and \
+            no DDP to obtain accurate results'
+    )
+
+    if not hasattr(cfg.model, 'test_ds'):
+        raise ValueError(f'model.test_ds was not found in the config, skipping evaluation')
+    else:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+
+    trainer = pl.Trainer(
+        gpus=gpu,
+        precision=cfg.trainer.precision,
+        amp_level=cfg.trainer.amp_level,
+        logger=False,
+        checkpoint_callback=False,
+    )
+    exp_dir = exp_manager(trainer, cfg.exp_manager)
+
+    if not cfg.pretrained_model:
+        raise ValueError(
+            'To run evaluation and inference script a pre-trained model or .nemo file must be provided.'
+            f'Choose from {PunctuationCapitalizationModel.list_available_models()} or "pretrained_model"="your_model.nemo"'
+        )
+
+    if os.path.exists(cfg.pretrained_model):
+        model = PunctuationCapitalizationModel.restore_from(cfg.pretrained_model)
+    elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names():
+        model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model)
+    else:
+        raise ValueError(
+            f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}'
+        )
+
+    data_dir = cfg.model.dataset.get('data_dir', None)
+
+    if data_dir is None:
+        logging.error(
+            'No dataset directory provided. Skipping evaluation. '
+            'To run evaluation on a file, specify path to the directory that contains test_ds.text_file and test_ds.labels_file with "model.dataset.data_dir" argument.'
+        )
+    elif not os.path.exists(data_dir):
+        logging.error(f'{data_dir} is not found, skipping evaluation on the test set.')
+    else:
+        model.update_data_dir(data_dir=data_dir)
+        model._cfg.dataset = cfg.model.dataset
+
+        if not hasattr(cfg.model, 'test_ds'):
+            logging.error(f'model.test_ds was not found in the config, skipping evaluation')
+        elif model.prepare_test(trainer):
+            model.setup_test_data(cfg.model.test_ds)
+            trainer.test(model)
+        else:
+            logging.error('Skipping the evaluation. The trainer is not setup properly.')
+
+    # run an inference on a few examples
+    queries = [
+        'we bought four shirts one pen and a mug from the nvidia gear store in santa clara',
+        'what can i do for you today',
+        'how are you',
+    ]
+
+    inference_results = model.add_punctuation_capitalization(queries, batch_size=len(queries), max_seq_length=512)
+
+    for query, result in zip(queries, inference_results):
+        logging.info(f'Query : {query}')
+        logging.info(f'Result: {result.strip()}\n')
+
+    logging.info(f'Results are saved at {exp_dir}')
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py

@@ -0,0 +1,101 @@
+# 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import PunctuationCapitalizationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+This script show how to train a Punctuation and Capitalization Model.
+More details on the task and data format could be found in tutorials/nlp/Punctuation_and_Capitalization.ipynb
+
+*** Setting the configs ***
+
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+This script uses the `/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml` config file
+by default. You may update the config file from the file directly. 
+The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb
+
+*** Model training ***
+
+To run this script and train the model from scratch, use:
+    python punctuation_and_capitalization_train.py \
+           model.dataset.data_dir=<PATH_TO_DATA_DIR>
+
+To use one of the pretrained versions of the model and finetune it, run:
+    python punctuation_and_capitalization.py \
+    pretrained_model=punctuation_en_bert \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>
+    
+    <PATH_TO_DATA_DIR> - a directory that contains test_ds.text_file and test_ds.labels_file (see the config)
+    pretrained_model   - pretrained PunctuationCapitalization model from list_available_models() or 
+                     path to a .nemo file, for example: punctuation_en_bert or model.nemo
+
+"""
+
+
+@hydra_runner(config_path="conf", config_name="punctuation_capitalization_config")
+def main(cfg: DictConfig) -> None:
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.pretrained_model:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = PunctuationCapitalizationModel(cfg.model, trainer=trainer)
+    else:
+        if os.path.exists(cfg.pretrained_model):
+            model = PunctuationCapitalizationModel.restore_from(cfg.pretrained_model)
+        elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names():
+            model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model)
+        else:
+            raise ValueError(
+                f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}'
+            )
+
+        data_dir = cfg.model.dataset.get('data_dir', None)
+        if data_dir:
+            if not os.path.exists(data_dir):
+                raise ValueError(f'{data_dir} is not found at')
+
+            # we can also do finetuning of the pretrained model but we would need to update the data dir
+            model.update_data_dir(data_dir)
+            # setup train and validation Pytorch DataLoaders
+            model.setup_training_data()
+            model.setup_validation_data()
+            logging.info(f'Using config file of the pretrained model')
+        else:
+            raise ValueError(
+                'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \
+                with "model.dataset.data_dir" argument'
+            )
+
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/token_classification/token_classification_evaluate.py

@@ -0,0 +1,137 @@
+# 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 pytorch_lightning as pl
+from omegaconf import DictConfig
+
+from nemo.collections.nlp.models import TokenClassificationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+This script shows how to perform evaluation and runs inference of a few examples.
+
+More details on Token Classification model could be found in tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb
+
+*** Setting the configs ***
+
+This script uses the `/examples/nlp/token_classification/conf/token_classification_config.yaml` config file
+by default. You may update the config file from the file directly. 
+The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb
+
+*** Model Evaluation ***
+
+The script runs two types of evaluation: 
+    * model.test() - this eval will use the config setting for evaluation such as model.dataset.max_seq_length
+    * model.evaluate_from_file():
+        * disregards model.dataset.max_seq_length and evaluates all the tokens, BERT max seq length - 512 tokens after tokenization
+        * creates confusion matrix
+        * saves predictions and labels (if provided)
+
+To run the script:
+
+    python token_classification_evaluate.py \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+    pretrained_model=ner_en_bert 
+
+<PATH_TO_DATA_DIR> - a directory that contains test_ds.text_file and test_ds.labels_file (see the config)
+pretrained_model   - pretrained TokenClassification model from list_available_models() or 
+                     path to a .nemo file, for example: ner_en_bert or your_model.nemo
+
+"""
+
+
+@hydra_runner(config_path="conf", config_name="token_classification_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(
+        'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU and \
+            no DDP to obtain accurate results'
+    )
+
+    if not hasattr(cfg.model, 'test_ds'):
+        raise ValueError(f'model.test_ds was not found in the config, skipping evaluation')
+    else:
+        gpu = 1 if cfg.trainer.gpus != 0 else 0
+
+    trainer = pl.Trainer(
+        gpus=gpu,
+        precision=cfg.trainer.precision,
+        amp_level=cfg.trainer.amp_level,
+        logger=False,
+        checkpoint_callback=False,
+    )
+    exp_dir = exp_manager(trainer, cfg.exp_manager)
+
+    if not cfg.pretrained_model:
+        raise ValueError(
+            'To run evaluation and inference script a pre-trained model or .nemo file must be provided.'
+            f'Choose from {TokenClassificationModel.list_available_models()} or "pretrained_model"="your_model.nemo"'
+        )
+
+    if os.path.exists(cfg.pretrained_model):
+        model = TokenClassificationModel.restore_from(cfg.pretrained_model)
+    elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names():
+        model = TokenClassificationModel.from_pretrained(cfg.pretrained_model)
+    else:
+        raise ValueError(
+            f'Provide path to the pre-trained .nemo checkpoint or choose from {TokenClassificationModel.list_available_models()}'
+        )
+
+    data_dir = cfg.model.dataset.get('data_dir', None)
+    if data_dir is None:
+        logging.error(
+            'No dataset directory provided. Skipping evaluation. '
+            'To run evaluation on a file, specify path to the directory that contains test_ds.text_file and test_ds.labels_file with "model.dataset.data_dir" argument.'
+        )
+    elif not os.path.exists(data_dir):
+        logging.error(f'{data_dir} is not found, skipping evaluation on the test set.')
+    else:
+        model.update_data_dir(data_dir=data_dir)
+        model._cfg.dataset = cfg.model.dataset
+
+        if not hasattr(cfg.model, 'test_ds'):
+            logging.error(f'model.test_ds was not found in the config, skipping evaluation')
+        elif model.prepare_test(trainer):
+            model.setup_test_data(cfg.model.test_ds)
+            trainer.test(model)
+
+            model.evaluate_from_file(
+                text_file=os.path.join(data_dir, cfg.model.test_ds.text_file),
+                labels_file=os.path.join(data_dir, cfg.model.test_ds.labels_file),
+                output_dir=exp_dir,
+                add_confusion_matrix=True,
+                normalize_confusion_matrix=True,
+            )
+        else:
+            logging.error('Skipping the evaluation. The trainer is not setup properly.')
+
+    # run an inference on a few examples
+    queries = ['we bought four shirts from the nvidia gear store in santa clara.', 'Nvidia is a company.']
+    results = model.add_predictions(queries, output_file='predictions.txt')
+
+    for query, result in zip(queries, results):
+        logging.info(f'Query : {query}')
+        logging.info(f'Result: {result.strip()}\n')
+
+    logging.info(f'Results are saved at {exp_dir}')
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/token_classification/token_classification_train.py

@@ -0,0 +1,147 @@
+# 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import TokenClassificationModel
+from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+This scripts shows how to train a Token Classification model.
+
+The Token Classification model supports Named Entity Recognition task and other token level classification tasks,
+as long as the data follows the format specified below.
+
+More details on how to use this script could be found in 
+tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb
+
+*** Data Format *** 
+Token Classification Model requires the data to be split into 2 files: text.txt and labels.txt.
+Each line of the text.txt file contains text sequences, where words are separated with spaces, i.e.:
+[WORD] [SPACE] [WORD] [SPACE] [WORD].
+The labels.txt file contains corresponding labels for each word in text.txt, the labels are separated with spaces, i.e.:
+[LABEL] [SPACE] [LABEL] [SPACE] [LABEL].
+
+Example of a text.txt file:
+Jennifer is from New York City .
+She likes ...
+...
+
+Corresponding labels.txt file:
+B-PER O O B-LOC I-LOC I-LOC O
+O O ...
+...
+
+*** Preparing the dataset ***
+
+To convert an IOB format data to the format required for training, run
+examples/nlp/token_classification/data/import_from_iob_format.py on your train and dev files, as follows:
+
+python examples/nlp/token_classification/data/import_from_iob_format.py --data_file PATH_TO_IOB_FORMAT_DATAFILE
+
+*** Setting the configs ***
+
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+This script uses the `/examples/nlp/token_classification/conf/token_classification_config.yaml` config file
+by default. You may update the config file from the file directly. 
+The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb
+
+*** Model Training ***
+
+To train TokenClassification model from scratch with the default config file, run:
+
+    python token_classification_train.py \
+           model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+           trainer.max_epochs=<NUM_EPOCHS> \
+           trainer.gpus=[<CHANGE_TO_GPU(s)_YOU_WANT_TO_USE>]
+
+To use one of the pretrained versions of the model specify a `pretrained_model` arg with either 
+TokenClassification model from list_available_models() or path to a .nemo file, for example: 
+ner_en_bert or model.nemo, run:
+
+    python token_classification_train.py pretrained_model=ner_en_bert
+
+To use one of the pretrained versions of the model and fine-tune it, run:
+
+    python token_classification_train.py \
+           model.dataset.data_dir=<PATH_TO_DATA_DIR>  \
+           pretrained_model=ner_en_bert
+
+<PATH_TO_DATA_DIR> - a directory that contains test_ds.text_file and test_ds.labels_file (see the config)
+pretrained_model   - pretrained TokenClassification model from list_available_models() or 
+                     path to a .nemo file, for example: ner_en_bert or model.nemo
+                     
+For more ways of restoring a pre-trained model, see tutorials/00_NeMo_Primer.ipynb
+"""
+
+
+@hydra_runner(config_path="conf", config_name="token_classification_config")
+def main(cfg: DictConfig) -> None:
+    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.pretrained_model:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = TokenClassificationModel(cfg.model, trainer=trainer)
+    else:
+        if os.path.exists(cfg.pretrained_model):
+            # TODO: can we drop strict=False?
+            model = TokenClassificationModel.restore_from(cfg.pretrained_model, trainer=trainer, strict=False)
+        elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names():
+            model = TokenClassificationModel.from_pretrained(cfg.pretrained_model)
+        else:
+            raise ValueError(
+                f'Provide path to the pre-trained .nemo file or choose from {TokenClassificationModel.list_available_models()}'
+            )
+
+        data_dir = cfg.model.dataset.get('data_dir', None)
+        if data_dir:
+            if not os.path.exists(data_dir):
+                raise ValueError(f'{data_dir} is not found at')
+
+            # we can also do finetuning of the pretrained model but it will require
+            # setup the data dir to get class weights statistics
+            model.update_data_dir(data_dir=data_dir)
+            # finally, setup train and validation Pytorch DataLoaders
+            model.setup_training_data()
+            model.setup_validation_data()
+            # then we're setting up loss, use model.dataset.class_balancing,
+            # if you want to add class weights to the CrossEntropyLoss
+            model.setup_loss(class_balancing=cfg.model.dataset.class_balancing)
+            logging.info(f'Using config file of the pretrained model')
+        else:
+            raise ValueError(
+                'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \
+                with "model.dataset.data_dir" argument'
+            )
+
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    main()