Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/nlp/models/__init__.py

@@ -0,0 +1,28 @@
+# 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 nemo.collections.nlp.models.dialogue_state_tracking.sgdqa_model import SGDQAModel
+from nemo.collections.nlp.models.entity_linking.entity_linking_model import EntityLinkingModel
+from nemo.collections.nlp.models.glue_benchmark.glue_benchmark_model import GLUEModel
+from nemo.collections.nlp.models.information_retrieval import BertDPRModel, BertJointIRModel
+from nemo.collections.nlp.models.intent_slot_classification import IntentSlotClassificationModel
+from nemo.collections.nlp.models.language_modeling.bert_lm_model import BERTLMModel
+from nemo.collections.nlp.models.language_modeling.transformer_lm_model import TransformerLMModel
+from nemo.collections.nlp.models.machine_translation import MTEncDecModel
+from nemo.collections.nlp.models.neural_machine_translation.neural_machine_translation_model import (
+    NeuralMachineTranslationModel,
+)
+from nemo.collections.nlp.models.question_answering.qa_model import QAModel
+from nemo.collections.nlp.models.text_classification import TextClassificationModel
+from nemo.collections.nlp.models.token_classification import PunctuationCapitalizationModel, TokenClassificationModel

tts/nemo/collections/nlp/models/dialogue_state_tracking/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.collections.nlp.models.dialogue_state_tracking.sgdqa_model import SGDQAModel

tts/nemo/collections/nlp/models/dialogue_state_tracking/sgdqa_model.py

@@ -0,0 +1,580 @@
+# 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 file contains code artifacts adapted from the original implementation:
+https://github.com/google-research/google-research/blob/master/schema_guided_dst/baseline/train_and_predict.py
+'''
+
+import os
+from typing import Dict, List, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.nlp.data.dialogue_state_tracking import Schema, SGDDataProcessor, SGDDataset
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.evaluate import evaluate, get_in_domain_services
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.prediction_utils import write_predictions_to_file
+from nemo.collections.nlp.losses import SGDDialogueStateLoss
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules import SGDDecoder, SGDEncoder
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import tensor2list
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.utils.get_rank import is_global_rank_zero
+
+__all__ = ['SGDQAModel']
+
+NUM_TASKS = 6  # number of multi-head tasks
+
+
+class SGDQAModel(NLPModel):
+    """Dialogue State Tracking Model SGD-QA"""
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.decoder.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        self.data_prepared = False
+        self.setup_tokenizer(cfg.tokenizer)
+        super().__init__(cfg=cfg, trainer=trainer)
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        self.encoder = SGDEncoder(hidden_size=self.bert_model.config.hidden_size, dropout=self._cfg.encoder.dropout)
+        self.decoder = SGDDecoder(embedding_dim=self.bert_model.config.hidden_size)
+        self.loss = SGDDialogueStateLoss(reduction="mean")
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        token_embeddings = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        encoded_utterance, token_embeddings = self.encoder(hidden_states=token_embeddings)
+        (
+            logit_intent_status,
+            logit_req_slot_status,
+            logit_cat_slot_status,
+            logit_cat_slot_value_status,
+            logit_noncat_slot_status,
+            logit_spans,
+        ) = self.decoder(
+            encoded_utterance=encoded_utterance, token_embeddings=token_embeddings, utterance_mask=attention_mask
+        )
+        return (
+            logit_intent_status,
+            logit_req_slot_status,
+            logit_cat_slot_status,
+            logit_cat_slot_value_status,
+            logit_noncat_slot_status,
+            logit_spans,
+        )
+
+    def training_step(self, batch, batch_idx):
+        (
+            example_id_num,
+            service_id,
+            utterance_ids,
+            token_type_ids,
+            attention_mask,
+            intent_status,
+            requested_slot_status,
+            categorical_slot_status,
+            categorical_slot_value_status,
+            noncategorical_slot_status,
+            noncategorical_slot_value_start,
+            noncategorical_slot_value_end,
+            start_char_idx,
+            end_char_idx,
+            task_mask,
+        ) = batch
+        (
+            logit_intent_status,
+            logit_req_slot_status,
+            logit_cat_slot_status,
+            logit_cat_slot_value_status,
+            logit_noncat_slot_status,
+            logit_spans,
+        ) = self(input_ids=utterance_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
+        loss = self.loss(
+            logit_intent_status=logit_intent_status,
+            intent_status=intent_status,
+            logit_req_slot_status=logit_req_slot_status,
+            requested_slot_status=requested_slot_status,
+            logit_cat_slot_status=logit_cat_slot_status,
+            categorical_slot_status=categorical_slot_status,
+            logit_cat_slot_value_status=logit_cat_slot_value_status,
+            categorical_slot_value_status=categorical_slot_value_status,
+            logit_noncat_slot_status=logit_noncat_slot_status,
+            noncategorical_slot_status=noncategorical_slot_status,
+            logit_spans=logit_spans,
+            noncategorical_slot_value_start=noncategorical_slot_value_start,
+            noncategorical_slot_value_end=noncategorical_slot_value_end,
+            task_mask=task_mask,
+        )
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': loss,
+            'lr': lr,
+        }
+
+    def validation_step(self, batch: List[torch.Tensor], batch_idx: int, dataloader_idx: int = 0) -> dict:
+        """
+        Called at every validation step to aggregate and postprocess outputs on each GPU
+        Args:
+            batch: input batch at validation step
+            batch_idx: batch index 
+            dataloader_idx: dataloader index
+        """
+        loss, tensors = self.eval_step_helper(batch=batch)
+        self.log(f'val_loss', loss)
+        return {f'val_loss': loss, f'tensors': tensors}
+
+    def test_step(self, batch: List[torch.Tensor], batch_idx: int, dataloader_idx: int = 0) -> dict:
+        """
+        Called at every test step to aggregate and postprocess outputs on each GPU
+        Args:
+            batch: input batch at test step
+            batch_idx: batch index 
+            dataloader_idx: dataloader index
+        """
+        loss, tensors = self.eval_step_helper(batch=batch)
+        return {f'test_loss': loss, f'tensors': tensors}
+
+    def eval_step_helper(self, batch: List[torch.Tensor]):
+        """
+        Helper called at every validation/test step to aggregate and postprocess outputs on each GPU
+        Args:
+            batch: input batch at step
+        Returns:
+            loss: averaged batch loss
+            tensors: collection of aggregated output tensors across all GPU workers
+        """
+        (
+            example_id_num,
+            service_id,
+            utterance_ids,
+            token_type_ids,
+            attention_mask,
+            intent_status,
+            requested_slot_status,
+            categorical_slot_status,
+            categorical_slot_value_status,
+            noncategorical_slot_status,
+            noncategorical_slot_value_start,
+            noncategorical_slot_value_end,
+            start_char_idx,
+            end_char_idx,
+            task_mask,
+        ) = batch
+        (
+            logit_intent_status,
+            logit_req_slot_status,
+            logit_cat_slot_status,
+            logit_cat_slot_value_status,
+            logit_noncat_slot_status,
+            logit_spans,
+        ) = self(input_ids=utterance_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
+        loss = self.loss(
+            logit_intent_status=logit_intent_status,
+            intent_status=intent_status,
+            logit_req_slot_status=logit_req_slot_status,
+            requested_slot_status=requested_slot_status,
+            logit_cat_slot_status=logit_cat_slot_status,
+            categorical_slot_status=categorical_slot_status,
+            logit_cat_slot_value_status=logit_cat_slot_value_status,
+            categorical_slot_value_status=categorical_slot_value_status,
+            logit_noncat_slot_status=logit_noncat_slot_status,
+            noncategorical_slot_status=noncategorical_slot_status,
+            logit_spans=logit_spans,
+            noncategorical_slot_value_start=noncategorical_slot_value_start,
+            noncategorical_slot_value_end=noncategorical_slot_value_end,
+            task_mask=task_mask,
+        )
+
+        all_example_id_num = []
+        all_service_id = []
+        all_logit_intent_status = []
+        all_logit_req_slot_status = []
+        all_logit_cat_slot_status = []
+        all_logit_cat_slot_value_status = []
+        all_logit_noncat_slot_status = []
+        all_logit_spans = []
+        all_start_char_idx = []
+        all_end_char_idx = []
+
+        if self.trainer.gpus and self.trainer.world_size > 1:
+            world_size = self.trainer.world_size
+            for ind in range(world_size):
+                all_example_id_num.append(torch.empty_like(example_id_num))
+                all_service_id.append(torch.empty_like(service_id))
+                all_logit_intent_status.append(torch.empty_like(logit_intent_status))
+                all_logit_req_slot_status.append(torch.empty_like(logit_req_slot_status))
+                all_logit_cat_slot_status.append(torch.empty_like(logit_cat_slot_status))
+                all_logit_cat_slot_value_status.append(torch.empty_like(logit_cat_slot_value_status))
+                all_logit_noncat_slot_status.append(torch.empty_like(logit_noncat_slot_status))
+                all_logit_spans.append(torch.empty_like(logit_spans))
+                all_start_char_idx.append(torch.empty_like(start_char_idx))
+                all_end_char_idx.append(torch.empty_like(end_char_idx))
+
+            torch.distributed.all_gather(all_example_id_num, example_id_num)
+            torch.distributed.all_gather(all_service_id, service_id)
+            torch.distributed.all_gather(all_logit_intent_status, logit_intent_status)
+            torch.distributed.all_gather(all_logit_req_slot_status, logit_req_slot_status)
+            torch.distributed.all_gather(all_logit_cat_slot_status, logit_cat_slot_status)
+            torch.distributed.all_gather(all_logit_cat_slot_value_status, logit_cat_slot_value_status)
+            torch.distributed.all_gather(all_logit_noncat_slot_status, logit_noncat_slot_status)
+            torch.distributed.all_gather(all_logit_spans, logit_spans)
+            torch.distributed.all_gather(all_start_char_idx, start_char_idx)
+            torch.distributed.all_gather(all_end_char_idx, end_char_idx)
+        else:
+            all_example_id_num.append(example_id_num)
+            all_service_id.append(service_id)
+            all_logit_intent_status.append(logit_intent_status)
+            all_logit_req_slot_status.append(logit_req_slot_status)
+            all_logit_cat_slot_status.append(logit_cat_slot_status)
+            all_logit_cat_slot_value_status.append(logit_cat_slot_value_status)
+            all_logit_noncat_slot_status.append(logit_noncat_slot_status)
+            all_logit_spans.append(logit_spans)
+            all_start_char_idx.append(start_char_idx)
+            all_end_char_idx.append(end_char_idx)
+
+        # after this: all_x is list of tensors, of length world_size
+        example_id_num = torch.cat(all_example_id_num)
+        service_id = torch.cat(all_service_id)
+        logit_intent_status = torch.cat(all_logit_intent_status)
+        logit_req_slot_status = torch.cat(all_logit_req_slot_status)
+        logit_cat_slot_status = torch.cat(all_logit_cat_slot_status)
+        logit_cat_slot_value_status = torch.cat(all_logit_cat_slot_value_status)
+        logit_noncat_slot_status = torch.cat(all_logit_noncat_slot_status)
+        logit_spans = torch.cat(all_logit_spans)
+        start_char_idx = torch.cat(all_start_char_idx)
+        end_char_idx = torch.cat(all_end_char_idx)
+
+        intent_status = torch.nn.Sigmoid()(logit_intent_status)
+
+        # Scores are output for each requested slot.
+        req_slot_status = torch.nn.Sigmoid()(logit_req_slot_status)
+
+        # For categorical slots, the status of each slot and the predicted value are output.
+        cat_slot_status_dist = torch.nn.Softmax(dim=-1)(logit_cat_slot_status)
+
+        cat_slot_status = torch.argmax(logit_cat_slot_status, axis=-1)
+        cat_slot_status_p = torch.max(cat_slot_status_dist, axis=-1)[0]
+        cat_slot_value_status = torch.nn.Sigmoid()(logit_cat_slot_value_status)
+
+        # For non-categorical slots, the status of each slot and the indices for spans are output.
+        noncat_slot_status_dist = torch.nn.Softmax(dim=-1)(logit_noncat_slot_status)
+
+        noncat_slot_status = torch.argmax(logit_noncat_slot_status, axis=-1)
+        noncat_slot_status_p = torch.max(noncat_slot_status_dist, axis=-1)[0]
+
+        softmax = torch.nn.Softmax(dim=1)
+
+        scores = softmax(logit_spans)
+        start_scores, end_scores = torch.unbind(scores, dim=-1)
+
+        batch_size, max_num_tokens = end_scores.size()
+        # Find the span with the maximum sum of scores for start and end indices.
+        total_scores = torch.unsqueeze(start_scores, axis=2) + torch.unsqueeze(end_scores, axis=1)
+        start_idx = torch.arange(max_num_tokens, device=total_scores.get_device()).view(1, -1, 1)
+        end_idx = torch.arange(max_num_tokens, device=total_scores.get_device()).view(1, 1, -1)
+        invalid_index_mask = (start_idx > end_idx).repeat(batch_size, 1, 1)
+        total_scores = torch.where(
+            invalid_index_mask,
+            torch.zeros(total_scores.size(), device=total_scores.get_device(), dtype=total_scores.dtype),
+            total_scores,
+        )
+        max_span_index = torch.argmax(total_scores.view(-1, max_num_tokens ** 2), axis=-1)
+        max_span_p = torch.max(total_scores.view(-1, max_num_tokens ** 2), axis=-1)[0]
+
+        span_start_index = torch.floor_divide(max_span_index, max_num_tokens)
+        span_end_index = torch.fmod(max_span_index, max_num_tokens)
+
+        tensors = {
+            'example_id_num': example_id_num,
+            'service_id': service_id,
+            'intent_status': intent_status,
+            'req_slot_status': req_slot_status,
+            'cat_slot_status': cat_slot_status,
+            'cat_slot_status_p': cat_slot_status_p,
+            'cat_slot_value_status': cat_slot_value_status,
+            'noncat_slot_status': noncat_slot_status,
+            'noncat_slot_status_p': noncat_slot_status_p,
+            'noncat_slot_p': max_span_p,
+            'noncat_slot_start': span_start_index,
+            'noncat_slot_end': span_end_index,
+            'noncat_alignment_start': start_char_idx,
+            'noncat_alignment_end': end_char_idx,
+        }
+        return loss, tensors
+
+    def multi_validation_epoch_end(self, outputs: List[dict], dataloader_idx: int = 0):
+        """
+        Called at the end of validation to post process outputs into human readable format
+        Args:
+            outputs: list of individual outputs of each validation step
+            dataloader_idx: dataloader index
+        """
+        avg_loss = torch.stack([x[f'val_loss'] for x in outputs]).mean()
+        split = self._validation_names[dataloader_idx][:-1]
+        dataloader = self._validation_dl[dataloader_idx]
+        metrics = self.multi_eval_epoch_end_helper(outputs=outputs, split=split, dataloader=dataloader)
+
+        for k, v in metrics.items():
+            self.log(f'{split}_{k}', v)
+
+        self.log(f'val_loss', avg_loss, prog_bar=True)
+
+    def multi_test_epoch_end(self, outputs: List[dict], dataloader_idx: int = 0):
+        """
+        Called at the end of test to post process outputs into human readable format
+        Args:
+            outputs: list of individual outputs of each test step
+            dataloader_idx: dataloader index
+        """
+        avg_loss = torch.stack([x[f'test_loss'] for x in outputs]).mean()
+        split = self._test_names[dataloader_idx][:-1]
+        dataloader = self._test_dl[dataloader_idx]
+        metrics = self.multi_eval_epoch_end_helper(outputs=outputs, split=split, dataloader=dataloader)
+
+        for k, v in metrics.items():
+            self.log(f'{split}_{k}', v)
+
+        self.log(f'test_loss', avg_loss, prog_bar=True)
+
+    def multi_eval_epoch_end_helper(
+        self, outputs: List[dict], split: str, dataloader: torch.utils.data.DataLoader
+    ) -> dict:
+        """
+        Helper called at the end of evaluation to post process outputs into human readable format
+        Args:
+            outputs: list of individual outputs of each step
+            split: data split
+            dataloader: dataloader
+        Returns:
+            metrics: metrics collection
+        """
+
+        def get_str_example_id(split: str, ids_to_service_names_dict: dict, example_id_num: torch.Tensor) -> str:
+            """
+            Constructs string representation of example ID
+            Args:
+                split: evaluation data split
+                ids_to_service_names_dict: id to service name mapping
+                example_id_num: tensor example id
+            """
+
+            def format_turn_id(ex_id_num):
+                dialog_id_1, dialog_id_2, turn_id, service_id, model_task_id, slot_intent_id, value_id = ex_id_num
+                return "{}-{}_{:05d}-{:02d}-{}-{}-{}-{}".format(
+                    split,
+                    dialog_id_1,
+                    dialog_id_2,
+                    turn_id,
+                    ids_to_service_names_dict[service_id],
+                    model_task_id,
+                    slot_intent_id,
+                    value_id,
+                )
+
+            return list(map(format_turn_id, tensor2list(example_id_num)))
+
+        def combine_predictions_in_example(predictions: dict, batch_size: int):
+            '''
+            Combines predicted values to a single example. 
+            Args:
+                predictions: predictions ordered by keys then batch
+                batch_size: batch size
+            Returns:
+                examples_preds: predictions ordered by batch then key
+            '''
+            examples_preds = [{} for _ in range(batch_size)]
+            for k, v in predictions.items():
+                if k != 'example_id':
+                    v = torch.chunk(v, batch_size)
+
+                for i in range(batch_size):
+                    if k == 'example_id':
+                        examples_preds[i][k] = v[i]
+                    else:
+                        examples_preds[i][k] = v[i].view(-1)
+            return examples_preds
+
+        example_id_num = torch.cat([x[f'tensors']['example_id_num'] for x in outputs])
+        service_id = torch.cat([x[f'tensors']['service_id'] for x in outputs])
+        intent_status = torch.cat([x[f'tensors']['intent_status'] for x in outputs])
+        req_slot_status = torch.cat([x[f'tensors']['req_slot_status'] for x in outputs])
+        cat_slot_status = torch.cat([x[f'tensors']['cat_slot_status'] for x in outputs])
+        cat_slot_status_p = torch.cat([x[f'tensors']['cat_slot_status_p'] for x in outputs])
+        cat_slot_value_status = torch.cat([x[f'tensors']['cat_slot_value_status'] for x in outputs])
+        noncat_slot_status = torch.cat([x[f'tensors']['noncat_slot_status'] for x in outputs])
+        noncat_slot_status_p = torch.cat([x[f'tensors']['noncat_slot_status_p'] for x in outputs])
+        noncat_slot_p = torch.cat([x[f'tensors']['noncat_slot_p'] for x in outputs])
+        noncat_slot_start = torch.cat([x[f'tensors']['noncat_slot_start'] for x in outputs])
+        noncat_slot_end = torch.cat([x[f'tensors']['noncat_slot_end'] for x in outputs])
+        noncat_alignment_start = torch.cat([x[f'tensors']['noncat_alignment_start'] for x in outputs])
+        noncat_alignment_end = torch.cat([x[f'tensors']['noncat_alignment_end'] for x in outputs])
+
+        ids_to_service_names_dict = self.dialogues_processor.schemas._services_id_to_vocab
+        example_id = get_str_example_id(dataloader.dataset, ids_to_service_names_dict, example_id_num)
+
+        metrics = {}
+        try:
+            prediction_dir = self.trainer.log_dir if self.trainer.log_dir is not None else ""
+        except:
+            prediction_dir = ""
+
+        if self.trainer.global_rank == 0:
+            prediction_dir = os.path.join(
+                prediction_dir, 'predictions', 'pred_res_{}_{}'.format(split, self._cfg.dataset.task_name)
+            )
+            os.makedirs(prediction_dir, exist_ok=True)
+
+            input_json_files = SGDDataProcessor.get_dialogue_files(
+                self._cfg.dataset.data_dir, split, self._cfg.dataset.task_name
+            )
+
+            predictions = {}
+            predictions['example_id'] = example_id
+            predictions['service_id'] = service_id
+            predictions['intent_status'] = intent_status
+            predictions['req_slot_status'] = req_slot_status
+            predictions['cat_slot_status'] = cat_slot_status
+            predictions['cat_slot_status_p'] = cat_slot_status_p
+            predictions['cat_slot_value_status'] = cat_slot_value_status
+            predictions['noncat_slot_status'] = noncat_slot_status
+            predictions['noncat_slot_status_p'] = noncat_slot_status_p
+            predictions['noncat_slot_p'] = noncat_slot_p
+            predictions['noncat_slot_start'] = noncat_slot_start
+            predictions['noncat_slot_end'] = noncat_slot_end
+            predictions['noncat_alignment_start'] = noncat_alignment_start
+            predictions['noncat_alignment_end'] = noncat_alignment_end
+
+            in_domain_services = get_in_domain_services(
+                os.path.join(self._cfg.dataset.data_dir, split, "schema.json"),
+                self.dialogues_processor.get_seen_services("train"),
+            )
+            predictions = combine_predictions_in_example(predictions, service_id.shape[0])
+
+            # write predictions to file in Dstc8/SGD format
+            write_predictions_to_file(
+                predictions,
+                input_json_files,
+                output_dir=prediction_dir,
+                schemas=self.dialogues_processor.schemas,
+                state_tracker=self._cfg.dataset.state_tracker,
+                eval_debug=False,
+                in_domain_services=in_domain_services,
+            )
+            metrics = evaluate(
+                prediction_dir,
+                self._cfg.dataset.data_dir,
+                split,
+                in_domain_services,
+                joint_acc_across_turn=self._cfg.dataset.joint_acc_across_turn,
+                use_fuzzy_match=self._cfg.dataset.use_fuzzy_match,
+            )
+
+        return metrics
+
+    def prepare_data(self):
+        """
+        Preprocessed schema and dialogues and caches this
+        """
+        if self.data_prepared:
+            return
+        schema_config = {
+            "MAX_NUM_CAT_SLOT": self._cfg.dataset.max_num_cat_slot,
+            "MAX_NUM_NONCAT_SLOT": self._cfg.dataset.max_num_noncat_slot,
+            "MAX_NUM_VALUE_PER_CAT_SLOT": self._cfg.dataset.max_value_per_cat_slot,
+            "MAX_NUM_INTENT": self._cfg.dataset.max_num_intent,
+            "NUM_TASKS": NUM_TASKS,
+            "MAX_SEQ_LENGTH": self._cfg.dataset.max_seq_length,
+        }
+        all_schema_json_paths = []
+        for dataset_split in ['train', 'test', 'dev']:
+            all_schema_json_paths.append(os.path.join(self._cfg.dataset.data_dir, dataset_split, "schema.json"))
+        schemas = Schema(all_schema_json_paths)
+
+        self.dialogues_processor = SGDDataProcessor(
+            task_name=self._cfg.dataset.task_name,
+            data_dir=self._cfg.dataset.data_dir,
+            dialogues_example_dir=self._cfg.dataset.dialogues_example_dir,
+            tokenizer=self.tokenizer,
+            schemas=schemas,
+            schema_config=schema_config,
+            subsample=self._cfg.dataset.subsample,
+        )
+
+        if is_global_rank_zero():
+            overwrite_dial_files = not self._cfg.dataset.use_cache
+            self.dialogues_processor.save_dialog_examples(overwrite_dial_files=overwrite_dial_files)
+
+        self.data_prepared = True
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig] = None):
+        self.prepare_data()
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config, split=train_data_config.ds_item)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig] = None):
+        self.prepare_data()
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config, split=val_data_config.ds_item)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig] = None):
+        self.prepare_data()
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config, split=test_data_config.ds_item)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig, split: str) -> DataLoader:
+        dataset_cfg = self._cfg.dataset
+        data_dir = dataset_cfg.data_dir
+
+        if not os.path.exists(data_dir):
+            raise FileNotFoundError(f"Data directory is not found at: {data_dir}.")
+
+        dataset = SGDDataset(dataset_split=split, dialogues_processor=self.dialogues_processor)
+
+        dl = torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            collate_fn=dataset.collate_fn,
+            drop_last=cfg.drop_last,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.num_workers,
+            pin_memory=cfg.pin_memory,
+        )
+        return dl
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        pass

tts/nemo/collections/nlp/models/enc_dec_nlp_model.py

@@ -0,0 +1,85 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import Any
+
+from omegaconf.omegaconf import MISSING
+from pytorch_lightning.trainer.trainer import Trainer
+
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common.decoder_module import DecoderModule
+from nemo.collections.nlp.modules.common.encoder_module import EncoderModule
+from nemo.collections.nlp.modules.common.tokenizer_utils import TokenizerConfig
+from nemo.core.config.modelPT import ModelConfig
+
+
+@dataclass
+class EncDecNLPModelConfig(ModelConfig):
+    encoder_tokenizer: TokenizerConfig = MISSING
+    decoder_tokenizer: TokenizerConfig = MISSING
+    encoder: Any = MISSING
+    decoder: Any = MISSING
+    head: Any = MISSING
+
+
+class EncDecNLPModel(NLPModel):
+    """Base class for encoder-decoder NLP models.
+    """
+
+    def __init__(self, cfg: EncDecNLPModelConfig, trainer: Trainer = None):
+        super().__init__(cfg=cfg, trainer=trainer)
+
+    @property
+    def encoder_vocab_size(self):
+        return self.encoder_tokenizer.vocab_size
+
+    @property
+    def decoder_vocab_size(self):
+        return self.decoder_tokenizer.vocab_size
+
+    @property
+    def encoder_tokenizer(self):
+        return self._encoder_tokenizer
+
+    @encoder_tokenizer.setter
+    def encoder_tokenizer(self, tokenizer):
+        self._encoder_tokenizer = tokenizer
+
+    @property
+    def decoder_tokenizer(self):
+        return self._decoder_tokenizer
+
+    @decoder_tokenizer.setter
+    def decoder_tokenizer(self, tokenizer):
+        self._decoder_tokenizer = tokenizer
+
+    @property
+    def encoder(self) -> EncoderModule:
+        return self._encoder
+
+    @encoder.setter
+    def encoder(self, encoder):
+        self._encoder = encoder
+
+    @property
+    def decoder(self) -> DecoderModule:
+        return self._decoder
+
+    @decoder.setter
+    def decoder(self, decoder):
+        self._decoder = decoder
+
+    def export(self, **kwargs):
+        raise NotImplementedError('For EncDecNLPModel, you must export encoder and decoder separately!')

tts/nemo/collections/nlp/models/entity_linking/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+from nemo.collections.nlp.models.entity_linking.entity_linking_model import EntityLinkingModel

tts/nemo/collections/nlp/models/entity_linking/entity_linking_model.py

@@ -0,0 +1,193 @@
+# 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.
+
+from typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig
+from pytorch_lightning import Trainer
+from transformers import AutoTokenizer
+
+from nemo.collections.common.losses import MultiSimilarityLoss
+from nemo.collections.nlp.data import EntityLinkingDataset
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+from nemo.core.neural_types import ChannelType, LogitsType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['EntityLinkingModel']
+
+
+class EntityLinkingModel(NLPModel, Exportable):
+    """
+    Second stage pretraining of BERT based language model
+    for entity linking task. An implementation of Liu et. al's
+    NAACL 2021 paper Self-Alignment Pretraining for Biomedical Entity Representations.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "attention_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "token_type_ids": NeuralType(('B', 'T'), ChannelType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"logits": NeuralType(('B', 'D'), LogitsType())}
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """Initializes the SAP-BERT model for entity linking."""
+
+        # tokenizer needed before super().__init__() so dataset and loader can process data
+        self._setup_tokenizer(cfg.tokenizer)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=cfg.language_model.config_file,
+            config_dict=cfg.language_model.config,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+        )
+
+        # Token to use for the self-alignment loss, typically the first token, [CLS]
+        self._idx_conditioned_on = 0
+        self.loss = MultiSimilarityLoss()
+
+    def _setup_tokenizer(self, cfg: DictConfig):
+        tokenizer = AutoTokenizer.from_pretrained(
+            cfg.tokenizer_name, vocab_file=cfg.vocab_file, do_lower_case=cfg.do_lower_case
+        )
+
+        self.tokenizer = tokenizer
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        hidden_states = self.model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
+
+        # normalize to unit sphere
+        logits = torch.nn.functional.normalize(hidden_states[:, self._idx_conditioned_on], p=2, dim=1)
+        return logits
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        input_ids, token_type_ids, attention_mask, concept_ids = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
+        train_loss = self.loss(logits=logits, labels=concept_ids)
+
+        # No hard examples found in batch,
+        # shouldn't use this batch to update model weights
+        if train_loss == 0:
+            train_loss = None
+            lr = None
+
+        else:
+            lr = self._optimizer.param_groups[0]["lr"]
+            self.log("train_loss", train_loss)
+            self.log("lr", lr, prog_bar=True)
+
+        return {"loss": train_loss, "lr": lr}
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, concept_ids = batch
+        with torch.no_grad():
+            logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+            val_loss = self.loss(logits=logits, labels=concept_ids)
+
+        # No hard examples found in batch,
+        # val loss not used to update model weights
+        if val_loss == 0:
+            val_loss = None
+        else:
+            self.log("val_loss", val_loss)
+            logging.info(f"val loss: {val_loss}")
+
+        return {"val_loss": val_loss}
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+
+        Args:
+            outputs: list of individual outputs of each validation step.
+        Returns:
+            
+        """
+        if outputs:
+            avg_loss = torch.stack([x["val_loss"] for x in outputs if x["val_loss"] != None]).mean()
+            self.log(f"val_loss", avg_loss, prog_bar=True)
+
+            return {"val_loss": avg_loss}
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        if not train_data_config or not train_data_config.data_file:
+            logging.info(
+                f"Dataloader config or file_path or processed data path for the train dataset is missing, \
+                        so no data loader for train is created!"
+            )
+
+            self._train_dl = None
+            return
+
+        self._train_dl = self.setup_dataloader(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        if not val_data_config or not val_data_config.data_file:
+            logging.info(
+                f"Dataloader config or file_path or processed data path for the val dataset is missing, \
+                        so no data loader for validation is created!"
+            )
+
+            self._validation_dl = None
+            return
+
+        self._validation_dl = self.setup_dataloader(cfg=val_data_config)
+
+    def setup_dataloader(self, cfg: Dict, is_index_data: bool = False) -> 'torch.utils.data.DataLoader':
+
+        dataset = EntityLinkingDataset(
+            tokenizer=self.tokenizer,
+            data_file=cfg.data_file,
+            max_seq_length=cfg.max_seq_length,
+            is_index_data=is_index_data,
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            collate_fn=dataset.collate_fn,
+            shuffle=cfg.get("shuffle", True),
+            num_workers=cfg.get("num_wokers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass
+
+    @classmethod
+    def from_pretrained(cls, name: str):
+        pass

tts/nemo/collections/nlp/models/glue_benchmark/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.collections.nlp.models.glue_benchmark.glue_benchmark_model import GLUEModel

tts/nemo/collections/nlp/models/glue_benchmark/glue_benchmark_model.py

@@ -0,0 +1,279 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 typing import Dict, Optional, Union
+
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import CrossEntropyLoss, MSELoss
+from nemo.collections.nlp.data.glue_benchmark.glue_benchmark_dataset import GLUE_TASKS_NUM_LABELS, GLUEDataset
+from nemo.collections.nlp.models.glue_benchmark.metrics_for_glue import compute_metrics
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import SequenceClassifier, SequenceRegression
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import list2str, tensor2list
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ['GLUEModel']
+
+'''
+Some transformer of this code were adapted from the HuggingFace library at
+https://github.com/huggingface/transformers
+Example of running a pretrained BERT model on the 9 GLUE tasks, read more
+about GLUE benchmark here: https://gluebenchmark.com
+Download the GLUE data by running the script:
+https://gist.github.com/W4ngatang/60c2bdb54d156a41194446737ce03e2e
+
+Some of these tasks have a small dataset and training can lead to high variance
+in the results between different runs. Below is the median on 5 runs
+(with different seeds) for each of the metrics on the dev set of the benchmark
+with an uncased BERT base model (the checkpoint bert-base-uncased)
+(source https://github.com/huggingface/transformers/tree/master/examples#glue).
+Task	Metric	                        Result
+CoLA	Matthew's corr	                48.87
+SST-2	Accuracy	                    91.74
+MRPC	F1/Accuracy	                 90.70/86.27
+STS-B	Person/Spearman corr.	     91.39/91.04
+QQP	    Accuracy/F1	                 90.79/87.66
+MNLI	Matched acc./Mismatched acc. 83.70/84.83
+QNLI	Accuracy	                    89.31
+RTE	    Accuracy	                    71.43
+WNLI	Accuracy	                    43.66
+
+'''
+
+
+class GLUEModel(NLPModel):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.pooler.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """
+        Initializes model to use BERT model for GLUE tasks.
+        """
+
+        if cfg.task_name not in cfg.supported_tasks:
+            raise ValueError(f'{cfg.task_name} not in supported task. Choose from {cfg.supported_tasks}')
+        self.task_name = cfg.task_name
+
+        # needed to setup validation on multiple datasets
+        # MNLI task has two separate dev sets: matched and mismatched
+        if not self._is_model_being_restored():
+            if self.task_name == "mnli":
+                cfg.validation_ds.ds_item = [
+                    os.path.join(cfg.dataset.data_dir, 'dev_matched.tsv'),
+                    os.path.join(cfg.dataset.data_dir, 'dev_mismatched.tsv'),
+                ]
+            else:
+                cfg.validation_ds.ds_item = os.path.join(cfg.dataset.data_dir, cfg.validation_ds.ds_item)
+            cfg.train_ds.ds_item = os.path.join(cfg.dataset.data_dir, cfg.train_ds.ds_item)
+            logging.info(f'Using {cfg.validation_ds.ds_item} for model evaluation.')
+
+        self.setup_tokenizer(cfg.tokenizer)
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        num_labels = GLUE_TASKS_NUM_LABELS[self.task_name]
+
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        # uses [CLS] token for classification (the first token)
+        if self.task_name == "sts-b":
+            self.pooler = SequenceRegression(hidden_size=self.bert_model.config.hidden_size)
+            self.loss = MSELoss()
+        else:
+            self.pooler = SequenceClassifier(
+                hidden_size=self.bert_model.config.hidden_size, num_classes=num_labels, log_softmax=False
+            )
+            self.loss = CrossEntropyLoss()
+
+    def update_data_dir(self, data_dir: str) -> None:
+        """
+        Update data directory and get data stats with Data Descriptor
+        Weights are later used to setup loss
+
+        Args:
+            data_dir: path to data directory
+        """
+        self._cfg.dataset.data_dir = data_dir
+        logging.info(f'Setting model.dataset.data_dir to {data_dir}.')
+        if self.task_name == "mnli":
+            self._cfg.validation_ds.ds_item = [
+                os.path.join(data_dir, 'dev_matched.tsv'),
+                os.path.join(data_dir, 'dev_mismatched.tsv'),
+            ]
+        else:
+            self._cfg.validation_ds.ds_item = os.path.join(data_dir, 'dev.tsv')
+
+        self._cfg.train_ds.ds_item = os.path.join(data_dir, 'train.tsv')
+        logging.info(f'Using {self._cfg.validation_ds.ds_item} for model evaluation.')
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        output = self.pooler(hidden_states=hidden_states)
+        return output
+
+    def training_step(self, batch, batch_idx):
+        input_ids, input_type_ids, input_mask, labels = batch
+        model_output = self(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        if self.task_name == "sts-b":
+            loss = self.loss(preds=model_output, labels=labels)
+        else:
+            loss = self.loss(logits=model_output, labels=labels)
+
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': loss,
+            'lr': lr,
+        }
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        input_ids, input_type_ids, input_mask, labels = batch
+        model_output = self(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        if self.task_name == "sts-b":
+            val_loss = self.loss(preds=model_output, labels=labels)
+        else:
+            val_loss = self.loss(logits=model_output, labels=labels)
+
+        if self.task_name != 'sts-b':
+            model_output = torch.argmax(model_output, 1)
+
+        eval_tensors = {'preds': model_output, 'labels': labels}
+        return {'val_loss': val_loss, 'eval_tensors': eval_tensors}
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        """
+        Called at the end of validation to aggregate outputs.
+        outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+        preds = torch.cat([x['eval_tensors']['preds'] for x in outputs])
+        labels = torch.cat([x['eval_tensors']['labels'] for x in outputs])
+
+        all_preds = []
+        all_labels = []
+        if torch.distributed.is_initialized():
+            world_size = torch.distributed.get_world_size()
+            for ind in range(world_size):
+                all_preds.append(torch.empty_like(preds))
+                all_labels.append(torch.empty_like(labels))
+            torch.distributed.all_gather(all_preds, preds)
+            torch.distributed.all_gather(all_labels, labels)
+        else:
+            all_preds.append(preds)
+            all_labels.append(labels)
+
+        tensorboard_logs = {}
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            preds = []
+            labels = []
+            for p in all_preds:
+                preds.extend(tensor2list(p))
+            for l in all_labels:
+                labels.extend(tensor2list(l))
+
+            tensorboard_logs = compute_metrics(self.task_name, np.array(preds), np.array(labels))
+            val_name = self._validation_names[dataloader_idx].upper()
+            logging.info(f'{val_name} evaluation: {tensorboard_logs}')
+
+            # writing labels and predictions to a file in output_dir is specified in the config
+            output_dir = self._cfg.output_dir
+            if output_dir:
+                os.makedirs(output_dir, exist_ok=True)
+                filename = os.path.join(output_dir, f'{self.task_name}_{val_name}.txt')
+                logging.info(f'Saving labels and predictions to {filename}')
+                with open(filename, 'w') as f:
+                    f.write('labels\t' + list2str(labels) + '\n')
+                    f.write('preds\t' + list2str(preds) + '\n')
+
+        tensorboard_logs['val_loss'] = avg_loss
+        for key in tensorboard_logs:
+            self.log(f'{key}', tensorboard_logs[key], prog_bar=True)
+
+        return {'val_loss': avg_loss, 'log': tensorboard_logs}
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig] = None):
+        if train_data_config is None:
+            train_data_config = self._cfg.train_ds
+
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig] = None):
+        if val_data_config is None:
+            val_data_config = self._cfg.validation_ds
+
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_multiple_validation_data(self, val_data_config: Union[DictConfig, Dict] = None):
+        if val_data_config is None:
+            val_data_config = self._cfg.validation_ds
+
+        return super().setup_multiple_validation_data(val_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+        file_name = cfg.ds_item
+        if not os.path.exists(file_name):
+            raise FileNotFoundError(
+                "GLUE datasets not found. For more details on how to get the data, see: "
+                "https://gist.github.com/W4ngatang/60c2bdb54d156a41194446737ce03e2e"
+            )
+
+        dataset = GLUEDataset(
+            file_name=file_name,
+            task_name=self.task_name,
+            tokenizer=self.tokenizer,
+            max_seq_length=self._cfg.dataset.max_seq_length,
+            use_cache=self._cfg.dataset.use_cache,
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=self._cfg.dataset.num_workers,
+            pin_memory=self._cfg.dataset.pin_memory,
+            drop_last=self._cfg.dataset.drop_last,
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass

tts/nemo/collections/nlp/models/glue_benchmark/metrics_for_glue.py

@@ -0,0 +1,66 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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 typing import Dict, List
+
+from scipy.stats import pearsonr, spearmanr
+from sklearn.metrics import f1_score, matthews_corrcoef
+
+__all__ = ['compute_metrics']
+
+
+def accuracy(preds: List[int], labels: List[int]):
+    return {"acc": (preds == labels).mean()}
+
+
+def acc_and_f1(preds: List[int], labels: List[int]):
+    accuracy = (preds == labels).mean()
+    f1 = f1_score(y_true=labels, y_pred=preds)
+    return {"acc": accuracy, "f1": f1}
+
+
+def mcc(preds: List[int], labels: List[int]):
+    return {"mcc": matthews_corrcoef(labels, preds)}
+
+
+def pearson_and_spearman(preds: List[int], labels: List[int]):
+    pearson_corr = pearsonr(preds, labels)[0]
+    spearman_corr = spearmanr(preds, labels)[0]
+    return {"pearson": pearson_corr, "spearmanr": spearman_corr, "pear+spear av": (pearson_corr + spearman_corr) / 2}
+
+
+def compute_metrics(task_name: str, preds: List[int], labels: List[int]) -> Dict[str, float]:
+    """
+    Computes metrics for GLUE tasks
+    Args:
+        task_name: GLUE task name
+        preds: model predictions
+        labels: golden labels
+    Returns:
+        metrics
+    """
+    if len(preds) != len(labels):
+        raise ValueError("Predictions and labels must have the same length")
+
+    metric_fn = accuracy
+    if task_name == 'cola':
+        metric_fn = mcc
+    elif task_name in ['mrpc', 'qqp']:
+        metric_fn = acc_and_f1
+    elif task_name == 'sts-b':
+        metric_fn = pearson_and_spearman
+
+    return metric_fn(preds, labels)

tts/nemo/collections/nlp/models/information_retrieval/__init__.py

@@ -0,0 +1,16 @@
+# 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 nemo.collections.nlp.models.information_retrieval.bert_dpr_model import BertDPRModel
+from nemo.collections.nlp.models.information_retrieval.bert_joint_ir_model import BertJointIRModel

tts/nemo/collections/nlp/models/information_retrieval/base_ir_model.py

@@ -0,0 +1,211 @@
+# 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 math
+from typing import Dict, Optional
+
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.nlp.data import BertInformationRetrievalDataset
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.core.classes.common import typecheck
+
+__all__ = ['BaseIRModel']
+
+
+class BaseIRModel(NLPModel):
+    """
+    Base class for information retrieval models.
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        self.setup_tokenizer(cfg.tokenizer)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+    @typecheck()
+    def forward(self, *args):
+        pass
+
+    def compute_scores_and_loss(self, inputs):
+        pass
+
+    @staticmethod
+    def get_lm_model_with_padded_embedding(cfg: DictConfig):
+        """
+        Function which ensures that vocabulary size is divisivble by 8
+        for faster mixed precision training.
+        """
+        model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=cfg.language_model.config_file,
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=cfg.tokenizer.vocab_file,
+        )
+        vocab_size, hidden_size = model.config.vocab_size, model.config.hidden_size
+        tokens_to_add = 8 * math.ceil(vocab_size / 8) - vocab_size
+        zeros = torch.zeros((tokens_to_add, hidden_size))
+        model.embeddings.word_embeddings.weight.data = torch.cat((model.embeddings.word_embeddings.weight.data, zeros))
+        return model
+
+    @staticmethod
+    def calculate_mean_reciprocal_rank(query2passages, query2rel):
+        """
+        Helper function which calculates mean reciprocal rank.
+        Args:
+            query2passages: dict which contains passage ids and corresponding
+                scores for each query
+            query2rel: dict which contains ids of relevant passages for each query
+        """
+        reciprocal_ranks = []
+
+        for query in query2passages:
+            indices = np.argsort(query2passages[query]["scores"])[::-1]
+            sorted_psgs = query2passages[query]["psg_ids"][indices]
+
+            reciprocal_ranks.append(0)
+            for i, psg_id in enumerate(sorted_psgs):
+                if psg_id in query2rel[query]:
+                    reciprocal_ranks[-1] = 1 / (i + 1)
+                    break
+        return np.mean(reciprocal_ranks)
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        scores, train_loss = self.compute_scores_and_loss(batch[:-2])
+        tensorboard_logs = {"train_loss": train_loss, "lr": self._optimizer.param_groups[0]["lr"]}
+        return {"loss": train_loss, "log": tensorboard_logs}
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        scores, val_loss = self.compute_scores_and_loss(batch[:-2])
+        query_ids, passage_ids = batch[-2:]
+        data_for_val = {
+            "val_loss": val_loss,
+            "scores": scores,
+            "query_ids": query_ids,
+            "passage_ids": passage_ids,
+        }
+        return data_for_val
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+
+        query_ids = torch.cat([x["query_ids"] for x in outputs])
+        passage_ids = torch.cat([x["passage_ids"] for x in outputs])
+        scores = torch.cat([x["scores"] for x in outputs])
+
+        all_query_ids, all_passage_ids, all_scores = [], [], []
+        if torch.distributed.is_initialized():
+            world_size = torch.distributed.get_world_size()
+            for ind in range(world_size):
+                all_query_ids.append(torch.empty_like(query_ids))
+                all_passage_ids.append(torch.empty_like(passage_ids))
+                all_scores.append(torch.empty_like(scores))
+            torch.distributed.all_gather(all_query_ids, query_ids)
+            torch.distributed.all_gather(all_passage_ids, passage_ids)
+            torch.distributed.all_gather(all_scores, scores)
+        else:
+            all_query_ids.append(query_ids)
+            all_passage_ids.append(passage_ids)
+            all_scores.append(scores)
+
+        val_mrr = 0
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            query2passages, query2rels = {}, {}
+            processed_queries = set()
+
+            for i in range(len(all_query_ids)):
+
+                query_ids = all_query_ids[i].detach().cpu().numpy()
+                passage_ids = all_passage_ids[i].detach().cpu().numpy()
+                scores = all_scores[i].detach().cpu().numpy()
+
+                for j, query_id in enumerate(query_ids):
+
+                    if query_id not in processed_queries:
+                        processed_queries.add(query_id)
+                        query2passages[query_id] = {
+                            "psg_ids": passage_ids[j],
+                            "scores": scores[j],
+                        }
+                        query2rels[query_id] = [passage_ids[j][0]]
+                    else:
+                        query2passages[query_id]["psg_ids"] = np.concatenate(
+                            (query2passages[query_id]["psg_ids"], passage_ids[j][1:])
+                        )
+                        query2passages[query_id]["scores"] = np.concatenate(
+                            (query2passages[query_id]["scores"], scores[j][1:])
+                        )
+
+            val_mrr = self.calculate_mean_reciprocal_rank(query2passages, query2rels)
+
+        val_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+
+        tensorboard_logs = {
+            "val_mrr": val_mrr,
+            "val_loss": val_loss,
+        }
+
+        return {"log": tensorboard_logs}
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+
+        dataset = BertInformationRetrievalDataset(
+            tokenizer=self.tokenizer,
+            passages=cfg.passages,
+            queries=cfg.queries,
+            query_to_passages=cfg.query_to_passages,
+            num_negatives=cfg.num_negatives,
+            psg_cache_format=cfg.get("psg_cache_format", "pkl"),
+            max_query_length=cfg.get("max_query_length", 31),
+            max_passage_length=cfg.get("max_passage_length", 190),
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass

tts/nemo/collections/nlp/models/information_retrieval/bert_dpr_model.py

@@ -0,0 +1,131 @@
+# 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 typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import SmoothedCrossEntropyLoss
+from nemo.collections.nlp.data import BertInformationRetrievalDataset
+from nemo.collections.nlp.models.information_retrieval.base_ir_model import BaseIRModel
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import typecheck
+from nemo.core.neural_types import ChannelType, LogitsType, MaskType, NeuralType
+
+__all__ = ["BertDPRModel"]
+
+
+class BertDPRModel(BaseIRModel):
+    """
+    Information retrieval model which encodes query and passage separately
+    with two different BERT encoders and computes their similarity score
+    as a dot-product between corresponding [CLS] token representations.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "q_input_ids": NeuralType(("B", "T"), ChannelType()),
+            "q_attention_mask": NeuralType(("B", "T"), MaskType()),
+            "q_token_type_ids": NeuralType(("B", "T"), ChannelType()),
+            "p_input_ids": NeuralType(("B", "T"), ChannelType()),
+            "p_attention_mask": NeuralType(("B", "T"), MaskType()),
+            "p_token_type_ids": NeuralType(("B", "T"), ChannelType()),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"logits": NeuralType(("B", "D"), LogitsType())}
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        model_name = cfg.language_model.pretrained_model_name
+        self.tokenizer = get_tokenizer(tokenizer_name=model_name)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.q_encoder = self.get_lm_model_with_padded_embedding(cfg)
+        self.p_encoder = self.get_lm_model_with_padded_embedding(cfg)
+        self.loss = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id)
+
+    @typecheck()
+    def forward(
+        self, q_input_ids, q_token_type_ids, q_attention_mask, p_input_ids, p_token_type_ids, p_attention_mask,
+    ):
+
+        q_vectors = self.q_encoder(
+            input_ids=q_input_ids, token_type_ids=q_token_type_ids, attention_mask=q_attention_mask,
+        )
+        q_vectors = q_vectors[:, 0]
+        batch_size, hidden_size = q_vectors.size()
+
+        p_vectors = self.p_encoder(
+            input_ids=p_input_ids, token_type_ids=p_token_type_ids, attention_mask=p_attention_mask,
+        )
+        num_passages = p_vectors.shape[0] // batch_size
+        p_vectors = p_vectors[:, 0].view(-1, num_passages, hidden_size)
+        p_positives, p_negatives = p_vectors[:, 0], p_vectors[:, 1:]
+        scores = torch.cat(
+            (torch.matmul(q_vectors, p_positives.T), torch.einsum("ij,ipj->ip", q_vectors, p_negatives),), dim=1,
+        )
+
+        return scores
+
+    def compute_scores_and_loss(self, inputs):
+        (q_input_ids, q_input_mask, q_input_type_ids, p_input_ids, p_input_mask, p_input_type_ids,) = inputs
+        batch_size, num_passages, p_seq_length = p_input_ids.size()
+        q_seq_length = q_input_ids.size()[-1]
+
+        scores = self(
+            q_input_ids=q_input_ids.view(-1, q_seq_length),
+            q_token_type_ids=q_input_type_ids.view(-1, q_seq_length),
+            q_attention_mask=q_input_mask.view(-1, q_seq_length),
+            p_input_ids=p_input_ids.view(-1, p_seq_length),
+            p_token_type_ids=p_input_type_ids.view(-1, p_seq_length),
+            p_attention_mask=p_input_mask.view(-1, p_seq_length),
+        ).view(batch_size, 1, batch_size + num_passages - 1)
+        normalized_scores = torch.log_softmax(scores, dim=-1)
+
+        labels = torch.arange(batch_size)[:, None].long().to(normalized_scores.device)
+        loss = self.loss(log_probs=normalized_scores, labels=labels, output_mask=torch.ones_like(labels),)
+
+        scores = scores[:, 0]
+        scores = torch.cat((torch.diag(scores)[:, None], scores[:, batch_size:]), dim=1,)
+
+        return scores, loss
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+
+        dataset = BertInformationRetrievalDataset(
+            tokenizer=self.tokenizer,
+            passages=cfg.passages,
+            queries=cfg.queries,
+            query_to_passages=cfg.query_to_passages,
+            num_negatives=cfg.num_negatives,
+            psg_cache_format=cfg.get("psg_cache_format", "pkl"),
+            max_query_length=cfg.get("max_query_length", 31),
+            max_passage_length=cfg.get("max_passage_length", 190),
+            preprocess_fn="preprocess_dpr",
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )

tts/nemo/collections/nlp/models/information_retrieval/bert_joint_ir_model.py

@@ -0,0 +1,84 @@
+# 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 typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import SmoothedCrossEntropyLoss
+from nemo.collections.nlp.models.information_retrieval.base_ir_model import BaseIRModel
+from nemo.collections.nlp.modules.common import SequenceRegression
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import typecheck
+from nemo.core.neural_types import NeuralType
+
+__all__ = ["BertJointIRModel"]
+
+
+class BertJointIRModel(BaseIRModel):
+    """
+    Information retrieval model which jointly encodes both query and passage
+    and passes them to BERT encoder followed by a fully-connected layer for
+    similarity score prediction.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.sim_score_regressor.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        model_name = cfg.language_model.pretrained_model_name
+        self.tokenizer = get_tokenizer(tokenizer_name=model_name)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.bert_model = self.get_lm_model_with_padded_embedding(cfg)
+        hidden_size = self.bert_model.config.hidden_size
+        self.sim_score_regressor = SequenceRegression(
+            hidden_size=hidden_size, num_layers=1, dropout=cfg.language_model.sim_score_dropout,
+        )
+        self.loss = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id)
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask,
+        )
+        scores = self.sim_score_regressor(hidden_states=hidden_states)
+
+        return scores
+
+    def compute_scores_and_loss(self, inputs):
+        input_ids, input_mask, input_type_ids = inputs
+        batch_size, num_passages, seq_length = input_ids.size()
+
+        unnormalized_scores = self(
+            input_ids=input_ids.view(-1, seq_length),
+            token_type_ids=input_type_ids.view(-1, seq_length),
+            attention_mask=input_mask.view(-1, seq_length),
+        ).view(batch_size, 1, num_passages)
+        scores = torch.log_softmax(unnormalized_scores, dim=-1)
+
+        labels = torch.zeros_like(input_ids[:, :1, 0])
+        loss = self.loss(log_probs=scores, labels=labels, output_mask=torch.ones_like(labels))
+
+        return unnormalized_scores[:, 0], loss

tts/nemo/collections/nlp/models/intent_slot_classification/__init__.py

@@ -0,0 +1,17 @@
+# 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 nemo.collections.nlp.models.intent_slot_classification.intent_slot_classification_model import (
+    IntentSlotClassificationModel,
+)

tts/nemo/collections/nlp/models/intent_slot_classification/intent_slot_classification_model.py

@@ -0,0 +1,487 @@
+# 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 typing import Dict, List, Optional
+
+import onnx
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.common.losses import AggregatorLoss, CrossEntropyLoss
+from nemo.collections.nlp.data.intent_slot_classification import (
+    IntentSlotClassificationDataset,
+    IntentSlotDataDesc,
+    IntentSlotInferenceDataset,
+)
+from nemo.collections.nlp.metrics.classification_report import ClassificationReport
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import SequenceTokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import tensor2list
+from nemo.core.classes import typecheck
+from nemo.core.classes.common import PretrainedModelInfo
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+
+class IntentSlotClassificationModel(NLPModel):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.classifier.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """ Initializes BERT Joint Intent and Slot model.
+        """
+        self.max_seq_length = cfg.language_model.max_seq_length
+
+        # Setup tokenizer.
+        self.setup_tokenizer(cfg.tokenizer)
+
+        # Check the presence of data_dir.
+        if not cfg.data_dir or not os.path.exists(cfg.data_dir):
+            # Disable setup methods.
+            IntentSlotClassificationModel._set_model_restore_state(is_being_restored=True)
+            # Set default values of data_desc.
+            self._set_defaults_data_desc(cfg)
+        else:
+            self.data_dir = cfg.data_dir
+            # Update configuration of data_desc.
+            self._set_data_desc_to_cfg(cfg, cfg.data_dir, cfg.train_ds, cfg.validation_ds)
+
+        # init superclass
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # Enable setup methods.
+        IntentSlotClassificationModel._set_model_restore_state(is_being_restored=False)
+
+        # Initialize Bert model
+        self.bert_model = get_lm_model(
+            pretrained_model_name=self.cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(self.cfg.language_model.config)
+            if self.cfg.language_model.config
+            else None,
+            checkpoint_file=self.cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        # Initialize Classifier.
+        self._reconfigure_classifier()
+
+    def _set_defaults_data_desc(self, cfg):
+        """
+        Method makes sure that cfg.data_desc params are set.
+        If not, set's them to "dummy" defaults.
+        """
+        if not hasattr(cfg, "data_desc"):
+            OmegaConf.set_struct(cfg, False)
+            cfg.data_desc = {}
+            # Intents.
+            cfg.data_desc.intent_labels = " "
+            cfg.data_desc.intent_label_ids = {" ": 0}
+            cfg.data_desc.intent_weights = [1]
+            # Slots.
+            cfg.data_desc.slot_labels = " "
+            cfg.data_desc.slot_label_ids = {" ": 0}
+            cfg.data_desc.slot_weights = [1]
+
+            cfg.data_desc.pad_label = "O"
+            OmegaConf.set_struct(cfg, True)
+
+    def _set_data_desc_to_cfg(self, cfg, data_dir, train_ds, validation_ds):
+        """ Method creates IntentSlotDataDesc and copies generated values to cfg.data_desc. """
+        # Save data from data desc to config - so it can be reused later, e.g. in inference.
+        data_desc = IntentSlotDataDesc(data_dir=data_dir, modes=[train_ds.prefix, validation_ds.prefix])
+        OmegaConf.set_struct(cfg, False)
+        if not hasattr(cfg, "data_desc") or cfg.data_desc is None:
+            cfg.data_desc = {}
+        # Intents.
+        cfg.data_desc.intent_labels = list(data_desc.intents_label_ids.keys())
+        cfg.data_desc.intent_label_ids = data_desc.intents_label_ids
+        cfg.data_desc.intent_weights = data_desc.intent_weights
+        # Slots.
+        cfg.data_desc.slot_labels = list(data_desc.slots_label_ids.keys())
+        cfg.data_desc.slot_label_ids = data_desc.slots_label_ids
+        cfg.data_desc.slot_weights = data_desc.slot_weights
+
+        cfg.data_desc.pad_label = data_desc.pad_label
+
+        # for older(pre - 1.0.0.b3) configs compatibility
+        if not hasattr(cfg, "class_labels") or cfg.class_labels is None:
+            cfg.class_labels = {}
+            cfg.class_labels = OmegaConf.create(
+                {'intent_labels_file': 'intent_labels.csv', 'slot_labels_file': 'slot_labels.csv'}
+            )
+
+        slot_labels_file = os.path.join(data_dir, cfg.class_labels.slot_labels_file)
+        intent_labels_file = os.path.join(data_dir, cfg.class_labels.intent_labels_file)
+        self._save_label_ids(data_desc.slots_label_ids, slot_labels_file)
+        self._save_label_ids(data_desc.intents_label_ids, intent_labels_file)
+
+        self.register_artifact('class_labels.intent_labels_file', intent_labels_file)
+        self.register_artifact('class_labels.slot_labels_file', slot_labels_file)
+        OmegaConf.set_struct(cfg, True)
+
+    def _save_label_ids(self, label_ids: Dict[str, int], filename: str) -> None:
+        """ Saves label ids map to a file """
+        with open(filename, 'w') as out:
+            labels, _ = zip(*sorted(label_ids.items(), key=lambda x: x[1]))
+            out.write('\n'.join(labels))
+            logging.info(f'Labels: {label_ids}')
+            logging.info(f'Labels mapping saved to : {out.name}')
+
+    def _reconfigure_classifier(self):
+        """ Method reconfigures the classifier depending on the settings of model cfg.data_desc """
+
+        self.classifier = SequenceTokenClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_intents=len(self.cfg.data_desc.intent_labels),
+            num_slots=len(self.cfg.data_desc.slot_labels),
+            dropout=self.cfg.head.fc_dropout,
+            num_layers=self.cfg.head.num_output_layers,
+            log_softmax=False,
+        )
+
+        # define losses
+        if self.cfg.class_balancing == 'weighted_loss':
+            # You may need to increase the number of epochs for convergence when using weighted_loss
+            self.intent_loss = CrossEntropyLoss(logits_ndim=2, weight=self.cfg.data_desc.intent_weights)
+            self.slot_loss = CrossEntropyLoss(logits_ndim=3, weight=self.cfg.data_desc.slot_weights)
+        else:
+            self.intent_loss = CrossEntropyLoss(logits_ndim=2)
+            self.slot_loss = CrossEntropyLoss(logits_ndim=3)
+
+        self.total_loss = AggregatorLoss(
+            num_inputs=2, weights=[self.cfg.intent_loss_weight, 1.0 - self.cfg.intent_loss_weight]
+        )
+
+        # setup to track metrics
+        self.intent_classification_report = ClassificationReport(
+            num_classes=len(self.cfg.data_desc.intent_labels),
+            label_ids=self.cfg.data_desc.intent_label_ids,
+            dist_sync_on_step=True,
+            mode='micro',
+        )
+        self.slot_classification_report = ClassificationReport(
+            num_classes=len(self.cfg.data_desc.slot_labels),
+            label_ids=self.cfg.data_desc.slot_label_ids,
+            dist_sync_on_step=True,
+            mode='micro',
+        )
+
+    def update_data_dir_for_training(self, data_dir: str, train_ds, validation_ds) -> None:
+        """
+        Update data directory and get data stats with Data Descriptor.
+        Also, reconfigures the classifier - to cope with data with e.g. different number of slots.
+
+        Args:
+            data_dir: path to data directory
+        """
+        logging.info(f'Setting data_dir to {data_dir}.')
+        self.data_dir = data_dir
+        # Update configuration with new data.
+        self._set_data_desc_to_cfg(self.cfg, data_dir, train_ds, validation_ds)
+        # Reconfigure the classifier for different settings (number of intents, slots etc.).
+        self._reconfigure_classifier()
+
+    def update_data_dir_for_testing(self, data_dir) -> None:
+        """
+        Update data directory.
+
+        Args:
+            data_dir: path to data directory
+        """
+        logging.info(f'Setting data_dir to {data_dir}.')
+        self.data_dir = data_dir
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        intent_logits, slot_logits = self.classifier(hidden_states=hidden_states)
+        return intent_logits, slot_logits
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        # forward pass
+        input_ids, input_type_ids, input_mask, loss_mask, subtokens_mask, intent_labels, slot_labels = batch
+        intent_logits, slot_logits = self(
+            input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask
+        )
+
+        # calculate combined loss for intents and slots
+        intent_loss = self.intent_loss(logits=intent_logits, labels=intent_labels)
+        slot_loss = self.slot_loss(logits=slot_logits, labels=slot_labels, loss_mask=loss_mask)
+        train_loss = self.total_loss(loss_1=intent_loss, loss_2=slot_loss)
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', train_loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': train_loss,
+            'lr': lr,
+        }
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, loss_mask, subtokens_mask, intent_labels, slot_labels = batch
+        intent_logits, slot_logits = self(
+            input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask
+        )
+
+        # calculate combined loss for intents and slots
+        intent_loss = self.intent_loss(logits=intent_logits, labels=intent_labels)
+        slot_loss = self.slot_loss(logits=slot_logits, labels=slot_labels, loss_mask=loss_mask)
+        val_loss = self.total_loss(loss_1=intent_loss, loss_2=slot_loss)
+
+        # calculate accuracy metrics for intents and slot reporting
+        # intents
+        preds = torch.argmax(intent_logits, axis=-1)
+        self.intent_classification_report.update(preds, intent_labels)
+        # slots
+        subtokens_mask = subtokens_mask > 0.5
+        preds = torch.argmax(slot_logits, axis=-1)[subtokens_mask]
+        slot_labels = slot_labels[subtokens_mask]
+        self.slot_classification_report.update(preds, slot_labels)
+
+        return {
+            'val_loss': val_loss,
+            'intent_tp': self.intent_classification_report.tp,
+            'intent_fn': self.intent_classification_report.fn,
+            'intent_fp': self.intent_classification_report.fp,
+            'slot_tp': self.slot_classification_report.tp,
+            'slot_fn': self.slot_classification_report.fn,
+            'slot_fp': self.slot_classification_report.fp,
+        }
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+
+        # calculate metrics and log classification report (separately for intents and slots)
+        intent_precision, intent_recall, intent_f1, intent_report = self.intent_classification_report.compute()
+        logging.info(f'Intent report: {intent_report}')
+
+        slot_precision, slot_recall, slot_f1, slot_report = self.slot_classification_report.compute()
+        logging.info(f'Slot report: {slot_report}')
+
+        self.log('val_loss', avg_loss)
+        self.log('intent_precision', intent_precision)
+        self.log('intent_recall', intent_recall)
+        self.log('intent_f1', intent_f1)
+        self.log('slot_precision', slot_precision)
+        self.log('slot_recall', slot_recall)
+        self.log('slot_f1', slot_f1)
+
+        self.intent_classification_report.reset()
+        self.slot_classification_report.reset()
+
+        return {
+            'val_loss': avg_loss,
+            'intent_precision': intent_precision,
+            'intent_recall': intent_recall,
+            'intent_f1': intent_f1,
+            'slot_precision': slot_precision,
+            'slot_recall': slot_recall,
+            'slot_f1': slot_f1,
+        }
+
+    def test_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the test loop with the data from the test dataloader
+        passed in as `batch`.
+        """
+        return self.validation_step(batch, batch_idx)
+
+    def test_epoch_end(self, outputs):
+        """
+        Called at the end of test to aggregate outputs.
+        :param outputs: list of individual outputs of each test step.
+        """
+        return self.validation_epoch_end(outputs)
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+        input_file = f'{self.data_dir}/{cfg.prefix}.tsv'
+        slot_file = f'{self.data_dir}/{cfg.prefix}_slots.tsv'
+
+        if not (os.path.exists(input_file) and os.path.exists(slot_file)):
+            raise FileNotFoundError(
+                f'{input_file} or {slot_file} not found. Please refer to the documentation for the right format \
+                 of Intents and Slots files.'
+            )
+
+        dataset = IntentSlotClassificationDataset(
+            input_file=input_file,
+            slot_file=slot_file,
+            tokenizer=self.tokenizer,
+            max_seq_length=self.max_seq_length,
+            num_samples=cfg.num_samples,
+            pad_label=self.cfg.data_desc.pad_label,
+            ignore_extra_tokens=self.cfg.ignore_extra_tokens,
+            ignore_start_end=self.cfg.ignore_start_end,
+        )
+
+        return DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.num_workers,
+            pin_memory=cfg.pin_memory,
+            drop_last=cfg.drop_last,
+            collate_fn=dataset.collate_fn,
+        )
+
+    def _setup_infer_dataloader(self, queries: List[str], test_ds) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a infer data loader.
+        Args:
+            queries: text
+            batch_size: batch size to use during inference
+        Returns:
+            A pytorch DataLoader.
+        """
+
+        dataset = IntentSlotInferenceDataset(
+            tokenizer=self.tokenizer, queries=queries, max_seq_length=-1, do_lower_case=False
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=test_ds.batch_size,
+            shuffle=test_ds.shuffle,
+            num_workers=test_ds.num_workers,
+            pin_memory=test_ds.pin_memory,
+            drop_last=test_ds.drop_last,
+        )
+
+    def predict_from_examples(self, queries: List[str], test_ds) -> List[List[str]]:
+        """
+        Get prediction for the queries (intent and slots)
+        Args:
+            queries: text sequences
+            test_ds: Dataset configuration section.
+        Returns:
+            predicted_intents, predicted_slots: model intent and slot predictions
+        """
+        predicted_intents = []
+        predicted_slots = []
+        mode = self.training
+        try:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+            # Retrieve intent and slot vocabularies from configuration.
+            intent_labels = self.cfg.data_desc.intent_labels
+            slot_labels = self.cfg.data_desc.slot_labels
+
+            # Initialize tokenizer.
+            # if not hasattr(self, "tokenizer"):
+            #    self._setup_tokenizer(self.cfg.tokenizer)
+            # Initialize modules.
+            # self._reconfigure_classifier()
+
+            # Switch model to evaluation mode
+            self.eval()
+            self.to(device)
+
+            # Dataset.
+            infer_datalayer = self._setup_infer_dataloader(queries, test_ds)
+
+            for batch in infer_datalayer:
+                input_ids, input_type_ids, input_mask, loss_mask, subtokens_mask = batch
+
+                intent_logits, slot_logits = self.forward(
+                    input_ids=input_ids.to(device),
+                    token_type_ids=input_type_ids.to(device),
+                    attention_mask=input_mask.to(device),
+                )
+
+                # predict intents and slots for these examples
+                # intents
+                intent_preds = tensor2list(torch.argmax(intent_logits, axis=-1))
+
+                # convert numerical outputs to Intent and Slot labels from the dictionaries
+                for intent_num in intent_preds:
+                    if intent_num < len(intent_labels):
+                        predicted_intents.append(intent_labels[int(intent_num)])
+                    else:
+                        # should not happen
+                        predicted_intents.append("Unknown Intent")
+
+                # slots
+                slot_preds = torch.argmax(slot_logits, axis=-1)
+
+                for slot_preds_query, mask_query in zip(slot_preds, subtokens_mask):
+                    query_slots = ''
+                    for slot, mask in zip(slot_preds_query, mask_query):
+                        if mask == 1:
+                            if slot < len(slot_labels):
+                                query_slots += slot_labels[int(slot)] + ' '
+                            else:
+                                query_slots += 'Unknown_slot '
+                    predicted_slots.append(query_slots.strip())
+
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+
+        return predicted_intents, predicted_slots
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        model = PretrainedModelInfo(
+            pretrained_model_name="Joint_Intent_Slot_Assistant",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemonlpmodels/versions/1.0.0a5/files/Joint_Intent_Slot_Assistant.nemo",
+            description="This models is trained on this https://github.com/xliuhw/NLU-Evaluation-Data dataset which includes 64 various intents and 55 slots. Final Intent accuracy is about 87%, Slot accuracy is about 89%.",
+        )
+        result.append(model)
+        return result

tts/nemo/collections/nlp/models/language_modeling/__init__.py

@@ -0,0 +1,16 @@
+# 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 nemo.collections.nlp.models.language_modeling.bert_lm_model import BERTLMModel
+from nemo.collections.nlp.models.language_modeling.transformer_lm_model import TransformerLMModel

tts/nemo/collections/nlp/models/language_modeling/bert_lm_model.py

@@ -0,0 +1,272 @@
+# 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 typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.common.losses import AggregatorLoss, CrossEntropyLoss, SmoothedCrossEntropyLoss
+from nemo.collections.common.metrics import Perplexity
+from nemo.collections.nlp.data.language_modeling.lm_bert_dataset import (
+    BertPretrainingDataset,
+    BertPretrainingPreprocessedDataloader,
+)
+from nemo.collections.nlp.modules.common import BertPretrainingTokenClassifier, SequenceClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ["BERTLMModel"]
+
+
+class BERTLMModel(ModelPT):
+    """
+    BERT language model pretraining.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        output_types_dict = {"mlm_log_probs": self.mlm_classifier.output_types["log_probs"]}
+        if not self.only_mlm_loss:
+            output_types_dict["nsp_logits"] = self.nsp_classifier.output_types["logits"]
+        return output_types_dict
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        if cfg.tokenizer is not None:
+            self._setup_tokenizer(cfg.tokenizer)
+        else:
+            self.tokenizer = None
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file)
+            if cfg.tokenizer
+            else None,
+        )
+
+        self.hidden_size = self.bert_model.config.hidden_size
+        self.vocab_size = self.bert_model.config.vocab_size
+        self.only_mlm_loss = cfg.only_mlm_loss
+
+        self.mlm_classifier = BertPretrainingTokenClassifier(
+            hidden_size=self.hidden_size,
+            num_classes=self.vocab_size,
+            num_layers=cfg.num_tok_classification_layers,
+            activation="gelu",
+            log_softmax=True,
+            use_transformer_init=True,
+        )
+
+        self.mlm_loss = SmoothedCrossEntropyLoss()
+
+        if not self.only_mlm_loss:
+            self.nsp_classifier = SequenceClassifier(
+                hidden_size=self.hidden_size,
+                num_classes=2,
+                num_layers=cfg.num_seq_classification_layers,
+                log_softmax=False,
+                activation="tanh",
+                use_transformer_init=True,
+            )
+
+            self.nsp_loss = CrossEntropyLoss()
+            self.agg_loss = AggregatorLoss(num_inputs=2)
+
+        # # tie weights of MLM softmax layer and embedding layer of the encoder
+        if (
+            self.mlm_classifier.mlp.last_linear_layer.weight.shape
+            != self.bert_model.embeddings.word_embeddings.weight.shape
+        ):
+            raise ValueError("Final classification layer does not match embedding layer.")
+        self.mlm_classifier.mlp.last_linear_layer.weight = self.bert_model.embeddings.word_embeddings.weight
+        # create extra bias
+
+        # setup to track metrics
+        self.validation_perplexity = Perplexity(compute_on_step=False)
+
+        self.setup_optimization(cfg.optim)
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask,
+        )
+        mlm_log_probs = self.mlm_classifier(hidden_states=hidden_states)
+        if self.only_mlm_loss:
+            return (mlm_log_probs,)
+        nsp_logits = self.nsp_classifier(hidden_states=hidden_states)
+        return mlm_log_probs, nsp_logits
+
+    def _compute_losses(self, mlm_log_probs, nsp_logits, output_ids, output_mask, labels):
+        mlm_loss = self.mlm_loss(log_probs=mlm_log_probs, labels=output_ids, output_mask=output_mask)
+        if self.only_mlm_loss:
+            loss, nsp_loss = mlm_loss, None
+        else:
+            nsp_loss = self.nsp_loss(logits=nsp_logits, labels=labels)
+            loss = self.agg_loss(loss_1=mlm_loss, loss_2=nsp_loss)
+        return mlm_loss, nsp_loss, loss
+
+    def _parse_forward_outputs(self, forward_outputs):
+        if self.only_mlm_loss:
+            return forward_outputs[0], None
+        else:
+            return forward_outputs
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, output_ids, output_mask, labels = batch
+        forward_outputs = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        mlm_log_probs, nsp_logits = self._parse_forward_outputs(forward_outputs)
+        _, _, loss = self._compute_losses(mlm_log_probs, nsp_logits, output_ids, output_mask, labels)
+        lr = self._optimizer.param_groups[0]['lr']
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+        return {"loss": loss, "lr": lr}
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, output_ids, output_mask, labels = batch
+        forward_outputs = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        mlm_log_probs, nsp_logits = self._parse_forward_outputs(forward_outputs)
+        _, _, loss = self._compute_losses(mlm_log_probs, nsp_logits, output_ids, output_mask, labels)
+        self.validation_perplexity(logits=mlm_log_probs)
+        return {'val_loss': loss}
+
+    def validation_epoch_end(self, outputs):
+        """Called at the end of validation to aggregate outputs.
+
+        Args:
+            outputs (list): The individual outputs of each validation step.
+
+        Returns:
+            dict: Validation loss and tensorboard logs.
+        """
+        if outputs:
+            avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+            perplexity = self.validation_perplexity.compute()
+            logging.info(f"evaluation perplexity {perplexity.cpu().item()}")
+            self.log(f'val_loss', avg_loss)
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = (
+            self._setup_preprocessed_dataloader(train_data_config)
+            if self.tokenizer is None
+            else self._setup_dataloader(train_data_config)
+        )
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = (
+            self._setup_preprocessed_dataloader(val_data_config)
+            if self.tokenizer is None
+            else self._setup_dataloader(val_data_config)
+        )
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        pass
+
+    def _setup_preprocessed_dataloader(self, cfg: Optional[DictConfig]):
+        dataset = cfg.data_file
+        max_predictions_per_seq = cfg.max_predictions_per_seq
+        batch_size = cfg.batch_size
+
+        if os.path.isdir(dataset):
+            files = [os.path.join(dataset, f) for f in os.listdir(dataset) if os.path.isfile(os.path.join(dataset, f))]
+        else:
+            files = [dataset]
+        files.sort()
+        dl = BertPretrainingPreprocessedDataloader(
+            data_files=files, max_predictions_per_seq=max_predictions_per_seq, batch_size=batch_size,
+        )
+        return dl
+
+    def _setup_tokenizer(self, cfg: DictConfig):
+        tokenizer = get_tokenizer(
+            tokenizer_name=cfg.tokenizer_name,
+            tokenizer_model=cfg.tokenizer_model,
+            special_tokens=OmegaConf.to_container(cfg.special_tokens) if cfg.special_tokens else None,
+            vocab_file=cfg.vocab_file,
+        )
+        self.tokenizer = tokenizer
+
+    def _setup_dataloader(self, cfg: DictConfig):
+        dataset = BertPretrainingDataset(
+            tokenizer=self.tokenizer,
+            data_file=cfg.data_file,
+            max_seq_length=cfg.max_seq_length,
+            mask_prob=cfg.mask_prob,
+            short_seq_prob=cfg.short_seq_prob,
+        )
+        dl = torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            collate_fn=dataset.collate_fn,
+            drop_last=cfg.get("drop_last", False),
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 0),
+        )
+        return dl
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="bertbaseuncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/bertbaseuncased/versions/1.0.0rc1/files/bertbaseuncased.nemo",
+                description="The model was trained EN Wikipedia and BookCorpus on a sequence length of 512.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="bertlargeuncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/bertlargeuncased/versions/1.0.0rc1/files/bertlargeuncased.nemo",
+                description="The model was trained EN Wikipedia and BookCorpus on a sequence length of 512.",
+            )
+        )
+        return result

tts/nemo/collections/nlp/models/language_modeling/transformer_lm_model.py

@@ -0,0 +1,276 @@
+# 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 math
+from typing import Dict, Optional
+
+import numpy as np
+import torch
+import torch.utils.data as pt_data
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import SmoothedCrossEntropyLoss
+from nemo.collections.common.metrics import GlobalAverageLossMetric
+from nemo.collections.common.parts import transformer_weights_init
+from nemo.collections.nlp.data import SentenceDataset, TarredSentenceDataset
+from nemo.collections.nlp.metrics import SequencePerplexity
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_transformer
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.modelPT import ModelPT
+from nemo.utils import logging, model_utils
+
+__all__ = ["TransformerLMModel"]
+
+
+class TransformerLMModel(ModelPT):
+    """
+    Left-to-right Transformer language model.
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        # Instantiates tokenizer and register to be saved with NeMo Model archive
+        # After this call, ther will be self.tokenizer which can convert between tokens and token_ids.
+        self.setup_tokenizer(
+            tokenizer_name=cfg.tokenizer.get("tokenizer_name", "yttm"),
+            tokenizer_model=cfg.tokenizer.get("tokenizer_model", None),
+            vocab_file=cfg.tokenizer.get("vocab_file", None),
+            bpe_dropout=cfg.tokenizer.get("bpe_dropout", 0.0),
+        )
+
+        # init superclass
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # make vocabulary size divisible by 8 for fast fp16 training
+        vocab_size = 8 * math.ceil(self.tokenizer.vocab_size / 8)
+
+        # encoder from NeMo, Megatron-LM, or HuggingFace
+        encoder_cfg_dict = OmegaConf.to_container(cfg.get('encoder'))
+        encoder_cfg_dict['vocab_size'] = vocab_size
+        library = encoder_cfg_dict.pop('library', 'nemo')
+        model_name = encoder_cfg_dict.pop('model_name', None)
+        pretrained = encoder_cfg_dict.pop('pretrained', False)
+        self.encoder = get_transformer(
+            library=library, model_name=model_name, pretrained=pretrained, config_dict=encoder_cfg_dict, encoder=True,
+        )
+
+        self.log_softmax = TokenClassifier(
+            hidden_size=self.encoder.hidden_size,
+            num_classes=vocab_size,
+            activation=cfg.head.activation,
+            log_softmax=cfg.head.log_softmax,
+            dropout=cfg.head.dropout,
+            use_transformer_init=cfg.head.use_transformer_init,
+        )
+
+        # tie weights of embedding and softmax matrices
+        self.log_softmax.mlp.layer0.weight = self.encoder.embedding.token_embedding.weight
+
+        std_init_range = 1 / self.encoder.hidden_size ** 0.5
+
+        # initialize weights if not using pretrained encoder
+        if not self._cfg.encoder.get('pretrained', False):
+            self.encoder.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.log_softmax.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.loss_fn = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id, label_smoothing=cfg.label_smoothing)
+        self.eval_loss_fn = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id)
+        self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True)
+        self.eval_ppl = SequencePerplexity()
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+
+        hidden_states = self.encoder(input_ids=input_ids, encoder_mask=attention_mask)
+        log_probs = self.log_softmax(hidden_states=hidden_states)
+
+        return log_probs
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        # forward pass
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1
+                # added by DataLoader is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+        ids, mask = batch
+        input_ids, labels = ids[:, :-1], ids[:, 1:]
+        input_mask = mask[:, :-1]
+        log_probs = self(input_ids=input_ids, attention_mask=input_mask)
+
+        train_loss = self.loss_fn(log_probs=log_probs, labels=labels)
+
+        tensorboard_logs = {
+            "train_loss": train_loss,
+            "lr": self._optimizer.param_groups[0]["lr"],
+        }
+        return {"loss": train_loss, "log": tensorboard_logs}
+
+    def eval_step(self, batch, batch_idx):
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1
+                # added by DataLoader is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+        ids, mask = batch
+        input_ids, labels = ids[:, :-1], ids[:, 1:]
+        input_mask, output_mask = mask[:, :-1], mask[:, 1:]
+        log_probs = self(input_ids=input_ids, attention_mask=input_mask)
+        eval_loss = self.eval_loss_fn(log_probs=log_probs, labels=labels)
+
+        self.eval_loss(loss=eval_loss, num_measurements=log_probs.shape[0] * log_probs.shape[1])
+        self.eval_ppl(log_probs=log_probs, labels=labels, mask=output_mask)
+        return {}
+
+    def test_step(self, batch, batch_idx):
+        return self.eval_step(batch, batch_idx)
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        return self.eval_step(batch, batch_idx)
+
+    def eval_epoch_end(self, outputs, mode):
+        eval_loss = self.eval_loss.compute()
+        eval_ppl = self.eval_ppl.compute()
+        ans = {f"{mode}_loss": eval_loss, f"{mode}_ppl": eval_ppl}
+        ans['log'] = dict(ans)
+        dataset_name = "Validation" if mode == 'val' else "Test"
+        logging.info(f"\n\n\n\n{dataset_name} PPL: {np.round(eval_ppl.item(), 2)}")
+        return ans
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+        self.log_dict(self.eval_epoch_end(outputs, 'val'), sync_dist=True)
+        self.eval_loss.reset()
+        self.eval_ppl.reset()
+
+    def test_epoch_end(self, outputs):
+        return self.eval_epoch_end(outputs, 'test')
+
+    def setup_tokenizer(
+        self, tokenizer_name=None, tokenizer_model=None, vocab_file=None, bpe_dropout=0.0,
+    ):
+
+        supported_tokenizers = ['yttm', 'huggingface', 'sentencepiece', 'word']
+        if tokenizer_name not in supported_tokenizers:
+            raise NotImplementedError(f"Currently we only support tokenizers in {supported_tokenizers}.")
+
+        self.tokenizer = get_tokenizer(
+            tokenizer_name=tokenizer_name,
+            tokenizer_model=self.register_artifact("cfg.tokenizer.tokenizer_model", tokenizer_model),
+            vocab_file=vocab_file,
+            bpe_dropout=bpe_dropout,
+            special_tokens=None,
+            use_fast=False,
+        )
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig, predict_last_k=0):
+
+        if cfg.get("use_tarred_dataset", False):
+            if cfg.get("metadata_file") is None:
+                raise FileNotFoundError("Trying to use tarred data set but could not find metadata path in config.")
+            else:
+                metadata_file = cfg.get('metadata_file')
+                with open(metadata_file) as metadata_reader:
+                    metadata = json.load(metadata_reader)
+                if cfg.get('tar_files') is None:
+                    tar_files = metadata.get('tar_files')
+                    if tar_files is not None:
+                        logging.info(f'Loading from tarred dataset {tar_files}')
+                    else:
+                        raise FileNotFoundError("Could not find tarred dataset in config or metadata.")
+                else:
+                    tar_files = cfg.get('tar_files')
+                    if metadata.get('tar_files') is not None:
+                        raise ValueError(
+                            'Tar files specified in config and in metadata file. Tar files should only be specified once.'
+                        )
+            dataset = TarredSentenceDataset(
+                text_tar_filepaths=tar_files,
+                metadata_path=metadata_file,
+                tokenizer=self.tokenizer,
+                shuffle_n=cfg.get("tar_shuffle_n", 100),
+                shard_strategy=cfg.get("shard_strategy", "scatter"),
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+            )
+            return torch.utils.data.DataLoader(
+                dataset=dataset,
+                batch_size=1,
+                num_workers=cfg.get("num_workers", 2),
+                pin_memory=cfg.get("pin_memory", False),
+                drop_last=cfg.get("drop_last", False),
+            )
+        else:
+            dataset = SentenceDataset(
+                tokenizer=self.tokenizer,
+                dataset=cfg.file_name,
+                tokens_in_batch=cfg.tokens_in_batch,
+                clean=cfg.get("clean", False),
+                max_seq_length=cfg.get("max_seq_length", 512),
+                min_seq_length=cfg.get("min_seq_length", 1),
+                cache_ids=cfg.get("cache_ids", False),
+            )
+        if cfg.shuffle:
+            sampler = pt_data.RandomSampler(dataset)
+        else:
+            sampler = pt_data.SequentialSampler(dataset)
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=1,
+            sampler=sampler,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass

tts/nemo/collections/nlp/models/machine_translation/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel

tts/nemo/collections/nlp/models/machine_translation/mt_enc_dec_config.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.
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+from omegaconf.omegaconf import MISSING
+
+from nemo.collections.nlp.data.machine_translation.machine_translation_dataset import TranslationDataConfig
+from nemo.collections.nlp.models.enc_dec_nlp_model import EncDecNLPModelConfig
+from nemo.collections.nlp.modules.common.token_classifier import TokenClassifierConfig
+from nemo.collections.nlp.modules.common.tokenizer_utils import TokenizerConfig
+from nemo.collections.nlp.modules.common.transformer.transformer import (
+    NeMoTransformerConfig,
+    NeMoTransformerEncoderConfig,
+)
+from nemo.core.config.modelPT import ModelConfig, OptimConfig, SchedConfig
+
+
+@dataclass
+class MTSchedConfig(SchedConfig):
+    name: str = 'InverseSquareRootAnnealing'
+    warmup_ratio: Optional[float] = None
+    last_epoch: int = -1
+
+
+# TODO: Refactor this dataclass to to support more optimizers (it pins the optimizer to Adam-like optimizers).
+@dataclass
+class MTOptimConfig(OptimConfig):
+    name: str = 'adam'
+    lr: float = 1e-3
+    betas: Tuple[float, float] = (0.9, 0.98)
+    weight_decay: float = 0.0
+    sched: Optional[MTSchedConfig] = MTSchedConfig()
+
+
+@dataclass
+class MTEncDecModelConfig(EncDecNLPModelConfig):
+    # machine translation configurations
+    num_val_examples: int = 3
+    num_test_examples: int = 3
+    max_generation_delta: int = 10
+    label_smoothing: Optional[float] = 0.0
+    beam_size: int = 4
+    len_pen: float = 0.0
+    src_language: Any = 'en'  # Any = str or List[str]
+    tgt_language: Any = 'en'  # Any = str or List[str]
+    find_unused_parameters: Optional[bool] = True
+    shared_tokenizer: Optional[bool] = True
+    multilingual: Optional[bool] = False
+    preproc_out_dir: Optional[str] = None
+
+    # network architecture configuration
+    encoder_tokenizer: Any = MISSING
+    encoder: Any = MISSING
+
+    decoder_tokenizer: Any = MISSING
+    decoder: Any = MISSING
+
+    head: TokenClassifierConfig = TokenClassifierConfig(log_softmax=True)
+
+    # dataset configurations
+    train_ds: Optional[TranslationDataConfig] = TranslationDataConfig(
+        src_file_name=MISSING,
+        tgt_file_name=MISSING,
+        tokens_in_batch=512,
+        clean=True,
+        shuffle=True,
+        cache_ids=False,
+        use_cache=False,
+    )
+    validation_ds: Optional[TranslationDataConfig] = TranslationDataConfig(
+        src_file_name=MISSING,
+        tgt_file_name=MISSING,
+        tokens_in_batch=512,
+        clean=False,
+        shuffle=False,
+        cache_ids=False,
+        use_cache=False,
+    )
+    test_ds: Optional[TranslationDataConfig] = TranslationDataConfig(
+        src_file_name=MISSING,
+        tgt_file_name=MISSING,
+        tokens_in_batch=512,
+        clean=False,
+        shuffle=False,
+        cache_ids=False,
+        use_cache=False,
+    )
+    optim: Optional[OptimConfig] = MTOptimConfig()
+
+
+@dataclass
+class AAYNBaseConfig(MTEncDecModelConfig):
+
+    # Attention is All You Need Base Configuration
+    encoder_tokenizer: TokenizerConfig = TokenizerConfig(library='yttm')
+    decoder_tokenizer: TokenizerConfig = TokenizerConfig(library='yttm')
+
+    encoder: NeMoTransformerEncoderConfig = NeMoTransformerEncoderConfig(
+        library='nemo',
+        model_name=None,
+        pretrained=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,
+    )
+
+    decoder: NeMoTransformerConfig = NeMoTransformerConfig(
+        library='nemo',
+        model_name=None,
+        pretrained=False,
+        inner_size=2048,
+        num_layers=6,
+        num_attention_heads=8,
+        ffn_dropout=0.1,
+        attn_score_dropout=0.1,
+        attn_layer_dropout=0.1,
+    )

tts/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py

@@ -0,0 +1,836 @@
+# 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 itertools
+import json
+import random
+from multiprocessing import Value
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.utils.data as pt_data
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.utilities import rank_zero_only
+from sacrebleu import corpus_bleu
+
+from nemo.collections.common.data import ConcatDataset
+from nemo.collections.common.losses import NLLLoss, SmoothedCrossEntropyLoss
+from nemo.collections.common.metrics import GlobalAverageLossMetric
+from nemo.collections.common.parts import transformer_weights_init
+from nemo.collections.common.tokenizers.chinese_tokenizers import ChineseProcessor
+from nemo.collections.common.tokenizers.en_ja_tokenizers import EnJaProcessor
+from nemo.collections.common.tokenizers.moses_tokenizers import MosesProcessor
+from nemo.collections.nlp.data import TarredTranslationDataset, TranslationDataset
+from nemo.collections.nlp.models.enc_dec_nlp_model import EncDecNLPModel
+from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTEncDecModelConfig
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_transformer
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
+from nemo.collections.nlp.modules.common.transformer import BeamSearchSequenceGenerator, TopKSequenceGenerator
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.utils import logging, model_utils
+
+__all__ = ['MTEncDecModel']
+
+
+class MTEncDecModel(EncDecNLPModel):
+    """
+    Encoder-decoder machine translation model.
+    """
+
+    def __init__(self, cfg: MTEncDecModelConfig, trainer: Trainer = None):
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        # Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
+
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        self.src_language = cfg.get("src_language", None)
+        self.tgt_language = cfg.get("tgt_language", None)
+
+        self.multilingual = cfg.get("multilingual", False)
+        self.multilingual_ids = []
+
+        # Instantiates tokenizers and register to be saved with NeMo Model archive
+        # After this call, ther will be self.encoder_tokenizer and self.decoder_tokenizer
+        # Which can convert between tokens and token_ids for SRC and TGT languages correspondingly.
+        self.setup_enc_dec_tokenizers(
+            encoder_tokenizer_library=cfg.encoder_tokenizer.get('library', 'yttm'),
+            encoder_tokenizer_model=cfg.encoder_tokenizer.get('tokenizer_model'),
+            encoder_bpe_dropout=cfg.encoder_tokenizer.get('bpe_dropout', 0.0)
+            if cfg.encoder_tokenizer.get('bpe_dropout', 0.0) is not None
+            else 0.0,
+            encoder_model_name=cfg.encoder.get('model_name') if hasattr(cfg.encoder, 'model_name') else None,
+            decoder_tokenizer_library=cfg.decoder_tokenizer.get('library', 'yttm'),
+            decoder_tokenizer_model=cfg.decoder_tokenizer.tokenizer_model,
+            decoder_bpe_dropout=cfg.decoder_tokenizer.get('bpe_dropout', 0.0)
+            if cfg.decoder_tokenizer.get('bpe_dropout', 0.0) is not None
+            else 0.0,
+            decoder_model_name=cfg.decoder.get('model_name') if hasattr(cfg.decoder, 'model_name') else None,
+        )
+
+        if self.multilingual:
+            if isinstance(self.src_language, ListConfig) and isinstance(self.tgt_language, ListConfig):
+                raise ValueError(
+                    "cfg.src_language and cfg.tgt_language cannot both be lists. We only support many-to-one or one-to-many multilingual models."
+                )
+            elif isinstance(self.src_language, ListConfig):
+                for lng in self.src_language:
+                    self.multilingual_ids.append(self.encoder_tokenizer.token_to_id("<" + lng + ">"))
+            elif isinstance(self.tgt_language, ListConfig):
+                for lng in self.tgt_language:
+                    self.multilingual_ids.append(self.encoder_tokenizer.token_to_id("<" + lng + ">"))
+            else:
+                raise ValueError(
+                    "Expect either cfg.src_language or cfg.tgt_language to be a list when multilingual=True."
+                )
+
+            if isinstance(self.src_language, ListConfig):
+                self.tgt_language = [self.tgt_language] * len(self.src_language)
+            else:
+                self.src_language = [self.src_language] * len(self.tgt_language)
+
+            self.source_processor_list = []
+            self.target_processor_list = []
+            for src_lng, tgt_lng in zip(self.src_language, self.tgt_language):
+                src_prcsr, tgt_prscr = self.setup_pre_and_post_processing_utils(
+                    source_lang=src_lng, target_lang=tgt_lng
+                )
+                self.source_processor_list.append(src_prcsr)
+                self.target_processor_list.append(tgt_prscr)
+
+        else:
+            # After this call, the model will have  self.source_processor and self.target_processor objects
+            self.setup_pre_and_post_processing_utils(source_lang=self.src_language, target_lang=self.tgt_language)
+            self.multilingual_ids = [None]
+
+        # TODO: Why is this base constructor call so late in the game?
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # encoder from NeMo, Megatron-LM, or HuggingFace
+        encoder_cfg_dict = OmegaConf.to_container(cfg.get('encoder'))
+        encoder_cfg_dict['vocab_size'] = self.encoder_vocab_size
+        library = encoder_cfg_dict.pop('library', 'nemo')
+        model_name = encoder_cfg_dict.pop('model_name', None)
+        pretrained = encoder_cfg_dict.pop('pretrained', False)
+        self.encoder = get_transformer(
+            library=library,
+            model_name=model_name,
+            pretrained=pretrained,
+            config_dict=encoder_cfg_dict,
+            encoder=True,
+            pre_ln_final_layer_norm=encoder_cfg_dict.get('pre_ln_final_layer_norm', False),
+        )
+
+        # decoder from NeMo, Megatron-LM, or HuggingFace
+        decoder_cfg_dict = OmegaConf.to_container(cfg.get('decoder'))
+        decoder_cfg_dict['vocab_size'] = self.decoder_vocab_size
+        library = decoder_cfg_dict.pop('library', 'nemo')
+        model_name = decoder_cfg_dict.pop('model_name', None)
+        pretrained = decoder_cfg_dict.pop('pretrained', False)
+        decoder_cfg_dict['hidden_size'] = self.encoder.hidden_size
+        self.decoder = get_transformer(
+            library=library,
+            model_name=model_name,
+            pretrained=pretrained,
+            config_dict=decoder_cfg_dict,
+            encoder=False,
+            pre_ln_final_layer_norm=decoder_cfg_dict.get('pre_ln_final_layer_norm', False),
+        )
+
+        self.log_softmax = TokenClassifier(
+            hidden_size=self.decoder.hidden_size,
+            num_classes=self.decoder_vocab_size,
+            activation=cfg.head.activation,
+            log_softmax=cfg.head.log_softmax,
+            dropout=cfg.head.dropout,
+            use_transformer_init=cfg.head.use_transformer_init,
+        )
+
+        self.beam_search = BeamSearchSequenceGenerator(
+            embedding=self.decoder.embedding,
+            decoder=self.decoder.decoder,
+            log_softmax=self.log_softmax,
+            max_sequence_length=self.decoder.max_sequence_length,
+            beam_size=cfg.beam_size,
+            bos=self.decoder_tokenizer.bos_id,
+            pad=self.decoder_tokenizer.pad_id,
+            eos=self.decoder_tokenizer.eos_id,
+            len_pen=cfg.len_pen,
+            max_delta_length=cfg.max_generation_delta,
+        )
+
+        # tie weights of embedding and softmax matrices
+        self.log_softmax.mlp.layer0.weight = self.decoder.embedding.token_embedding.weight
+
+        # TODO: encoder and decoder with different hidden size?
+        std_init_range = 1 / self.encoder.hidden_size ** 0.5
+
+        # initialize weights if not using pretrained encoder/decoder
+        if not self._cfg.encoder.get('pretrained', False):
+            self.encoder.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        if not self._cfg.decoder.get('pretrained', False):
+            self.decoder.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.log_softmax.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.loss_fn = SmoothedCrossEntropyLoss(
+            pad_id=self.decoder_tokenizer.pad_id, label_smoothing=cfg.label_smoothing
+        )
+        self.eval_loss_fn = NLLLoss(ignore_index=self.decoder_tokenizer.pad_id)
+
+    def filter_predicted_ids(self, ids):
+        ids[ids >= self.decoder_tokenizer.vocab_size] = self.decoder_tokenizer.unk_id
+        return ids
+
+    @typecheck()
+    def forward(self, src, src_mask, tgt, tgt_mask):
+        src_hiddens = self.encoder(input_ids=src, encoder_mask=src_mask)
+        tgt_hiddens = self.decoder(
+            input_ids=tgt, decoder_mask=tgt_mask, encoder_embeddings=src_hiddens, encoder_mask=src_mask
+        )
+        log_probs = self.log_softmax(hidden_states=tgt_hiddens)
+        return log_probs
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        # forward pass
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1 added by DataLoader
+                # is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+        src_ids, src_mask, tgt_ids, tgt_mask, labels = batch
+        log_probs = self(src_ids, src_mask, tgt_ids, tgt_mask)
+        train_loss = self.loss_fn(log_probs=log_probs, labels=labels)
+        tensorboard_logs = {
+            'train_loss': train_loss,
+            'lr': self._optimizer.param_groups[0]['lr'],
+        }
+        return {'loss': train_loss, 'log': tensorboard_logs}
+
+    def eval_step(self, batch, batch_idx, mode, dataloader_idx=0):
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1 added by DataLoader
+                # is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+
+        if self.multilingual:
+            self.source_processor = self.source_processor_list[dataloader_idx]
+            self.target_processor = self.target_processor_list[dataloader_idx]
+
+        src_ids, src_mask, tgt_ids, tgt_mask, labels = batch
+        log_probs = self(src_ids, src_mask, tgt_ids, tgt_mask)
+        eval_loss = self.eval_loss_fn(log_probs=log_probs, labels=labels)
+        # this will run encoder twice -- TODO: potentially fix
+        _, translations = self.batch_translate(src=src_ids, src_mask=src_mask)
+        if dataloader_idx == 0:
+            getattr(self, f'{mode}_loss')(loss=eval_loss, num_measurements=log_probs.shape[0] * log_probs.shape[1])
+        else:
+            getattr(self, f'{mode}_loss_{dataloader_idx}')(
+                loss=eval_loss, num_measurements=log_probs.shape[0] * log_probs.shape[1]
+            )
+        np_tgt = tgt_ids.detach().cpu().numpy()
+        ground_truths = [self.decoder_tokenizer.ids_to_text(tgt) for tgt in np_tgt]
+        ground_truths = [self.target_processor.detokenize(tgt.split(' ')) for tgt in ground_truths]
+        num_non_pad_tokens = np.not_equal(np_tgt, self.decoder_tokenizer.pad_id).sum().item()
+        return {
+            'translations': translations,
+            'ground_truths': ground_truths,
+            'num_non_pad_tokens': num_non_pad_tokens,
+        }
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        return self.eval_step(batch, batch_idx, 'test', dataloader_idx)
+
+    @rank_zero_only
+    def log_param_stats(self):
+        for name, p in self.named_parameters():
+            if p.requires_grad:
+                self.trainer.logger.experiment.add_histogram(name + '_hist', p, global_step=self.global_step)
+                self.trainer.logger.experiment.add_scalars(
+                    name,
+                    {'mean': p.mean(), 'stddev': p.std(), 'max': p.max(), 'min': p.min()},
+                    global_step=self.global_step,
+                )
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        return self.eval_step(batch, batch_idx, 'val', dataloader_idx)
+
+    def eval_epoch_end(self, outputs, mode):
+        # if user specifies one validation dataloader, then PTL reverts to giving a list of dictionary instead of a list of list of dictionary
+        if isinstance(outputs[0], dict):
+            outputs = [outputs]
+
+        loss_list = []
+        sb_score_list = []
+        for dataloader_idx, output in enumerate(outputs):
+            if dataloader_idx == 0:
+                eval_loss = getattr(self, f'{mode}_loss').compute()
+            else:
+                eval_loss = getattr(self, f'{mode}_loss_{dataloader_idx}').compute()
+
+            translations = list(itertools.chain(*[x['translations'] for x in output]))
+            ground_truths = list(itertools.chain(*[x['ground_truths'] for x in output]))
+            assert len(translations) == len(ground_truths)
+
+            # Gather translations and ground truths from all workers
+            tr_and_gt = [None for _ in range(self.world_size)]
+            # we also need to drop pairs where ground truth is an empty string
+            dist.all_gather_object(
+                tr_and_gt, [(t, g) for (t, g) in zip(translations, ground_truths) if g.strip() != '']
+            )
+            if self.global_rank == 0:
+                _translations = []
+                _ground_truths = []
+                for rank in range(0, self.world_size):
+                    _translations += [t for (t, g) in tr_and_gt[rank]]
+                    _ground_truths += [g for (t, g) in tr_and_gt[rank]]
+
+                if self.tgt_language in ['ja']:
+                    sacre_bleu = corpus_bleu(_translations, [_ground_truths], tokenize="ja-mecab")
+                elif self.tgt_language in ['zh']:
+                    sacre_bleu = corpus_bleu(_translations, [_ground_truths], tokenize="zh")
+                else:
+                    sacre_bleu = corpus_bleu(_translations, [_ground_truths], tokenize="13a")
+
+                # because the reduction op later is average (over word_size)
+                sb_score = sacre_bleu.score * self.world_size
+
+                dataset_name = "Validation" if mode == 'val' else "Test"
+                logging.info(
+                    f"Dataset name: {dataset_name}, Dataloader index: {dataloader_idx}, Set size: {len(translations)}"
+                )
+                logging.info(
+                    f"Dataset name: {dataset_name}, Dataloader index: {dataloader_idx}, Val Loss = {eval_loss}"
+                )
+                logging.info(
+                    f"Dataset name: {dataset_name}, Dataloader index: {dataloader_idx}, Sacre BLEU = {sb_score / self.world_size}"
+                )
+                logging.info(
+                    f"Dataset name: {dataset_name}, Dataloader index: {dataloader_idx}, Translation Examples:"
+                )
+                for i in range(0, 3):
+                    ind = random.randint(0, len(translations) - 1)
+                    logging.info("    " + '\u0332'.join(f"Example {i}:"))
+                    logging.info(f"    Prediction:   {translations[ind]}")
+                    logging.info(f"    Ground Truth: {ground_truths[ind]}")
+            else:
+                sb_score = 0.0
+
+            loss_list.append(eval_loss.cpu().numpy())
+            sb_score_list.append(sb_score)
+            if dataloader_idx == 0:
+                self.log(f"{mode}_loss", eval_loss, sync_dist=True)
+                self.log(f"{mode}_sacreBLEU", sb_score, sync_dist=True)
+                getattr(self, f'{mode}_loss').reset()
+            else:
+                self.log(f"{mode}_loss_dl_index_{dataloader_idx}", eval_loss, sync_dist=True)
+                self.log(f"{mode}_sacreBLEU_dl_index_{dataloader_idx}", sb_score, sync_dist=True)
+                getattr(self, f'{mode}_loss_{dataloader_idx}').reset()
+
+        if len(loss_list) > 1:
+            self.log(f"{mode}_loss_avg", np.mean(loss_list), sync_dist=True)
+            self.log(f"{mode}_sacreBLEU_avg", np.mean(sb_score_list), sync_dist=True)
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+        self.eval_epoch_end(outputs, 'val')
+
+    def test_epoch_end(self, outputs):
+        self.eval_epoch_end(outputs, 'test')
+
+    def setup_enc_dec_tokenizers(
+        self,
+        encoder_tokenizer_library=None,
+        encoder_tokenizer_model=None,
+        encoder_bpe_dropout=0.0,
+        encoder_model_name=None,
+        decoder_tokenizer_library=None,
+        decoder_tokenizer_model=None,
+        decoder_bpe_dropout=0.0,
+        decoder_model_name=None,
+    ):
+
+        supported_tokenizers = ['yttm', 'huggingface', 'sentencepiece']
+        if (
+            encoder_tokenizer_library not in supported_tokenizers
+            or decoder_tokenizer_library not in supported_tokenizers
+        ):
+            raise NotImplementedError(f"Currently we only support tokenizers in {supported_tokenizers}.")
+
+        self.encoder_tokenizer = get_nmt_tokenizer(
+            library=encoder_tokenizer_library,
+            tokenizer_model=self.register_artifact("encoder_tokenizer.tokenizer_model", encoder_tokenizer_model),
+            bpe_dropout=encoder_bpe_dropout,
+            model_name=encoder_model_name,
+            vocab_file=None,
+            special_tokens=None,
+            use_fast=False,
+        )
+        self.decoder_tokenizer = get_nmt_tokenizer(
+            library=decoder_tokenizer_library,
+            tokenizer_model=self.register_artifact("decoder_tokenizer.tokenizer_model", decoder_tokenizer_model),
+            bpe_dropout=decoder_bpe_dropout,
+            model_name=decoder_model_name,
+            vocab_file=None,
+            special_tokens=None,
+            use_fast=False,
+        )
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_multiple_validation_data(self, val_data_config: Union[DictConfig, Dict]):
+        self.setup_validation_data(self._cfg.get('validation_ds'))
+
+    def setup_multiple_test_data(self, test_data_config: Union[DictConfig, Dict]):
+        self.setup_test_data(self._cfg.get('test_ds'))
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self._setup_eval_dataloader_from_config(cfg=val_data_config)
+        # instantiate Torchmetric for each val dataloader
+        if self._validation_dl is not None:
+            for dataloader_idx in range(len(self._validation_dl)):
+                if dataloader_idx == 0:
+                    setattr(
+                        self, f'val_loss', GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
+                    )
+                else:
+                    setattr(
+                        self,
+                        f'val_loss_{dataloader_idx}',
+                        GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
+                    )
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self._setup_eval_dataloader_from_config(cfg=test_data_config)
+        # instantiate Torchmetric for each test dataloader
+        if self._test_dl is not None:
+            for dataloader_idx in range(len(self._test_dl)):
+                if dataloader_idx == 0:
+                    setattr(
+                        self, f'test_loss', GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
+                    )
+                else:
+                    setattr(
+                        self,
+                        f'test_loss_{dataloader_idx}',
+                        GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
+                    )
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+        if cfg.get("use_tarred_dataset", False):
+            if cfg.get("metadata_file") is None:
+                raise FileNotFoundError("Trying to use tarred data set but could not find metadata path in config.")
+            else:
+                if not self.multilingual:
+                    metadata_file_list = [cfg.get('metadata_file')]
+                else:
+                    metadata_file_list = cfg.get('metadata_file')
+
+                datasets = []
+                for idx, metadata_file in enumerate(metadata_file_list):
+                    with open(metadata_file) as metadata_reader:
+                        metadata = json.load(metadata_reader)
+                    if cfg.get('tar_files') is None:
+                        tar_files = metadata.get('tar_files')
+                        if tar_files is not None:
+                            logging.info(f'Loading from tarred dataset {tar_files}')
+                        else:
+                            raise FileNotFoundError("Could not find tarred dataset in config or metadata.")
+                    else:
+                        tar_files = cfg.get('tar_files')
+                        if self.multilingual:
+                            tar_files = tar_files[idx]
+                        if metadata.get('tar_files') is not None:
+                            logging.info(
+                                f'Tar file paths found in both cfg and metadata using one in cfg by default - {tar_files}'
+                            )
+
+                    dataset = TarredTranslationDataset(
+                        text_tar_filepaths=tar_files,
+                        metadata_path=metadata_file,
+                        encoder_tokenizer=self.encoder_tokenizer,
+                        decoder_tokenizer=self.decoder_tokenizer,
+                        shuffle_n=cfg.get("tar_shuffle_n", 100),
+                        shard_strategy=cfg.get("shard_strategy", "scatter"),
+                        global_rank=self.global_rank,
+                        world_size=self.world_size,
+                        reverse_lang_direction=cfg.get("reverse_lang_direction", False),
+                        prepend_id=self.multilingual_ids[idx],
+                    )
+                    datasets.append(dataset)
+
+                if self.multilingual:
+                    dataset = ConcatDataset(
+                        datasets=datasets,
+                        sampling_technique=cfg.get('concat_sampling_technique'),
+                        sampling_temperature=cfg.get('concat_sampling_temperature'),
+                        sampling_probabilities=cfg.get('concat_sampling_probabilities'),
+                        global_rank=self.global_rank,
+                        world_size=self.world_size,
+                    )
+                else:
+                    dataset = datasets[0]
+
+            return torch.utils.data.DataLoader(
+                dataset=dataset,
+                batch_size=1,
+                num_workers=cfg.get("num_workers", 2),
+                pin_memory=cfg.get("pin_memory", False),
+                drop_last=cfg.get("drop_last", False),
+            )
+        else:
+            if not self.multilingual:
+                src_file_list = [cfg.src_file_name]
+                tgt_file_list = [cfg.tgt_file_name]
+            else:
+                src_file_list = cfg.src_file_name
+                tgt_file_list = cfg.tgt_file_name
+
+            if len(src_file_list) != len(tgt_file_list):
+                raise ValueError(
+                    'The same number of filepaths must be passed in for source and target while training multilingual.'
+                )
+
+            datasets = []
+            for idx, src_file in enumerate(src_file_list):
+                dataset = TranslationDataset(
+                    dataset_src=str(Path(src_file).expanduser()),
+                    dataset_tgt=str(Path(tgt_file_list[idx]).expanduser()),
+                    tokens_in_batch=cfg.tokens_in_batch,
+                    clean=cfg.get("clean", False),
+                    max_seq_length=cfg.get("max_seq_length", 512),
+                    min_seq_length=cfg.get("min_seq_length", 1),
+                    max_seq_length_diff=cfg.get("max_seq_length_diff", 512),
+                    max_seq_length_ratio=cfg.get("max_seq_length_ratio", 512),
+                    cache_ids=cfg.get("cache_ids", False),
+                    cache_data_per_node=cfg.get("cache_data_per_node", False),
+                    use_cache=cfg.get("use_cache", False),
+                    reverse_lang_direction=cfg.get("reverse_lang_direction", False),
+                    prepend_id=self.multilingual_ids[idx],
+                )
+                dataset.batchify(self.encoder_tokenizer, self.decoder_tokenizer)
+                datasets.append(dataset)
+
+            if self.multilingual:
+                dataset = ConcatDataset(
+                    datasets=datasets,
+                    shuffle=cfg.get('shuffle'),
+                    sampling_technique=cfg.get('concat_sampling_technique'),
+                    sampling_temperature=cfg.get('concat_sampling_temperature'),
+                    sampling_probabilities=cfg.get('concat_sampling_probabilities'),
+                    global_rank=self.global_rank,
+                    world_size=self.world_size,
+                )
+                return torch.utils.data.DataLoader(
+                    dataset=dataset,
+                    batch_size=1,
+                    num_workers=cfg.get("num_workers", 2),
+                    pin_memory=cfg.get("pin_memory", False),
+                    drop_last=cfg.get("drop_last", False),
+                )
+            else:
+                dataset = datasets[0]
+
+        if cfg.shuffle:
+            sampler = pt_data.RandomSampler(dataset)
+        else:
+            sampler = pt_data.SequentialSampler(dataset)
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=1,
+            sampler=sampler,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )
+
+    def replace_beam_with_sampling(self, topk=500):
+        self.beam_search = TopKSequenceGenerator(
+            embedding=self.decoder.embedding,
+            decoder=self.decoder.decoder,
+            log_softmax=self.log_softmax,
+            max_sequence_length=self.beam_search.max_seq_length,
+            beam_size=topk,
+            bos=self.decoder_tokenizer.bos_id,
+            pad=self.decoder_tokenizer.pad_id,
+            eos=self.decoder_tokenizer.eos_id,
+        )
+
+    def _setup_eval_dataloader_from_config(self, cfg: DictConfig):
+        src_file_name = cfg.get('src_file_name')
+        tgt_file_name = cfg.get('tgt_file_name')
+
+        if src_file_name is None or tgt_file_name is None:
+            raise ValueError(
+                'Validation dataloader needs both cfg.src_file_name and cfg.tgt_file_name to not be None.'
+            )
+        else:
+            # convert src_file_name and tgt_file_name to list of strings
+            if isinstance(src_file_name, str):
+                src_file_list = [src_file_name]
+            elif isinstance(src_file_name, ListConfig):
+                src_file_list = src_file_name
+            else:
+                raise ValueError("cfg.src_file_name must be string or list of strings")
+            if isinstance(tgt_file_name, str):
+                tgt_file_list = [tgt_file_name]
+            elif isinstance(tgt_file_name, ListConfig):
+                tgt_file_list = tgt_file_name
+            else:
+                raise ValueError("cfg.tgt_file_name must be string or list of strings")
+        if len(src_file_list) != len(tgt_file_list):
+            raise ValueError('The same number of filepaths must be passed in for source and target validation.')
+
+        dataloaders = []
+        prepend_idx = 0
+        for idx, src_file in enumerate(src_file_list):
+            if self.multilingual:
+                prepend_idx = idx
+            dataset = TranslationDataset(
+                dataset_src=str(Path(src_file).expanduser()),
+                dataset_tgt=str(Path(tgt_file_list[idx]).expanduser()),
+                tokens_in_batch=cfg.tokens_in_batch,
+                clean=cfg.get("clean", False),
+                max_seq_length=cfg.get("max_seq_length", 512),
+                min_seq_length=cfg.get("min_seq_length", 1),
+                max_seq_length_diff=cfg.get("max_seq_length_diff", 512),
+                max_seq_length_ratio=cfg.get("max_seq_length_ratio", 512),
+                cache_ids=cfg.get("cache_ids", False),
+                cache_data_per_node=cfg.get("cache_data_per_node", False),
+                use_cache=cfg.get("use_cache", False),
+                reverse_lang_direction=cfg.get("reverse_lang_direction", False),
+                prepend_id=self.multilingual_ids[prepend_idx],
+            )
+            dataset.batchify(self.encoder_tokenizer, self.decoder_tokenizer)
+
+            if cfg.shuffle:
+                sampler = pt_data.RandomSampler(dataset)
+            else:
+                sampler = pt_data.SequentialSampler(dataset)
+
+            dataloader = torch.utils.data.DataLoader(
+                dataset=dataset,
+                batch_size=1,
+                sampler=sampler,
+                num_workers=cfg.get("num_workers", 2),
+                pin_memory=cfg.get("pin_memory", False),
+                drop_last=cfg.get("drop_last", False),
+            )
+            dataloaders.append(dataloader)
+
+        return dataloaders
+
+    def setup_pre_and_post_processing_utils(self, source_lang, target_lang):
+        """
+        Creates source and target processor objects for input and output pre/post-processing.
+        """
+        self.source_processor, self.target_processor = None, None
+        if (source_lang == 'en' and target_lang == 'ja') or (source_lang == 'ja' and target_lang == 'en'):
+            self.source_processor = EnJaProcessor(source_lang)
+            self.target_processor = EnJaProcessor(target_lang)
+        else:
+            if source_lang == 'zh':
+                self.source_processor = ChineseProcessor()
+            if target_lang == 'zh':
+                self.target_processor = ChineseProcessor()
+            if source_lang is not None and source_lang not in ['ja', 'zh']:
+                self.source_processor = MosesProcessor(source_lang)
+            if target_lang is not None and target_lang not in ['ja', 'zh']:
+                self.target_processor = MosesProcessor(target_lang)
+
+        return self.source_processor, self.target_processor
+
+    @torch.no_grad()
+    def batch_translate(
+        self, src: torch.LongTensor, src_mask: torch.LongTensor,
+    ):
+        """	
+        Translates a minibatch of inputs from source language to target language.	
+        Args:	
+            src: minibatch of inputs in the src language (batch x seq_len)	
+            src_mask: mask tensor indicating elements to be ignored (batch x seq_len)	
+        Returns:	
+            translations: a list strings containing detokenized translations	
+            inputs: a list of string containing detokenized inputs	
+        """
+        mode = self.training
+        try:
+            self.eval()
+            src_hiddens = self.encoder(input_ids=src, encoder_mask=src_mask)
+            beam_results = self.beam_search(encoder_hidden_states=src_hiddens, encoder_input_mask=src_mask)
+            beam_results = self.filter_predicted_ids(beam_results)
+
+            translations = [self.decoder_tokenizer.ids_to_text(tr) for tr in beam_results.cpu().numpy()]
+            inputs = [self.encoder_tokenizer.ids_to_text(inp) for inp in src.cpu().numpy()]
+            if self.target_processor is not None:
+                translations = [
+                    self.target_processor.detokenize(translation.split(' ')) for translation in translations
+                ]
+
+            if self.source_processor is not None:
+                inputs = [self.source_processor.detokenize(item.split(' ')) for item in inputs]
+        finally:
+            self.train(mode=mode)
+        return inputs, translations
+
+    # TODO: We should drop source/target_lang arguments in favor of using self.src/tgt_language
+    @torch.no_grad()
+    def translate(self, text: List[str], source_lang: str = None, target_lang: str = None) -> List[str]:
+        """
+        Translates list of sentences from source language to target language.
+        Should be regular text, this method performs its own tokenization/de-tokenization
+        Args:
+            text: list of strings to translate
+            source_lang: if not None, corresponding MosesTokenizer and MosesPunctNormalizer will be run
+            target_lang: if not None, corresponding MosesDecokenizer will be run
+        Returns:
+            list of translated strings
+        """
+        # __TODO__: This will reset both source and target processors even if you want to reset just one.
+        if source_lang is not None or target_lang is not None:
+            self.setup_pre_and_post_processing_utils(source_lang, target_lang)
+
+        mode = self.training
+        prepend_ids = []
+        if self.multilingual:
+            if source_lang is None or target_lang is None:
+                raise ValueError("Expect source_lang and target_lang to infer for multilingual model.")
+            src_symbol = self.encoder_tokenizer.token_to_id('<' + source_lang + '>')
+            tgt_symbol = self.encoder_tokenizer.token_to_id('<' + target_lang + '>')
+            prepend_ids = [src_symbol if src_symbol in self.multilingual_ids else tgt_symbol]
+        try:
+            self.eval()
+            inputs = []
+            for txt in text:
+                if self.source_processor is not None:
+                    txt = self.source_processor.normalize(txt)
+                    txt = self.source_processor.tokenize(txt)
+                ids = self.encoder_tokenizer.text_to_ids(txt)
+                ids = prepend_ids + [self.encoder_tokenizer.bos_id] + ids + [self.encoder_tokenizer.eos_id]
+                inputs.append(ids)
+            max_len = max(len(txt) for txt in inputs)
+            src_ids_ = np.ones((len(inputs), max_len)) * self.encoder_tokenizer.pad_id
+            for i, txt in enumerate(inputs):
+                src_ids_[i][: len(txt)] = txt
+
+            src_mask = torch.FloatTensor((src_ids_ != self.encoder_tokenizer.pad_id)).to(self.device)
+            src = torch.LongTensor(src_ids_).to(self.device)
+            _, translations = self.batch_translate(src, src_mask)
+        finally:
+            self.train(mode=mode)
+        return translations
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_en_de_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_de_transformer12x2/versions/1.0.0rc1/files/nmt_en_de_transformer12x2.nemo",
+            description="En->De translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_de_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_de_en_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_de_en_transformer12x2/versions/1.0.0rc1/files/nmt_de_en_transformer12x2.nemo",
+            description="De->En translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_de_en_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_en_es_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_es_transformer12x2/versions/1.0.0rc1/files/nmt_en_es_transformer12x2.nemo",
+            description="En->Es translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_es_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_es_en_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_es_en_transformer12x2/versions/1.0.0rc1/files/nmt_es_en_transformer12x2.nemo",
+            description="Es->En translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_es_en_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_en_fr_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_fr_transformer12x2/versions/1.0.0rc1/files/nmt_en_fr_transformer12x2.nemo",
+            description="En->Fr translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_fr_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_fr_en_transformer12x2",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_fr_en_transformer12x2/versions/1.0.0rc1/files/nmt_fr_en_transformer12x2.nemo",
+            description="Fr->En translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_fr_en_transformer12x2",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_en_ru_transformer6x6",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_ru_transformer6x6/versions/1.0.0rc1/files/nmt_en_ru_transformer6x6.nemo",
+            description="En->Ru translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_ru_transformer6x6",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_ru_en_transformer6x6",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_ru_en_transformer6x6/versions/1.0.0rc1/files/nmt_ru_en_transformer6x6.nemo",
+            description="Ru->En translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_ru_en_transformer6x6",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_zh_en_transformer6x6",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_zh_en_transformer6x6/versions/1.0.0rc1/files/nmt_zh_en_transformer6x6.nemo",
+            description="Zh->En translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_zh_en_transformer6x6",
+        )
+        result.append(model)
+
+        model = PretrainedModelInfo(
+            pretrained_model_name="nmt_en_zh_transformer6x6",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/nmt_en_zh_transformer6x6/versions/1.0.0rc1/files/nmt_en_zh_transformer6x6.nemo",
+            description="En->Zh translation model. See details here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_zh_transformer6x6",
+        )
+        result.append(model)
+
+        return result

tts/nemo/collections/nlp/models/neural_machine_translation/__init__.py

@@ -0,0 +1,17 @@
+# 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 nemo.collections.nlp.models.neural_machine_translation.neural_machine_translation_model import (
+    NeuralMachineTranslationModel,
+)

tts/nemo/collections/nlp/models/neural_machine_translation/neural_machine_translation_model.py

@@ -0,0 +1,245 @@
+# Copyright 2018 The Google AI Language Team Authors and
+# The HuggingFace Inc. team.
+# 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.
+
+
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from transformers import AutoModel, BartForConditionalGeneration, EncoderDecoderModel
+
+from nemo.collections.common.metrics import Perplexity
+from nemo.collections.nlp.data.neural_machine_translation import NeuralMachineTranslationDataset
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.neural_types import ChannelType, LossType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ["NeuralMachineTranslationModel"]
+
+
+class NeuralMachineTranslationModel(ModelPT):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "attention_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "decoder_input_ids": NeuralType(('B', 'T'), ChannelType(), optional=True),
+            "labels": NeuralType(('B', 'T'), ChannelType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "loss": NeuralType((), LossType()),
+            "decoder_hidden_states": NeuralType(("B", "T", "D"), ChannelType(), optional=True),
+            "encoder_hidden_states": NeuralType(("B", "T", "D"), ChannelType(), optional=True),
+        }
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        # must assign tokenizers before init
+        if cfg.language_model.pretrained_model_name:
+            if cfg.language_model.pretrained_encoder_model_name or cfg.language_model.pretrained_decoder_model_name:
+                raise ValueError(
+                    "Must have either pretrained_model_name or both pretrained_encoder_model name and "
+                    "pretrained_decoder_model_name."
+                )
+            # setup tokenizer
+            self.encoder_tokenizer = self.setup_tokenizer(cfg.encoder_tokenizer)
+            self.encoder_add_special_tokens = cfg.encoder_tokenizer.add_special_tokens
+
+            # set decoder to encoder
+            self.decoder_tokenizer = self.encoder_tokenizer
+            self.decoder_add_special_tokens = self.encoder_add_special_tokens
+        else:
+            if not (
+                cfg.language_model.pretrained_encoder_model_name and cfg.language_model.pretrained_decoder_model_name
+            ):
+                raise ValueError("Both encoder and decoder must be specified")
+
+            # setup tokenizers
+            self.encoder_tokenizer = self.setup_tokenizer(cfg.encoder_tokenizer)
+            self.encoder_add_special_tokens = cfg.encoder_tokenizer.add_special_tokens
+
+            self.decoder_tokenizer = self.setup_tokenizer(cfg.decoder_tokenizer)
+            self.decoder_add_special_tokens = cfg.decoder_tokenizer.add_special_tokens
+
+        if not self.encoder_tokenizer:
+            raise TypeError("encoder_tokenizer failed to initialize")
+        if not self.decoder_tokenizer:
+            raise TypeError("decoder_tokenizer failed to initialize")
+
+        # init superclass
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # must assign modules after init
+        if cfg.language_model.pretrained_model_name:
+            # Setup end-to-end model
+            if "bart" in cfg.language_model.pretrained_model_name:
+                self.model = BartForConditionalGeneration.from_pretrained(cfg.language_model.pretrained_model_name)
+            else:
+                self.model = AutoModel.from_pretrained(cfg.language_model.pretrained_model_name)
+        else:
+            if not (
+                cfg.language_model.pretrained_encoder_model_name and cfg.language_model.pretrained_decoder_model_name
+            ):
+                raise ValueError("Both encoder and decoder must be specified")
+
+            # Setup encoder/decoder model
+            self.model = EncoderDecoderModel.from_encoder_decoder_pretrained(
+                encoder=cfg.language_model.pretrained_encoder_model_name,
+                decoder=cfg.language_model.pretrained_decoder_model_name,
+            )
+
+        self.validation_perplexity = Perplexity(compute_on_step=False)
+
+        self.setup_optimization(cfg.optim)
+
+    @typecheck()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+        decoder_input_ids: torch.Tensor = None,
+        labels: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            labels=labels,
+            return_dict=False,
+        )
+        return outputs
+
+    @typecheck.disable_checks()
+    def generate(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Wraps huggingface EncoderDecoder.generate()."""
+        outputs = self.model.generate(
+            input_ids=input_ids,
+            pad_token_id=self.encoder_tokenizer.pad_id,
+            bos_token_id=self.encoder_tokenizer.bos_id,
+            eos_token_id=self.encoder_tokenizer.eos_id,
+            decoder_start_token_id=self.decoder_tokenizer.bos_id,
+            **self._cfg.generate,
+        )
+        return outputs
+
+    def training_step(self, batch: Tuple, batch_idx: int) -> Dict:
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`. Loss calculation from HuggingFace's BartForConditionalGeneration.
+        """
+        input_ids, input_mask, decoder_input_ids, labels = batch
+        loss = self.forward(
+            input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=decoder_input_ids, labels=labels,
+        )[0]
+
+        tensorboard_logs = {"train_loss": loss, "lr": self._optimizer.param_groups[0]["lr"]}
+
+        return {"loss": loss, "log": tensorboard_logs}
+
+    def validation_step(self, batch: Tuple, batch_idx: int) -> Dict:
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`. Loss calculation from HuggingFace's BartForConditionalGeneration.
+        """
+        input_ids, input_mask, decoder_input_ids, labels = batch
+        loss, logits = self.forward(
+            input_ids=input_ids, attention_mask=input_mask, decoder_input_ids=decoder_input_ids, labels=labels,
+        )[:2]
+
+        self.validation_perplexity(logits=logits)
+
+        tensorboard_logs = {"val_loss": loss}
+
+        return {"val_loss": loss, "log": tensorboard_logs}
+
+    def validation_epoch_end(self, outputs: List[Dict]) -> Dict:
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+        perplexity = self.validation_perplexity.compute()
+        tensorboard_logs = {"val_loss": avg_loss, "perplexity": perplexity}
+        logging.info(f"evaluation perplexity {perplexity.item()}")
+        return {"val_loss": avg_loss, "log": tensorboard_logs}
+
+    @typecheck.disable_checks()
+    def test_step(self, batch: Tuple, batch_idx: int) -> torch.Tensor:
+        """Lightning calls this inside the test loop with data from the test dataloader."""
+        input_ids, input_mask, decoder_input_ids, labels = batch
+        sequences = self.generate(input_ids=input_ids)
+        return sequences
+
+    @typecheck.disable_checks()
+    def test_epoch_end(self, outputs: List[torch.Tensor]) -> Dict[str, List[str]]:
+        """Called at the end of test to aggregate outputs and decode them."""
+        texts = [self.encoder_tokenizer.ids_to_text(seq) for batch in outputs for seq in batch]
+        return {"texts": texts}
+
+    def setup_tokenizer(self, cfg: DictConfig):
+        tokenizer = get_tokenizer(
+            tokenizer_name=cfg.tokenizer_name,
+            tokenizer_model=cfg.tokenizer_model,
+            special_tokens=OmegaConf.to_container(cfg.special_tokens) if cfg.special_tokens else None,
+            vocab_file=cfg.vocab_file,
+        )
+        return tokenizer
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self.setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self.setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self.setup_dataloader_from_config(cfg=test_data_config)
+
+    def setup_dataloader_from_config(self, cfg: DictConfig):
+        dataset = NeuralMachineTranslationDataset(
+            filepath=cfg.filepath,
+            encoder_tokenizer=self.encoder_tokenizer,
+            decoder_tokenizer=self.decoder_tokenizer,
+            encoder_add_special_tokens=self.encoder_add_special_tokens,
+            decoder_add_special_tokens=self.decoder_add_special_tokens,
+            max_seq_length=self._cfg.max_seq_length,
+            num_samples=cfg.get("num_samples", -1),
+            convert_labels=self._cfg.convert_labels,
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=self._cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+            collate_fn=dataset.collate_fn,
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass

tts/nemo/collections/nlp/models/nlp_model.py

@@ -0,0 +1,463 @@
+# 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 glob
+import hashlib
+import json
+import os
+import shutil
+import tarfile
+import tempfile
+from typing import Any, Dict, Optional, Union
+
+import torch
+from megatron import mpu
+from megatron.checkpointing import get_checkpoint_version, set_checkpoint_version
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.utilities import rank_zero_only
+from transformers import TRANSFORMERS_CACHE
+
+from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+from nemo.collections.nlp.modules import BertModule, MegatronBertEncoder
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import compute_model_parallel_rank
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.collections.nlp.parts.nlp_overrides import NLPCheckpointConnector
+from nemo.core.classes import ModelPT
+from nemo.core.classes.exportable import Exportable
+from nemo.utils import AppState, logging
+from nemo.utils.exp_manager import configure_checkpointing
+from nemo.utils.get_rank import is_global_rank_zero
+
+__all__ = ['NLPModel']
+
+NEMO_NLP_TMP = os.path.join(os.path.dirname(str(TRANSFORMERS_CACHE)), "nemo_nlp_tmp")
+
+os.makedirs(NEMO_NLP_TMP, exist_ok=True)
+
+
+class NLPModel(ModelPT, Exportable):
+    """Base class for NLP Models.
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        super().__init__(cfg, trainer)
+        self.set_world_size(trainer)
+
+    def register_artifact(self, config_path: str, src: str, verify_src_exists: bool = False):
+        """ Overrides ModelPT register_artifact default behavior. NLP models usually need artifacts that are optional."""
+        return super().register_artifact(config_path, src, verify_src_exists=verify_src_exists)
+
+    @rank_zero_only
+    def register_bert_model(self):
+        """Adds encoder config to .nemo archive for Jarvis.
+        """
+        # check if there is an encoder, warn if not
+        if self.bert_model is None:
+            raise ValueError('Instantiate self.bert_model before registering it.')
+        else:
+            # get encoder config and create source for artifact
+            if isinstance(self.bert_model, MegatronBertEncoder):
+                pretrained_model_name = self.bert_model._model_name
+                encoder_config_path = pretrained_model_name + '_encoder_config'
+                encoder_config_src = os.path.join(NEMO_NLP_TMP, encoder_config_path + '.json')
+                config_for_json = OmegaConf.to_container(self.bert_model.config)
+                with open(encoder_config_src, 'w', encoding='utf-8') as f:
+                    f.write(json.dumps(config_for_json, indent=2, sort_keys=True) + '\n')
+                self.register_artifact('language_model.config_file', encoder_config_src)  # for .nemo
+            elif isinstance(self.bert_model, BertModule):
+                # HuggingFace Transformer Config
+                pretrained_model_name = self.bert_model.name_or_path
+                # Some HF names have "/" in them so we replace with _
+                pretrained_model_name = pretrained_model_name.replace("/", "_")
+                encoder_config_path = pretrained_model_name + '_encoder_config'
+                encoder_config_src = os.path.join(NEMO_NLP_TMP, encoder_config_path + '.json')
+                self.bert_model.config.to_json_file(encoder_config_src)  # name requested by jarvis team
+                self.register_artifact('language_model.config_file', encoder_config_src)  # for .nemo
+            else:
+                logging.info(
+                    f'Registering BERT model config for {self.bert_model} is not yet supported. Please override this method if needed.'
+                )
+
+    def setup_tokenizer(self, cfg: DictConfig):
+        """Instantiates tokenizer based on config and registers tokenizer artifacts.
+
+           If model is being restored from .nemo file then the tokenizer.vocab_file will
+           be used (if it exists).
+
+           Otherwise, we will use the vocab file provided in the config (if it exists).
+
+           Finally, if no vocab file is given (this happens frequently when using HF),
+           we will attempt to extract the vocab from the tokenizer object and then register it.
+
+        Args:
+            cfg (DictConfig): Tokenizer config
+        """
+        vocab_file = None
+        if cfg.vocab_file:
+            vocab_file = self.register_artifact(config_path='tokenizer.vocab_file', src=cfg.vocab_file)
+        tokenizer = get_tokenizer(
+            tokenizer_name=cfg.tokenizer_name,
+            vocab_file=vocab_file,
+            special_tokens=OmegaConf.to_container(cfg.special_tokens) if cfg.special_tokens else None,
+            tokenizer_model=self.register_artifact(config_path='tokenizer.tokenizer_model', src=cfg.tokenizer_model),
+        )
+        self.tokenizer = tokenizer
+
+        if vocab_file is None:
+            # when there is no vocab file we try to get the vocab from the tokenizer and register it
+            self._register_vocab_from_tokenizer(vocab_file_config_path='tokenizer.vocab_file', cfg=cfg)
+
+    @rank_zero_only
+    def _register_vocab_from_tokenizer(
+        self,
+        vocab_file_config_path: str = 'tokenizer.vocab_file',
+        vocab_dict_config_path: str = 'tokenizer_vocab_dict',
+        cfg: DictConfig = None,
+    ):
+        """Creates vocab file from tokenizer if vocab file is None.
+
+        Args:
+            vocab_file_config_path: path to the vocab_file in the config
+            vocab_dict_config_path: path to the vocab_dict in the config
+            cfg: tokenizer config
+        """
+        if self.tokenizer is None:
+            raise ValueError('Instantiate self.tokenizer before registering vocab from it.')
+        else:
+            if isinstance(self.tokenizer, AutoTokenizer):
+                # extract vocab from tokenizer
+                vocab_dict = self.tokenizer.tokenizer.get_vocab()
+
+                # for fast and slow tokenizer vocabularies compatibility
+                vocab_dict = dict(sorted(vocab_dict.items(), key=lambda item: item[1]))
+
+                # get hash of vocab_dict to create a unique directory to write vocab_dict and vocab_file
+                m = hashlib.md5()
+                if 'tokenizer_name' in cfg:
+                    if cfg.tokenizer_name is not None:
+                        # different pretrained models with the same vocab will have different hash
+                        m.update(cfg.tokenizer_name.encode())
+                # get string representation of vocab_dict
+                vocab_dict_str = json.dumps(vocab_dict, sort_keys=True).encode()
+                m.update(vocab_dict_str)
+                vocab_dict_hash = m.hexdigest()
+
+                hash_path = os.path.join(NEMO_NLP_TMP, vocab_dict_hash)
+                os.makedirs(hash_path, exist_ok=True)
+
+                vocab_json_src = os.path.join(hash_path, vocab_dict_config_path)
+
+                with open(vocab_json_src, 'w', encoding='utf-8') as f:
+                    f.write(json.dumps(vocab_dict, indent=2, sort_keys=True) + '\n')
+                self.register_artifact(config_path=vocab_dict_config_path, src=vocab_json_src)
+                # create vocab file
+                vocab_file_src = os.path.join(hash_path, vocab_file_config_path)
+                with open(vocab_file_src, 'w', encoding='utf-8') as f:
+                    for key in vocab_dict:
+                        f.write(key + '\n')
+
+                cfg.vocab_file = vocab_file_src
+                self.register_artifact(config_path=vocab_file_config_path, src=vocab_file_src)
+            else:
+                logging.info(
+                    f'Registering tokenizer vocab for {self.tokenizer} is not yet supported. Please override this method if needed.'
+                )
+
+    def _clip_gradients(self, optimizer, clip_val=None):
+        """ Override of PTL Gradient Clipping.
+            Enables model parallel gradient clipping from Megatron-LM.
+
+        Args:
+            optimizer ([type]): [description]
+            clip_val ([type], optional): [description]. Defaults to None.
+        """
+        app_state = AppState()
+
+        # get clip_val from trainer if None is provided
+        if clip_val is None:
+            clip_val = float(self._trainer.gradient_clip_val)
+
+        if app_state.model_parallel_size is not None:
+            model = self._trainer.get_model()
+            parameters = model.parameters()
+            if mpu.model_parallel_is_initialized():
+                mpu.grads.clip_grad_norm(parameters=parameters, max_norm=clip_val)
+            else:
+                raise ValueError('Model parallel groups must be intialized to use model parallel gradient clipping.')
+
+        else:
+            return Accelerator._clip_gradients(self, optimizer, clip_val)
+
+    def setup(self, stage: str) -> None:
+        """ PTL hook that is called on all DDP processes. """
+
+        if stage == 'fit':
+
+            # adds self.bert_model config to .nemo file
+            if hasattr(self, 'bert_model') and self.bert_model is not None:
+                self.register_bert_model()
+
+            app_state = AppState()
+
+            if app_state.model_parallel_size is not None:
+
+                self._trainer.checkpoint_connector = NLPCheckpointConnector(self._trainer)
+
+                # # Configure checkpointing for model parallel
+                # if app_state.create_checkpoint_callback:
+                #     # global rank 0 is configured by exp_manager
+                #     if not is_global_rank_zero() and app_state.data_parallel_rank == 0:
+                #         configure_checkpointing(
+                #             self._trainer,
+                #             app_state.log_dir,
+                #             app_state.checkpoint_name,
+                #             app_state.checkpoint_callback_params,
+                #         )
+
+    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        """ LightningModule hook that's used to save things in addition to model weights. """
+
+        if hasattr(self, "bert_model") and isinstance(self.bert_model, MegatronBertEncoder):
+            checkpoint['checkpoint_version'] = get_checkpoint_version()
+        return None
+
+    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        """ LightningModule hook that's used to restore things saved with on_save_checkpoint."""
+
+        if hasattr(self, "bert_model") and isinstance(self.bert_model, MegatronBertEncoder):
+            if get_checkpoint_version():
+                assert (
+                    checkpoint['checkpoint_version'] == get_checkpoint_version()
+                ), 'checkpoint version found on_load_checkpoint different than get_checkpoint_version'
+            else:
+                set_checkpoint_version(checkpoint['checkpoint_version'])
+                logging.info(f"Setting Megatron checkpoint version: {checkpoint['checkpoint_version']}")
+        return None
+
+    # no rank check as model parallel models need to be saved on data parallel rank 0
+    def save_to(self, save_path: str):
+        """
+        Saves model instance (weights and configuration) into .nemo file
+         You can use "restore_from" method to fully restore instance from .nemo file.
+
+        .nemo file is an archive (tar.gz) with the following:
+            model_config.yaml - model configuration in .yaml format. You can deserialize this into cfg argument for model's constructor
+            model_weights.ckpt - model checkpoint
+
+        Args:
+            save_path: Path to .nemo file where model instance should be saved
+        """
+        save_path = os.path.abspath(save_path)
+        app_state = AppState()
+        if app_state.model_parallel_size is not None:
+            self._default_save_to(save_path)
+        else:
+            # super.save_to only runs on global rank 0
+            return super().save_to(save_path)
+
+    def _default_save_to(self, save_path: str):
+        app_state = AppState()
+        if app_state.model_parallel_size is not None:
+            # each model parallel rank creates a .nemo file
+            # after all .nemo files are created, each rank
+            # will add their checkpoint to global rank 0
+
+            base_dir = os.path.dirname(save_path)  # use the directory to merge mp_rank .nemo files into one
+
+            # update save_path based on model parallel_rank
+            base_path = os.path.splitext(save_path)[0]  # everything except the extension
+
+            mp_save_path = f'{base_path}_mp_rank_{app_state.model_parallel_rank:02d}.nemo'
+
+            if app_state.data_parallel_rank == 0:
+                super()._default_save_to(mp_save_path)
+
+            # barrier so that all processes have finished writing their weights before creating .nemo file
+            torch.distributed.barrier()
+
+            if is_global_rank_zero():
+                # extract all tar files
+                for mp_rank in range(app_state.model_parallel_size):
+                    mp_tar_path = f'{base_path}_mp_rank_{mp_rank:02d}.nemo'
+                    mp_tar = tarfile.open(mp_tar_path, 'r:gz')
+                    mp_tar.extractall(path=os.path.join(base_dir, f'mp_rank_{mp_rank:02d}'))
+                    mp_tar.close()
+                    os.remove(mp_tar_path)
+
+                # move rank 0 .nemo extract to base_path
+                shutil.move(os.path.join(base_dir, 'mp_rank_00'), base_path)
+
+                # move mp_rank_00 checkpoint to mp_rank_00 directory inside base_path
+                os.mkdir(os.path.join(base_path, 'mp_rank_00'))
+                shutil.move(os.path.join(base_path, 'model_weights.ckpt'), os.path.join(base_path, 'mp_rank_00'))
+
+                # move other mp_rank checkpoints from base_dir to base_path
+                for mp_rank in range(1, app_state.model_parallel_size):
+                    os.mkdir(os.path.join(base_path, f'mp_rank_{mp_rank:02d}'))
+                    shutil.move(
+                        os.path.join(base_dir, f'mp_rank_{mp_rank:02d}', 'model_weights.ckpt'),
+                        os.path.join(base_path, f'mp_rank_{mp_rank:02d}'),
+                    )
+                    # clean up leftover directory
+                    shutil.rmtree(os.path.join(base_dir, f'mp_rank_{mp_rank:02d}'))
+
+                # create tar file from base_path
+                self._make_nemo_file_from_folder(save_path, base_path)
+
+                # clean up base_path
+                shutil.rmtree(base_path)
+
+        elif is_global_rank_zero():
+            return super()._default_save_to(save_path)
+        else:
+            return
+
+    @classmethod
+    def restore_from(
+        cls,
+        restore_path: str,
+        override_config_path: Optional[Union[OmegaConf, str]] = None,
+        map_location: Optional[torch.device] = None,
+        strict: bool = True,
+        return_config: bool = False,
+        trainer: Trainer = None,
+    ):
+        """
+        Restores model instance (weights and configuration) from .nemo file.
+
+        Args:
+            restore_path: path to .nemo file from which model should be instantiated
+            override_config_path: path to a yaml config that will override the internal
+                config file or an OmegaConf / DictConfig object representing the model config.
+            map_location: Optional torch.device() to map the instantiated model to a device.
+                By default (None), it will select a GPU if available, falling back to CPU otherwise.
+            strict: Passed to load_state_dict. Set to True by default.
+            return_config: If set to true, will return just the underlying config of the restored
+                model as an OmegaConf DictConfig object without instantiating the model.
+            trainer: PyTorch Lightning trainer. Must be passed in order to use model parallel .nemo
+
+            Example:
+                ```
+                model = nemo.collections.nlp.models.TokenClassificationModel.restore_from('token_classification.nemo')
+                assert isinstance(model, nemo.collections.nlp.models.TokenClassificationModel)
+                ```
+
+        Returns:
+            An instance of type cls or its underlying config (if return_config is set).
+        """
+        if not os.path.exists(restore_path):
+            raise FileNotFoundError(f"Can't find {restore_path}")
+
+        app_state = AppState()
+        app_state.model_restore_path = os.path.abspath(os.path.expanduser(restore_path))
+
+        # detect if we have a model parallel .nemo file
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cwd = os.getcwd()
+            os.chdir(tmpdir)
+            # detect if model parallel from tarfile
+            tar = tarfile.open(app_state.model_restore_path, "r:gz")
+            names = tar.getnames()
+            mp_ranks = []
+            for name in names:
+                if 'mp_rank' in name:
+                    mp_ranks.append(name)
+            if mp_ranks:
+                app_state.model_parallel_size = len(mp_ranks) // 2  # directory and file are included in getnames()
+
+                # get checkpoint version
+                checkpoint_version_member = None
+                for member in tar.getmembers():
+                    if 'megatron_checkpoint_version.json' in member.name:
+                        checkpoint_version_member = member
+                tar.extract(checkpoint_version_member, tmpdir)
+                with open(checkpoint_version_member.name, 'r') as f:
+                    checkpoint_version = json.load(f).get('checkpoint_version', None)
+                logging.info(
+                    (
+                        f'Detected model parallel .nemo file: {restore_path}. '
+                        f'Assuming megatron model parallelism with '
+                        f'model_parallel_size: {app_state.model_parallel_size} '
+                        f'and checkpoint version: {checkpoint_version}'
+                    )
+                )
+            tar.close()
+            os.chdir(cwd)
+
+        if app_state.model_parallel_size is not None:
+            if not isinstance(trainer, Trainer):
+                raise ValueError("trainer must be a PyTorch Lightning Trainer to restore model parallel .nemo files.")
+
+            if checkpoint_version is None:
+                raise ValueError(
+                    "Restoring from megatron model parallel .nemo but could not find megatron checkpoint version."
+                )
+            else:
+                logging.info(f"Setting megatron checkpoint version: {checkpoint_version}")
+                set_checkpoint_version(checkpoint_version)
+
+            app_state.world_size = trainer.num_gpus * trainer.num_nodes
+
+            if trainer.local_rank is not None:
+                app_state.local_rank = trainer.local_rank
+            else:
+                raise ValueError("trainer.local_rank is None. local_rank needed to restore model parallel models.")
+
+            model_parallel_rank = compute_model_parallel_rank(trainer.local_rank, app_state.model_parallel_size)
+            app_state.model_parallel_rank = model_parallel_rank
+
+            restored_model = cls._default_restore_from(
+                restore_path, override_config_path, map_location, strict, return_config
+            )
+            restored_model.set_trainer(trainer)
+            return restored_model
+        else:
+            return super().restore_from(restore_path, override_config_path, map_location, strict, return_config)
+
+    @rank_zero_only
+    def register_megatron_checkpoint_version(self):
+        """ Adds checkpoint version to .nemo archive """
+        if self.bert_model is None:
+            raise ValueError('Instantiate self.bert_model before registering megatron checkpoint version.')
+        else:
+            # get encoder config and create source for artifact
+            if isinstance(self.bert_model, MegatronBertEncoder):
+                checkpoint_version = get_checkpoint_version()
+                if checkpoint_version is None:
+                    raise ValueError('Unable to get megatron checkpoint version.')
+                else:
+                    checkpoint_version_dict = {'checkpoint_version': checkpoint_version}
+                    checkpoint_version_path = 'megatron_checkpoint_version.json'
+                    checkpoint_version_src = os.path.join(NEMO_NLP_TMP, checkpoint_version_path)
+                    with open(checkpoint_version_src, 'w') as f:
+                        f.write(json.dumps(checkpoint_version_dict))
+                    self.register_artifact(checkpoint_version_path, checkpoint_version_src)
+
+    @staticmethod
+    def _unpack_nemo_file(path2file: str, out_folder: str) -> str:
+        return super(NLPModel, NLPModel)._unpack_nemo_file(path2file, out_folder)
+
+    @staticmethod
+    def _make_nemo_file_from_folder(filename, source_dir):
+        return super(NLPModel, NLPModel)._make_nemo_file_from_folder(filename, source_dir)
+
+    @property
+    def input_module(self):
+        return self.bert_model
+
+    @property
+    def output_module(self):
+        return self.classifier

tts/nemo/collections/nlp/models/question_answering/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.collections.nlp.models.question_answering.qa_model import QAModel

tts/nemo/collections/nlp/models/question_answering/qa_model.py

@@ -0,0 +1,404 @@
+# 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
+from typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.common.losses import SpanningLoss
+from nemo.collections.nlp.data import SquadDataset
+from nemo.collections.nlp.data.question_answering_squad.qa_squad_processing import (
+    EVALUATION_MODE,
+    INFERENCE_MODE,
+    TRAINING_MODE,
+)
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import tensor2list
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ['QAModel']
+
+
+class QAModel(NLPModel):
+    """
+    BERT encoder with QA head training.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.classifier.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        self.setup_tokenizer(cfg.tokenizer)
+        super().__init__(cfg=cfg, trainer=trainer)
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file)
+            if cfg.tokenizer is not None
+            else None,
+        )
+
+        self.classifier = TokenClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_classes=cfg.token_classifier.num_classes,
+            num_layers=cfg.token_classifier.num_layers,
+            activation=cfg.token_classifier.activation,
+            log_softmax=cfg.token_classifier.log_softmax,
+            dropout=cfg.token_classifier.dropout,
+            use_transformer_init=cfg.token_classifier.use_transformer_init,
+        )
+
+        self.loss = SpanningLoss()
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        logits = self.classifier(hidden_states=hidden_states)
+        return logits
+
+    def training_step(self, batch, batch_idx):
+        input_ids, input_type_ids, input_mask, unique_ids, start_positions, end_positions = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        loss, _, _ = self.loss(logits=logits, start_positions=start_positions, end_positions=end_positions)
+
+        lr = self._optimizer.param_groups[0]['lr']
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+        return {'loss': loss, 'lr': lr}
+
+    def validation_step(self, batch, batch_idx):
+        if self.trainer.testing:
+            prefix = 'test'
+        else:
+            prefix = 'val'
+
+        input_ids, input_type_ids, input_mask, unique_ids, start_positions, end_positions = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        loss, start_logits, end_logits = self.loss(
+            logits=logits, start_positions=start_positions, end_positions=end_positions
+        )
+
+        tensors = {
+            'unique_ids': unique_ids,
+            'start_logits': start_logits,
+            'end_logits': end_logits,
+        }
+        return {f'{prefix}_loss': loss, f'{prefix}_tensors': tensors}
+
+    def test_step(self, batch, batch_idx):
+        return self.validation_step(batch, batch_idx)
+
+    def validation_epoch_end(self, outputs):
+        if self.trainer.testing:
+            prefix = 'test'
+        else:
+            prefix = 'val'
+
+        avg_loss = torch.stack([x[f'{prefix}_loss'] for x in outputs]).mean()
+
+        unique_ids = torch.cat([x[f'{prefix}_tensors']['unique_ids'] for x in outputs])
+        start_logits = torch.cat([x[f'{prefix}_tensors']['start_logits'] for x in outputs])
+        end_logits = torch.cat([x[f'{prefix}_tensors']['end_logits'] for x in outputs])
+
+        all_unique_ids = []
+        all_start_logits = []
+        all_end_logits = []
+        if torch.distributed.is_initialized():
+            world_size = torch.distributed.get_world_size()
+            for ind in range(world_size):
+                all_unique_ids.append(torch.empty_like(unique_ids))
+                all_start_logits.append(torch.empty_like(start_logits))
+                all_end_logits.append(torch.empty_like(end_logits))
+            torch.distributed.all_gather(all_unique_ids, unique_ids)
+            torch.distributed.all_gather(all_start_logits, start_logits)
+            torch.distributed.all_gather(all_end_logits, end_logits)
+        else:
+            all_unique_ids.append(unique_ids)
+            all_start_logits.append(start_logits)
+            all_end_logits.append(end_logits)
+
+        exact_match, f1, all_predictions, all_nbest = -1, -1, [], []
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+
+            unique_ids = []
+            start_logits = []
+            end_logits = []
+            for u in all_unique_ids:
+                unique_ids.extend(tensor2list(u))
+            for u in all_start_logits:
+                start_logits.extend(tensor2list(u))
+            for u in all_end_logits:
+                end_logits.extend(tensor2list(u))
+
+            eval_dataset = self._test_dl.dataset if self.trainer.testing else self._validation_dl.dataset
+            exact_match, f1, all_predictions, all_nbest = eval_dataset.evaluate(
+                unique_ids=unique_ids,
+                start_logits=start_logits,
+                end_logits=end_logits,
+                n_best_size=self._cfg.dataset.n_best_size,
+                max_answer_length=self._cfg.dataset.max_answer_length,
+                version_2_with_negative=self._cfg.dataset.version_2_with_negative,
+                null_score_diff_threshold=self._cfg.dataset.null_score_diff_threshold,
+                do_lower_case=self._cfg.dataset.do_lower_case,
+            )
+
+        logging.info(f"{prefix} exact match {exact_match}")
+        logging.info(f"{prefix} f1 {f1}")
+
+        self.log(f'{prefix}_loss', avg_loss)
+        self.log(f'{prefix}_exact_match', exact_match)
+        self.log(f'{prefix}_f1', f1)
+
+    def test_epoch_end(self, outputs):
+        return self.validation_epoch_end(outputs)
+
+    @torch.no_grad()
+    def inference(
+        self,
+        file: str,
+        batch_size: int = 1,
+        num_samples: int = -1,
+        output_nbest_file: Optional[str] = None,
+        output_prediction_file: Optional[str] = None,
+    ):
+        """
+        Get prediction for unlabeled inference data
+
+        Args:
+            file: inference data
+            batch_size: batch size to use during inference
+            num_samples: number of samples to use of inference data. Default: -1 if all data should be used.
+            output_nbest_file: optional output file for writing out nbest list
+            output_prediction_file: optional output file for writing out predictions
+            
+        Returns:
+            model predictions, model nbest list
+        """
+        # store predictions for all queries in a single list
+        all_predictions = []
+        all_nbest = []
+        mode = self.training
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        try:
+            # Switch model to evaluation mode
+            self.eval()
+            self.to(device)
+            logging_level = logging.get_verbosity()
+            logging.set_verbosity(logging.WARNING)
+            dataloader_cfg = {
+                "batch_size": batch_size,
+                "file": file,
+                "shuffle": False,
+                "num_samples": num_samples,
+                'num_workers': 2,
+                'pin_memory': False,
+                'drop_last': False,
+            }
+            dataloader_cfg = OmegaConf.create(dataloader_cfg)
+            infer_datalayer = self._setup_dataloader_from_config(cfg=dataloader_cfg, mode=INFERENCE_MODE)
+
+            all_logits = []
+            all_unique_ids = []
+            for i, batch in enumerate(infer_datalayer):
+                input_ids, token_type_ids, attention_mask, unique_ids = batch
+                logits = self.forward(
+                    input_ids=input_ids.to(device),
+                    token_type_ids=token_type_ids.to(device),
+                    attention_mask=attention_mask.to(device),
+                )
+                all_logits.append(logits)
+                all_unique_ids.append(unique_ids)
+            logits = torch.cat(all_logits)
+            unique_ids = tensor2list(torch.cat(all_unique_ids))
+            s, e = logits.split(dim=-1, split_size=1)
+            start_logits = tensor2list(s.squeeze())
+            end_logits = tensor2list(e.squeeze())
+            (all_predictions, all_nbest, scores_diff) = infer_datalayer.dataset.get_predictions(
+                unique_ids=unique_ids,
+                start_logits=start_logits,
+                end_logits=end_logits,
+                n_best_size=self._cfg.dataset.n_best_size,
+                max_answer_length=self._cfg.dataset.max_answer_length,
+                version_2_with_negative=self._cfg.dataset.version_2_with_negative,
+                null_score_diff_threshold=self._cfg.dataset.null_score_diff_threshold,
+                do_lower_case=self._cfg.dataset.do_lower_case,
+            )
+
+            with open(file, 'r') as test_file_fp:
+                test_data = json.load(test_file_fp)["data"]
+                id_to_question_mapping = {}
+                for title in test_data:
+                    for par in title["paragraphs"]:
+                        for question in par["qas"]:
+                            id_to_question_mapping[question["id"]] = question["question"]
+
+            for question_id in all_predictions:
+                all_predictions[question_id] = (id_to_question_mapping[question_id], all_predictions[question_id])
+
+            if output_nbest_file is not None:
+                with open(output_nbest_file, "w") as writer:
+                    writer.write(json.dumps(all_nbest, indent=4) + "\n")
+            if output_prediction_file is not None:
+                with open(output_prediction_file, "w") as writer:
+                    writer.write(json.dumps(all_predictions, indent=4) + "\n")
+
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+            logging.set_verbosity(logging_level)
+
+        return all_predictions, all_nbest
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        if not train_data_config or not train_data_config.file:
+            logging.info(
+                f"Dataloader config or file_path for the train is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config, mode=TRAINING_MODE)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        if not val_data_config or not val_data_config.file:
+            logging.info(
+                f"Dataloader config or file_path for the validation is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config, mode=EVALUATION_MODE)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        if not test_data_config or test_data_config.file is None:
+            logging.info(
+                f"Dataloader config or file_path for the test is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config, mode=EVALUATION_MODE)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig, mode: str):
+        dataset = SquadDataset(
+            tokenizer=self.tokenizer,
+            data_file=cfg.file,
+            doc_stride=self._cfg.dataset.doc_stride,
+            max_query_length=self._cfg.dataset.max_query_length,
+            max_seq_length=self._cfg.dataset.max_seq_length,
+            version_2_with_negative=self._cfg.dataset.version_2_with_negative,
+            num_samples=cfg.num_samples,
+            mode=mode,
+            use_cache=self._cfg.dataset.use_cache,
+        )
+
+        dl = torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            collate_fn=dataset.collate_fn,
+            drop_last=cfg.drop_last,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.num_workers,
+            pin_memory=cfg.pin_memory,
+        )
+        return dl
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv1.1_bertbase",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv1_1_bertbase/versions/1.0.0rc1/files/qa_squadv1.1_bertbase.nemo",
+                description="Question answering model finetuned from NeMo BERT Base Uncased on SQuAD v1.1 dataset which obtains an exact match (EM) score of 82.78% and an F1 score of 82.78%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv2.0_bertbase",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv2_0_bertbase/versions/1.0.0rc1/files/qa_squadv2.0_bertbase.nemo",
+                description="Question answering model finetuned from NeMo BERT Base Uncased on SQuAD v2.0 dataset which obtains an exact match (EM) score of 75.04% and an F1 score of 78.08%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv1_1_bertlarge",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv1_1_bertlarge/versions/1.0.0rc1/files/qa_squadv1.1_bertlarge.nemo",
+                description="Question answering model finetuned from NeMo BERT Large Uncased on SQuAD v1.1 dataset which obtains an exact match (EM) score of 85.44% and an F1 score of 92.06%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv2.0_bertlarge",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv2_0_bertlarge/versions/1.0.0rc1/files/qa_squadv2.0_bertlarge.nemo",
+                description="Question answering model finetuned from NeMo BERT Large Uncased on SQuAD v2.0 dataset which obtains an exact match (EM) score of 80.22% and an F1 score of 83.05%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv1_1_megatron_cased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv1_1_megatron_cased/versions/1.0.0rc1/files/qa_squadv1.1_megatron_cased.nemo",
+                description="Question answering model finetuned from Megatron Cased on SQuAD v1.1 dataset which obtains an exact match (EM) score of 88.18% and an F1 score of 94.07%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv2.0_megatron_cased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv2_0_megatron_cased/versions/1.0.0rc1/files/qa_squadv2.0_megatron_cased.nemo",
+                description="Question answering model finetuned from Megatron Cased on SQuAD v2.0 dataset which obtains an exact match (EM) score of 84.73% and an F1 score of 87.89%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv1.1_megatron_uncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv1_1_megatron_uncased/versions/1.0.0rc1/files/qa_squadv1.1_megatron_uncased.nemo",
+                description="Question answering model finetuned from Megatron Unased on SQuAD v1.1 dataset which obtains an exact match (EM) score of 87.61% and an F1 score of 94.00%.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="qa_squadv2.0_megatron_uncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/qa_squadv2_0_megatron_uncased/versions/1.0.0rc1/files/qa_squadv2.0_megatron_uncased.nemo",
+                description="Question answering model finetuned from Megatron Uncased on SQuAD v2.0 dataset which obtains an exact match (EM) score of 84.48% and an F1 score of 87.65%.",
+            )
+        )
+        return result

tts/nemo/collections/nlp/models/text_classification/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.collections.nlp.models.text_classification.text_classification_model import TextClassificationModel

tts/nemo/collections/nlp/models/text_classification/text_classification_model.py

@@ -0,0 +1,324 @@
+# 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 typing import Dict, List, Optional
+
+import torch
+from omegaconf import DictConfig
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import CrossEntropyLoss
+from nemo.collections.nlp.data.text_classification import TextClassificationDataset, calc_class_weights
+from nemo.collections.nlp.metrics.classification_report import ClassificationReport
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import SequenceClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import tensor2list
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.exportable import Exportable
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ['TextClassificationModel']
+
+
+class TextClassificationModel(NLPModel, Exportable):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.classifier.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """Initializes the BERTTextClassifier model."""
+
+        # shared params for dataset and data loaders
+        self.dataset_cfg = cfg.dataset
+        # tokenizer needs to get initialized before the super.__init__()
+        # as dataloaders and datasets need it to process the data
+        self.setup_tokenizer(cfg.tokenizer)
+
+        self.class_weights = None
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=cfg.language_model.config,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        self.classifier = SequenceClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_classes=cfg.dataset.num_classes,
+            num_layers=cfg.classifier_head.num_output_layers,
+            activation='relu',
+            log_softmax=False,
+            dropout=cfg.classifier_head.fc_dropout,
+            use_transformer_init=True,
+            idx_conditioned_on=0,
+        )
+
+        self.create_loss_module()
+
+        # setup to track metrics
+        self.classification_report = ClassificationReport(
+            num_classes=cfg.dataset.num_classes, mode='micro', dist_sync_on_step=True
+        )
+
+        # register the file containing the labels into the artifacts to get stored in the '.nemo' file later
+        if 'class_labels' in cfg and 'class_labels_file' in cfg.class_labels and cfg.class_labels.class_labels_file:
+            self.register_artifact('class_labels', cfg.class_labels.class_labels_file)
+
+    def create_loss_module(self):
+        # create the loss module if it is not yet created by the training data loader
+        if not hasattr(self, 'loss'):
+            if hasattr(self, 'class_weights') and self.class_weights:
+                # You may need to increase the number of epochs for convergence when using weighted_loss
+                self.loss = CrossEntropyLoss(weight=self.class_weights)
+            else:
+                self.loss = CrossEntropyLoss()
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        logits = self.classifier(hidden_states=hidden_states)
+        return logits
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        # forward pass
+        input_ids, input_type_ids, input_mask, labels = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        train_loss = self.loss(logits=logits, labels=labels)
+
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', train_loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': train_loss,
+            'lr': lr,
+        }
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, labels = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        val_loss = self.loss(logits=logits, labels=labels)
+
+        preds = torch.argmax(logits, axis=-1)
+
+        tp, fn, fp, _ = self.classification_report(preds, labels)
+
+        return {'val_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp}
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        :param outputs: list of individual outputs of each validation step.
+        """
+        if not outputs:
+            return {}
+        if self.trainer.testing:
+            prefix = 'test'
+        else:
+            prefix = 'val'
+
+        avg_loss = torch.stack([x[f'val_loss'] for x in outputs]).mean()
+
+        # calculate metrics and classification report
+        precision, recall, f1, report = self.classification_report.compute()
+
+        logging.info(f'{prefix}_report: {report}')
+
+        self.log(f'{prefix}_loss', avg_loss, prog_bar=True)
+        self.log(f'{prefix}_precision', precision)
+        self.log(f'{prefix}_f1', f1)
+        self.log(f'{prefix}_recall', recall)
+
+        self.classification_report.reset()
+
+    def test_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the test loop with the data from the test dataloader
+        passed in as `batch`.
+        """
+        return self.validation_step(batch, batch_idx)
+
+    def test_epoch_end(self, outputs):
+        """
+        Called at the end of test to aggregate outputs.
+        :param outputs: list of individual outputs of each test step.
+        """
+        return self.validation_epoch_end(outputs)
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        if not train_data_config or not train_data_config.file_path:
+            logging.info(
+                f"Dataloader config or file_path for the train is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+        # calculate the class weights to be used in the loss function
+        if self.cfg.dataset.class_balancing == 'weighted_loss':
+            self.class_weights = calc_class_weights(train_data_config.file_path, self.cfg.dataset.num_classes)
+        else:
+            self.class_weights = None
+        # we need to create/update the loss module by using the weights calculated from the training data
+        self.create_loss_module()
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        if not val_data_config or not val_data_config.file_path:
+            logging.info(
+                f"Dataloader config or file_path for the validation is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        if not test_data_config or not test_data_config.file_path:
+            logging.info(
+                f"Dataloader config or file_path for the test is missing, so no data loader for test is created!"
+            )
+            self._test_dl = None
+            return
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: Dict) -> 'torch.utils.data.DataLoader':
+        input_file = cfg.file_path
+        if not os.path.exists(input_file):
+            raise FileNotFoundError(
+                f'{input_file} not found! The data should be be stored in TAB-separated files \n\
+                "validation_ds.file_path" and "train_ds.file_path" for train and evaluation respectively. \n\
+                Each line of the files contains text sequences, where words are separated with spaces. \n\
+                The label of the example is separated with TAB at the end of each line. \n\
+                Each line of the files should follow the format: \n\
+                [WORD][SPACE][WORD][SPACE][WORD][...][TAB][LABEL]'
+            )
+
+        dataset = TextClassificationDataset(
+            tokenizer=self.tokenizer,
+            input_file=input_file,
+            max_seq_length=self.dataset_cfg.max_seq_length,
+            num_samples=cfg.get("num_samples", -1),
+            shuffle=cfg.shuffle,
+            use_cache=self.dataset_cfg.use_cache,
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 0),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+            collate_fn=dataset.collate_fn,
+        )
+
+    @torch.no_grad()
+    def classifytext(self, queries: List[str], batch_size: int = 1, max_seq_length: int = -1) -> List[int]:
+        """
+        Get prediction for the queries
+        Args:
+            queries: text sequences
+            batch_size: batch size to use during inference
+            max_seq_length: sequences longer than max_seq_length will get truncated. default -1 disables truncation.
+        Returns:
+            all_preds: model predictions
+        """
+        # store predictions for all queries in a single list
+        all_preds = []
+        mode = self.training
+        device = next(self.parameters()).device
+        try:
+            # Switch model to evaluation mode
+            self.eval()
+            logging_level = logging.get_verbosity()
+            logging.set_verbosity(logging.WARNING)
+            dataloader_cfg = {"batch_size": batch_size, "num_workers": 3, "pin_memory": False}
+            infer_datalayer = self._setup_infer_dataloader(dataloader_cfg, queries, max_seq_length)
+
+            for i, batch in enumerate(infer_datalayer):
+                input_ids, input_type_ids, input_mask, subtokens_mask = batch
+
+                logits = self.forward(
+                    input_ids=input_ids.to(device),
+                    token_type_ids=input_type_ids.to(device),
+                    attention_mask=input_mask.to(device),
+                )
+
+                preds = tensor2list(torch.argmax(logits, axis=-1))
+                all_preds.extend(preds)
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+            logging.set_verbosity(logging_level)
+        return all_preds
+
+    def _setup_infer_dataloader(
+        self, cfg: Dict, queries: List[str], max_seq_length: int = -1
+    ) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a infer data loader.
+
+        Args:
+            cfg: config dictionary containing data loader params like batch_size, num_workers and pin_memory
+            queries: text
+            max_seq_length: maximum length of queries, default is -1 for no limit
+        Returns:
+            A pytorch DataLoader.
+        """
+        dataset = TextClassificationDataset(tokenizer=self.tokenizer, queries=queries, max_seq_length=max_seq_length)
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg["batch_size"],
+            shuffle=False,
+            num_workers=cfg.get("num_workers", 0),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=False,
+            collate_fn=dataset.collate_fn,
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass
+
+    @classmethod
+    def from_pretrained(cls, name: str):
+        pass

tts/nemo/collections/nlp/models/token_classification/__init__.py

@@ -0,0 +1,19 @@
+# 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 nemo.collections.nlp.models.token_classification.punctuation_capitalization_model import (
+    PunctuationCapitalizationModel,
+)
+from nemo.collections.nlp.models.token_classification.token_classification_model import TokenClassificationModel

tts/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py

@@ -0,0 +1,490 @@
+# 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 typing import Dict, List, Optional
+
+import onnx
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import AggregatorLoss, CrossEntropyLoss
+from nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset import (
+    BertPunctuationCapitalizationDataset,
+    BertPunctuationCapitalizationInferDataset,
+)
+from nemo.collections.nlp.metrics.classification_report import ClassificationReport
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import tensor2list
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.classes.exportable import Exportable, ExportFormat
+from nemo.core.neural_types import LogitsType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['PunctuationCapitalizationModel']
+
+
+class PunctuationCapitalizationModel(NLPModel, Exportable):
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "punct_logits": NeuralType(('B', 'T', 'C'), LogitsType()),
+            "capit_logits": NeuralType(('B', 'T', 'C'), LogitsType()),
+        }
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """
+        Initializes BERT Punctuation and Capitalization model.
+        """
+        self.setup_tokenizer(cfg.tokenizer)
+
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        self.punct_classifier = TokenClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_classes=len(self._cfg.punct_label_ids),
+            activation=cfg.punct_head.activation,
+            log_softmax=False,
+            dropout=cfg.punct_head.fc_dropout,
+            num_layers=cfg.punct_head.punct_num_fc_layers,
+            use_transformer_init=cfg.punct_head.use_transformer_init,
+        )
+
+        self.capit_classifier = TokenClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_classes=len(self._cfg.capit_label_ids),
+            activation=cfg.capit_head.activation,
+            log_softmax=False,
+            dropout=cfg.capit_head.fc_dropout,
+            num_layers=cfg.capit_head.capit_num_fc_layers,
+            use_transformer_init=cfg.capit_head.use_transformer_init,
+        )
+
+        self.loss = CrossEntropyLoss(logits_ndim=3)
+        self.agg_loss = AggregatorLoss(num_inputs=2)
+
+        # setup to track metrics
+        self.punct_class_report = ClassificationReport(
+            num_classes=len(self._cfg.punct_label_ids),
+            label_ids=self._cfg.punct_label_ids,
+            mode='macro',
+            dist_sync_on_step=True,
+        )
+        self.capit_class_report = ClassificationReport(
+            num_classes=len(self._cfg.capit_label_ids),
+            label_ids=self._cfg.capit_label_ids,
+            mode='macro',
+            dist_sync_on_step=True,
+        )
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        punct_logits = self.punct_classifier(hidden_states=hidden_states)
+        capit_logits = self.capit_classifier(hidden_states=hidden_states)
+        return punct_logits, capit_logits
+
+    def _make_step(self, batch):
+        input_ids, input_type_ids, input_mask, subtokens_mask, loss_mask, punct_labels, capit_labels = batch
+        punct_logits, capit_logits = self(
+            input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask
+        )
+
+        punct_loss = self.loss(logits=punct_logits, labels=punct_labels, loss_mask=loss_mask)
+        capit_loss = self.loss(logits=capit_logits, labels=capit_labels, loss_mask=loss_mask)
+        loss = self.agg_loss(loss_1=punct_loss, loss_2=capit_loss)
+        return loss, punct_logits, capit_logits
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        loss, _, _ = self._make_step(batch)
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('lr', lr, prog_bar=True)
+        self.log('train_loss', loss)
+
+        return {'loss': loss, 'lr': lr}
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        _, _, _, subtokens_mask, _, punct_labels, capit_labels = batch
+        val_loss, punct_logits, capit_logits = self._make_step(batch)
+
+        subtokens_mask = subtokens_mask > 0.5
+        punct_preds = torch.argmax(punct_logits, axis=-1)[subtokens_mask]
+        punct_labels = punct_labels[subtokens_mask]
+        self.punct_class_report.update(punct_preds, punct_labels)
+
+        capit_preds = torch.argmax(capit_logits, axis=-1)[subtokens_mask]
+        capit_labels = capit_labels[subtokens_mask]
+        self.capit_class_report.update(capit_preds, capit_labels)
+
+        return {
+            'val_loss': val_loss,
+            'punct_tp': self.punct_class_report.tp,
+            'punct_fn': self.punct_class_report.fn,
+            'punct_fp': self.punct_class_report.fp,
+            'capit_tp': self.capit_class_report.tp,
+            'capit_fn': self.capit_class_report.fn,
+            'capit_fp': self.capit_class_report.fp,
+        }
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        _, _, _, subtokens_mask, _, punct_labels, capit_labels = batch
+        test_loss, punct_logits, capit_logits = self._make_step(batch)
+
+        subtokens_mask = subtokens_mask > 0.5
+        punct_preds = torch.argmax(punct_logits, axis=-1)[subtokens_mask]
+        punct_labels = punct_labels[subtokens_mask]
+        self.punct_class_report.update(punct_preds, punct_labels)
+
+        capit_preds = torch.argmax(capit_logits, axis=-1)[subtokens_mask]
+        capit_labels = capit_labels[subtokens_mask]
+        self.capit_class_report.update(capit_preds, capit_labels)
+
+        return {
+            'test_loss': test_loss,
+            'punct_tp': self.punct_class_report.tp,
+            'punct_fn': self.punct_class_report.fn,
+            'punct_fp': self.punct_class_report.fp,
+            'capit_tp': self.capit_class_report.tp,
+            'capit_fn': self.capit_class_report.fn,
+            'capit_fp': self.capit_class_report.fp,
+        }
+
+    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
+        """
+        Called at the end of validation to aggregate outputs.
+        outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+
+        # calculate metrics and log classification report for Punctuation task
+        punct_precision, punct_recall, punct_f1, punct_report = self.punct_class_report.compute()
+        logging.info(f'Punctuation report: {punct_report}')
+
+        # calculate metrics and log classification report for Capitalization task
+        capit_precision, capit_recall, capit_f1, capit_report = self.capit_class_report.compute()
+        logging.info(f'Capitalization report: {capit_report}')
+
+        self.log('val_loss', avg_loss, prog_bar=True)
+        self.log('punct_precision', punct_precision)
+        self.log('punct_f1', punct_f1)
+        self.log('punct_recall', punct_recall)
+        self.log('capit_precision', capit_precision)
+        self.log('capit_f1', capit_f1)
+        self.log('capit_recall', capit_recall)
+
+        self.punct_class_report.reset()
+        self.capit_class_report.reset()
+
+    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
+        """
+            Called at the end of test to aggregate outputs.
+            outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
+
+        # calculate metrics and log classification report for Punctuation task
+        punct_precision, punct_recall, punct_f1, punct_report = self.punct_class_report.compute()
+        logging.info(f'Punctuation report: {punct_report}')
+
+        # calculate metrics and log classification report for Capitalization task
+        capit_precision, capit_recall, capit_f1, capit_report = self.capit_class_report.compute()
+        logging.info(f'Capitalization report: {capit_report}')
+
+        self.log('test_loss', avg_loss, prog_bar=True)
+        self.log('punct_precision', punct_precision)
+        self.log('punct_f1', punct_f1)
+        self.log('punct_recall', punct_recall)
+        self.log('capit_precision', capit_precision)
+        self.log('capit_f1', capit_f1)
+        self.log('capit_recall', capit_recall)
+
+    def update_data_dir(self, data_dir: str) -> None:
+        """
+        Update data directory
+
+        Args:
+            data_dir: path to data directory
+        """
+        if os.path.exists(data_dir):
+            logging.info(f'Setting model.dataset.data_dir to {data_dir}.')
+            self._cfg.dataset.data_dir = data_dir
+        else:
+            raise ValueError(f'{data_dir} not found')
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig] = None):
+        """Setup training data"""
+        if train_data_config is None:
+            train_data_config = self._cfg.train_ds
+
+        # for older(pre - 1.0.0.b3) configs compatibility
+        if not hasattr(self._cfg, "class_labels") or self._cfg.class_labels is None:
+            OmegaConf.set_struct(self._cfg, False)
+            self._cfg.class_labels = {}
+            self._cfg.class_labels = OmegaConf.create(
+                {'punct_labels_file': 'punct_label_ids.csv', 'capit_labels_file': 'capit_label_ids.csv'}
+            )
+
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            self.register_artifact('class_labels.punct_labels_file', self._train_dl.dataset.punct_label_ids_file)
+            self.register_artifact('class_labels.capit_labels_file', self._train_dl.dataset.capit_label_ids_file)
+
+            # save label maps to the config
+            self._cfg.punct_label_ids = OmegaConf.create(self._train_dl.dataset.punct_label_ids)
+            self._cfg.capit_label_ids = OmegaConf.create(self._train_dl.dataset.capit_label_ids)
+
+    def setup_validation_data(self, val_data_config: Optional[Dict] = None):
+        """
+        Setup validaton data
+
+        val_data_config: validation data config
+        """
+        if val_data_config is None:
+            val_data_config = self._cfg.validation_ds
+
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[Dict] = None):
+        if test_data_config is None:
+            test_data_config = self._cfg.test_ds
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig):
+        # use data_dir specified in the ds_item to run evaluation on multiple datasets
+        if 'ds_item' in cfg and cfg.ds_item is not None:
+            data_dir = cfg.ds_item
+        else:
+            data_dir = self._cfg.dataset.data_dir
+
+        text_file = os.path.join(data_dir, cfg.text_file)
+        label_file = os.path.join(data_dir, cfg.labels_file)
+
+        dataset = BertPunctuationCapitalizationDataset(
+            tokenizer=self.tokenizer,
+            text_file=text_file,
+            label_file=label_file,
+            pad_label=self._cfg.dataset.pad_label,
+            punct_label_ids=self._cfg.punct_label_ids,
+            capit_label_ids=self._cfg.capit_label_ids,
+            max_seq_length=self._cfg.dataset.max_seq_length,
+            ignore_extra_tokens=self._cfg.dataset.ignore_extra_tokens,
+            ignore_start_end=self._cfg.dataset.ignore_start_end,
+            use_cache=self._cfg.dataset.use_cache,
+            num_samples=cfg.num_samples,
+            punct_label_ids_file=self._cfg.class_labels.punct_labels_file
+            if 'class_labels' in self._cfg
+            else 'punct_label_ids.csv',
+            capit_label_ids_file=self._cfg.class_labels.capit_labels_file
+            if 'class_labels' in self._cfg
+            else 'capit_label_ids.csv',
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=self._cfg.dataset.num_workers,
+            pin_memory=self._cfg.dataset.pin_memory,
+            drop_last=self._cfg.dataset.drop_last,
+        )
+
+    def _setup_infer_dataloader(
+        self, queries: List[str], batch_size: int, max_seq_length: int = None
+    ) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a infer data loader.
+
+        Args:
+            queries: lower cased text without punctuation
+            batch_size: batch size to use during inference
+            max_seq_length: maximum sequence length after tokenization
+        Returns:
+            A pytorch DataLoader.
+        """
+        if max_seq_length is None:
+            max_seq_length = self._cfg.dataset.max_seq_length
+
+        dataset = BertPunctuationCapitalizationInferDataset(
+            tokenizer=self.tokenizer, queries=queries, max_seq_length=max_seq_length
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=self._cfg.dataset.num_workers,
+            pin_memory=self._cfg.dataset.pin_memory,
+            drop_last=False,
+        )
+
+    def add_punctuation_capitalization(
+        self, queries: List[str], batch_size: int = None, max_seq_length: int = 512
+    ) -> List[str]:
+        """
+        Adds punctuation and capitalization to the queries. Use this method for debugging and prototyping.
+        Args:
+            queries: lower cased text without punctuation
+            batch_size: batch size to use during inference
+            max_seq_length: maximum sequence length after tokenization
+        Returns:
+            result: text with added capitalization and punctuation
+        """
+        if queries is None or len(queries) == 0:
+            return []
+        if batch_size is None:
+            batch_size = len(queries)
+            logging.info(f'Using batch size {batch_size} for inference')
+
+        # We will store the output here
+        result = []
+
+        # Model's mode and device
+        mode = self.training
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        try:
+            # Switch model to evaluation mode
+            self.eval()
+            self = self.to(device)
+
+            infer_datalayer = self._setup_infer_dataloader(queries, batch_size, max_seq_length)
+
+            # store predictions for all queries in a single list
+            all_punct_preds = []
+            all_capit_preds = []
+
+            for batch in infer_datalayer:
+                input_ids, input_type_ids, input_mask, subtokens_mask = batch
+
+                punct_logits, capit_logits = self.forward(
+                    input_ids=input_ids.to(device),
+                    token_type_ids=input_type_ids.to(device),
+                    attention_mask=input_mask.to(device),
+                )
+
+                subtokens_mask = subtokens_mask > 0.5
+
+                punct_preds = [
+                    tensor2list(p_l[subtokens_mask[i]]) for i, p_l in enumerate(torch.argmax(punct_logits, axis=-1))
+                ]
+                capit_preds = [
+                    tensor2list(c_l[subtokens_mask[i]]) for i, c_l in enumerate(torch.argmax(capit_logits, axis=-1))
+                ]
+
+                all_punct_preds.extend(punct_preds)
+                all_capit_preds.extend(capit_preds)
+
+            punct_ids_to_labels = {v: k for k, v in self._cfg.punct_label_ids.items()}
+            capit_ids_to_labels = {v: k for k, v in self._cfg.capit_label_ids.items()}
+
+            queries = [q.strip().split() for q in queries]
+            for i, query in enumerate(queries):
+                punct_preds = all_punct_preds[i]
+                capit_preds = all_capit_preds[i]
+                if len(query) != len(punct_preds):
+                    logging.warning(
+                        f'Max sequence length of query {query} is set to {max_seq_length}. Truncating the input.'
+                    )
+
+                    # removing the end of phrase punctuation of the truncated segment
+                    punct_preds[-1] = 0
+                    max_len = len(punct_preds)
+                    query = query[:max_len]
+
+                query_with_punct_and_capit = ''
+                for j, word in enumerate(query):
+                    punct_label = punct_ids_to_labels[punct_preds[j]]
+                    capit_label = capit_ids_to_labels[capit_preds[j]]
+
+                    if capit_label != self._cfg.dataset.pad_label:
+                        word = word.capitalize()
+                    query_with_punct_and_capit += word
+                    if punct_label != self._cfg.dataset.pad_label:
+                        query_with_punct_and_capit += punct_label
+                    query_with_punct_and_capit += ' '
+
+                result.append(query_with_punct_and_capit.strip())
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+        return result
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="punctuation_en_bert",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/punctuation_en_bert/versions/1.0.0rc1/files/punctuation_en_bert.nemo",
+                description="The model was trained with NeMo BERT base uncased checkpoint on a subset of data from the following sources: Tatoeba sentences, books from Project Gutenberg, Fisher transcripts.",
+            )
+        )
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="punctuation_en_distilbert",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/punctuation_en_distilbert/versions/1.0.0rc1/files/punctuation_en_distilbert.nemo",
+                description="The model was trained with DiltilBERT base uncased checkpoint from HuggingFace on a subset of data from the following sources: Tatoeba sentences, books from Project Gutenberg, Fisher transcripts.",
+            )
+        )
+        return result
+
+    @property
+    def input_module(self):
+        return self.bert_model
+
+    @property
+    def output_module(self):
+        return self

tts/nemo/collections/nlp/models/token_classification/token_classification_model.py

@@ -0,0 +1,522 @@
+# 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 typing import Dict, List, Optional, Union
+
+import onnx
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.common.losses import CrossEntropyLoss
+from nemo.collections.nlp.data.data_utils.data_preprocessing import get_labels_to_labels_id_mapping
+from nemo.collections.nlp.data.token_classification.token_classification_dataset import (
+    BertTokenClassificationDataset,
+    BertTokenClassificationInferDataset,
+)
+from nemo.collections.nlp.data.token_classification.token_classification_utils import get_label_ids
+from nemo.collections.nlp.metrics.classification_report import ClassificationReport
+from nemo.collections.nlp.models.nlp_model import NLPModel
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.parts.utils_funcs import get_classification_report, plot_confusion_matrix, tensor2list
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ['TokenClassificationModel']
+
+
+class TokenClassificationModel(NLPModel):
+    """Token Classification Model with BERT, applicable for tasks such as Named Entity Recognition"""
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.classifier.output_types
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+        """Initializes Token Classification Model."""
+        # extract str to int labels mapping if a mapping file provided
+        if isinstance(cfg.label_ids, str):
+            if os.path.exists(cfg.label_ids):
+                logging.info(f'Reusing label_ids file found at {cfg.label_ids}.')
+                label_ids = get_labels_to_labels_id_mapping(cfg.label_ids)
+                # update the config to store name to id mapping
+                cfg.label_ids = OmegaConf.create(label_ids)
+            else:
+                raise ValueError(f'{cfg.label_ids} not found.')
+
+        self.setup_tokenizer(cfg.tokenizer)
+        self.class_weights = None
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        self.bert_model = get_lm_model(
+            pretrained_model_name=cfg.language_model.pretrained_model_name,
+            config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file),
+            config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None,
+            checkpoint_file=cfg.language_model.lm_checkpoint,
+            vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file),
+        )
+
+        self.classifier = TokenClassifier(
+            hidden_size=self.bert_model.config.hidden_size,
+            num_classes=len(self._cfg.label_ids),
+            num_layers=self._cfg.head.num_fc_layers,
+            activation=self._cfg.head.activation,
+            log_softmax=False,
+            dropout=self._cfg.head.fc_dropout,
+            use_transformer_init=self._cfg.head.use_transformer_init,
+        )
+
+        self.loss = self.setup_loss(class_balancing=self._cfg.dataset.class_balancing)
+
+        # setup to track metrics
+        self.classification_report = ClassificationReport(
+            len(self._cfg.label_ids), label_ids=self._cfg.label_ids, dist_sync_on_step=True
+        )
+
+    def update_data_dir(self, data_dir: str) -> None:
+        """
+        Update data directory and get data stats with Data Descriptor
+        Weights are later used to setup loss
+
+        Args:
+            data_dir: path to data directory
+        """
+        self._cfg.dataset.data_dir = data_dir
+        logging.info(f'Setting model.dataset.data_dir to {data_dir}.')
+
+    def setup_loss(self, class_balancing: str = None):
+        """Setup loss
+           Setup or update loss.
+
+        Args:
+            class_balancing: whether to use class weights during training
+        """
+        if class_balancing not in ['weighted_loss', None]:
+            raise ValueError(f'Class balancing {class_balancing} is not supported. Choose from: [null, weighted_loss]')
+        if class_balancing == 'weighted_loss' and self.class_weights:
+            # you may need to increase the number of epochs for convergence when using weighted_loss
+            loss = CrossEntropyLoss(logits_ndim=3, weight=self.class_weights)
+            logging.debug(f'Using {class_balancing} class balancing.')
+        else:
+            loss = CrossEntropyLoss(logits_ndim=3)
+            logging.debug(f'Using CrossEntropyLoss class balancing.')
+        return loss
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        hidden_states = self.bert_model(
+            input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+        )
+        logits = self.classifier(hidden_states=hidden_states)
+        return logits
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, subtokens_mask, loss_mask, labels = batch
+        logits = self(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        loss = self.loss(logits=logits, labels=labels, loss_mask=loss_mask)
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': loss,
+            'lr': lr,
+        }
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        input_ids, input_type_ids, input_mask, subtokens_mask, loss_mask, labels = batch
+        logits = self(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        val_loss = self.loss(logits=logits, labels=labels, loss_mask=loss_mask)
+
+        subtokens_mask = subtokens_mask > 0.5
+
+        preds = torch.argmax(logits, axis=-1)[subtokens_mask]
+        labels = labels[subtokens_mask]
+        tp, fn, fp, _ = self.classification_report(preds, labels)
+
+        return {'val_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp}
+
+    def validation_epoch_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs.
+        outputs: list of individual outputs of each validation step.
+        """
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+
+        # calculate metrics and classification report
+        precision, recall, f1, report = self.classification_report.compute()
+
+        logging.info(report)
+
+        self.log('val_loss', avg_loss, prog_bar=True)
+        self.log('precision', precision)
+        self.log('f1', f1)
+        self.log('recall', recall)
+
+        self.classification_report.reset()
+
+    def test_step(self, batch, batch_idx):
+        input_ids, input_type_ids, input_mask, subtokens_mask, loss_mask, labels = batch
+        logits = self(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        val_loss = self.loss(logits=logits, labels=labels, loss_mask=loss_mask)
+
+        subtokens_mask = subtokens_mask > 0.5
+
+        preds = torch.argmax(logits, axis=-1)[subtokens_mask]
+        labels = labels[subtokens_mask]
+        tp, fn, fp, _ = self.classification_report(preds, labels)
+
+        return {'test_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp}
+
+    def test_epoch_end(self, outputs):
+        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
+        # calculate metrics and classification report
+        precision, recall, f1, report = self.classification_report.compute()
+        logging.info(report)
+
+        self.log('test_loss', avg_loss, prog_bar=True)
+        self.log('precision', precision)
+        self.log('f1', f1)
+        self.log('recall', recall)
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig] = None):
+        if train_data_config is None:
+            train_data_config = self._cfg.train_ds
+
+        labels_file = os.path.join(self._cfg.dataset.data_dir, train_data_config.labels_file)
+
+        # for older(pre - 1.0.0.b3) configs compatibility
+        if not hasattr(self._cfg, "class_labels") or self._cfg.class_labels is None:
+            OmegaConf.set_struct(self._cfg, False)
+            self._cfg.class_labels = {}
+            self._cfg.class_labels = OmegaConf.create({'class_labels_file': 'label_ids.csv'})
+            OmegaConf.set_struct(self._cfg, True)
+
+        label_ids, label_ids_filename, self.class_weights = get_label_ids(
+            label_file=labels_file,
+            is_training=True,
+            pad_label=self._cfg.dataset.pad_label,
+            label_ids_dict=self._cfg.label_ids,
+            get_weights=True,
+            class_labels_file_artifact=self._cfg.class_labels.class_labels_file,
+        )
+        # save label maps to the config
+        self._cfg.label_ids = OmegaConf.create(label_ids)
+
+        self.register_artifact('class_labels.class_labels_file', label_ids_filename)
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig] = None):
+        if val_data_config is None:
+            val_data_config = self._cfg.validation_ds
+
+        labels_file = os.path.join(self._cfg.dataset.data_dir, val_data_config.labels_file)
+        get_label_ids(
+            label_file=labels_file,
+            is_training=False,
+            pad_label=self._cfg.dataset.pad_label,
+            label_ids_dict=self._cfg.label_ids,
+            get_weights=False,
+        )
+
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig] = None):
+        if test_data_config is None:
+            test_data_config = self._cfg.test_ds
+
+        labels_file = os.path.join(self._cfg.dataset.data_dir, test_data_config.labels_file)
+        get_label_ids(
+            label_file=labels_file,
+            is_training=False,
+            pad_label=self._cfg.dataset.pad_label,
+            label_ids_dict=self._cfg.label_ids,
+            get_weights=False,
+        )
+
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig) -> DataLoader:
+        """
+        Setup dataloader from config
+        Args:
+            cfg: config for the dataloader
+        Return:
+            Pytorch Dataloader
+        """
+        dataset_cfg = self._cfg.dataset
+        data_dir = dataset_cfg.data_dir
+
+        if not os.path.exists(data_dir):
+            raise FileNotFoundError(f"Data directory is not found at: {data_dir}.")
+
+        text_file = os.path.join(data_dir, cfg.text_file)
+        labels_file = os.path.join(data_dir, cfg.labels_file)
+
+        if not (os.path.exists(text_file) and os.path.exists(labels_file)):
+            raise FileNotFoundError(
+                f'{text_file} or {labels_file} not found. The data should 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).'
+            )
+        dataset = BertTokenClassificationDataset(
+            text_file=text_file,
+            label_file=labels_file,
+            max_seq_length=dataset_cfg.max_seq_length,
+            tokenizer=self.tokenizer,
+            num_samples=cfg.num_samples,
+            pad_label=dataset_cfg.pad_label,
+            label_ids=self._cfg.label_ids,
+            ignore_extra_tokens=dataset_cfg.ignore_extra_tokens,
+            ignore_start_end=dataset_cfg.ignore_start_end,
+            use_cache=dataset_cfg.use_cache,
+        )
+        return DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=cfg.batch_size,
+            shuffle=cfg.shuffle,
+            num_workers=dataset_cfg.num_workers,
+            pin_memory=dataset_cfg.pin_memory,
+            drop_last=dataset_cfg.drop_last,
+        )
+
+    def _setup_infer_dataloader(self, queries: List[str], batch_size: int) -> 'torch.utils.data.DataLoader':
+        """
+        Setup function for a infer data loader.
+
+        Args:
+            queries: text
+            batch_size: batch size to use during inference
+
+        Returns:
+            A pytorch DataLoader.
+        """
+
+        dataset = BertTokenClassificationInferDataset(tokenizer=self.tokenizer, queries=queries, max_seq_length=-1)
+
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            collate_fn=dataset.collate_fn,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=self._cfg.dataset.num_workers,
+            pin_memory=self._cfg.dataset.pin_memory,
+            drop_last=False,
+        )
+
+    @torch.no_grad()
+    def _infer(self, queries: List[str], batch_size: int = None) -> List[int]:
+        """
+        Get prediction for the queries
+        Args:
+            queries: text sequences
+            batch_size: batch size to use during inference.
+        Returns:
+            all_preds: model predictions
+        """
+        # store predictions for all queries in a single list
+        all_preds = []
+        mode = self.training
+        try:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            # Switch model to evaluation mode
+            self.eval()
+            self.to(device)
+            infer_datalayer = self._setup_infer_dataloader(queries, batch_size)
+
+            for batch in infer_datalayer:
+                input_ids, input_type_ids, input_mask, subtokens_mask = batch
+
+                logits = self.forward(
+                    input_ids=input_ids.to(device),
+                    token_type_ids=input_type_ids.to(device),
+                    attention_mask=input_mask.to(device),
+                )
+
+                subtokens_mask = subtokens_mask > 0.5
+                preds = tensor2list(torch.argmax(logits, axis=-1)[subtokens_mask])
+                all_preds.extend(preds)
+        finally:
+            # set mode back to its original value
+            self.train(mode=mode)
+        return all_preds
+
+    def add_predictions(
+        self, queries: Union[List[str], str], batch_size: int = 32, output_file: Optional[str] = None
+    ) -> List[str]:
+        """
+        Add predicted token labels to the queries. Use this method for debugging and prototyping.
+        Args:
+            queries: text
+            batch_size: batch size to use during inference.
+            output_file: file to save models predictions
+        Returns:
+            result: text with added entities
+        """
+        if queries is None or len(queries) == 0:
+            return []
+
+        if isinstance(queries, str):
+            logging.info(f'Reading from {queries} file')
+            with open(queries, 'r') as f:
+                queries = f.readlines()
+
+        result = []
+        all_preds = self._infer(queries, batch_size)
+
+        queries = [q.strip().split() for q in queries]
+        num_words = [len(q) for q in queries]
+        if sum(num_words) != len(all_preds):
+            raise ValueError('Pred and words must have the same length')
+
+        ids_to_labels = {v: k for k, v in self._cfg.label_ids.items()}
+        start_idx = 0
+        end_idx = 0
+        for query in queries:
+            end_idx += len(query)
+
+            # extract predictions for the current query from the list of all predictions
+            preds = all_preds[start_idx:end_idx]
+            start_idx = end_idx
+
+            query_with_entities = ''
+            for j, word in enumerate(query):
+                # strip out the punctuation to attach the entity tag to the word not to a punctuation mark
+                # that follows the word
+                if word[-1].isalpha():
+                    punct = ''
+                else:
+                    punct = word[-1]
+                    word = word[:-1]
+
+                query_with_entities += word
+                label = ids_to_labels[preds[j]]
+
+                if label != self._cfg.dataset.pad_label:
+                    query_with_entities += '[' + label + ']'
+                query_with_entities += punct + ' '
+            result.append(query_with_entities.strip())
+
+        if output_file is not None:
+            with open(output_file, 'w') as f:
+                for r in result:
+                    f.write(r + '\n')
+            logging.info(f'Predictions saved to {output_file}')
+        return result
+
+    def evaluate_from_file(
+        self,
+        output_dir: str,
+        text_file: str,
+        labels_file: Optional[str] = None,
+        add_confusion_matrix: Optional[bool] = False,
+        normalize_confusion_matrix: Optional[bool] = True,
+        batch_size: int = 1,
+    ) -> None:
+        """
+        Run inference on data from a file, plot confusion matrix and calculate classification report.
+        Use this method for final evaluation.
+
+        Args:
+            output_dir: path to output directory to store model predictions, confusion matrix plot (if set to True)
+            text_file: path to file with text. Each line of the text.txt file contains text sequences, where words
+                are separated with spaces: [WORD] [SPACE] [WORD] [SPACE] [WORD]
+            labels_file (Optional): path to file with labels. Each line of the labels_file should contain
+                labels corresponding to each word in the text_file, the labels are separated with spaces:
+                [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).'
+            add_confusion_matrix: whether to generate confusion matrix
+            normalize_confusion_matrix: whether to normalize confusion matrix
+            batch_size: batch size to use during inference.
+        """
+        output_dir = os.path.abspath(output_dir)
+
+        with open(text_file, 'r') as f:
+            queries = f.readlines()
+
+        all_preds = self._infer(queries, batch_size)
+        with_labels = labels_file is not None
+        if with_labels:
+            with open(labels_file, 'r') as f:
+                all_labels_str = f.readlines()
+                all_labels_str = ' '.join([labels.strip() for labels in all_labels_str])
+
+        # writing labels and predictions to a file in output_dir is specified in the config
+        os.makedirs(output_dir, exist_ok=True)
+        filename = os.path.join(output_dir, 'infer_' + os.path.basename(text_file))
+        try:
+            with open(filename, 'w') as f:
+                if with_labels:
+                    f.write('labels\t' + all_labels_str + '\n')
+                    logging.info(f'Labels save to {filename}')
+
+                # convert labels from string label to ids
+                ids_to_labels = {v: k for k, v in self._cfg.label_ids.items()}
+                all_preds_str = [ids_to_labels[pred] for pred in all_preds]
+                f.write('preds\t' + ' '.join(all_preds_str) + '\n')
+                logging.info(f'Predictions saved to {filename}')
+
+            if with_labels and add_confusion_matrix:
+                all_labels = all_labels_str.split()
+                # convert labels from string label to ids
+                label_ids = self._cfg.label_ids
+                all_labels = [label_ids[label] for label in all_labels]
+
+                plot_confusion_matrix(
+                    all_labels, all_preds, output_dir, label_ids=label_ids, normalize=normalize_confusion_matrix
+                )
+                logging.info(get_classification_report(all_labels, all_preds, label_ids))
+        except Exception:
+            logging.error(
+                f'When providing a file with labels, check that all labels in {labels_file} were'
+                f'seen during training.'
+            )
+            raise
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+        model = PretrainedModelInfo(
+            pretrained_model_name="ner_en_bert",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/ner_en_bert/versions/1.0.0rc1/files/ner_en_bert.nemo",
+            description="The model was trained on GMB (Groningen Meaning Bank) corpus for entity recognition and achieves 74.61 F1 Macro score.",
+        )
+        result.append(model)
+        return result