Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/nlp/__init__.py

@@ -0,0 +1,25 @@
+# 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 import data, losses, models, modules
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Natural Language Processing collection"

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

@@ -0,0 +1,40 @@
+# 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.data.data_utils import *
+from nemo.collections.nlp.data.entity_linking.entity_linking_dataset import EntityLinkingDataset
+from nemo.collections.nlp.data.information_retrieval.information_retrieval_dataset import (
+    BertInformationRetrievalDataset,
+)
+from nemo.collections.nlp.data.language_modeling.l2r_lm_dataset import (
+    L2RLanguageModelingDataset,
+    TarredL2RLanguageModelingDataset,
+)
+from nemo.collections.nlp.data.language_modeling.lm_bert_dataset import (
+    BertPretrainingDataset,
+    BertPretrainingPreprocessedDataloader,
+)
+from nemo.collections.nlp.data.language_modeling.sentence_dataset import SentenceDataset, TarredSentenceDataset
+from nemo.collections.nlp.data.machine_translation.machine_translation_dataset import (
+    TarredTranslationDataset,
+    TranslationDataset,
+)
+from nemo.collections.nlp.data.neural_machine_translation.neural_machine_translation_dataset import (
+    NeuralMachineTranslationDataset,
+)
+from nemo.collections.nlp.data.question_answering_squad.qa_dataset import SquadDataset
+from nemo.collections.nlp.data.token_classification.token_classification_dataset import (
+    BertTokenClassificationDataset,
+    BertTokenClassificationInferDataset,
+)

tts/nemo/collections/nlp/data/data_utils/__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.data.data_utils.data_preprocessing import *

tts/nemo/collections/nlp/data/data_utils/data_preprocessing.py

@@ -0,0 +1,479 @@
+# 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 csv
+import json
+import os
+import pickle
+import random
+import re
+import string
+from collections import Counter
+
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+
+from nemo.utils import logging
+from nemo.utils.env_var_parsing import get_envint
+
+__all__ = [
+    'DataProcessor',
+    'get_label_stats',
+    'partition_data',
+    'write_files',
+    'write_data',
+    'create_dataset',
+    'read_csv',
+    'get_dataset',
+    'partition',
+    'map_entities',
+    'get_entities',
+    'get_data',
+    'reverse_dict',
+    'get_intent_labels',
+    'get_stats',
+    'DATABASE_EXISTS_TMP',
+    'MODE_EXISTS_TMP',
+    'is_whitespace',
+    'write_vocab',
+    'if_exist',
+    'remove_punctuation_from_sentence',
+    'dataset_to_ids',
+    'get_freq_weights',
+    'fill_class_weights',
+    'normalize_answer',
+    'get_labels_to_labels_id_mapping',
+    'get_vocab',
+    'find_newlines',
+    'load_data_indices',
+]
+
+DATABASE_EXISTS_TMP = '{} dataset has already been processed and stored at {}'
+MODE_EXISTS_TMP = '{} mode of {} dataset has already been processed and stored at {}'
+
+
+class DataProcessor(object):
+    """Base class for data converters for sequence classification data sets."""
+
+    def get_train_examples(self, data_dir):
+        """Gets a collection of `InputExample`s for the train set."""
+        raise NotImplementedError()
+
+    def get_dev_examples(self, data_dir):
+        """Gets a collection of `InputExample`s for the dev set."""
+        raise NotImplementedError()
+
+    def get_labels(self):
+        """Gets the list of labels for this data set."""
+        raise NotImplementedError()
+
+    @classmethod
+    def _read_tsv(cls, input_file, quotechar=None):
+        """Reads a tab separated value file."""
+        with open(input_file, "r", encoding="utf-8-sig") as f:
+            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
+            lines = []
+            for line in reader:
+                # if sys.version_info[0] == 2:
+                #     line = list(unicode(cell, 'utf-8') for cell in line)
+                lines.append(line)
+            return lines
+
+
+def normalize_answer(s):
+    """Lower text and remove punctuation, articles and extra whitespace."""
+
+    def remove_articles(text):
+        return re.sub(r'\b(a|an|the)\b', ' ', text)
+
+    def white_space_fix(text):
+        return ' '.join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return ''.join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+
+def get_label_stats(labels, outfile='stats.tsv', verbose=True):
+    '''
+
+    Args:
+        labels: list of all labels
+        outfile: path to the file where to save label stats
+
+    Returns:
+        total (int): total number of labels
+        label_frequencies (list of tuples): each tuple represent (label, label frequency)
+        max id of the labels
+    '''
+    labels = Counter(labels)
+    total = sum(labels.values())
+    out = open(outfile, 'w')
+    i = 0
+    freq_dict = {}
+    label_frequencies = labels.most_common()
+    for k, v in label_frequencies:
+        out.write(f'{k}\t\t{round(v/total,5)}\t\t{v}\n')
+        if verbose and i < 3:
+            logging.info(f'label: {k}, {v} out of {total} ({(v / total)*100.0:.2f}%).')
+        i += 1
+        freq_dict[k] = v
+
+    return total, freq_dict, max(labels.keys())
+
+
+def partition_data(intent_queries, slot_tags, split=0.1):
+    n = len(intent_queries)
+    n_dev = int(n * split)
+    dev_idx = set(random.sample(range(n), n_dev))
+    dev_intents, dev_slots, train_intents, train_slots = [], [], [], []
+
+    dev_intents.append('sentence\tlabel\n')
+    train_intents.append('sentence\tlabel\n')
+
+    for i, item in enumerate(intent_queries):
+        if i in dev_idx:
+            dev_intents.append(item)
+            dev_slots.append(slot_tags[i])
+        else:
+            train_intents.append(item)
+            train_slots.append(slot_tags[i])
+    return train_intents, train_slots, dev_intents, dev_slots
+
+
+def write_files(data, outfile):
+    with open(outfile, 'w') as f:
+        for item in data:
+            item = f'{item.strip()}\n'
+            f.write(item)
+
+
+def write_data(data, slot_dict, intent_dict, outfold, mode, uncased):
+    intent_file = open(f'{outfold}/{mode}.tsv', 'w')
+    intent_file.write('sentence\tlabel\n')
+    slot_file = open(f'{outfold}/{mode}_slots.tsv', 'w')
+    for tokens, slots, intent in data:
+        text = ' '.join(tokens)
+        if uncased:
+            text = text.lower()
+        intent_file.write(f'{text}\t{intent_dict[intent]}\n')
+        slots = [str(slot_dict[slot]) for slot in slots]
+        slot_file.write(' '.join(slots) + '\n')
+    intent_file.close()
+    slot_file.close()
+
+
+def create_dataset(train, dev, slots, intents, uncased, outfold):
+    os.makedirs(outfold, exist_ok=True)
+    if 'O' in slots:
+        slots.remove('O')
+    slots = sorted(list(slots)) + ['O']
+    intents = sorted(list(intents))
+    slots = write_vocab(slots, f'{outfold}/dict.slots.csv')
+    intents = write_vocab(intents, f'{outfold}/dict.intents.csv')
+    write_data(train, slots, intents, outfold, 'train', uncased)
+    write_data(dev, slots, intents, outfold, 'test', uncased)
+
+
+def read_csv(file_path):
+    rows = []
+    with open(file_path, 'r') as csvfile:
+        read_csv = csv.reader(csvfile, delimiter=',')
+        for row in read_csv:
+            rows.append(row)
+    return rows
+
+
+def get_dataset(files, dev_split=0.1):
+    # entity2value, value2entity = get_entities(files)
+    data, slots, intents = get_data(files)
+    if len(data) == 1:
+        train, dev = partition(data[0], split=dev_split)
+    else:
+        train, dev = data[0], data[1]
+    return train, dev, slots, intents
+
+
+def partition(data, split=0.1):
+    n = len(data)
+    n_dev = int(n * split)
+    dev_idx = set(random.sample(range(n), n_dev))
+    dev, train = [], []
+
+    for i, item in enumerate(data):
+        if i in dev_idx:
+            dev.append(item)
+        else:
+            train.append(item)
+    return train, dev
+
+
+def map_entities(entity2value, entities):
+    for key in entities:
+        if 'data' in entities[key]:
+            if key not in entity2value:
+                entity2value[key] = set([])
+
+            values = []
+            for value in entities[key]['data']:
+                values.append(value['value'])
+                values.extend(value['synonyms'])
+            entity2value[key] = entity2value[key] | set(values)
+
+    return entity2value
+
+
+def get_entities(files):
+    entity2value = {}
+    for file in files:
+        with open(file, 'r') as json_file:
+            data = json.load(json_file)
+            entity2value = map_entities(entity2value, data['entities'])
+
+    value2entity = reverse_dict(entity2value)
+    return entity2value, value2entity
+
+
+def get_data(files):
+    all_data, all_slots, all_intents = [], set(['O']), set()
+    for file in files:
+        file_data = []
+        with open(file, 'r') as json_file:
+            data = json.load(json_file)
+            for intent in data['intents']:
+                all_intents.add(intent)
+                utterances = data['intents'][intent]['utterances']
+                for utterance in utterances:
+                    tokens, slots = [], []
+                    for frag in utterance['data']:
+                        frag_tokens = frag['text'].strip().split()
+                        tokens.extend(frag_tokens)
+                        if 'slot_name' not in frag:
+                            slot = 'O'
+                        else:
+                            slot = frag['slot_name']
+                            all_slots.add(slot)
+                        slots.extend([slot] * len(frag_tokens))
+                    file_data.append((tokens, slots, intent))
+        all_data.append(file_data)
+    return all_data, all_slots, all_intents
+
+
+def reverse_dict(entity2value):
+    value2entity = {}
+    for entity in entity2value:
+        for value in entity2value[entity]:
+            value2entity[value] = entity
+    return value2entity
+
+
+def get_intent_labels(intent_file):
+    labels = {}
+    label = 0
+    with open(intent_file, 'r') as f:
+        for line in f:
+            intent = line.strip()
+            labels[intent] = label
+            label += 1
+    return labels
+
+
+def get_stats(lengths):
+    logging.info('Some stats of the lengths of the sequences:')
+    lengths = np.asarray(lengths)
+    logging.info(
+        f'Min: {np.min(lengths)} | \
+                 Max: {np.max(lengths)} | \
+                 Mean: {np.mean(lengths)} | \
+                 Median: {np.median(lengths)}'
+    )
+    logging.info(f'75 percentile: {np.percentile(lengths, 75):.2f}')
+    logging.info(f'99 percentile: {np.percentile(lengths, 99):.2f}')
+
+
+def is_whitespace(c):
+    if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
+        return True
+    return False
+
+
+def write_vocab(items, outfile):
+    vocab = {}
+    idx = 0
+    with open(outfile, 'w') as f:
+        for item in items:
+            f.write(item + '\n')
+            vocab[item] = idx
+            idx += 1
+    return vocab
+
+
+def get_labels_to_labels_id_mapping(file):
+    '''
+    Reads labels from the file and returns labels to id mapping dictionary
+    Args:
+        file: path to file
+    Returns:
+        labels to id mapping dictionary
+    '''
+    lines = open(file, 'r').readlines()
+    lines = [line.strip() for line in lines if line.strip()]
+    label_ids = {lines[i]: i for i in range(len(lines))}
+    return label_ids
+
+
+def if_exist(outfold, files):
+    if not os.path.exists(outfold):
+        return False
+    for file in files:
+        if not os.path.exists(f'{outfold}/{file}'):
+            return False
+    return True
+
+
+def remove_punctuation_from_sentence(sentence):
+    sentence = re.sub('[' + string.punctuation + ']', '', sentence)
+    sentence = sentence.lower()
+    return sentence
+
+
+def dataset_to_ids(dataset, tokenizer, cache_ids=False, add_bos_eos=True, cache_data_per_node=False, use_cache=False):
+    """
+    Reads dataset from file line by line, tokenizes each line with tokenizer,
+    and returns list of lists which corresponds to ids of tokenized strings.
+
+    Args:
+        dataset (str): path to dataset
+        tokenizer: tokenizer to convert text into ids
+        cache_ids (bool): if True, ids are saved to disk as pickle file
+            with similar name (e.g., data.txt --> data.txt.pkl)
+        add_bos_eos (bool): whether to add <s> and </s> symbols (e.g., for NMT)
+        cache_data_per_node (bool): Cache data on local_rank 0. Use when there is not a shared-filesystem.
+        use_cache (bool): Use cached ids if they exist.
+    Returns:
+        ids: list of ids which correspond to tokenized strings of the dataset
+    """
+
+    cached_ids_dataset = dataset + str(".pkl")
+    if use_cache and os.path.isfile(cached_ids_dataset):
+        logging.info("Loading cached tokenized dataset ...")
+        ids = pickle.load(open(cached_ids_dataset, "rb"))
+    else:
+        logging.info(f"Tokenizing dataset {dataset}...")
+        data = open(dataset, "rb").readlines()
+        ids = []
+        for sentence in tqdm(data, desc='Tokenizing sentence'):
+            sent_ids = tokenizer.text_to_ids(sentence.decode("utf-8"))
+            if add_bos_eos:
+                sent_ids = [tokenizer.bos_id] + sent_ids + [tokenizer.eos_id]
+            ids.append(sent_ids)
+        if cache_ids and (
+            not torch.distributed.is_initialized() or (cache_data_per_node and get_envint("LOCAL_RANK", 0) == 0)
+        ):
+            logging.info("Caching tokenized dataset ...")
+            pickle.dump(ids, open(cached_ids_dataset, "wb"))
+    return ids
+
+
+def get_freq_weights(label_freq):
+    """
+    Goal is to give more weight to the classes with less samples
+    so as to match the ones with the higher frequencies. We achieve this by
+    dividing the total frequency by the freq of each label to calculate its weight.
+    """
+    total_size = 0
+    for lf in label_freq.values():
+        total_size += lf
+    weighted_slots = {label: (total_size / (len(label_freq) * freq)) for label, freq in label_freq.items()}
+    return weighted_slots
+
+
+def fill_class_weights(weights, max_id=-1):
+    """
+    Gets a dictionary of labels with their weights and creates a list with size of the labels filled with those weights.
+    Missing labels in the dictionary would get value 1.
+
+    Args:
+        weights: dictionary of weights for labels, labels as keys and weights are their values
+        max_id: the largest label id in the dataset, default=-1 would consider the largest label in the weights dictionary as max_id
+    Returns:
+        weights_list: list of weights for labels
+    """
+    if max_id < 0:
+        max_id = 0
+        for l in weights.keys():
+            max_id = max(max_id, l)
+
+    all_weights = [1.0] * (max_id + 1)
+    for i in range(len(all_weights)):
+        if i in weights:
+            all_weights[i] = weights[i]
+    return all_weights
+
+
+def get_vocab(file):
+    lines = open(file, 'r').readlines()
+    lines = [line.strip() for line in lines if line.strip()]
+    labels = {i: lines[i] for i in range(len(lines))}
+    return labels
+
+
+def find_newlines(contents):
+    """
+    Finds all of the newline positions in a text file.
+    """
+    start = 0
+
+    while True:
+        try:
+            # index and split are much faster than Python for loops
+            new_start = contents.index(b"\n", start)
+            line = (
+                contents[start:new_start]
+                .replace(b"\xc2\x99", b" ")
+                .replace(b"\xc2\xa0", b" ")
+                .decode("utf-8", errors="ignore")
+            )
+
+            if len(line.split()) > 0:
+                yield start
+
+            start = new_start + 1
+
+        except ValueError:
+            break
+
+
+def load_data_indices(idx_file: str, data_file: str, savename: str):
+    """
+    Loads dataset index file if it exsits
+    """
+    data_dir = data_file[: data_file.rfind('/')]
+    mode = data_file[data_file.rfind('/') + 1 : data_file.rfind('.')]
+    idx_file = f"{data_dir}/{mode}_{savename}.pkl"
+
+    if os.path.isfile(idx_file):
+        # If the sentence indices file already exists, load from it
+        with open(idx_file, "rb") as f:
+            indices = pickle.load(f)
+
+            return indices, idx_file, data_dir
+
+    return None, idx_file, data_dir

tts/nemo/collections/nlp/data/dialogue_state_tracking/__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.data.dialogue_state_tracking.sgd.data_processor import SGDDataProcessor
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.dataset import SGDDataset
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.schema import Schema

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/data_processor.py

@@ -0,0 +1,676 @@
+# 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/data_utils.py
+"""
+
+import collections
+import json
+import os
+import pickle
+import re
+from typing import Dict, List
+
+import numpy as np
+
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.input_example import InputExample
+from nemo.utils import logging
+
+__all__ = ['SGDDataProcessor']
+
+FILE_RANGES = {
+    "sgd_single_domain": {"train": range(1, 44), "dev": range(1, 8), "test": range(1, 12)},
+    "sgd_multi_domain": {"train": range(44, 128), "dev": range(8, 21), "test": range(12, 35)},
+    "sgd_all": {"train": range(1, 128), "dev": range(1, 21), "test": range(1, 35)},
+    "sgd_all_single": {"train": range(1, 128), "dev": range(1, 8), "test": range(1, 12)},
+    "multiwoz": {"train": range(1, 18), "dev": range(1, 3), "test": range(1, 3)},
+    "debug_sample": {"train": range(1, 2), "dev": range(1, 2), "test": range(1, 2)},
+}
+
+
+class SGDDataProcessor(object):
+    """Data generator for SGD dialogues."""
+
+    def __init__(
+        self,
+        task_name: str,
+        data_dir: str,
+        dialogues_example_dir: str,
+        tokenizer: object,
+        schemas: object,
+        schema_config: Dict[str, int],
+        subsample: bool = False,
+    ):
+        """
+        Constructs SGDDataProcessor
+        Args:
+            task_name: task name, e.g. "sgd_single_domain"
+            data_dir: path to data directory
+            dialogues_example_dir: path to store processed dialogue examples
+            tokenizer: tokenizer object
+            schemas: schema object 
+            schema_config: schema configuration
+            subsample: whether to balance positive and negative samples in dataset
+        """
+        self.data_dir = data_dir
+
+        self._task_name = task_name
+        self.schemas = schemas
+        self.schema_config = schema_config
+
+        train_file_range = FILE_RANGES[task_name]["train"]
+        dev_file_range = FILE_RANGES[task_name]["dev"]
+        test_file_range = FILE_RANGES[task_name]["test"]
+
+        self._file_ranges = {
+            "train": train_file_range,
+            "dev": dev_file_range,
+            "test": test_file_range,
+        }
+
+        self._seen_services = {
+            "train": set(),
+            "dev": set(),
+            "test": set(),
+        }
+
+        self._tokenizer = tokenizer
+        self._subsample = subsample
+        self._dialogues_example_dir = dialogues_example_dir
+
+        self.dial_files = {}
+
+        # slots_relation_list.np would contain the candidate list of slots for each (service, slot) which would be
+        # looked into when a switch between two services happens in the dialogue and we can not find any value for a slot in the current user utterance.
+        # This file would get generated from the dialogues in the training set.
+        self.slots_relation_file = os.path.join(dialogues_example_dir, f"{task_name}_train_slots_relation_list.np")
+        for dataset in ["train", "dev", "test"]:
+            # Process dialogue files
+            dial_file = f"{task_name}_{dataset}_examples.processed"
+            dial_file = os.path.join(dialogues_example_dir, dial_file)
+            self.dial_files[(task_name, dataset)] = dial_file
+
+            dialog_paths = SGDDataProcessor.get_dialogue_files(data_dir, dataset, task_name)
+            dialogs = SGDDataProcessor.load_dialogues(dialog_paths)
+            for dialog in dialogs:
+                self._seen_services[dataset].update(set(dialog['services']))
+
+    def save_dialog_examples(self, overwrite_dial_files: bool):
+        """
+        Preprocesses dialogues and saves to disk.
+        Args:
+            overwrite_dial_files: whether or not to overwrite saved file if already exists
+        """
+        for dataset in ["train", "dev", "test"]:
+            dial_file = self.dial_files[(self._task_name, dataset)]
+            if not os.path.exists(dial_file) or overwrite_dial_files:
+                logging.info(f"Start generating the dialogue examples for {dataset} dataset.")
+                if not os.path.exists(self._dialogues_example_dir):
+                    os.makedirs(self._dialogues_example_dir)
+                dial_examples, slots_relation_list = self._generate_dialog_examples(
+                    dataset, self.schemas, self._subsample
+                )
+                for ex in dial_examples:
+                    del ex.service_schema
+                    del ex._tokenizer
+                    del ex.schema_config
+                    del ex.user_utterance
+                    del ex.categorical_slot_id
+                    del ex.system_utterance
+                    del ex.noncategorical_slot_id
+                    del ex.categorical_slot_value_id
+                    del ex.requested_slot_id
+                    del ex.intent_id
+                with open(dial_file, "wb") as f:
+                    np.save(f, dial_examples)
+
+                if dataset == "train":
+                    with open(self.slots_relation_file, "wb") as f:
+                        pickle.dump(slots_relation_list, f)
+                    logging.info(f"The slot carry-over list for train set is stored at {self.slots_relation_file}")
+
+                logging.info(f"The dialogue examples for {dataset} dataset saved at {dial_file}")
+                logging.info(f"Finish generating the dialogue examples for {dataset} dataset.")
+
+    def get_dialog_examples(self, dataset_split: str) -> List[object]:
+        """
+        Loads preprocessed dialogue examples from disk. 
+        Args:
+            dataset_split: dataset split
+        Returns:
+            dial_examples:  list of InputExample's.
+        """
+        if (self._task_name, dataset_split) not in self.dial_files or not os.path.exists(
+            self.dial_files[(self._task_name, dataset_split)]
+        ):
+            raise ValueError(
+                f"{dataset_split} dialogue examples were not processed for {self._task_name} task. Re-initialize SGDDataProcessor and add {dataset_split} dataset split to datasets arg."
+            )
+        dial_file = self.dial_files[(self._task_name, dataset_split)]
+        logging.info(f"Loading dialogue examples from {dial_file}.")
+
+        with open(dial_file, "rb") as f:
+            dial_examples = np.load(f, allow_pickle=True)
+            f.close()
+
+        if not os.path.exists(self.slots_relation_file):
+            raise ValueError(
+                f"Slots relation file {self.slots_relation_file} does not exist. It is needed for the carry-over mechanism of state tracker for switches between services."
+            )
+        if os.path.getsize(self.slots_relation_file) > 0:
+            with open(self.slots_relation_file, "rb") as f:
+                self.schemas._slots_relation_list = pickle.load(f)
+            logging.info(
+                f"Loaded the slot relation list for value carry-over between services from {self.slots_relation_file}."
+            )
+
+        return dial_examples
+
+    def get_seen_services(self, dataset_split: str):
+        """
+        Returns list of seen services, i.e. both in given and training split
+        Args:
+            dataset_split: data split
+        Returns:
+            seen_services: list of seen services
+        """
+        seen_services = self._seen_services[dataset_split]
+        return seen_services
+
+    def _generate_dialog_examples(self, dataset_split: str, schemas: object, subsample: bool):
+        """
+        Returns a list of `InputExample`s of the data splits' dialogues.
+        Args:
+            dataset_split: data split, can be "train", "dev", or "test".
+            schemas: schema for all services of all datasets 
+            subsample: whether to balance postive and negative samples in the dataset
+        Returns:
+            examples: a list of `InputExample`s.
+        """
+        logging.info(f'Creating examples and slot relation list from the dialogues started...')
+        dialog_paths = [
+            os.path.join(self.data_dir, dataset_split, "dialogues_{:03d}.json".format(i))
+            for i in self._file_ranges[dataset_split]
+        ]
+        dialogs = SGDDataProcessor.load_dialogues(dialog_paths)
+
+        examples = []
+        slot_carryover_candlist = collections.defaultdict(int)
+        for dialog_idx, dialog in enumerate(dialogs):
+            if dialog_idx % 1000 == 0:
+                logging.info(f'Processed {dialog_idx} dialogues.')
+            examples.extend(
+                self._create_examples_from_dialog(dialog, schemas, dataset_split, slot_carryover_candlist, subsample)
+            )
+
+        slots_relation_list = collections.defaultdict(list)
+        for slots_relation, relation_size in slot_carryover_candlist.items():
+            if relation_size > 0:
+                slots_relation_list[(slots_relation[0], slots_relation[1])].append(
+                    (slots_relation[2], slots_relation[3], relation_size)
+                )
+                slots_relation_list[(slots_relation[2], slots_relation[3])].append(
+                    (slots_relation[0], slots_relation[1], relation_size)
+                )
+
+        return examples, slots_relation_list
+
+    def _create_examples_from_dialog(
+        self, dialog: dict, schemas: object, dataset_split: str, slot_carryover_candlist: dict, subsample: bool
+    ):
+        """
+        Create examples for every turn in the dialogue.
+        Args:
+            dialog: dialogue example
+            schemas: schema for all services of all datasets
+            dataset_split: data split
+            slot_carryover_candlist: a dictionary to keep and count the number of carry-over cases between two slots from two different services
+            subsample: whether to balance postive and negative samples in the dataset
+        Returns:
+            examples: a list of `InputExample`s.
+        """
+        dialog_id = dialog["dialogue_id"]
+        prev_states = {}
+        examples = []
+        for turn_idx, turn in enumerate(dialog["turns"]):
+            # Generate an example for every frame in every user turn.
+            if turn["speaker"] == "USER":
+                user_utterance = turn["utterance"]
+                user_frames = {f["service"]: f for f in turn["frames"]}
+                if turn_idx > 0:
+                    system_turn = dialog["turns"][turn_idx - 1]
+                    system_utterance = system_turn["utterance"]
+                    system_frames = {f["service"]: f for f in system_turn["frames"]}
+                else:
+                    system_utterance = ""
+                    system_frames = {}
+
+                turn_id = "{}-{}-{:02d}".format(dataset_split, dialog_id, turn_idx)
+                turn_examples, prev_states, slot_carryover_values = self._create_examples_from_turn(
+                    turn_id,
+                    system_utterance,
+                    user_utterance,
+                    system_frames,
+                    user_frames,
+                    prev_states,
+                    schemas,
+                    subsample,
+                )
+                examples.extend(turn_examples)
+
+                for value, slots_list in slot_carryover_values.items():
+                    if value in ["True", "False"]:
+                        continue
+                    if len(slots_list) > 1:
+                        for service1, slot1 in slots_list:
+                            for service2, slot2 in slots_list:
+                                if service1 == service2:
+                                    continue
+                                if service1 > service2:
+                                    service1, service2 = service2, service1
+                                    slot1, slot2 = slot2, slot1
+                                slot_carryover_candlist[(service1, slot1, service2, slot2)] += 1
+        return examples
+
+    def _get_state_update(self, current_state: dict, prev_state: dict) -> dict:
+        """
+        Updates dialogue state
+        Args:
+            current_state: slot values pairs for the current dialogue turn
+            prev_state: slot values pairs for the previous dialogue turns
+        Returns:
+            state_update: slot values pairs that are added/updated during the current dialogue turn
+        """
+        state_update = dict(current_state)
+        for slot, values in current_state.items():
+            if slot in prev_state and prev_state[slot][0] in values:
+                # Remove the slot from state if its value didn't change.
+                state_update.pop(slot)
+        return state_update
+
+    def _create_examples_from_turn(
+        self,
+        turn_id: int,
+        system_utterance: str,
+        user_utterance: str,
+        system_frames: dict,
+        user_frames: dict,
+        prev_states: dict,
+        schemas: object,
+        subsample: bool,
+    ):
+        """
+        Creates an example for each frame in the user turn.
+        Args:
+            turn_id: turn number
+            system_utterance: last system utterance
+            user_utterance: lst user utterance
+            system_frames: all system utterances and slot - slot value pairs
+            user_frames: all user utterances and slot - slot value pairs
+            prev_states: slot - slot value pairs from the previous turns
+            schemas: schema for all services of all datasets
+            subsample: whether to balance postive and negative samples in the dataset
+        Returns:
+            examples: a list of `InputExample`s.
+            prev_states: updated dialogue state e.g. {'Restaurants_1': {'city': ['San Jose'], 'cuisine': ['American']}}
+        """
+        user_tokens, user_alignments, user_inv_alignments = self._tokenize(user_utterance)
+        system_tokens, system_alignments, system_inv_alignments = self._tokenize(system_utterance)
+        system_user_utterance = system_utterance + ' ' + user_utterance
+        system_user_tokens, system_user_alignments, system_user_inv_alignments = self._tokenize(system_user_utterance)
+        states = {}
+
+        examples = []
+        slot_carryover_values = collections.defaultdict(list)
+        for service, user_frame in user_frames.items():
+
+            base_example = InputExample(schema_config=self.schema_config, tokenizer=self._tokenizer,)
+            base_example.service_schema = schemas.get_service_schema(service)
+            base_example.service_id = schemas.get_service_schema(service).service_id
+            system_frame = system_frames.get(service, None)
+            state = user_frame["state"]["slot_values"]
+            state_update = self._get_state_update(state, prev_states.get(service, {}))
+            states[service] = state
+            dataset_split, dialog_id, turn_id_ = turn_id.split('-')
+            dialog_id_1, dialog_id_2 = dialog_id.split('_')
+
+            base_example.example_id = f"{turn_id}-{service}"
+            base_example.example_id_num = [
+                int(dialog_id_1),
+                int(dialog_id_2),
+                int(turn_id_),
+                schemas.get_service_id(service),
+            ]
+
+            for model_task in range(self.schema_config["NUM_TASKS"]):
+                if model_task == 0:
+                    for intent_id, intent in enumerate(schemas.get_service_schema(service).intents):
+                        task_example = base_example.make_copy()
+                        task_example.task_mask[model_task] = 1
+                        task_example.intent_id = intent_id
+                        task_example.example_id += f"-{model_task}-{intent_id}-0"
+                        task_example.example_id_num.extend([model_task, intent_id, 0])
+                        intent_description = (
+                            intent + " " + schemas.get_service_schema(service).intent_descriptions[intent]
+                        )
+                        intent_tokens, intent_alignments, intent_inv_alignments = self._tokenize(intent_description)
+                        task_example.add_utterance_features(
+                            intent_tokens,
+                            intent_inv_alignments,
+                            system_user_tokens,
+                            system_user_inv_alignments,
+                            intent_description,
+                            system_user_utterance,
+                        )
+                        task_example.add_intents(user_frame)
+                        examples.append(task_example)
+
+                if model_task == 1:
+                    for slot_id, slot in enumerate(schemas.get_service_schema(service).slots):
+                        task_example = base_example.make_copy()
+                        task_example.task_mask[model_task] = 1
+                        task_example.requested_slot_id = slot_id
+                        task_example.example_id += f"-{model_task}-{slot_id}-0"
+                        task_example.example_id_num.extend([model_task, slot_id, 0])
+                        slot_description = slot + " " + schemas.get_service_schema(service).slot_descriptions[slot]
+                        slot_tokens, slot_alignments, slot_inv_alignments = self._tokenize(slot_description)
+                        task_example.add_utterance_features(
+                            slot_tokens,
+                            slot_inv_alignments,
+                            user_tokens,
+                            user_inv_alignments,
+                            slot_description,
+                            user_utterance,
+                        )
+                        task_example.add_requested_slots(user_frame)
+                        examples.append(task_example)
+                if model_task == 2:
+                    off_slots = []
+                    on_slots = []
+                    for slot_id, slot in enumerate(schemas.get_service_schema(service).categorical_slots):
+                        task_example = base_example.make_copy()
+                        task_example.task_mask[model_task] = 1
+
+                        # assert task_example.task_mask == [0, 0, 1, 0, 0, 0]
+                        task_example.categorical_slot_id = slot_id
+                        task_example.example_id += f"-{model_task}-{slot_id}-0"
+                        task_example.example_id_num.extend([model_task, slot_id, 0])
+                        slot_description = slot + " " + schemas.get_service_schema(service).slot_descriptions[slot]
+                        slot_tokens, slot_alignments, slot_inv_alignments = self._tokenize(slot_description)
+                        task_example.add_utterance_features(
+                            slot_tokens,
+                            slot_inv_alignments,
+                            system_user_tokens,
+                            system_user_inv_alignments,
+                            slot_description,
+                            system_user_utterance,
+                        )
+                        task_example.add_categorical_slots(state_update)
+                        if task_example.categorical_slot_status == 0:
+                            off_slots.append(task_example)
+                        else:
+                            on_slots.append(task_example)
+                            examples.append(task_example)
+                        old_example = task_example
+
+                        for value_id, value in enumerate(
+                            schemas.get_service_schema(service).get_categorical_slot_values(slot)
+                        ):
+                            if dataset_split != 'train' or task_example.categorical_slot_status == 1:
+                                task_example = old_example.make_copy_of_categorical_features()
+                                task_example.task_mask[3] = 1
+                                # assert task_example.task_mask == [0, 0, 0, 1, 0, 0]
+                                task_example.categorical_slot_id = slot_id
+                                task_example.categorical_slot_value_id = value_id
+                                task_example.example_id = base_example.example_id + f"-3-{slot_id}-{value_id}"
+                                task_example.example_id_num = base_example.example_id_num + [3, slot_id, value_id]
+                                slot_description = slot + " " + value  # add slot description
+                                slot_tokens, slot_alignments, slot_inv_alignments = self._tokenize(slot_description)
+                                task_example.add_utterance_features(
+                                    slot_tokens,
+                                    slot_inv_alignments,
+                                    system_user_tokens,
+                                    system_user_inv_alignments,
+                                    slot_description,
+                                    system_user_utterance,
+                                )
+                                task_example.add_categorical_slots(state_update)
+                                assert task_example.categorical_slot_status == old_example.categorical_slot_status
+                                examples.append(task_example)
+
+                    if dataset_split == 'train' and subsample:
+                        num_on_slots = len(on_slots)
+                        examples.extend(
+                            np.random.choice(off_slots, replace=False, size=min(max(num_on_slots, 1), len(off_slots)))
+                        )
+                    else:
+                        examples.extend(off_slots)
+
+                if model_task == 4:  # noncat slot status
+                    off_slots = []
+                    on_slots = []
+                    for slot_id, slot in enumerate(schemas.get_service_schema(service).non_categorical_slots):
+                        task_example = base_example.make_copy()
+                        task_example.task_mask[model_task] = 1
+                        # assert task_example.task_mask == [0, 0, 0, 0, 1, 0]
+                        task_example.noncategorical_slot_id = slot_id
+                        task_example.example_id += f"-{model_task}-{slot_id}-0"
+                        task_example.example_id_num.extend([model_task, slot_id, 0])
+                        slot_description = slot + " " + schemas.get_service_schema(service).slot_descriptions[slot]
+                        slot_tokens, slot_alignments, slot_inv_alignments = self._tokenize(slot_description)
+                        task_example.add_utterance_features(
+                            slot_tokens,
+                            slot_inv_alignments,
+                            system_user_tokens,
+                            system_user_inv_alignments,
+                            slot_description,
+                            system_user_utterance,
+                        )
+
+                        user_span_boundaries = self._find_subword_indices(
+                            state_update,
+                            user_utterance,
+                            user_frame["slots"],
+                            user_alignments,
+                            user_tokens,
+                            2 + len(slot_tokens) + len(system_tokens),
+                        )
+                        if system_frame is not None:
+                            system_span_boundaries = self._find_subword_indices(
+                                state_update,
+                                system_utterance,
+                                system_frame["slots"],
+                                system_alignments,
+                                system_tokens,
+                                2 + len(slot_tokens),
+                            )
+                        else:
+                            system_span_boundaries = {}
+                        task_example.add_noncategorical_slots(
+                            state_update, user_span_boundaries, system_span_boundaries
+                        )
+                        if task_example.noncategorical_slot_status == 0:
+                            off_slots.append(task_example)
+                        else:
+                            on_slots.append(task_example)
+                            examples.append(task_example)
+
+                        if dataset_split != 'train' or task_example.noncategorical_slot_status == 1:
+                            task_example = task_example.make_copy_of_non_categorical_features()
+                            task_example.task_mask[5] = 1
+                            # assert task_example.task_mask == [0, 0, 0, 0, 0, 1]
+                            task_example.example_id = base_example.example_id + f"-5-{slot_id}-0"
+                            task_example.example_id_num = base_example.example_id_num + [5, slot_id, 0]
+                            examples.append(task_example)
+                    if dataset_split == 'train' and subsample:
+                        num_on_slots = len(on_slots)
+                        examples.extend(
+                            np.random.choice(off_slots, replace=False, size=min(max(num_on_slots, 1), len(off_slots)))
+                        )
+                    else:
+                        examples.extend(off_slots)
+
+            if service not in prev_states and int(turn_id_) > 0:
+                for slot_name, values in state_update.items():
+                    for value in values:
+                        slot_carryover_values[value].append((service, slot_name))
+                for prev_service, prev_slot_value_list in prev_states.items():
+                    if prev_service == service:
+                        continue
+                    if prev_service in state:
+                        prev_slot_value_list = state[prev_service]
+                    for prev_slot_name, prev_values in prev_slot_value_list.items():
+                        for prev_value in prev_values:
+                            slot_carryover_values[prev_value].append((prev_service, prev_slot_name))
+
+        return examples, states, slot_carryover_values
+
+    def _find_subword_indices(
+        self,
+        slot_values: dict,
+        utterance: str,
+        char_slot_spans: dict,
+        alignments: List[int],
+        subwords: List[str],
+        bias: int,
+    ) -> dict:
+        """
+        Find indices for subwords corresponding to slot values.
+        Args:
+            slot_values: slot - slot value pairs
+            utterance: utterance
+            char_slot_spans: char - slot spans
+            alignments: alignments
+            subwords: subtokens mapping
+            bias: offset
+        Returns:
+            span_boundaries: span boundaries
+        """
+        span_boundaries = {}
+        for slot, values in slot_values.items():
+            # Get all values present in the utterance for the specified slot.
+            value_char_spans = {}
+            for slot_span in char_slot_spans:
+                if slot_span["slot"] == slot:
+                    value = utterance[slot_span["start"] : slot_span["exclusive_end"]]
+                    start_tok_idx = alignments[slot_span["start"]]
+                    end_tok_idx = alignments[slot_span["exclusive_end"] - 1]
+                    if 0 <= start_tok_idx < len(subwords):
+                        end_tok_idx = min(end_tok_idx, len(subwords) - 1)
+                        value_char_spans[value] = (start_tok_idx + bias, end_tok_idx + bias)
+            for v in values:
+                if v in value_char_spans:
+                    span_boundaries[slot] = value_char_spans[v]
+                    break
+        return span_boundaries
+
+    def _tokenize(self, utterance: str):
+        """
+        Tokenize the utterance
+
+        Args:
+            utterance: A string containing the utterance to be tokenized.
+
+        Returns:
+            bert_tokens: A list of tokens obtained by word-piece tokenization of the
+                utterance.
+            alignments: A dict mapping indices of characters corresponding to start
+                and end positions of words (not subwords) to corresponding indices in
+                bert_tokens list.
+            inverse_alignments: A list of size equal to bert_tokens. Each element is a
+                tuple containing the index of the starting and inclusive ending
+                character of the word corresponding to the subword. This list is used
+                during inference to map word-piece indices to spans in the original
+                utterance.
+        """
+        # utterance = tokenization.convert_to_unicode(utterance)
+
+        # After _naive_tokenize, spaces and punctuation marks are all retained, i.e.
+        # direct concatenation of all the tokens in the sequence will be the
+        # original string.
+        tokens = SGDDataProcessor._naive_tokenize(utterance)
+        # ['I', ' ', 'am', ' ', 'feeling', ' ', 'hungry', ' ', 'so', ' ', 'I', ' ', 'would', ' ', 'like', ' ', 'to', ' ', 'find', ' ', 'a', ' ', 'place', ' ', 'to', ' ', 'eat', '.']
+        # Filter out empty tokens and obtain aligned character index for each token.
+        alignments = {}
+        char_index = 0
+        bert_tokens = (
+            []
+        )  # ['I', 'am', 'feeling', 'hungry', 'so', 'I', 'would', 'like', 'to', 'find', 'a', 'place', 'to', 'eat', '.']
+        # These lists store inverse alignments to be used during inference.
+        bert_tokens_start_chars = []
+        bert_tokens_end_chars = []
+        for token in tokens:
+            if token.strip():
+                subwords = self._tokenizer.text_to_tokens(token)
+                # Store the alignment for the index of starting character and the
+                # inclusive ending character of the token.
+                alignments[char_index] = len(bert_tokens)
+                bert_tokens_start_chars.extend([char_index] * len(subwords))
+                bert_tokens.extend(subwords)
+                # The inclusive ending character index corresponding to the word.
+                inclusive_char_end = char_index + len(token) - 1
+                alignments[inclusive_char_end] = len(bert_tokens) - 1
+                bert_tokens_end_chars.extend([inclusive_char_end] * len(subwords))
+            char_index += len(token)
+        inverse_alignments = list(zip(bert_tokens_start_chars, bert_tokens_end_chars))
+        return bert_tokens, alignments, inverse_alignments
+
+    @classmethod
+    def _naive_tokenize(cls, s: str):
+        """
+        Tokenizes a string, separating words, spaces and punctuations.
+        Args:
+            s: a string
+        Returns:
+            seq_tok: list of words, spaces and punctuations from the string
+        """
+        # Spaces and punctuation marks are all retained, i.e. direct concatenation
+        # of all the tokens in the sequence will be the original string.
+        seq_tok = [tok for tok in re.split(r"([^a-zA-Z0-9])", s) if tok]
+        return seq_tok
+
+    @classmethod
+    def load_dialogues(cls, dialog_json_filepaths: List[str]) -> List[dict]:
+        """
+        Obtain the list of all dialogues from specified json files.
+        Args:
+            dialog_json_filepaths: list of json files
+        Returns:
+            dialogs: the list of all dialogues
+        """
+        dialogs = []
+        for dialog_json_filepath in sorted(dialog_json_filepaths):
+            with open(dialog_json_filepath, 'r') as f:
+                dialogs.extend(json.load(f))
+                f.close()
+        return dialogs
+
+    @classmethod
+    def get_dialogue_files(cls, data_dir: str, dataset_split: str, task_name: str):
+        """
+        Obtain the list of all dialogue json files
+        Args:
+            data_dir: path to the data folder
+            dataset_split: data split
+            task_name: SGD task name, see keys of the FILE_RANGES
+        Returns:
+            dialog: the list of all dialogue json files paths
+        """
+        return [
+            os.path.join(data_dir, dataset_split, 'dialogues_{:03d}.json'.format(fid))
+            for fid in FILE_RANGES[task_name][dataset_split]
+        ]

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/dataset.py

@@ -0,0 +1,85 @@
+# 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
+"""
+from typing import Dict, Optional
+
+import numpy as np
+
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, LabelsType, NeuralType
+
+__all__ = ['SGDDataset']
+
+
+class SGDDataset(Dataset):
+    """Processes SGD dataset"""
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            "example_id_num": NeuralType(('B', 'T'), ChannelType()),
+            "service_id": NeuralType(('B'), ChannelType()),
+            "utterance_ids": NeuralType(('B', 'T'), ChannelType()),
+            "token_type_ids": NeuralType(('B', 'T'), ChannelType()),  # utterance segment
+            "attention_mask": NeuralType(('B', 'T'), ChannelType()),  # utterance mask
+            "intent_status": NeuralType(('B'), LabelsType()),
+            "requested_slot_status": NeuralType(('B'), LabelsType()),
+            "categorical_slot_status": NeuralType(('B'), LabelsType()),
+            "categorical_slot_value_status": NeuralType(('B'), LabelsType()),
+            "noncategorical_slot_status": NeuralType(('B'), LabelsType()),
+            "noncategorical_slot_value_start": NeuralType(('B'), LabelsType()),
+            "noncategorical_slot_value_end": NeuralType(('B'), LabelsType()),
+            "start_char_idx": NeuralType(('B', 'T'), LabelsType()),
+            "end_char_idx": NeuralType(('B', 'T'), LabelsType()),
+            "task_mask": NeuralType(('B', 'T'), ChannelType()),
+        }
+
+    def __init__(self, dataset_split: str, dialogues_processor: object):
+        """ Constructor
+        Args:
+            dataset_split: dataset split
+            dialogues_processor: Data generator for SGD dialogues
+        """
+        self.features = dialogues_processor.get_dialog_examples(dataset_split)
+
+    def __len__(self):
+        return len(self.features)
+
+    def __getitem__(self, idx: int):
+        ex = self.features[idx]
+        service_id = ex.service_id
+
+        return (
+            np.array(ex.example_id_num),
+            np.array(service_id),
+            np.array(ex.utterance_ids),
+            np.array(ex.utterance_segment),
+            np.array(ex.utterance_mask, dtype=np.long),
+            np.array(ex.intent_status, dtype=np.float32),
+            np.array(ex.requested_slot_status, dtype=np.float32),
+            np.array(ex.categorical_slot_status),
+            np.array(ex.categorical_slot_value_status, dtype=np.float32),
+            np.array(ex.noncategorical_slot_status),
+            np.array(ex.noncategorical_slot_value_start),
+            np.array(ex.noncategorical_slot_value_end),
+            np.array(ex.start_char_idx),  # noncat_alignment_start
+            np.array(ex.end_char_idx),  # noncat_alignment_end
+            np.array(ex.task_mask),  # noncat_alignment_end
+        )

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/evaluate.py

@@ -0,0 +1,294 @@
+# 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.
+
+"""
+Evaluate predictions JSON file, w.r.t. ground truth file.
+This file contains code artifacts adapted from the original implementation:
+https://github.com/google-research/google-research/blob/master/schema_guided_dst/evaluate.py
+"""
+
+import collections
+import glob
+import json
+import os
+
+import numpy as np
+
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.metrics import (
+    ACTIVE_INTENT_ACCURACY,
+    JOINT_CAT_ACCURACY,
+    JOINT_GOAL_ACCURACY,
+    JOINT_NONCAT_ACCURACY,
+    NAN_VAL,
+    REQUESTED_SLOTS_F1,
+    REQUESTED_SLOTS_PRECISION,
+    REQUESTED_SLOTS_RECALL,
+    SLOT_TAGGING_F1,
+    SLOT_TAGGING_PRECISION,
+    SLOT_TAGGING_RECALL,
+    get_active_intent_accuracy,
+    get_average_and_joint_goal_accuracy,
+    get_requested_slots_f1,
+    get_slot_tagging_f1,
+)
+from nemo.utils import logging
+
+__all__ = ['get_in_domain_services']
+
+ALL_SERVICES = "#ALL_SERVICES"
+SEEN_SERVICES = "#SEEN_SERVICES"
+UNSEEN_SERVICES = "#UNSEEN_SERVICES"
+
+# Name of the file containing all predictions and their corresponding frame metrics.
+PER_FRAME_OUTPUT_FILENAME = "dialogues_and_metrics.json"
+
+
+def get_service_set(schema_path: str) -> set:
+    """
+    Get the set of all services present in a schema.
+    Args:
+        schema_path: schema file path
+    Returns:
+        service_set: set of services in file
+    """
+    service_set = set()
+    with open(schema_path) as f:
+        schema = json.load(f)
+        for service in schema:
+            service_set.add(service["service_name"])
+        f.close()
+    return service_set
+
+
+def get_in_domain_services(schema_path: str, service_set: set) -> set:
+    """Get the set of common services between a schema and set of services.
+    Args:
+        schema_path: path to schema file
+        service_set: set of services
+    Returns: 
+        joint_services: joint services between schema path file and service set
+    """
+    joint_services = get_service_set(schema_path) & service_set
+    return joint_services
+
+
+def get_dataset_as_dict(file_path_patterns) -> dict:
+    """Read the DSTC8/SGD json dialogue data as dictionary with dialog ID as keys.
+    Args:
+        file_path_patterns: list or directory of files 
+    Returns:
+        dataset_dict: dataset dictionary with dialog ID as keys
+    """
+    dataset_dict = {}
+    if isinstance(file_path_patterns, list):
+        list_fp = file_path_patterns
+    else:
+        list_fp = sorted(glob.glob(file_path_patterns))
+    for fp in list_fp:
+        if PER_FRAME_OUTPUT_FILENAME in fp:
+            continue
+        logging.debug("Loading file: %s", fp)
+        with open(fp) as f:
+            data = json.load(f)
+            if isinstance(data, list):
+                for dial in data:
+                    dataset_dict[dial["dialogue_id"]] = dial
+            elif isinstance(data, dict):
+                dataset_dict.update(data)
+            f.close()
+    return dataset_dict
+
+
+def get_metrics(
+    dataset_ref: dict,
+    dataset_hyp: dict,
+    service_schemas: dict,
+    in_domain_services: set,
+    joint_acc_across_turn: bool,
+    use_fuzzy_match: bool,
+):
+    """Calculate the DSTC8/SGD metrics.
+    Args:
+        dataset_ref: The ground truth dataset represented as a dict mapping dialogue id to the corresponding dialogue.
+        dataset_hyp: The predictions in the same format as `dataset_ref`.
+        service_schemas: A dict mapping service name to the schema for the service.
+        in_domain_services: The set of services which are present in the training set.
+        joint_acc_across_turn: Whether to compute joint accuracy across turn instead of across service. Should be set to True when conducting multiwoz style evaluation.
+        use_fuzzy_match: Whether to use fuzzy string matching when comparing non-categorical slot values. Should be set to False when conducting multiwoz style evaluation.
+
+    Returns:
+        all_metric_aggregate: A dict mapping a metric collection name to a dict containing the values
+            for various metrics. Each metric collection aggregates the metrics across a specific set of frames in the dialogues.
+        per_frame_metric: metrics aggregated for each frame
+    """
+    # Metrics can be aggregated in various ways, eg over all dialogues, only for
+    # dialogues containing unseen services or for dialogues corresponding to a
+    # single service. This aggregation is done through metric_collections, which
+    # is a dict mapping a collection name to a dict, which maps a metric to a list
+    # of values for that metric. Each value in this list is the value taken by
+    # the metric on a frame.
+    metric_collections = collections.defaultdict(lambda: collections.defaultdict(list))
+
+    # Ensure the dialogs in dataset_hyp also occur in dataset_ref.
+    assert set(dataset_hyp.keys()).issubset(set(dataset_ref.keys()))
+    logging.debug("len(dataset_hyp)=%d, len(dataset_ref)=%d", len(dataset_hyp), len(dataset_ref))
+
+    # Store metrics for every frame for debugging.
+    per_frame_metric = {}
+
+    for dial_id, dial_hyp in dataset_hyp.items():
+        dial_ref = dataset_ref[dial_id]
+
+        if set(dial_ref["services"]) != set(dial_hyp["services"]):
+            raise ValueError(
+                "Set of services present in ground truth and predictions don't match "
+                "for dialogue with id {}".format(dial_id)
+            )
+
+        joint_metrics = [JOINT_GOAL_ACCURACY, JOINT_CAT_ACCURACY, JOINT_NONCAT_ACCURACY]
+        for turn_id, (turn_ref, turn_hyp) in enumerate(zip(dial_ref["turns"], dial_hyp["turns"])):
+            metric_collections_per_turn = collections.defaultdict(lambda: collections.defaultdict(lambda: 1.0))
+            if turn_ref["speaker"] != turn_hyp["speaker"]:
+                raise ValueError("Speakers don't match in dialogue with id {}".format(dial_id))
+
+            # Skip system turns because metrics are only computed for user turns.
+            if turn_ref["speaker"] != "USER":
+                continue
+
+            if turn_ref["utterance"] != turn_hyp["utterance"]:
+                logging.error("Ref utt: %s", turn_ref["utterance"])
+                logging.error("Hyp utt: %s", turn_hyp["utterance"])
+                raise ValueError("Utterances don't match for dialogue with id {}".format(dial_id))
+
+            hyp_frames_by_service = {frame["service"]: frame for frame in turn_hyp["frames"]}
+
+            # Calculate metrics for each frame in each user turn.
+            for frame_ref in turn_ref["frames"]:
+                service_name = frame_ref["service"]
+                if service_name not in hyp_frames_by_service:
+                    raise ValueError(
+                        "Frame for service {} not found in dialogue with id {}".format(service_name, dial_id)
+                    )
+                service = service_schemas[service_name]
+                frame_hyp = hyp_frames_by_service[service_name]
+
+                active_intent_acc = get_active_intent_accuracy(frame_ref, frame_hyp)
+                slot_tagging_f1_scores = get_slot_tagging_f1(frame_ref, frame_hyp, turn_ref["utterance"], service)
+                requested_slots_f1_scores = get_requested_slots_f1(frame_ref, frame_hyp)
+                goal_accuracy_dict = get_average_and_joint_goal_accuracy(
+                    frame_ref, frame_hyp, service, use_fuzzy_match
+                )
+
+                frame_metric = {
+                    ACTIVE_INTENT_ACCURACY: active_intent_acc,
+                    REQUESTED_SLOTS_F1: requested_slots_f1_scores.f1,
+                    REQUESTED_SLOTS_PRECISION: requested_slots_f1_scores.precision,
+                    REQUESTED_SLOTS_RECALL: requested_slots_f1_scores.recall,
+                }
+                if slot_tagging_f1_scores is not None:
+                    frame_metric[SLOT_TAGGING_F1] = slot_tagging_f1_scores.f1
+                    frame_metric[SLOT_TAGGING_PRECISION] = slot_tagging_f1_scores.precision
+                    frame_metric[SLOT_TAGGING_RECALL] = slot_tagging_f1_scores.recall
+                frame_metric.update(goal_accuracy_dict)
+
+                frame_id = "{:s}-{:03d}-{:s}".format(dial_id, turn_id, frame_hyp["service"])
+                per_frame_metric[frame_id] = frame_metric
+                # Add the frame-level metric result back to dialogues.
+                frame_hyp["metrics"] = frame_metric
+
+                # Get the domain name of the service.
+                domain_name = frame_hyp["service"].split("_")[0]
+                domain_keys = [ALL_SERVICES, frame_hyp["service"], domain_name]
+                if frame_hyp["service"] in in_domain_services:
+                    domain_keys.append(SEEN_SERVICES)
+
+                else:
+                    domain_keys.append(UNSEEN_SERVICES)
+                for domain_key in domain_keys:
+                    for metric_key, metric_value in frame_metric.items():
+                        if metric_value != NAN_VAL:
+                            if joint_acc_across_turn and metric_key in joint_metrics:
+                                metric_collections_per_turn[domain_key][metric_key] *= metric_value
+                            else:
+                                metric_collections[domain_key][metric_key].append(metric_value)
+            if joint_acc_across_turn:
+                # Conduct multiwoz style evaluation that computes joint goal accuracy
+                # across all the slot values of all the domains for each turn.
+                for domain_key in metric_collections_per_turn:
+                    for metric_key, metric_value in metric_collections_per_turn[domain_key].items():
+                        metric_collections[domain_key][metric_key].append(metric_value)
+
+    all_metric_aggregate = {}
+    for domain_key, domain_metric_vals in metric_collections.items():
+        domain_metric_aggregate = {}
+        for metric_key, value_list in domain_metric_vals.items():
+            if value_list:
+                # Metrics are macro-averaged across all frames.
+                domain_metric_aggregate[metric_key] = round(float(np.mean(value_list)) * 100.0, 2)
+            else:
+                domain_metric_aggregate[metric_key] = NAN_VAL
+        all_metric_aggregate[domain_key] = domain_metric_aggregate
+    return all_metric_aggregate, per_frame_metric
+
+
+def evaluate(
+    prediction_dir: str,
+    data_dir: str,
+    eval_dataset: str,
+    in_domain_services: set,
+    joint_acc_across_turn: bool,
+    use_fuzzy_match: bool,
+) -> dict:
+    """Calculate the DSTC8/SGD metrics for given data.
+
+    Args:
+        prediction_dir: prediction location
+        data_dir: ground truth data location.
+        eval_dataset: evaluation data split
+        in_domain_services: The set of services which are present in the training set.
+        joint_acc_across_turn: Whether to compute joint goal accuracy across turn instead of across service. Should be set to True when conducting multiwoz style evaluation.
+        use_fuzzy_match: Whether to use fuzzy string matching when comparing non-categorical slot values. Should be set to False when conducting multiwoz style evaluation.
+
+    Returns:
+        A dict mapping a metric collection name to a dict containing the values
+        for various metrics for all dialogues and all services
+    """
+
+    with open(os.path.join(data_dir, eval_dataset, "schema.json")) as f:
+        eval_services = {}
+        list_services = json.load(f)
+        for service in list_services:
+            eval_services[service["service_name"]] = service
+        f.close()
+
+    dataset_ref = get_dataset_as_dict(os.path.join(data_dir, eval_dataset, "dialogues_*.json"))
+    dataset_hyp = get_dataset_as_dict(os.path.join(prediction_dir, "*.json"))
+
+    # has ALLSERVICE, SEEN_SERVICES, UNSEEN_SERVICES, SERVICE, DOMAIN
+    all_metric_aggregate, _ = get_metrics(
+        dataset_ref, dataset_hyp, eval_services, in_domain_services, joint_acc_across_turn, use_fuzzy_match
+    )
+    if SEEN_SERVICES in all_metric_aggregate:
+        logging.info(f'Dialog metrics for {SEEN_SERVICES}  : {sorted(all_metric_aggregate[SEEN_SERVICES].items())}')
+    if UNSEEN_SERVICES in all_metric_aggregate:
+        logging.info(f'Dialog metrics for {UNSEEN_SERVICES}: {sorted(all_metric_aggregate[UNSEEN_SERVICES].items())}')
+    if ALL_SERVICES in all_metric_aggregate:
+        logging.info(f'Dialog metrics for {ALL_SERVICES}   : {sorted(all_metric_aggregate[ALL_SERVICES].items())}')
+
+    # Write the per-frame metrics values with the corrresponding dialogue frames.
+    with open(os.path.join(prediction_dir, PER_FRAME_OUTPUT_FILENAME), "w") as f:
+        json.dump(dataset_hyp, f, indent=2, separators=(",", ": "))
+        f.close()
+    return all_metric_aggregate[ALL_SERVICES]

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/input_example.py

@@ -0,0 +1,399 @@
+# 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/data_utils.py
+"""
+
+from typing import List
+
+from nemo.utils import logging
+
+__all__ = ['InputExample', 'STR_DONTCARE', 'STATUS_OFF', 'STATUS_ACTIVE', 'STATUS_DONTCARE']
+
+STR_DONTCARE = "dontcare"
+
+# These are used to represent the status of slots (off, active, dontcare) and
+# intents (off, active) in dialogue state tracking.
+STATUS_OFF = 0
+STATUS_ACTIVE = 1
+STATUS_DONTCARE = 2
+
+
+class InputExample(object):
+    """An example for training/inference."""
+
+    def __init__(
+        self,
+        schema_config: dict,
+        tokenizer: object,
+        service_schema: object = None,
+        example_id: str = "NONE",
+        example_id_num: List[int] = [],
+    ):
+        """
+        Constructs an InputExample.
+        Args:
+            schema_config: configuration
+            tokenizer: tokenizer object
+            service_schema: A ServiceSchema object wrapping the schema for the service
+                corresponding to this example.
+            example_id: Unique identifier for the example, like: 'train-1_00000-00-Restaurants_1'
+            example_id_num: dialogue_id and turn_id combined and service id combined into a list of ints,
+                like: [1, 0, 0, 18]
+        """
+        self.schema_config = schema_config
+        self.service_schema = service_schema
+        self.service_id = None
+        if service_schema:
+            self.service_id = service_schema.service_id
+        self.example_id = example_id
+        self.example_id_num = example_id_num
+        self._max_seq_length = schema_config["MAX_SEQ_LENGTH"]
+        self._tokenizer = tokenizer
+        if self._tokenizer is None:
+            raise ValueError("Must specify tokenizer")
+
+        self.user_utterance = ''
+        self.system_utterance = ''
+        # The id of each subword in the vocabulary for BERT.
+        self.utterance_ids = [0] * self._max_seq_length
+        # Denotes the identity of the sequence. Takes values 0 (schema description) and 1 (system and user utterance).
+        self.utterance_segment = [0] * self._max_seq_length
+        # Mask which takes the value 0 for padded tokens and 1 otherwise.
+        self.utterance_mask = [0] * self._max_seq_length
+        # Start and inclusive end character indices in the original utterance
+        # corresponding to the tokens. This is used to obtain the character indices
+        # from the predicted subword indices during inference.
+        # NOTE: A positive value indicates the character indices in the schema description
+        # whereas a negative value indicates the character indices in the
+        # utterance. The indices are offset by 1 to prevent ambiguity in the
+        # 0 index, which could be in either the schema description or utterance by the
+        # above convention. Now the 0 index corresponds to padded tokens.
+        self.start_char_idx = [0] * self._max_seq_length
+        self.end_char_idx = [0] * self._max_seq_length
+
+        # Id of categorical slot present in the example or 0 if not present.
+        self.categorical_slot_id = 0
+        # Id of non categorical slot present in the example or 0 if not present.
+        self.noncategorical_slot_id = 0
+        # The status of categorical slot in the example.
+        self.categorical_slot_status = STATUS_OFF
+        # The status of non categorical slot in the example.
+        self.noncategorical_slot_status = STATUS_OFF
+        # Masks out tasks not represented by example
+        self.task_mask = [0] * schema_config["NUM_TASKS"]
+
+        # The index of the starting subword corresponding to the slot span
+        # for a non-categorical slot value.
+        self.noncategorical_slot_value_start = 0
+        # The index of the ending (inclusive) subword corresponding to the slot span
+        # for a non-categorical slot value.
+        self.noncategorical_slot_value_end = 0
+
+        # Id of categorical slot value present in the example or 0 if not present.
+        self.categorical_slot_value_id = 0
+        # The status of categorical slot value in the example.
+        self.categorical_slot_value_status = STATUS_OFF
+        # Id of requested slot present in the example or 0 if not present.
+        self.requested_slot_id = 0
+        # Takes value 1 if the corresponding slot is requested, 0 otherwise.
+        self.requested_slot_status = STATUS_OFF
+
+        # ID of intent present in the example.
+        self.intent_id = 0
+        # Takes value 1 if the intent is active, 0 otherwise.
+        self.intent_status = STATUS_OFF
+
+    @property
+    def readable_summary(self):
+        """Get a readable dict that summarizes the attributes of an InputExample."""
+        seq_length = sum(self.utterance_mask)
+        utt_toks = self._tokenizer.ids_to_tokens(self.utterance_ids[:seq_length])
+        utt_tok_mask_pairs = list(zip(utt_toks, self.utterance_segment[:seq_length]))
+        active_intent = (
+            self.service_schema.get_intent_from_id(self.intent_id) if self.intent_status == STATUS_ACTIVE else ""
+        )
+        slot_values_in_state = {}
+        if self.categorical_slot_status == STATUS_ACTIVE:
+            slot_values_in_state[
+                self.service_schema.get_categorical_slot_from_id(self.categorical_slot_id)
+            ] = self.service_schema.get_categorical_slot_value_from_id(
+                self.categorical_slot_id, self.categorical_slot_value_id
+            )
+        elif self.categorical_slot_status == STATUS_DONTCARE:
+            slot_values_in_state[
+                self.service_schema.get_categorical_slot_from_id(self.categorical_slot_id)
+            ] = STR_DONTCARE
+        if self.noncategorical_slot_status == STATUS_ACTIVE:
+            slot = self.service_schema.get_non_categorical_slot_from_id(self.noncategorical_slot_id)
+            start_id = self.noncategorical_slot_value_start[idx]
+            end_id = self.noncategorical_slot_value_end[idx]
+            # Token list is consisted of the subwords that may start with "##". We
+            # remove "##" to reconstruct the original value. Note that it's not a
+            # strict restoration of the original string. It's primarily used for
+            # debugging.
+            # ex. ["san", "j", "##ose"] --> "san jose"
+            readable_value = " ".join(utt_toks[start_id : end_id + 1]).replace(" ##", "")
+            slot_values_in_state[slot] = readable_value
+        elif self.noncategorical_slot_status == STATUS_DONTCARE:
+            slot = self.service_schema.get_non_categorical_slot_from_id(self.noncategorical_slot_id)
+            slot_values_in_state[slot] = STR_DONTCARE
+
+        summary_dict = {
+            "utt_tok_mask_pairs": utt_tok_mask_pairs,
+            "utt_len": seq_length,
+            "categorical_slot_id": self.categorical_slot_id,
+            "noncategorical_slot_id": self.noncategorical_slot_id,
+            "intent_id": self.intent_id,
+            "service_name": self.service_schema.service_name,
+            "active_intent": active_intent,
+            "slot_values_in_state": slot_values_in_state,
+        }
+        return summary_dict
+
+    def add_utterance_features(
+        self, system_tokens, system_inv_alignments, user_tokens, user_inv_alignments, system_utterance, user_utterance
+    ):
+        """Add utterance related features input to InputExample.
+
+        Note: this method modifies the system tokens and user_tokens in place to
+        make their total length <= the maximum input length for BERT model.
+
+        Args:
+          system_tokens: a list of strings which represents schema description.
+          system_inv_alignments: a list of tuples which denotes the start and end
+            charater of the tpken that a bert token originates from in the original
+            schema description.
+          user_tokens: a list of strings which represents utterance.
+          user_inv_alignments: a list of tuples which denotes the start and end
+            charater of the token that a bert token originates from in the original
+            system and user utterance.
+        """
+        # Input sequence length for utterance BERT encoder
+        max_utt_len = self._max_seq_length
+
+        # Modify lengths of schema description & utterance so that length of total utt
+        # (including cls_token, setp_token, sep_token) is no more than max_utt_len
+        is_too_long = truncate_seq_pair(system_tokens, user_tokens, max_utt_len - 3)
+        if is_too_long:
+            logging.debug(
+                f'Utterance sequence truncated in example id - {self.example_id} from {len(system_tokens) + len(user_tokens)}.'
+            )
+
+        # Construct the tokens, segment mask and valid token mask which will be
+        # input to BERT, using the tokens for schema description (sequence A) and
+        # system and user utterance (sequence B).
+        utt_subword = []
+        utt_seg = []
+        utt_mask = []
+        start_char_idx = []
+        end_char_idx = []
+
+        utt_subword.append(self._tokenizer.cls_token)
+        utt_seg.append(0)
+        utt_mask.append(1)
+        start_char_idx.append(0)
+        end_char_idx.append(0)
+
+        for subword_idx, subword in enumerate(system_tokens):
+            utt_subword.append(subword)
+            utt_seg.append(0)
+            utt_mask.append(1)
+            st, en = system_inv_alignments[subword_idx]
+            start_char_idx.append(-(st + 1))
+            end_char_idx.append(-(en + 1))
+
+        utt_subword.append(self._tokenizer.sep_token)
+        utt_seg.append(0)
+        utt_mask.append(1)
+        start_char_idx.append(0)
+        end_char_idx.append(0)
+
+        for subword_idx, subword in enumerate(user_tokens):
+            utt_subword.append(subword)
+            utt_seg.append(1)
+            utt_mask.append(1)
+            st, en = user_inv_alignments[subword_idx]
+            start_char_idx.append(st + 1)
+            end_char_idx.append(en + 1)
+
+        utt_subword.append(self._tokenizer.sep_token)
+        utt_seg.append(1)
+        utt_mask.append(1)
+        start_char_idx.append(0)
+        end_char_idx.append(0)
+
+        utterance_ids = self._tokenizer.tokens_to_ids(utt_subword)
+
+        # Zero-pad up to the BERT input sequence length.
+        while len(utterance_ids) < max_utt_len:
+            utterance_ids.append(0)
+            utt_seg.append(0)
+            utt_mask.append(0)
+            start_char_idx.append(0)
+            end_char_idx.append(0)
+        self.utterance_ids = utterance_ids
+        self.utterance_segment = utt_seg
+        self.utterance_mask = utt_mask
+        self.start_char_idx = start_char_idx
+        self.end_char_idx = end_char_idx
+
+        self.user_utterance = user_utterance
+        self.system_utterance = system_utterance
+
+    def make_copy(self):
+        """Make a copy of the current example with utterance features."""
+        new_example = InputExample(
+            schema_config=self.schema_config,
+            service_schema=self.service_schema,
+            example_id=self.example_id,
+            example_id_num=self.example_id_num.copy(),
+            tokenizer=self._tokenizer,
+        )
+        return new_example
+
+    def make_copy_of_categorical_features(self):
+        """Make a copy of the current example with utterance and categorical features."""
+        new_example = self.make_copy()
+
+        new_example.categorical_slot_status = self.categorical_slot_status
+        return new_example
+
+    def make_copy_of_non_categorical_features(self):
+        """Make a copy of the current example with utterance features and non categorical features."""
+        new_example = self.make_copy()
+        new_example.noncategorical_slot_id = self.noncategorical_slot_id
+        new_example.noncategorical_slot_status = self.noncategorical_slot_status
+        new_example.utterance_ids = list(self.utterance_ids)
+        new_example.utterance_segment = list(self.utterance_segment)
+        new_example.utterance_mask = list(self.utterance_mask)
+        new_example.start_char_idx = list(self.start_char_idx)
+        new_example.end_char_idx = list(self.end_char_idx)
+        new_example.user_utterance = self.user_utterance
+        new_example.system_utterance = self.system_utterance
+        new_example.noncategorical_slot_status = self.noncategorical_slot_status
+        new_example.noncategorical_slot_value_start = self.noncategorical_slot_value_start
+        new_example.noncategorical_slot_value_end = self.noncategorical_slot_value_end
+        return new_example
+
+    def add_categorical_slots(self, state_update: dict):
+        """Add features for categorical slots.
+        Args:
+            state_update: slot value pairs of the state update
+        """
+
+        categorical_slots = self.service_schema.categorical_slots
+        if not categorical_slots:
+            return
+        slot = categorical_slots[self.categorical_slot_id]
+        values = state_update.get(slot, [])
+
+        if not values:
+            self.categorical_slot_status = STATUS_OFF
+        elif values[0] == STR_DONTCARE:
+            self.categorical_slot_status = STATUS_DONTCARE
+        else:
+            self.categorical_slot_status = STATUS_ACTIVE
+            self.categorical_slot_value_status = (
+                self.categorical_slot_value_id == self.service_schema.get_categorical_slot_value_id(slot, values[0])
+            )
+
+    def add_noncategorical_slots(self, state_update: dict, system_span_boundaries: dict, user_span_boundaries: dict):
+        """Add features for non-categorical slots.
+        Args:
+            state_update: slot value pairs of state update
+            system_span_boundaries: span boundaries of schema description
+            user_span_boundaries: span boundaries of utterance 
+        """
+
+        noncategorical_slots = self.service_schema.non_categorical_slots
+        slot = noncategorical_slots[self.noncategorical_slot_id]
+
+        values = state_update.get(slot, [])
+        if not values:
+            self.noncategorical_slot_status = STATUS_OFF
+        elif values[0] == STR_DONTCARE:
+            self.noncategorical_slot_status = STATUS_DONTCARE
+        else:
+            self.noncategorical_slot_status = STATUS_ACTIVE
+            # Add indices of the start and end tokens for the first encountered
+            # value. Spans in user utterance are prioritized over the system
+            # utterance. If a span is not found, the slot value is ignored.
+            if slot in user_span_boundaries:
+                start, end = user_span_boundaries[slot]
+            elif slot in system_span_boundaries:
+                start, end = system_span_boundaries[slot]
+            else:
+                # A span may not be found because the value was cropped out or because
+                # the value was mentioned earlier in the dialogue. Since this model
+                # only makes use of the last two utterances to predict state updates,
+                # it will fail in such cases.
+                logging.debug(
+                    f'"Slot values {str(values)} not found in user or system utterance in example with id - {self.example_id}.'
+                )
+                start = 0
+                end = 0
+            self.noncategorical_slot_value_start = start
+            self.noncategorical_slot_value_end = end
+
+    def add_requested_slots(self, frame: dict):
+        """Add requested slots to InputExample
+        Args:
+            frame: frame object from which requested slots are extracted
+        """
+        all_slots = self.service_schema.slots
+        slot = all_slots[self.requested_slot_id]
+        if slot in frame["state"]["requested_slots"]:
+            self.requested_slot_status = STATUS_ACTIVE
+
+    def add_intents(self, frame):
+        """Add intents to InputExample
+        Args:
+            frame: frame object from which intents are extracted
+        """
+        all_intents = self.service_schema.intents
+        intent = all_intents[self.intent_id]
+        if intent == frame["state"]["active_intent"]:
+            self.intent_status = STATUS_ACTIVE
+
+
+# Modified from run_classifier._truncate_seq_pair in the public bert model repo.
+# https://github.com/google-research/bert/blob/master/run_classifier.py.
+def truncate_seq_pair(tokens_a: List[int], tokens_b: List[int], max_length: int) -> bool:
+    """Truncate a seq pair in place so that their total length <= max_length.
+    Args:
+        tokens_a: first token sequence
+        tokens_b: second token sequence
+        max_length: truncated sequence length
+    Returns:
+        is_too_long: whether combined sequences exceed maximum sequence length
+    """
+    is_too_long = False
+    # This is a simple heuristic which will always truncate the longer sequence
+    # one token at a time. This makes more sense than truncating an equal percent
+    # of tokens from each, since if one sequence is very short then each token
+    # that's truncated likely contains more information than a longer sequence.
+    while True:
+        total_length = len(tokens_a) + len(tokens_b)
+        if total_length <= max_length:
+            break
+        is_too_long = True
+        if len(tokens_a) > len(tokens_b):
+            tokens_a.pop()
+        else:
+            tokens_b.pop()
+    return is_too_long

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/metrics.py

@@ -0,0 +1,341 @@
+# 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.
+
+"""Evaluation metrics for Schema-guided dialogue.
+
+This library provides functions for calculating the evaluation metrics for a
+single dialogue. The following metrics are defined:
+
+(1) Active intent accuracy: The fraction of user turns for which the active
+  intent has been correctly predicted.
+(2) Slot tagging F1: The macro-averaged F1 score for tagging slot values for
+  non-categorical slots. This metric is optional to report in the final paper
+  if participants decide not to use slot tagging.
+(3) Requested slots F1: The macro-averaged F1 score for requested slots over the
+  turns. For a turn, if there are no requested slots in both the ground truth
+  and the prediction, that turn is skipped. The reported number is the average
+  F1 score for all un-skipped user turns. This metric is optional to report in
+  the final paper.
+(4) Average goal accuracy: For each turn, participants must predict a single
+  value for each slot present in the dialogue state. The slots which have a
+  non-empty assignment in the ground truth dialogue state are only considered.
+  This is the average accuracy of predicting the value of a slot correctly. A
+  fuzzy matching based score is used for non-categorical slots.
+(5) Joint goal accuracy: This is the average accuracy of predicting all slot
+  assignments for a turn correctly. A fuzzy matching based score is used for
+  non-categorical slots. This is the primary evaluation metric used for ranking
+  submissions. More details to follow with the evaluation script.
+
+This file contains code artifacts adapted from the original implementation:
+https://github.com/google-research/google-research/blob/master/schema_guided_dst/metrics.py
+"""
+
+import collections
+
+import numpy as np
+from rapidfuzz import fuzz
+
+F1Scores = collections.namedtuple("F1Scores", ["f1", "precision", "recall"])
+
+# Evaluation and other relevant metrics for DSTC8/SGD Schema-guided DST.
+# (1) Active intent accuracy.
+ACTIVE_INTENT_ACCURACY = "active_intent_accuracy"
+# (2) Slot tagging F1.
+SLOT_TAGGING_F1 = "slot_tagging_f1"
+SLOT_TAGGING_PRECISION = "slot_tagging_precision"
+SLOT_TAGGING_RECALL = "slot_tagging_recall"
+# (3) Requested slots F1.
+REQUESTED_SLOTS_F1 = "requested_slots_f1"
+REQUESTED_SLOTS_PRECISION = "requested_slots_precision"
+REQUESTED_SLOTS_RECALL = "requested_slots_recall"
+# (4) Average goal accuracy.
+AVERAGE_GOAL_ACCURACY = "average_goal_accuracy"
+AVERAGE_CAT_ACCURACY = "average_cat_accuracy"
+AVERAGE_NONCAT_ACCURACY = "average_noncat_accuracy"
+# (5) Joint goal accuracy.
+JOINT_GOAL_ACCURACY = "joint_goal_accuracy"
+JOINT_CAT_ACCURACY = "joint_cat_accuracy"
+JOINT_NONCAT_ACCURACY = "joint_noncat_accuracy"
+
+
+AVERAGE_CAT_STATUS_ACCURACY = "average_cat_status_accuracy"
+AVERAGE_CAT_VALUE_ACCURACY = "average_cat_value_accuracy"
+AVERAGE_NONCAT_STATUS_ACCURACY = "average_noncat_status_accuracy"
+AVERAGE_NONCAT_VALUE_ACCURACY = "average_noncat_value_accuracy"
+
+JOINT_CAT_STATUS_ACCURACY = "joint_cat_status_accuracy"
+JOINT_CAT_VALUE_ACCURACY = "joint_cat_value_accuracy"
+JOINT_NONCAT_STATUS_ACCURACY = "joint_noncat_status_accuracy"
+JOINT_NONCAT_VALUE_ACCURACY = "joint_noncat_value_accuracy"
+
+
+NAN_VAL = "NA"
+
+
+def compute_f1(list_ref, list_hyp):
+    """Compute F1 score from reference (grouth truth) list and hypothesis list.
+    Args:
+      list_ref: List of true elements.
+      list_hyp: List of postive (retrieved) elements.
+    Returns:
+      A F1Scores object containing F1, precision, and recall scores.
+    """
+
+    ref = collections.Counter(list_ref)
+    hyp = collections.Counter(list_hyp)
+    true = sum(ref.values())
+    positive = sum(hyp.values())
+    true_positive = sum((ref & hyp).values())
+    precision = float(true_positive) / positive if positive else 1.0
+    recall = float(true_positive) / true if true else 1.0
+    if precision + recall > 0.0:
+        f1 = 2.0 * precision * recall / (precision + recall)
+    else:  # The F1-score is defined to be 0 if both precision and recall are 0.
+        f1 = 0.0
+
+    return F1Scores(f1=f1, precision=precision, recall=recall)
+
+
+def fuzzy_string_match(str_ref, str_hyp):
+    """Returns fuzzy string similarity score in range [0.0, 1.0].
+    Args:
+      str_ref: reference string
+      str_hyp: hypothesis string
+    Returns:
+      fuzzy string similarity
+    """
+
+    # The higher the score, the higher the similarity between the two strings.
+    return fuzz.token_sort_ratio(str_ref, str_hyp) / 100.0
+
+
+def noncat_slot_value_match(str_ref_list, str_hyp, use_fuzzy_match):
+    """Calculate non-categorical slots correctness.
+    Args:
+      str_ref_list: a list of reference strings.
+      str_hyp: the hypothesis string.
+      use_fuzzy_match: whether to use fuzzy string matching.
+    Returns:
+      score: The highest fuzzy string match score of the references and hypotheis.
+    """
+    score = 0.0
+    for str_ref in str_ref_list:
+        if use_fuzzy_match:
+            match_score = fuzzy_string_match(str_ref, str_hyp)
+        else:
+            match_score = float(str_ref == str_hyp)
+        score = max(score, match_score)
+    return score
+
+
+def compare_slot_values(slot_values_ref, slot_values_hyp, service, use_fuzzy_match):
+    """Compare and get correctness of goal state's slot_values.
+
+    Args:
+      slot_values_ref: goal state slot_values from reference (ground truth).
+      slot_values_hyp: goal state slot_values from hypothesis (prediction).
+      service: a service data structure in the schema. We use it to obtain the
+        list of slots in the service and infer whether a slot is categorical.
+      use_fuzzy_match: whether to use fuzzy string matching for non-categorical
+        slot values
+
+    Returns:
+      list_cor: list of corectness scores, each corresponding to one slot in the
+          service. The score is a float either 0.0 or 1.0 for categorical slot,
+          and in range [0.0, 1.0] for non-categorical slot.
+      slot_active: list indicating whether the element in list_cor corresponds to
+          an active ground-truth slot.
+      slot_cat: list indicating whether the element in list_cor corresponds to a
+          categorical slot.
+      list_cor_status: list of correct slot statuses 
+      list_cor_value: list of correctness score only for active slots. Monactive slots are assigned -1.
+    """
+    list_cor = []
+    list_cor_status = []
+    list_cor_value = []
+    slot_active = []
+    slot_cat = []
+
+    for slot in service["slots"]:
+        slot_name = slot["name"]
+        slot_cat.append(slot["is_categorical"])
+
+        if slot_name in slot_values_ref:  # REF=active
+            slot_active.append(True)
+            if slot_name in slot_values_hyp:  # HYP=active, apply matching
+                value_ref_list = slot_values_ref[slot_name]
+                value_hyp = slot_values_hyp[slot_name][0]
+                if slot["is_categorical"]:
+                    cor = float(value_ref_list[0] == value_hyp)
+                else:
+                    cor = noncat_slot_value_match(value_ref_list, value_hyp, use_fuzzy_match)
+                list_cor.append(cor)
+                list_cor_status.append(1.0)
+                list_cor_value.append(cor)
+            else:  # HYP=off
+                list_cor.append(0.0)
+                list_cor_status.append(0.0)
+                list_cor_value.append(-1.0)
+        else:  # REF=off
+            slot_active.append(False)
+            if slot_name in slot_values_hyp:  # HYP=active
+                list_cor.append(0.0)
+                list_cor_status.append(0.0)
+            else:  # HYP=off
+                list_cor.append(1.0)
+                list_cor_status.append(1.0)
+            list_cor_value.append(-1.0)
+
+    assert len(list_cor) == len(service["slots"])
+    assert len(slot_active) == len(service["slots"])
+    assert len(slot_cat) == len(service["slots"])
+    return list_cor, slot_active, slot_cat, list_cor_status, list_cor_value
+
+
+def get_active_intent_accuracy(frame_ref, frame_hyp):
+    """Get active intent accuracy of a frame.
+
+    Args:
+      frame_ref: single semantic frame from reference (ground truth) file.
+      frame_hyp: single semantic frame from hypothesis (prediction) file.
+
+    Returns:
+      1.0 if the intent prediction is correct, otherwise 0.0.
+    """
+    return float(frame_ref["state"]["active_intent"] == frame_hyp["state"]["active_intent"])
+
+
+def get_slot_tagging_f1(frame_ref, frame_hyp, utt, service):
+    """Get slot tagging (non-categorical slots only) F1 scores of a frame.
+
+    Args:
+      frame_ref: single semantic frame from reference (ground truth) file.
+      frame_hyp: single semantic frame from hypothesis (prediction) file.
+      utt: user utterance. Slot tagging annotations are the character positions in
+        the utterance.
+      service: a service data structure in the schema. We use it to infer whether
+        a slot is non-categorical.
+
+    Returns:
+      A F1Scores object containing F1, precision, and recall scores.
+    """
+    list_noncat_slots = [s["name"] for s in service["slots"] if not s["is_categorical"]]
+    if "slots" not in frame_hyp:
+        return None
+    else:
+        list_ref = [
+            (s["slot"], utt[s["start"] : s["exclusive_end"]])
+            for s in frame_ref["slots"]
+            if s["slot"] in list_noncat_slots
+        ]
+        list_hyp = [
+            (s["slot"], utt[s["start"] : s["exclusive_end"]])
+            for s in frame_hyp["slots"]
+            if s["slot"] in list_noncat_slots
+        ]
+        return compute_f1(list_ref, list_hyp)
+
+
+def get_requested_slots_f1(frame_ref, frame_hyp):
+    """Get requested slots F1 scores of a frame.
+
+    Args:
+      frame_ref: single semantic frame from reference (ground truth) file.
+      frame_hyp: single semantic frame from hypothesis (prediction) file.
+
+    Returns:
+      A F1Scores object containing F1, precision, and recall scores.
+    """
+    return compute_f1(frame_ref["state"]["requested_slots"], frame_hyp["state"]["requested_slots"])
+
+
+def get_average_and_joint_goal_accuracy(frame_ref, frame_hyp, service, use_fuzzy_match):
+    """Get average and joint goal accuracies of a frame.
+
+    Args:
+      frame_ref: single semantic frame from reference (ground truth) file.
+      frame_hyp: single semantic frame from hypothesis (prediction) file.
+      service: a service data structure in the schema. We use it to obtain the
+        list of slots in the service and infer whether a slot is categorical.
+      use_fuzzy_match: whether to use fuzzy string matching for comparing
+        non-categorical slot values.
+
+    Returns:
+      goal_acc: a dict whose values are average / joint
+          all-goal / categorical-goal / non-categorical-goal accuracies.
+    """
+    goal_acc = {}
+
+    list_acc, slot_active, slot_cat, list_status_acc, list_value_acc = compare_slot_values(
+        frame_ref["state"]["slot_values"], frame_hyp["state"]["slot_values"], service, use_fuzzy_match
+    )
+
+    # (4) Average goal accuracy.
+    active_acc = [acc for acc, active in zip(list_acc, slot_active) if active]
+    goal_acc[AVERAGE_GOAL_ACCURACY] = np.mean(active_acc) if active_acc else NAN_VAL
+    # (4-a) categorical.
+    active_cat_acc = [acc for acc, active, cat in zip(list_acc, slot_active, slot_cat) if active and cat]
+    goal_acc[AVERAGE_CAT_ACCURACY] = np.mean(active_cat_acc) if active_cat_acc else NAN_VAL
+    # (4-b) non-categorical.
+    active_noncat_acc = [acc for acc, active, cat in zip(list_acc, slot_active, slot_cat) if active and not cat]
+    goal_acc[AVERAGE_NONCAT_ACCURACY] = np.mean(active_noncat_acc) if active_noncat_acc else NAN_VAL
+
+    # (5) Joint goal accuracy.
+    goal_acc[JOINT_GOAL_ACCURACY] = np.prod(list_acc) if list_acc else NAN_VAL
+    # (5-a) categorical.
+    cat_acc = [acc for acc, cat in zip(list_acc, slot_cat) if cat]
+    goal_acc[JOINT_CAT_ACCURACY] = np.prod(cat_acc) if cat_acc else NAN_VAL
+    # (5-b) non-categorical.
+    noncat_acc = [acc for acc, cat in zip(list_acc, slot_cat) if not cat]
+    goal_acc[JOINT_NONCAT_ACCURACY] = np.prod(noncat_acc) if noncat_acc else NAN_VAL
+
+    # !!!!!!!!!!DEBUG!!!!!!!!!!!!!
+    # cat status acc for both active and non active
+    active_cat_status_acc = [acc for acc, active, cat in zip(list_status_acc, slot_active, slot_cat) if cat and active]
+    goal_acc[AVERAGE_CAT_STATUS_ACCURACY] = np.mean(active_cat_status_acc) if active_cat_status_acc else NAN_VAL
+    # joint cat status acc for both active and non active
+    cat_status_acc = [acc for acc, cat in zip(list_status_acc, slot_cat) if cat]
+    goal_acc[JOINT_CAT_STATUS_ACCURACY] = np.prod(cat_status_acc) if cat_status_acc else NAN_VAL
+
+    # non cat status acc for both active and non active
+    active_noncat_status_acc = [
+        acc for acc, active, cat in zip(list_status_acc, slot_active, slot_cat) if not cat and active
+    ]
+    goal_acc[AVERAGE_NONCAT_STATUS_ACCURACY] = (
+        np.mean(active_noncat_status_acc) if active_noncat_status_acc else NAN_VAL
+    )
+    # joint non cat status acc for both active and non active
+    noncat_status_acc = [acc for acc, cat in zip(list_status_acc, slot_cat) if not cat]
+    goal_acc[JOINT_NONCAT_STATUS_ACCURACY] = np.prod(noncat_status_acc) if noncat_status_acc else NAN_VAL
+
+    # cat value acc for both active and non active
+    active_cat_val_acc = [
+        acc for acc, active, cat in zip(list_value_acc, slot_active, slot_cat) if cat and acc > -0.5 and active
+    ]
+    goal_acc[AVERAGE_CAT_VALUE_ACCURACY] = np.mean(active_cat_val_acc) if active_cat_val_acc else NAN_VAL
+    # joint cat value acc for both active and non active
+    cat_val_acc = [acc for acc, cat in zip(list_value_acc, slot_cat) if cat and acc > -0.5]
+    goal_acc[JOINT_CAT_VALUE_ACCURACY] = np.prod(cat_val_acc) if cat_val_acc else NAN_VAL
+
+    # cat non value acc for both active and non active
+    active_noncat_val_acc = [
+        acc for acc, active, cat in zip(list_value_acc, slot_active, slot_cat) if not cat and acc > -0.5 and active
+    ]
+    goal_acc[AVERAGE_NONCAT_VALUE_ACCURACY] = np.mean(active_noncat_val_acc) if active_noncat_val_acc else NAN_VAL
+    # joint non cat value acc for both active and non active
+    noncat_val_acc = [acc for acc, cat in zip(list_value_acc, slot_cat) if not cat and acc > -0.5]
+    goal_acc[JOINT_NONCAT_VALUE_ACCURACY] = np.prod(noncat_val_acc) if noncat_val_acc else NAN_VAL
+
+    return goal_acc

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/prediction_utils.py

@@ -0,0 +1,251 @@
+# 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.
+
+"""
+Prediction and evaluation-related utility functions.
+This file contains code artifacts adapted from the original implementation:
+https://github.com/google-research/google-research/blob/master/schema_guided_dst/baseline/pred_utils.py
+"""
+
+import json
+import os
+from collections import OrderedDict, defaultdict
+from typing import Dict, List, Optional
+
+from nemo.collections.nlp.data.dialogue_state_tracking.sgd.input_example import (
+    STATUS_ACTIVE,
+    STATUS_DONTCARE,
+    STR_DONTCARE,
+)
+from nemo.utils import logging
+
+REQ_SLOT_THRESHOLD = 0.5
+
+
+__all__ = ['write_predictions_to_file']
+
+
+def set_cat_slot(predictions_status: dict, predictions_value: dict, cat_slot_values: Dict[str, List[str]]) -> dict:
+    """
+    Extract predicted categorical slot information 
+    Args:
+        predictions_status: predicted statuses
+        predictions_value: predicted slot values
+        cat_slot_values: possible categorical slots and their potential values for this service
+    Returns:
+        out_dict: predicted slot value pairs
+    """
+    out_dict = {}
+    for slot_idx, slot in enumerate(cat_slot_values):
+        slot_status = predictions_status[slot_idx][0]["cat_slot_status"]
+        if slot_status == STATUS_DONTCARE:
+            out_dict[slot] = STR_DONTCARE
+        elif slot_status == STATUS_ACTIVE:
+            tmp = predictions_value[slot_idx]
+            value_idx = max(tmp, key=lambda k: tmp[k]['cat_slot_value_status'][0].item())
+            out_dict[slot] = cat_slot_values[slot][value_idx]
+    return out_dict
+
+
+def set_noncat_slot(
+    predictions_status: dict,
+    predictions_value: dict,
+    non_cat_slots: List[str],
+    user_utterance: str,
+    sys_slots_agg: Optional[dict] = None,
+) -> dict:
+    """
+    Extract predicted non categorical slot information 
+    Args:
+        predictions_status: predicted statuses
+        predictions_value: predicted slot values
+        non_cat_slots: list of possible non categorical slots for this service
+        user_utterance: system and user utterance
+        sys_slots_agg: system retrieval lookup table. Contains for each slot the most recent value seen in the history
+    Returns:
+        out_dict: predicted slot value pairs
+    """
+    out_dict = {}
+    for slot_idx, slot in enumerate(non_cat_slots):
+        slot_status = predictions_status[slot_idx][0]["noncat_slot_status"]
+        if slot_status == STATUS_DONTCARE:
+            out_dict[slot] = STR_DONTCARE
+        elif slot_status == STATUS_ACTIVE:
+            tok_start_idx = predictions_value[slot_idx][0]["noncat_slot_start"]
+            tok_end_idx = predictions_value[slot_idx][0]["noncat_slot_end"]
+            ch_start_idx = predictions_value[slot_idx][0]["noncat_alignment_start"][tok_start_idx]
+            ch_end_idx = predictions_value[slot_idx][0]["noncat_alignment_end"][tok_end_idx]
+            if ch_start_idx > 0 and ch_end_idx > 0:
+                # Add span from the utterance.
+                out_dict[slot] = user_utterance[ch_start_idx - 1 : ch_end_idx]
+            elif sys_slots_agg and slot in sys_slots_agg:
+                # system retrieval
+                out_dict[slot] = sys_slots_agg[slot]
+    return out_dict
+
+
+def get_predicted_dialog(dialog: dict, all_predictions: dict, schemas: object, state_tracker: str) -> dict:
+    """Overwrite the labels in the turn with the predictions from the model. For test set, these labels are missing from the data and hence they are added. 
+    Args:
+        dialog: ground truth dialog
+        all_predictions: predictions
+        schemas: schema object of all services of all datasets
+        state_tracker: state tracker option, e.g. nemotracker
+    Returns:
+        dialog: dialog overwritten with prediction information
+    """
+    dialog_id = dialog["dialogue_id"]
+    if state_tracker == "baseline":
+        sys_slots_agg = {}
+    else:
+        sys_slots_agg = defaultdict(OrderedDict)
+    all_slot_values = defaultdict(dict)
+    for turn_idx, turn in enumerate(dialog["turns"]):
+        if turn["speaker"] == "SYSTEM" and state_tracker == 'nemotracker':
+            for frame in turn["frames"]:
+                if frame["service"] not in sys_slots_agg:
+                    sys_slots_agg[frame["service"]] = OrderedDict()
+                for action in frame["actions"]:
+                    if action["slot"] and len(action["values"]) > 0:
+                        sys_slots_agg[frame["service"]][action["slot"]] = action["values"][0]
+        if turn["speaker"] == "USER":
+            user_utterance = turn["utterance"]
+            system_utterance = dialog["turns"][turn_idx - 1]["utterance"] if turn_idx else ""
+            system_user_utterance = system_utterance + ' ' + user_utterance
+            turn_id = "{:02d}".format(turn_idx)
+            for frame in turn["frames"]:
+
+                predictions = all_predictions[(dialog_id, turn_id, frame["service"])]
+                slot_values = all_slot_values[frame["service"]]
+                service_schema = schemas.get_service_schema(frame["service"])
+                # Remove the slot spans and state if present.
+                frame.pop("slots", None)
+                frame.pop("state", None)
+
+                # The baseline model doesn't predict slot spans. Only state predictions
+                # are added.
+                state = {}
+
+                # Add prediction for active intent. No Offset is subtracted since schema has now NONE intent at index 0
+                state["active_intent"] = get_predicted_intent(
+                    predictions=predictions[0], intents=service_schema.intents
+                )
+                # Add prediction for requested slots.
+                state["requested_slots"] = get_requested_slot(predictions=predictions[1], slots=service_schema.slots)
+
+                # Add prediction for user goal (slot values).
+                # Categorical slots.
+                cat_out_dict = set_cat_slot(
+                    predictions_status=predictions[2],
+                    predictions_value=predictions[3],
+                    cat_slot_values=service_schema.categorical_slot_values,
+                )
+                for k, v in cat_out_dict.items():
+                    slot_values[k] = v
+
+                # Non-categorical slots.
+                noncat_out_dict = set_noncat_slot(
+                    predictions_status=predictions[4],
+                    predictions_value=predictions[5],
+                    non_cat_slots=service_schema.non_categorical_slots,
+                    user_utterance=system_user_utterance,
+                    sys_slots_agg=sys_slots_agg.get(frame["service"], None),
+                )
+                for k, v in noncat_out_dict.items():
+                    slot_values[k] = v
+                # Create a new dict to avoid overwriting the state in previous turns
+                # because of use of same objects.
+                state["slot_values"] = {s: [v] for s, v in slot_values.items()}
+                frame["state"] = state
+    return dialog
+
+
+def get_predicted_intent(predictions: dict, intents: List[str]) -> str:
+    """
+    Returns intent name with maximum score
+    Args:
+        predictions: predictions
+        intents: list of possible intents for this service
+    Returns:
+        intent: predicted intent
+    """
+    assert len(predictions) == len(intents)
+    active_intent_id = max(predictions, key=lambda k: predictions[k][0]['intent_status'])
+    intent = intents[active_intent_id]
+    return intent
+
+
+def get_requested_slot(predictions: dict, slots: List[str]) -> List[str]:
+    """
+    Returns list of slots which are predicted to be requested
+    Args:
+        predictions: predictions
+        slots: list of possible slots
+    Returns:
+        requested_slots: list of requested slots
+    """
+    active_indices = [k for k in predictions if predictions[k][0]["req_slot_status"] > REQ_SLOT_THRESHOLD]
+    requested_slots = list(map(lambda k: slots[k], active_indices))
+    return requested_slots
+
+
+def write_predictions_to_file(
+    predictions: List[dict],
+    input_json_files: List[str],
+    output_dir: str,
+    schemas: object,
+    state_tracker: str,
+    eval_debug: bool,
+    in_domain_services: set,
+):
+    """Save predicted dialogues as json files.
+
+    Args:
+        predictions: An iterator containing model predictions. This is the output of
+            the predict method in the estimator.
+        input_json_files: A list of json paths containing the dialogues to run
+            inference on.
+        output_dir: The directory where output json files will be created.
+        schemas: Schemas to all services in the dst dataset
+        state_tracker: state tracker option
+        eval_debug: output evaluation debugging information
+        in_domain_services: in domain services
+    """
+    logging.info(f"Writing predictions to {output_dir} started.")
+
+    # Index all predictions.
+    all_predictions = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))
+    for idx, prediction in enumerate(predictions):
+        eval_dataset, dialog_id, turn_id, service_name, model_task, slot_intent_id, value_id = prediction[
+            'example_id'
+        ].split('-')
+        all_predictions[(dialog_id, turn_id, service_name)][int(model_task)][int(slot_intent_id)][
+            int(value_id)
+        ] = prediction
+    logging.info(f'Predictions for {idx} examples in {eval_dataset} dataset are getting processed.')
+
+    # Read each input file and write its predictions.
+    for input_file_path in input_json_files:
+        with open(input_file_path) as f:
+            dialogs = json.load(f)
+            logging.debug(f'{input_file_path} file is loaded')
+            pred_dialogs = []
+            for d in dialogs:
+                pred_dialog = get_predicted_dialog(d, all_predictions, schemas, state_tracker)
+                pred_dialogs.append(pred_dialog)
+        input_file_name = os.path.basename(input_file_path)
+        output_file_path = os.path.join(output_dir, input_file_name)
+        with open(output_file_path, "w") as f:
+            json.dump(pred_dialogs, f, indent=2, separators=(",", ": "), sort_keys=True)

tts/nemo/collections/nlp/data/dialogue_state_tracking/sgd/schema.py

@@ -0,0 +1,222 @@
+# 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.
+
+"""
+Wrappers for schemas of different services.
+This file contains code artifacts adapted from the original implementation:
+https://github.com/google-research/google-research/blob/master/schema_guided_dst/schema.py
+"""
+
+import json
+from typing import List, Optional, Union
+
+from nemo.utils import logging
+
+__all__ = ['Schema']
+
+
+class ServiceSchema(object):
+    """A wrapper for schema for a service."""
+
+    def __init__(self, schema_json: dict, service_id: Optional[int] = None):
+        """
+        Constructor for ServiceSchema.
+        Args:
+            schema_json: schema json dict
+            service_id: service ID
+        """
+        self._service_name = schema_json["service_name"]
+        self._description = schema_json["description"]
+        self._schema_json = schema_json
+        self._service_id = service_id
+
+        # Construct the vocabulary for intents, slots, categorical slots,
+        # non-categorical slots and categorical slot values.
+        self._intents = ["NONE"] + sorted(i["name"] for i in schema_json["intents"])
+        self._intent_descriptions = {i["name"]: i["description"] for i in schema_json["intents"]}
+        self._intent_descriptions["NONE"] = "none"
+        self._slots = sorted(s["name"] for s in schema_json["slots"])
+        self._slots_descriptions = {s["name"]: s["description"] for s in schema_json["slots"]}
+        self._categorical_slots = sorted(
+            s["name"] for s in schema_json["slots"] if s["is_categorical"] and s["name"] in self.state_slots
+        )
+        self._non_categorical_slots = sorted(
+            s["name"] for s in schema_json["slots"] if not s["is_categorical"] and s["name"] in self.state_slots
+        )
+        slot_schemas = {s["name"]: s for s in schema_json["slots"]}
+        categorical_slot_values = {}
+        categorical_slot_value_ids = {}
+        categorical_slot_ids = {}
+        non_categorical_slot_ids = {}
+        for slot_id, slot in enumerate(self._categorical_slots):
+            slot_schema = slot_schemas[slot]
+            values = sorted(slot_schema["possible_values"])
+            categorical_slot_values[slot] = values
+            value_ids = {value: idx for idx, value in enumerate(values)}
+            categorical_slot_value_ids[slot] = value_ids
+            categorical_slot_ids[slot] = slot_id
+
+        for slot_id, slot in enumerate(self._non_categorical_slots):
+            non_categorical_slot_ids[slot] = slot_id
+
+        self._categorical_slot_values = categorical_slot_values
+        self._categorical_slot_value_ids = categorical_slot_value_ids
+
+        self._categorical_slot_ids = categorical_slot_ids
+        self._non_categorical_slot_ids = non_categorical_slot_ids
+
+    @property
+    def schema_json(self) -> dict:
+        """Returns schema json dictionary"""
+        return self._schema_json
+
+    @property
+    def state_slots(self) -> set:
+        """Set of slots which are permitted to be in the dialogue state."""
+        state_slots = set()
+        for intent in self._schema_json["intents"]:
+            state_slots.update(intent["required_slots"])
+            state_slots.update(intent["optional_slots"])
+        return state_slots
+
+    @property
+    def service_name(self):
+        return self._service_name
+
+    @property
+    def service_id(self):
+        return self._service_id
+
+    @property
+    def description(self):
+        return self._description
+
+    @property
+    def slots(self):
+        return self._slots
+
+    @property
+    def intents(self):
+        return self._intents
+
+    @property
+    def intent_descriptions(self):
+        return self._intent_descriptions
+
+    @property
+    def slot_descriptions(self):
+        return self._slots_descriptions
+
+    @property
+    def categorical_slots(self):
+        return self._categorical_slots
+
+    @property
+    def non_categorical_slots(self):
+        return self._non_categorical_slots
+
+    @property
+    def categorical_slot_values(self):
+        return self._categorical_slot_values
+
+    def get_categorical_slot_values(self, slot):
+        return self._categorical_slot_values[slot]
+
+    def get_slot_from_id(self, slot_id):
+        return self._slots[slot_id]
+
+    def get_intent_from_id(self, intent_id):
+        return self._intents[intent_id]
+
+    def get_categorical_slot_from_id(self, slot_id):
+        return self._categorical_slots[slot_id]
+
+    def get_non_categorical_slot_from_id(self, slot_id):
+        return self._non_categorical_slots[slot_id]
+
+    def get_categorical_slot_value_from_id(self, slot_id, value_id):
+        slot = self._categorical_slots[slot_id]
+        return self._categorical_slot_values[slot][value_id]
+
+    def get_categorical_slot_value_id(self, slot, value):
+        return self._categorical_slot_value_ids[slot][value]
+
+    def get_categorical_slot_id(self, slot):
+        return self._categorical_slot_ids[slot]
+
+    def get_non_categorical_slot_id(self, slot):
+        return self._non_categorical_slot_ids[slot]
+
+
+class Schema(object):
+    """Wrapper for schemas for all services in a dataset."""
+
+    def __init__(self, schema_json_paths: Union[str, List[str]]):
+        """
+        schema_json_paths: list of .json path to schema files of a single str with path to the json file.
+        """
+        # Load the schema from the json file.
+        if isinstance(schema_json_paths, str):
+            with open(schema_json_paths, "r") as f:
+                all_schemas = json.load(f)
+                f.close()
+        else:
+            # load multiple schemas from the list of the json files
+            all_schemas = []
+            completed_services = []
+            for schema_json_path in schema_json_paths:
+                with open(schema_json_path, "r") as f:
+                    schemas = json.load(f)
+                    f.close()
+                    logging.debug("Num of services in %s: %s", schema_json_path, len(schemas))
+
+                for service in schemas:
+                    if service['service_name'] not in completed_services:
+                        completed_services.append(service['service_name'])
+                        all_schemas.append(service)
+
+        self._services = sorted(schema["service_name"] for schema in all_schemas)
+        self._services_vocab = {v: k for k, v in enumerate(self._services)}
+        self._services_id_to_vocab = {v: k for k, v in self._services_vocab.items()}
+        service_schemas = {}
+        for schema in all_schemas:
+            service = schema["service_name"]
+            service_schemas[service] = ServiceSchema(schema, service_id=self.get_service_id(service))
+
+        self._service_schemas = service_schemas
+        self._schemas = all_schemas
+        self._slots_relation_list = {}
+
+    def get_service_id(self, service: str):
+        return self._services_vocab[service]
+
+    def get_service_from_id(self, service_id: int):
+        return self._services[service_id]
+
+    def get_service_schema(self, service: str):
+        return self._service_schemas[service]
+
+    @property
+    def services(self):
+        return self._services
+
+    def save_to_file(self, file_path):
+        """
+        Saves schema object to file
+        Args:
+            file_path: path to store schema object at
+        """
+        with open(file_path, "w") as f:
+            json.dump(self._schemas, f, indent=2)

tts/nemo/collections/nlp/data/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.data.entity_linking.entity_linking_dataset import EntityLinkingDataset

tts/nemo/collections/nlp/data/entity_linking/entity_linking_dataset.py

@@ -0,0 +1,135 @@
+# 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 array
+import pickle as pkl
+from typing import Optional
+
+import torch
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import find_newlines, load_data_indices
+from nemo.core.classes import Dataset
+from nemo.utils import logging
+
+__all__ = ['EntityLinkingDataset']
+
+
+class EntityLinkingDataset(Dataset):
+    """
+    Parent class for entity linking encoder training and index
+    datasets
+
+    Args:
+        tokenizer (obj): huggingface tokenizer,
+        data_file (str): path to tab separated column file where data 
+            pairs apear in the format 
+            concept_ID\tconcept_synonym1\tconcept_synonym2\n
+        newline_idx_file (str): path to pickle file containing location
+            of data_file newline characters
+        max_seq_length (int): maximum length of a concept in tokens
+        is_index_data (bool): Whether dataset will be used for building
+                            a nearest neighbors index
+    """
+
+    def __init__(
+        self,
+        tokenizer: object,
+        data_file: str,
+        newline_idx_file: Optional[str] = None,
+        max_seq_length: Optional[int] = 512,
+        is_index_data: bool = False,
+    ):
+
+        self.tokenizer = tokenizer
+
+        # Try and load pair indices file if already exists
+        newline_indices, newline_idx_file, _ = load_data_indices(newline_idx_file, data_file, "newline_indices")
+
+        # If pair indices file doesn't exists, generate and store them
+        if newline_indices is None:
+            logging.info("Getting datafile newline indices")
+
+            with open(data_file, "rb") as f:
+                contents = f.read()
+                newline_indices = find_newlines(contents)
+                newline_indices = array.array("I", newline_indices)
+
+            # Store data file indicies to avoid generating them again
+            with open(newline_idx_file, "wb") as f:
+                pkl.dump(newline_indices, f)
+
+        self.newline_indices = newline_indices
+        self.data_file = data_file
+        self.num_lines = len(newline_indices)
+        self.max_seq_length = max_seq_length
+        self.is_index_data = is_index_data
+
+        logging.info(f"Loaded dataset with {self.num_lines} examples")
+
+    def __len__(self):
+        return self.num_lines
+
+    def __getitem__(self, idx):
+
+        concept_offset = self.newline_indices[idx]
+
+        with open(self.data_file, "r", encoding='utf-8-sig') as f:
+            # Find data pair within datafile using byte offset
+            f.seek(concept_offset)
+            concept = f.readline()[:-1]
+            concept = concept.strip().split("\t")
+
+            if self.is_index_data:
+                concept_id, concept = concept
+                return (int(concept_id), concept)
+
+            else:
+                concept_id, concept1, concept2 = concept
+                return (int(concept_id), concept1, concept2)
+
+    def _collate_fn(self, batch):
+        """collate batch of input_ids, segment_ids, input_mask, and label
+
+        Args:
+            batch:  A list of tuples of format (concept_ID, concept_synonym1, concept_synonym2).
+        """
+        if self.is_index_data:
+            concept_ids, concepts = zip(*batch)
+            concept_ids = list(concept_ids)
+            concepts = list(concepts)
+
+        else:
+            concept_ids, concepts1, concepts2 = zip(*batch)
+            concept_ids = list(concept_ids)
+            concept_ids.extend(concept_ids)  # Need to double label list to match each concept
+            concepts = list(concepts1)
+            concepts.extend(concepts2)
+
+        batch = self.tokenizer(
+            concepts,
+            add_special_tokens=True,
+            padding=True,
+            truncation=True,
+            max_length=self.max_seq_length,
+            return_token_type_ids=True,
+            return_attention_mask=True,
+            return_length=True,
+        )
+
+        return (
+            torch.LongTensor(batch["input_ids"]),
+            torch.LongTensor(batch["token_type_ids"]),
+            torch.LongTensor(batch["attention_mask"]),
+            torch.LongTensor(concept_ids),
+        )

tts/nemo/collections/nlp/data/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.data.glue_benchmark.glue_benchmark_dataset import GLUEDataset

tts/nemo/collections/nlp/data/glue_benchmark/data_processors.py

@@ -0,0 +1,323 @@
+# 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 nemo.collections.nlp.data.data_utils.data_preprocessing import DataProcessor
+from nemo.utils import logging
+
+__all__ = [
+    'ColaProcessor',
+    'MnliProcessor',
+    'MnliMismatchedProcessor',
+    'MrpcProcessor',
+    'Sst2Processor',
+    'StsbProcessor',
+    'QqpProcessor',
+    'QnliProcessor',
+    'RteProcessor',
+    'WnliProcessor',
+]
+
+
+class MrpcProcessor(DataProcessor):
+    """Processor for the MRPC data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        logging.info(f'LOOKING AT {os.path.join(data_dir, "train.tsv")}')
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, i)
+            text_a = line[3]
+            text_b = line[4]
+            label = line[0]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class MnliProcessor(DataProcessor):
+    """Processor for the MultiNLI data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev_matched.tsv")), "dev_matched")
+
+    def get_labels(self):
+        """See base class."""
+        return ["contradiction", "entailment", "neutral"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            text_a = line[8]
+            text_b = line[9]
+            label = line[-1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class MnliMismatchedProcessor(MnliProcessor):
+    """Processor for the MultiNLI Mismatched data set (GLUE version)."""
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev_mismatched.tsv")), "dev_matched")
+
+
+class ColaProcessor(DataProcessor):
+    """Processor for the CoLA data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            guid = "%s-%s" % (set_type, i)
+            text_a = line[3]
+            label = line[1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
+        return examples
+
+
+class Sst2Processor(DataProcessor):
+    """Processor for the SST-2 data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, i)
+            text_a = line[0]
+            label = line[1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
+        return examples
+
+
+class StsbProcessor(DataProcessor):
+    """Processor for the STS-B data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return [None]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            text_a = line[7]
+            text_b = line[8]
+            label = line[-1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class QqpProcessor(DataProcessor):
+    """Processor for the QQP data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            try:
+                text_a = line[3]
+                text_b = line[4]
+                label = line[5]
+            except IndexError:
+                continue
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class QnliProcessor(DataProcessor):
+    """Processor for the QNLI data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["entailment", "not_entailment"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            text_a = line[1]
+            text_b = line[2]
+            label = line[-1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class RteProcessor(DataProcessor):
+    """Processor for the RTE data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["entailment", "not_entailment"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            text_a = line[1]
+            text_b = line[2]
+            label = line[-1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class WnliProcessor(DataProcessor):
+    """Processor for the WNLI data set (GLUE version)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1"]
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (i, line) in enumerate(lines):
+            if i == 0:
+                continue
+            guid = "%s-%s" % (set_type, line[0])
+            text_a = line[1]
+            text_b = line[2]
+            label = line[-1]
+            examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
+        return examples
+
+
+class InputExample(object):
+    """A single training/test example for simple sequence classification.
+
+    Args:
+        guid: Unique id for the example.
+        text_a: The untokenized text of the first sequence.
+        For single sequence tasks, only this sequence must be specified.
+        text_b: The untokenized text of the second
+        sequence. Only must be specified for sequence pair tasks.
+        label:The label of the example. This should be
+        specified for train and dev examples, but not for test examples.
+    """
+
+    def __init__(self, guid: int, text_a: str, text_b: str = None, label: str = None):
+        """Constructs a InputExample."""
+        self.guid = guid
+        self.text_a = text_a
+        self.text_b = text_b
+        self.label = label

tts/nemo/collections/nlp/data/glue_benchmark/glue_benchmark_dataset.py

@@ -0,0 +1,365 @@
+# 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.
+
+# Some code of this file was adapted from the HuggingFace library available at
+# https://github.com/huggingface/transformers
+
+import os
+import pickle
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.glue_benchmark.data_processors import (
+    ColaProcessor,
+    MnliMismatchedProcessor,
+    MnliProcessor,
+    MrpcProcessor,
+    QnliProcessor,
+    QqpProcessor,
+    RteProcessor,
+    Sst2Processor,
+    StsbProcessor,
+    WnliProcessor,
+)
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import CategoricalValuesType, ChannelType, MaskType, NeuralType, RegressionValuesType
+from nemo.utils import logging
+
+__all__ = ['GLUEDataset']
+
+processors = {
+    "cola": ColaProcessor,
+    "mnli": MnliProcessor,
+    "mnli-mm": MnliMismatchedProcessor,
+    "mrpc": MrpcProcessor,
+    "sst-2": Sst2Processor,
+    "sts-b": StsbProcessor,
+    "qqp": QqpProcessor,
+    "qnli": QnliProcessor,
+    "rte": RteProcessor,
+    "wnli": WnliProcessor,
+}
+output_modes = {
+    "cola": "classification",
+    "mnli": "classification",
+    "mnli-mm": "classification",
+    "mrpc": "classification",
+    "sst-2": "classification",
+    "sts-b": "regression",
+    "qqp": "classification",
+    "qnli": "classification",
+    "rte": "classification",
+    "wnli": "classification",
+}
+GLUE_TASKS_NUM_LABELS = {
+    "cola": 2,
+    "mnli": 3,
+    "mrpc": 2,
+    "sst-2": 2,
+    "sts-b": 1,
+    "qqp": 2,
+    "qnli": 2,
+    "rte": 2,
+    "wnli": 2,
+}
+
+
+class GLUEDataset(Dataset):
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            "labels": NeuralType(
+                tuple('B'), RegressionValuesType() if self.task_name == 'sts-b' else CategoricalValuesType()
+            ),
+        }
+
+    def __init__(
+        self, file_name: str, task_name: str, tokenizer: TokenizerSpec, max_seq_length: str, use_cache: bool = True,
+    ):
+        """
+        Processes GLUE datasets
+        Args:
+            file_name: path to file
+            task_name: GLUE task name
+            tokenizer: such as AutoTokenizer
+            max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+            use_cache: whether to use data cache
+        """
+        logging.info(f'Processing {file_name}')
+        data_dir, file_name = os.path.split(file_name)
+        file_name = file_name[:-4]
+        self.tokenizer = tokenizer
+        evaluate = False if 'train' in file_name else True
+
+        if task_name not in processors:
+            raise ValueError(f'{task_name} not supported. Choose from {processors.keys()}')
+
+        if task_name == 'mnli' and 'dev_mismatched' in file_name:
+            self.task_name = 'mnli-mm'
+        else:
+            self.task_name = task_name
+
+        processor = processors[self.task_name]()
+        output_mode = output_modes[self.task_name]
+        self.label_list = processor.get_labels()
+
+        self.examples = processor.get_dev_examples(data_dir) if evaluate else processor.get_train_examples(data_dir)
+        processor_name = type(processor).__name__
+        vocab_size = getattr(tokenizer, "vocab_size", 0)
+        cached_features_file = os.path.join(
+            data_dir,
+            "cached_{}_{}_{}_{}_{}".format(
+                processor_name, file_name, tokenizer.name, str(max_seq_length), str(vocab_size)
+            ),
+        )
+
+        if use_cache and os.path.exists(cached_features_file):
+            logging.info(f"loading from {cached_features_file}")
+            with open(cached_features_file, "rb") as reader:
+                self.features = pickle.load(reader)
+        else:
+            token_params = {
+                'bos_token': None,
+                'eos_token': tokenizer.eos_token,
+                'pad_token': tokenizer.pad_token,
+                'cls_token': tokenizer.cls_token,
+                'sep_token_extra': tokenizer.eos_token if 'roberta' in tokenizer.name.lower() else None,
+            }
+
+            self.features = self.convert_examples_to_features(
+                self.examples, self.label_list, max_seq_length, tokenizer, output_mode, **token_params
+            )
+            master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+            if master_device:
+                logging.info(f'Saving train features into {cached_features_file}')
+                with open(cached_features_file, "wb") as writer:
+                    pickle.dump(self.features, writer)
+
+    def __len__(self):
+        return len(self.features)
+
+    def __getitem__(self, idx):
+        feature = self.features[idx]
+        return (
+            np.array(feature.input_ids),
+            np.array(feature.segment_ids),
+            np.array(feature.input_mask, dtype=np.long),
+            np.array(feature.label_id),
+        )
+
+    def convert_examples_to_features(
+        self,
+        examples: List[str],
+        label_list: List[int],
+        max_seq_length: int,
+        tokenizer: TokenizerSpec,
+        output_mode: str,
+        bos_token: str = None,
+        eos_token: str = '[SEP]',
+        pad_token: str = '[PAD]',
+        cls_token: str = '[CLS]',
+        sep_token_extra: str = None,
+        cls_token_at_end: bool = False,
+        cls_token_segment_id: int = 0,
+        pad_token_segment_id: int = 0,
+        pad_on_left: bool = False,
+        mask_padding_with_zero: bool = True,
+        sequence_a_segment_id: int = 0,
+        sequence_b_segment_id: int = 1,
+    ):
+        """
+        Loads a data file into a list of `InputBatch`s.
+        The `cls_token_at_end` defines the location of the CLS token:
+
+            * False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]
+            * True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
+
+        The `cls_token_segment_id` defines the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
+        
+        The convention in BERT is:
+        
+            a. For sequence pairs:
+                * tokens:   [CLS] is this jack ##ville ? [SEP] no it is not . [SEP]
+                * type_ids:   0   0  0    0    0       0   0   1  1  1  1   1   1
+            b. For single sequences:
+                * tokens:   [CLS] the dog is hairy . [SEP]
+                * type_ids:   0   0   0   0  0     0   0
+
+        Where "type_ids" are used to indicate whether this is the first
+        sequence or the second sequence. The embedding vectors for `type=0`
+        and `type=1` were learned during pre-training and are added to the
+        wordpiece embedding vector (and position vector). This is
+        not *strictly* necessarysince the [SEP] token unambiguously separates
+        the sequences, but it makes it easier for the model to learn
+        the concept of sequences.
+        For classification tasks, the first vector (corresponding to [CLS])
+        is used as as the "sentence vector". Note that this only makes sense
+        because the entire model is fine-tuned.
+        
+        The convention for NMT is:
+        
+            a. For sequence pairs:
+                * tokens:<BOS> is this jack ##ville ? <EOS> <BOS> no it is not . <EOS>
+                * type_ids:0   0  0    0    0       0   0     1   1  1  1  1   1   1
+            b. For single sequences:
+                * tokens:   <BOS> the dog is hairy . <EOS>
+                * type_ids:   0   0   0   0  0     0   0
+
+        """
+        label_map = {label: i for i, label in enumerate(label_list)}
+
+        features = []
+        for ex_index, example in enumerate(examples):
+            if ex_index % 10000 == 0:
+                logging.info("Writing example %d of %d" % (ex_index, len(examples)))
+
+            tokens_a = tokenizer.text_to_tokens(example.text_a)
+
+            tokens_b = None
+            if example.text_b:
+                tokens_b = tokenizer.text_to_tokens(example.text_b)
+
+                special_tokens_count = 2 if eos_token else 0
+                special_tokens_count += 1 if sep_token_extra else 0
+                special_tokens_count += 2 if bos_token else 0
+                special_tokens_count += 1 if cls_token else 0
+                self._truncate_seq_pair(tokens_a, tokens_b, max_seq_length - special_tokens_count)
+            else:
+                special_tokens_count = 1 if eos_token else 0
+                special_tokens_count += 1 if sep_token_extra else 0
+                special_tokens_count += 1 if bos_token else 0
+                if len(tokens_a) > max_seq_length - special_tokens_count:
+                    tokens_a = tokens_a[: max_seq_length - special_tokens_count]
+            # Add special tokens to sequence_a
+            tokens = tokens_a
+            if bos_token:
+                tokens = [bos_token] + tokens
+            if eos_token:
+                tokens += [eos_token]
+            segment_ids = [sequence_a_segment_id] * len(tokens)
+
+            # Add sequence separator between sequences
+            if tokens_b and sep_token_extra:
+                tokens += [sep_token_extra]
+                segment_ids += [sequence_a_segment_id]
+
+            # Add special tokens to sequence_b
+            if tokens_b:
+                if bos_token:
+                    tokens += [bos_token]
+                    segment_ids += [sequence_b_segment_id]
+                tokens += tokens_b
+                segment_ids += [sequence_b_segment_id] * (len(tokens_b))
+                if eos_token:
+                    tokens += [eos_token]
+                    segment_ids += [sequence_b_segment_id]
+
+            # Add classification token - for BERT models
+            if cls_token:
+                if cls_token_at_end:
+                    tokens += [cls_token]
+                    segment_ids += [cls_token_segment_id]
+                else:
+                    tokens = [cls_token] + tokens
+                    segment_ids = [cls_token_segment_id] + segment_ids
+            input_ids = tokenizer.tokens_to_ids(tokens)
+
+            # The mask has 1 for real tokens and 0 for padding tokens. Only real
+            # tokens are attended to.
+            input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)
+
+            # Zero-pad up to the sequence length.
+            padding_length = max_seq_length - len(input_ids)
+            pad_token_id = tokenizer.tokens_to_ids([pad_token])[0]
+            if pad_on_left:
+                input_ids = ([pad_token_id] * padding_length) + input_ids
+                input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
+                segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids
+            else:
+                input_ids = input_ids + ([pad_token_id] * padding_length)
+                input_mask = input_mask + ([0 if mask_padding_with_zero else 1] * padding_length)
+                segment_ids = segment_ids + ([pad_token_segment_id] * padding_length)
+            if len(input_ids) != max_seq_length:
+                raise ValueError("input_ids must be of length max_seq_length")
+            if len(input_mask) != max_seq_length:
+                raise ValueError("input_mask must be of length max_seq_length")
+            if len(segment_ids) != max_seq_length:
+                raise ValueError("segment_ids must be of length max_seq_length")
+            if output_mode == "classification":
+                label_id = label_map[example.label]
+            elif output_mode == "regression":
+                label_id = np.float32(example.label)
+            else:
+                raise KeyError(output_mode)
+
+            if ex_index < 5:
+                logging.info("*** Example ***")
+                logging.info("guid: %s" % (example.guid))
+                logging.info("tokens: %s" % " ".join(list(map(str, tokens))))
+                logging.info("input_ids: %s" % " ".join(list(map(str, input_ids))))
+                logging.info("input_mask: %s" % " ".join(list(map(str, input_mask))))
+                logging.info("segment_ids: %s" % " ".join(list(map(str, segment_ids))))
+                logging.info("label: %s (id = %d)" % (example.label, label_id))
+
+            features.append(
+                InputFeatures(input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids, label_id=label_id)
+            )
+        return features
+
+    def _truncate_seq_pair(self, tokens_a: str, tokens_b: str, max_length: int):
+        """Truncates a sequence pair in place to the maximum length.
+
+        This will always truncate the longer sequence one token at a time.
+        This makes more sense than truncating an equal percent
+        of tokens from each, since if one sequence is very short then each token
+        that's truncated likely contains more information than a longer sequence.
+        """
+        while True:
+            total_length = len(tokens_a) + len(tokens_b)
+            if total_length <= max_length:
+                break
+            if len(tokens_a) > len(tokens_b):
+                tokens_a.pop()
+            else:
+                tokens_b.pop()
+
+
+class InputFeatures(object):
+    """A single set of features of data.
+
+    Args:
+        input_ids: input/token ids
+        input_mask: masks out subword tokens
+        segment_ids: distinguish one sentence from the other one (if present)
+        label_ids: label for the current example
+    """
+
+    def __init__(
+        self, input_ids: List[int], input_mask: List[int], segment_ids: List[int], label_id: Union[float, int]
+    ):
+        """Initialized InputFeatures."""
+        self.input_ids = input_ids
+        self.input_mask = input_mask
+        self.segment_ids = segment_ids
+        self.label_id = label_id

tts/nemo/collections/nlp/data/information_retrieval/__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.data.information_retrieval.information_retrieval_dataset import (
+    BertInformationRetrievalDataset,
+)

tts/nemo/collections/nlp/data/information_retrieval/information_retrieval_dataset.py

@@ -0,0 +1,278 @@
+# 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 multiprocessing as mp
+import os
+import pickle
+import random
+from typing import Optional
+
+import numpy as np
+from torch.utils.data import Dataset
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+
+__all__ = ["BertInformationRetrievalDataset"]
+
+
+class BaseInformationRetrievalDataset(Dataset):
+    """
+    Base information retrieval dataset on which other datasets are built.
+
+    Args:
+        tokenizer: tokenizer
+        max_query_length: maximum length of query in tokens
+        max_passage_length: maximum length of passage in tokens
+    """
+
+    def __init__(
+        self, tokenizer: TokenizerSpec, max_query_length: Optional[int] = 31, max_passage_length: Optional[int] = 190,
+    ):
+        self.tokenizer = tokenizer
+        self.max_query_length = max_query_length
+        self.max_passage_length = max_passage_length
+
+    def parse_npz(self, file, max_seq_length):
+        """
+        Function which parses passages (documents) in npz format.
+        After pre-processing and tokenization, the dataset will be saved
+        as numpy matrix, i_th entry of which corresponds to i_th passage (document)
+        and has the following form:
+        [n, token_1, ..., token_n, 0, ..., 0]
+        where n is the passage length (in tokens) and 0s correspond to pad tokens.
+
+        Args:
+            file: str, path to file with passages (documents)
+            max_seq_length: maximum length of sequence in tokens
+        """
+        cached_collection = file + ".npz"
+        if os.path.isfile(cached_collection):
+            dataset_npz = np.load(cached_collection)["data"]
+        else:
+            dataset_dict = self.tokenize_dataset(file, max_seq_length)
+            dataset_npz = np.zeros((len(dataset_dict), max_seq_length + 1))
+            for key in dataset_dict:
+                dataset_npz[key][0] = len(dataset_dict[key])
+                dataset_npz[key][1 : len(dataset_dict[key]) + 1] = dataset_dict[key]
+            np.savez(cached_collection, data=dataset_npz)
+        return dataset_npz
+
+    def parse_pkl(self, file, max_seq_length):
+        """
+        Function which parses passages (documents, queries) in pkl format.
+        After pre-processing and tokenization, the dataset will be saved
+        as pkl dict, i_th entry of which corresponds to i_th passage (document, query)
+        and has the following form:
+        {passage_id: [token_1, ..., token_n]}
+        where n is the passage length (in tokens).
+
+        Args:
+            file: str, path to file with passages (documents)
+            max_seq_length: maximum length of sequence in tokens
+        """
+        cached_collection = file + ".pkl"
+        if os.path.isfile(cached_collection):
+            dataset_dict = pickle.load(open(cached_collection, "rb"))
+        else:
+            dataset_dict = self.tokenize_dataset(file, max_seq_length)
+            pickle.dump(dataset_dict, open(cached_collection, "wb"))
+        return dataset_dict
+
+    def tokenize_dataset(self, file, max_seq_length):
+        """
+        Function which pre-tokenizes the dataset.
+        """
+        lines = open(file, "r").readlines()
+        with mp.Pool() as pool:
+            dataset_dict = pool.map(self.preprocess_line, lines)
+        dataset_dict = {id_: tokens[:max_seq_length] for (id_, tokens) in dataset_dict}
+        return dataset_dict
+
+    def preprocess_line(self, line):
+        """
+        Parse a single entry (line) of tsv file.
+        """
+        if "\t" not in line:
+            raise ValueError(f"Provided dataset does not have a form of tsv file")
+        id_, text = line.split("\t")
+        token_ids = self.tokenizer.text_to_ids(text.strip())
+        return int(id_), token_ids
+
+    def construct_input(self, token_ids1, max_seq_length, token_ids2=None):
+        """
+        Function which constructs a valid input to BERT from tokens.
+
+        If only one list of tokens (token_ids1) is passed, the input will be
+        [CLS] token_ids1 [SEP]
+
+        if two lists of tokens are passed, the input will be
+        [CLS] token_ids1 [SEP] token_ids2 [SEP]
+        """
+
+        input_ids = [self.tokenizer.pad_id] * max_seq_length
+        bert_input = [self.tokenizer.cls_id] + token_ids1 + [self.tokenizer.sep_id]
+        sentence1_length = len(bert_input)
+        if token_ids2 is not None:
+            bert_input = bert_input + token_ids2 + [self.tokenizer.sep_id]
+
+        bert_input = bert_input[:max_seq_length]
+
+        num_nonpad_tokens = len(bert_input)
+
+        input_ids[:num_nonpad_tokens] = bert_input
+        input_ids = np.array(input_ids, dtype=np.long)
+        input_mask = input_ids != self.tokenizer.pad_id
+        input_type_ids = np.ones_like(input_ids)
+        input_type_ids[:sentence1_length] = 0
+
+        return input_ids, input_mask, input_type_ids
+
+    def preprocess_bert(self, query_id, psg_ids):
+        """
+        Transforms query id (Q) and a list of passages ids (P1, ..., Pk)
+        into a tensor of size [k, max_length] with the following rows:
+        [CLS] Q_text [SEP] Pi_text [SEP], i = 1, ..., k
+        """
+
+        max_seq_length = self.max_query_length + self.max_passage_length + 3
+        input_ids, input_mask, input_type_ids = [], [], []
+        for psg_id in psg_ids:
+            inputs = self.construct_input(self.queries[query_id], max_seq_length, self._psgid2tokens(psg_id))
+            input_ids.append(inputs[0])
+            input_mask.append(inputs[1])
+            input_type_ids.append(inputs[2])
+
+        input_ids = np.stack(input_ids)
+        input_mask = np.stack(input_mask)
+        input_type_ids = np.stack(input_type_ids)
+
+        return input_ids, input_mask, input_type_ids
+
+    def preprocess_dpr(self, query_id, psg_ids):
+        """
+        Transforms query id (Q) and a list of passages ids (P1, ..., Pk)
+        into two tensors of sizes [1, max_q_length] and [k, max_p_length]
+        with the following rows:
+        1) [CLS] Q_text [SEP]
+        2) [CLS] Pi_text [SEP], i = 1, ..., k
+        """
+
+        q_input_ids, q_input_mask, q_type_ids = self.construct_input(self.queries[query_id], self.max_query_length + 2)
+        input_ids, input_mask, input_type_ids = [], [], []
+        for psg_id in psg_ids:
+            inputs = self.construct_input(self._psgid2tokens(psg_id), self.max_passage_length + 2)
+            input_ids.append(inputs[0])
+            input_mask.append(inputs[1])
+            input_type_ids.append(inputs[2])
+        input_ids = np.stack(input_ids)
+        input_mask = np.stack(input_mask)
+        input_type_ids = np.stack(input_type_ids)
+        return (
+            q_input_ids[None, ...],
+            q_input_mask[None, ...],
+            q_type_ids[None, ...],
+            input_ids,
+            input_mask,
+            input_type_ids,
+        )
+
+    def _psgid2tokens(self, psg_id):
+        """
+        Internal function which maps passage id to its tokens.
+        """
+        pass
+
+    def psgid2tokens_npz(self, psg_id):
+        """
+        Mapping from passage id to its tokens in case of npz cache format.
+        """
+        seq_len = self.passages[psg_id][0]
+        return self.passages[psg_id][1 : seq_len + 1].tolist()
+
+    def psgid2tokens_pkl(self, psg_id):
+        """
+        Mapping from passage id to its tokens in case of pkl cache format.
+        """
+        return self.passages[psg_id]
+
+
+class BertInformationRetrievalDataset(BaseInformationRetrievalDataset):
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        passages: str,
+        queries: str,
+        query_to_passages: str,
+        max_query_length: Optional[int] = 31,
+        max_passage_length: Optional[int] = 190,
+        num_negatives: Optional[int] = 10,
+        preprocess_fn: Optional[str] = "preprocess_bert",
+        psg_cache_format: Optional[str] = "npz",
+    ):
+        """
+        Dataset for training information retrieval models.
+        
+        Args:
+            tokenizer: tokenizer
+            passages: path to tsv with [psg_id, psg_text] entries
+            queries: path to tsv with [query_id, query_text] entries
+            query_to_passages: path to tsv with
+                [query_id, pos_psg_id, neg_psg_id_1, ..., neg_psg_id_k] entries
+            max_query_length: maximum length of query in tokens
+            max_passage_length: maximum length of passage in tokens
+            num_negatives: number of negative passages per positive to use for training
+            preprocess_fn: either preprocess_bert or preprocess_dpr
+                preprocess_bert: joint input: [CLS] query [SEP] passage [SEP]
+                preprocess_dpr: separate inputs: [CLS] query [SEP], [CLS] passage [SEP]
+            psg_cache_format: either pkl or npz
+        """
+
+        super().__init__(tokenizer, max_query_length, max_passage_length)
+        self.num_negatives = num_negatives
+
+        self.passages = getattr(self, f"parse_{psg_cache_format}")(passages, max_passage_length)
+        self._psgid2tokens = getattr(self, f"psgid2tokens_{psg_cache_format}")
+        self.queries = self.parse_pkl(queries, max_query_length)
+        self.idx2psgs = self.parse_query_to_passages(query_to_passages)
+        self._preprocess_fn = getattr(self, preprocess_fn)
+
+    def __getitem__(self, idx):
+        query_and_psgs = self.idx2psgs[idx]
+        query_id, psg_ids = query_and_psgs[0], query_and_psgs[1:]
+        inputs = self._preprocess_fn(query_id, psg_ids)
+        return [*inputs, query_id, np.array(psg_ids)]
+
+    def __len__(self):
+        return len(self.idx2psgs)
+
+    def parse_query_to_passages(self, file):
+        """
+        Function which parses query to passages correspondence file.
+        """
+        idx2psgs = {}
+        idx = 0
+        for line in open(file, "r").readlines():
+            if "\t" not in line:
+                raise ValueError(f"Provided dataset does not have a form of tsv file")
+            query_and_psgs = line.split("\t")
+            query_and_psgs_ids = [int(id_) for id_ in query_and_psgs]
+            query_and_rel_psg_ids, irrel_psgs_ids = query_and_psgs_ids[:2], query_and_psgs_ids[2:]
+            random.shuffle(irrel_psgs_ids)
+            num_samples = len(irrel_psgs_ids) // self.num_negatives
+            for j in range(num_samples):
+                left = self.num_negatives * j
+                right = self.num_negatives * (j + 1)
+                idx2psgs[idx] = query_and_rel_psg_ids + irrel_psgs_ids[left:right]
+                idx += 1
+        return idx2psgs

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

@@ -0,0 +1,22 @@
+# 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.data.intent_slot_classification.intent_slot_classification_dataset import (
+    IntentSlotClassificationDataset,
+    IntentSlotInferenceDataset,
+)
+from nemo.collections.nlp.data.intent_slot_classification.intent_slot_classification_descriptor import (
+    IntentSlotDataDesc,
+)

tts/nemo/collections/nlp/data/intent_slot_classification/intent_slot_classification_dataset.py

@@ -0,0 +1,292 @@
+# 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 Any, Dict, Optional
+
+import numpy as np
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.data_utils import get_stats
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['IntentSlotClassificationDataset', 'IntentSlotInferenceDataset']
+
+
+def get_features(
+    queries,
+    max_seq_length,
+    tokenizer,
+    pad_label=128,
+    raw_slots=None,
+    ignore_extra_tokens=False,
+    ignore_start_end=False,
+):
+    all_subtokens = []
+    all_loss_mask = []
+    all_subtokens_mask = []
+    all_segment_ids = []
+    all_input_ids = []
+    all_input_mask = []
+    sent_lengths = []
+    all_slots = []
+
+    with_label = False
+    if raw_slots is not None:
+        with_label = True
+
+    for i, query in enumerate(queries):
+        words = query.strip().split()
+        subtokens = [tokenizer.cls_token]
+        loss_mask = [1 - ignore_start_end]
+        subtokens_mask = [0]
+        if with_label:
+            slots = [pad_label]
+
+        for j, word in enumerate(words):
+            word_tokens = tokenizer.text_to_tokens(word)
+            subtokens.extend(word_tokens)
+
+            loss_mask.append(1)
+            loss_mask.extend([int(not ignore_extra_tokens)] * (len(word_tokens) - 1))
+
+            subtokens_mask.append(1)
+            subtokens_mask.extend([0] * (len(word_tokens) - 1))
+
+            if with_label:
+                slots.extend([raw_slots[i][j]] * len(word_tokens))
+
+        subtokens.append(tokenizer.sep_token)
+        loss_mask.append(1 - ignore_start_end)
+        subtokens_mask.append(0)
+        sent_lengths.append(len(subtokens))
+        all_subtokens.append(subtokens)
+        all_loss_mask.append(loss_mask)
+        all_subtokens_mask.append(subtokens_mask)
+        all_input_mask.append([1] * len(subtokens))
+        if with_label:
+            slots.append(pad_label)
+            all_slots.append(slots)
+
+    max_seq_length_data = max(sent_lengths)
+    max_seq_length = min(max_seq_length, max_seq_length_data) if max_seq_length > 0 else max_seq_length_data
+    logging.info(f'Setting max length to: {max_seq_length}')
+    get_stats(sent_lengths)
+    too_long_count = 0
+
+    for i, subtokens in enumerate(all_subtokens):
+        if len(subtokens) > max_seq_length:
+            subtokens = [tokenizer.cls_token] + subtokens[-max_seq_length + 1 :]
+            all_input_mask[i] = [1] + all_input_mask[i][-max_seq_length + 1 :]
+            all_loss_mask[i] = [1 - ignore_start_end] + all_loss_mask[i][-max_seq_length + 1 :]
+            all_subtokens_mask[i] = [0] + all_subtokens_mask[i][-max_seq_length + 1 :]
+
+            if with_label:
+                all_slots[i] = [pad_label] + all_slots[i][-max_seq_length + 1 :]
+            too_long_count += 1
+
+        all_input_ids.append([tokenizer.tokens_to_ids(t) for t in subtokens])
+
+        if len(subtokens) < max_seq_length:
+            extra = max_seq_length - len(subtokens)
+            all_input_ids[i] = all_input_ids[i] + [0] * extra
+            all_loss_mask[i] = all_loss_mask[i] + [0] * extra
+            all_subtokens_mask[i] = all_subtokens_mask[i] + [0] * extra
+            all_input_mask[i] = all_input_mask[i] + [0] * extra
+
+            if with_label:
+                all_slots[i] = all_slots[i] + [pad_label] * extra
+
+        all_segment_ids.append([0] * max_seq_length)
+
+    logging.info(f'{too_long_count} are longer than {max_seq_length}')
+
+    # May be useful for debugging
+    logging.debug("*** Some Examples of Processed Data ***")
+    for i in range(min(len(all_input_ids), 5)):
+        logging.debug("i: %s" % (i))
+        logging.debug("subtokens: %s" % " ".join(list(map(str, all_subtokens[i]))))
+        logging.debug("loss_mask: %s" % " ".join(list(map(str, all_loss_mask[i]))))
+        logging.debug("input_mask: %s" % " ".join(list(map(str, all_input_mask[i]))))
+        logging.debug("subtokens_mask: %s" % " ".join(list(map(str, all_subtokens_mask[i]))))
+        if with_label:
+            logging.debug("slots_label: %s" % " ".join(list(map(str, all_slots[i]))))
+
+    return (all_input_ids, all_segment_ids, all_input_mask, all_loss_mask, all_subtokens_mask, all_slots)
+
+
+class IntentSlotClassificationDataset(Dataset):
+    """
+    Creates dataset to use for the task of joint intent
+    and slot classification with pretrained model.
+
+    Converts from raw data to an instance that can be used by
+    NMDataLayer.
+
+    For dataset to use during inference without labels, see
+    IntentSlotDataset.
+
+    Args:
+        input_file: file to sequence + label. the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        slot_file: file to slot labels, each line corresponding to slot labels for a sentence in input_file. No header.
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer: such as NemoBertTokenizer
+        num_samples: number of samples you want to use for the dataset. If -1, use all dataset. Useful for testing.
+        pad_label: pad value use for slot labels. by default, it's the neutral label.
+        ignore_extra_tokens: whether to ignore extra tokens in the loss_mask.
+        ignore_start_end: whether to ignore bos and eos tokens in the loss_mask.
+        do_lower_case: convert query to lower case or not
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'loss_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+            'intent_labels': NeuralType(('B'), LabelsType()),
+            'slot_labels': NeuralType(('B', 'T'), LabelsType()),
+        }
+
+    def __init__(
+        self,
+        input_file: str,
+        slot_file: str,
+        max_seq_length: int,
+        tokenizer: TokenizerSpec,
+        num_samples: int = -1,
+        pad_label: int = 128,
+        ignore_extra_tokens: bool = False,
+        ignore_start_end: bool = False,
+        do_lower_case: bool = False,
+    ):
+        if num_samples == 0:
+            raise ValueError("num_samples has to be positive", num_samples)
+
+        with open(slot_file, 'r') as f:
+            slot_lines = f.readlines()
+
+        with open(input_file, 'r') as f:
+            input_lines = f.readlines()[1:]
+
+        assert len(slot_lines) == len(input_lines)
+
+        dataset = list(zip(slot_lines, input_lines))
+
+        if num_samples > 0:
+            dataset = dataset[:num_samples]
+
+        raw_slots, queries, raw_intents = [], [], []
+        for slot_line, input_line in dataset:
+            raw_slots.append([int(slot) for slot in slot_line.strip().split()])
+            parts = input_line.strip().split()
+            raw_intents.append(int(parts[-1]))
+            query = ' '.join(parts[:-1])
+            if do_lower_case:
+                query = query.lower()
+            queries.append(query)
+
+        features = get_features(
+            queries,
+            max_seq_length,
+            tokenizer,
+            pad_label=pad_label,
+            raw_slots=raw_slots,
+            ignore_extra_tokens=ignore_extra_tokens,
+            ignore_start_end=ignore_start_end,
+        )
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_loss_mask = features[3]
+        self.all_subtokens_mask = features[4]
+        self.all_slots = features[5]
+        self.all_intents = raw_intents
+
+    def __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.long),
+            np.array(self.all_loss_mask[idx]),
+            np.array(self.all_subtokens_mask[idx]),
+            self.all_intents[idx],
+            np.array(self.all_slots[idx]),
+        )
+
+
+class IntentSlotInferenceDataset(Dataset):
+    """
+    Creates dataset to use for the task of joint intent
+    and slot classification with pretrained model.
+    This is to be used during inference only.
+    It uses list of queries as the input.
+
+    Args:
+        queries (list): list of queries to run inference on
+        max_seq_length (int): max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer (Tokenizer): such as NemoBertTokenizer
+        pad_label (int): pad value use for slot labels.
+            by default, it's the neutral label.
+
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+            Returns definitions of module output ports.
+        """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'loss_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+        }
+
+    def __init__(self, queries, max_seq_length, tokenizer, do_lower_case):
+        if do_lower_case:
+            for idx, query in enumerate(queries):
+                queries[idx] = queries[idx].lower()
+
+        features = get_features(queries, max_seq_length, tokenizer)
+
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_loss_mask = features[3]
+        self.all_subtokens_mask = features[4]
+
+    def __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.long),
+            np.array(self.all_loss_mask[idx]),
+            np.array(self.all_subtokens_mask[idx]),
+        )

tts/nemo/collections/nlp/data/intent_slot_classification/intent_slot_classification_descriptor.py

@@ -0,0 +1,163 @@
+# 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
+from typing import List
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import (
+    fill_class_weights,
+    get_freq_weights,
+    get_label_stats,
+    if_exist,
+)
+from nemo.utils import logging
+
+
+class IntentSlotDataDesc:
+    """ Convert the raw data to the standard format supported by
+    IntentSlotDataDesc.
+
+    By default, the None label for slots is 'O'.
+
+    IntentSlotDataDesc requires two files:
+
+        input_file: file to sequence + label.
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+
+        slot_file: file to slot labels, each line corresponding to
+            slot labels for a sentence in input_file. No header.
+
+    To keep the mapping from label index to label consistent during
+    training and inferencing we require the following files:
+        dicts.intents.csv: each line is an intent. The first line
+            corresponding to the 0 intent label, the second line
+            corresponding to the 1 intent label, and so on.
+
+        dicts.slots.csv: each line is a slot. The first line
+            corresponding to the 0 slot label, the second line
+            corresponding to the 1 slot label, and so on.
+
+    Args:
+        data_dir: the directory of the dataset
+        modes: ['train', 'test', 'dev'],
+        none_slot_label: the label for slots that aren't identified defaulted to 'O'
+        pad_label: the int used for padding. If set to -1, it'll be set to the whatever the None label is.
+    """
+
+    def __init__(
+        self,
+        data_dir: str,
+        modes: List[str] = ['train', 'test', 'dev'],
+        none_slot_label: str = 'O',
+        pad_label: int = -1,
+    ):
+        if not if_exist(data_dir, ['dict.intents.csv', 'dict.slots.csv']):
+            raise FileNotFoundError(
+                "Make sure that your data follows the standard format "
+                "supported by JointIntentSlotDataset. Your data must "
+                "contain dict.intents.csv and dict.slots.csv."
+            )
+
+        self.data_dir = data_dir
+        self.intent_dict_file = self.data_dir + '/dict.intents.csv'
+        self.slot_dict_file = self.data_dir + '/dict.slots.csv'
+
+        self.intents_label_ids = IntentSlotDataDesc.label2idx(self.intent_dict_file)
+        self.num_intents = len(self.intents_label_ids)
+        self.slots_label_ids = IntentSlotDataDesc.label2idx(self.slot_dict_file)
+        self.num_slots = len(self.slots_label_ids)
+
+        infold = self.data_dir
+        for mode in modes:
+            if not if_exist(self.data_dir, [f'{mode}.tsv']):
+                logging.info(f' Stats calculation for {mode} mode' f' is skipped as {mode}.tsv was not found.')
+                continue
+            logging.info(f' Stats calculating for {mode} mode...')
+            slot_file = f'{self.data_dir}/{mode}_slots.tsv'
+            with open(slot_file, 'r') as f:
+                slot_lines = f.readlines()
+
+            input_file = f'{self.data_dir}/{mode}.tsv'
+            with open(input_file, 'r') as f:
+                input_lines = f.readlines()[1:]  # Skipping headers at index 0
+
+            if len(slot_lines) != len(input_lines):
+                raise ValueError(
+                    "Make sure that the number of slot lines match the "
+                    "number of intent lines. There should be a 1-1 "
+                    "correspondence between every slot and intent lines."
+                )
+
+            dataset = list(zip(slot_lines, input_lines))
+
+            raw_slots, raw_intents = [], []
+            for slot_line, input_line in dataset:
+                slot_list = [int(slot) for slot in slot_line.strip().split()]
+                raw_slots.append(slot_list)
+                parts = input_line.strip().split()
+                raw_intents.append(int(parts[-1]))
+
+            logging.info(f'Three most popular intents in {mode} mode:')
+            total_intents, intent_label_freq, max_id = get_label_stats(
+                raw_intents, infold + f'/{mode}_intent_stats.tsv'
+            )
+
+            merged_slots = itertools.chain.from_iterable(raw_slots)
+            logging.info(f'Three most popular slots in {mode} mode:')
+            slots_total, slots_label_freq, max_id = get_label_stats(merged_slots, infold + f'/{mode}_slot_stats.tsv')
+
+            logging.info(f'Total Number of Intents: {total_intents}')
+            logging.info(f'Intent Label Frequencies: {intent_label_freq}')
+            logging.info(f'Total Number of Slots: {slots_total}')
+            logging.info(f'Slots Label Frequencies: {slots_label_freq}')
+
+            if mode == 'train':
+                intent_weights_dict = get_freq_weights(intent_label_freq)
+                logging.info(f'Intent Weights: {intent_weights_dict}')
+                slot_weights_dict = get_freq_weights(slots_label_freq)
+                logging.info(f'Slot Weights: {slot_weights_dict}')
+
+        self.intent_weights = fill_class_weights(intent_weights_dict, self.num_intents - 1)
+        self.slot_weights = fill_class_weights(slot_weights_dict, self.num_slots - 1)
+
+        if pad_label != -1:
+            self.pad_label = pad_label
+        else:
+            if none_slot_label not in self.slots_label_ids:
+                raise ValueError(f'none_slot_label {none_slot_label} not ' f'found in {self.slot_dict_file}.')
+            self.pad_label = self.slots_label_ids[none_slot_label]
+
+    @staticmethod
+    def label2idx(file):
+        lines = open(file, 'r').readlines()
+        lines = [line.strip() for line in lines if line.strip()]
+        labels = {lines[i]: i for i in range(len(lines))}
+        return labels
+
+    @staticmethod
+    def intent_slot_dicts(data_dir):
+        '''
+        Return Intent and slot dictionaries
+        '''
+        intent_dict_file = data_dir + '/dict.intents.csv'
+        slot_dict_file = data_dir + '/dict.slots.csv'
+
+        intents_labels = open(intent_dict_file, 'r').readlines()
+        intents_labels = [line.strip() for line in intents_labels if line.strip()]
+
+        slots_labels = open(slot_dict_file, 'r').readlines()
+        slots_labels = [line.strip() for line in slots_labels if line.strip()]
+
+        return intents_labels, slots_labels

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

@@ -0,0 +1,20 @@
+# 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.data.language_modeling.l2r_lm_dataset import L2RLanguageModelingDataset
+from nemo.collections.nlp.data.language_modeling.lm_bert_dataset import (
+    BertPretrainingDataset,
+    BertPretrainingPreprocessedDataloader,
+)
+from nemo.collections.nlp.data.language_modeling.sentence_dataset import SentenceDataset, TarredSentenceDataset

tts/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py

@@ -0,0 +1,240 @@
+# 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 io
+import json
+from typing import Optional
+
+import braceexpand
+import numpy as np
+import webdataset as wd
+from torch.utils.data import Dataset, IterableDataset
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.data_utils import dataset_to_ids
+from nemo.utils import logging
+
+__all__ = ['L2RLanguageModelingDataset', 'TarredL2RLanguageModelingDataset']
+
+
+class L2RLanguageModelingDataset(Dataset):
+    """
+    Dataset for training and evaluating left-to-right language models.
+    
+    Args:
+        tokenizer: tokenizer, such as WordTokenizer or CharTokenizer
+        dataset: path to data
+        max_seq_length: maximum sequence length (in tokens) of input tensors
+        batch_step: distance (in tokens) between two successive sequences of
+            the text. By default, it is equal to max_seq_length which corresponds
+            to splitting text into disjoint segments covering full dataset
+        use_cache: bool value, defaults to False. Determines whether the preprocessed,
+            tokenized dataset should be cached into a pickle file. If true, cache is saved
+            at the path provided in `dataset`.
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        dataset: str,
+        max_seq_length: Optional[int] = 512,
+        batch_step: Optional[int] = None,
+        use_cache: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.batch_step = batch_step or self.max_seq_length
+        ids = dataset_to_ids(dataset, tokenizer, cache_ids=use_cache, add_bos_eos=False)
+        self.ids = np.array([j for i in ids for j in i])
+
+    def __len__(self):
+        return (len(self.ids) - self.max_seq_length) // self.batch_step
+
+    def __getitem__(self, idx):
+        left = idx * self.batch_step
+        right = left + self.max_seq_length
+        src_ids = self.ids[left:right]
+        labels = self.ids[left + 1 : right + 1]
+        src_mask = (src_ids != self.tokenizer.pad_id).astype(np.float32)
+        return src_ids, src_mask, labels
+
+
+class TarredL2RLanguageModelingDataset(IterableDataset):
+    """
+    A similar Dataset to the L2RLanguageModelingDataset, but which loads tarred tokenized numpy files.
+    Accepts a single JSON metadata manifest file as well as the path(s) to the tarball(s) containing the wav files. 
+    The manifest should contain information such as the number of shards, the number of tokens in the corpus,
+    and the number of tokens contained within each shard of the tarfile(s).
+
+    Valid formats for the text_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar.gz', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
+
+    Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
+    This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
+    Supported opening braces - { <=> (, [, < and the special tag _OP_.
+    Supported closing braces - } <=> ), ], > and the special tag _CL_.
+    For SLURM based tasks, we suggest the use of the special tags for ease of use.
+    See the WebDataset documentation for more information about accepted data and input formats.
+
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+
+    Additionally, please note that the len() of this DataLayer is assumed to be the number of tokens
+    of the text data. An incorrect manifest length may lead to some DataLoader issues down the line.
+
+    Args:
+        text_tar_filepaths: Either a list of tokenized text tarball filepaths, or a
+            string (can be brace-expandable).
+        metadata_path (str): Path to the metadata manifest.
+        tokenizer: tokenizer, such as WordTokenizer or CharTokenizer
+        dataset: path to data
+        max_seq_length: maximum sequence length (in tokens) of input tensors
+        batch_step: distance (in tokens) between two successive sequences of
+            the text. By default, it is equal to max_seq_length which corresponds
+            to splitting text into disjoint segments covering full dataset
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        text_tar_filepaths: str,
+        metadata_path: str,
+        tokenizer,
+        max_seq_length: int = 512,
+        batch_step: int = None,
+        shuffle_n: int = 1,
+        shard_strategy: str = "scatter",
+        global_rank: int = 0,
+        world_size: int = 0,
+    ):
+        super(TarredL2RLanguageModelingDataset, self).__init__()
+
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.batch_step = batch_step or self.max_seq_length
+
+        valid_shard_strategies = ['scatter', 'replicate']
+        if shard_strategy not in valid_shard_strategies:
+            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+
+        self.metadata = metadata
+
+        if isinstance(text_tar_filepaths, str):
+            # Replace '(', '[', '<' and '_OP_' with '{'
+            brace_keys_open = ['(', '[', '<', '_OP_']
+            for bkey in brace_keys_open:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "{")
+
+            # Replace ')', ']', '>' and '_CL_' with '}'
+            brace_keys_close = [')', ']', '>', '_CL_']
+            for bkey in brace_keys_close:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "}")
+
+        if isinstance(text_tar_filepaths, str):
+            # Brace expand
+            text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
+
+        if shard_strategy == 'scatter':
+            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+
+            if len(text_tar_filepaths) % world_size != 0:
+                logging.warning(
+                    f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
+                    f"by number of distributed workers ({world_size})."
+                )
+
+            begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
+            end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
+            text_tar_filepaths = text_tar_filepaths[begin_idx:end_idx]
+            logging.info(
+                "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
+            )
+
+        elif shard_strategy == 'replicate':
+            logging.info("All tarred dataset shards will be replicated across all nodes.")
+
+        else:
+            raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
+
+        self.tarpath = text_tar_filepaths
+
+        # Put together WebDataset
+        self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
+
+        if shuffle_n > 0:
+            self._dataset = self._dataset.shuffle(shuffle_n)
+        else:
+            logging.info("WebDataset will not shuffle files within the tar files.")
+
+        self._dataset = self._dataset.rename(npy='npy', key='__key__').to_tuple('npy', 'key').map(f=self._build_sample)
+
+    def _build_sample(self, tup):
+        # Load file
+        npy, filepath = tup
+        npy = io.BytesIO(npy)
+        data = np.load(npy)  # loads np.int64 vector
+        npy.close()
+
+        # Select random contiguous subsegment
+        idx = np.random.randint(0, (len(data) - self.max_seq_length) // self.batch_step)
+
+        # Slice of data chunk
+        left = idx * self.batch_step
+        right = left + self.max_seq_length
+        data = data[left : right + 1]
+
+        # Create batch
+        src_ids = data[:-1]
+        labels = data[1:]
+        src_mask = (src_ids != self.tokenizer.pad_id).astype(np.float32)
+        return src_ids, src_mask, labels
+
+    def __iter__(self):
+        # We need to wrap an infinite generator since the actual files
+        # within the tar files contains large chunks of contiguous data.
+        # This prevents PTL from early exiting the train loop after exhausting
+        # all of the files in one iteration (though the actual dataset is many
+        # times larger due to each file containing a large chunk of data).
+        dl_iter = iter(self._dataset)
+        while True:
+            try:
+                batch = next(dl_iter)
+                yield batch
+            except StopIteration:
+                dl_iter = iter(self._dataset)
+                continue
+
+    def __len__(self):
+        return (self.metadata['num_text'] - self.max_seq_length) // self.batch_step

tts/nemo/collections/nlp/data/language_modeling/lm_bert_dataset.py

@@ -0,0 +1,406 @@
+# 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 array
+import os
+import pickle
+import random
+from typing import Dict, List, Optional
+
+import h5py
+import numpy as np
+from torch.utils.data import DataLoader, DistributedSampler
+from tqdm import tqdm
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import find_newlines, load_data_indices
+from nemo.core.classes import Dataset
+
+__all__ = ['BertPretrainingDataset', 'BertPretrainingPreprocessedDataloader']
+
+
+def load_h5(input_file: str):
+    return h5py.File(input_file, "r")
+
+
+class BertPretrainingDataset(Dataset):
+    """
+    Dataset for bert pretraining when using data preprocessing including tokenization
+    """
+
+    def __init__(
+        self,
+        tokenizer: object,
+        data_file: str,
+        max_seq_length: Optional[int] = 128,
+        mask_prob: Optional[float] = 0.15,
+        short_seq_prob: Optional[float] = 0.1,
+        seq_a_ratio: Optional[float] = 0.6,
+        sentence_idx_file: Optional[str] = None,
+    ):
+        """
+        Args:
+            tokenizer: tokenizer
+            data_file: path to data
+            max_seq_length: maximum sequence length of input tensors
+            mask_probability: proability to mask token
+            short_seq_prob: probability to create a sequence shorter than max_seq_length
+            seq_a_ratio: ratio between lengths of first and second sequence
+            sentence_idx_file: sentence indices file for caching
+        """
+        self.tokenizer = tokenizer
+
+        # Loading enormous datasets into RAM isn't always feasible -- for
+        # example, the pubmed corpus is 200+ GB, which doesn't fit into RAM on
+        # most computers. To get around this, we store the indices of newlines
+        # in each file so we can seek to and retrieve sentences immediately
+        # from main memory when needed during training.
+
+        # Try and load sentence indices file if already exists
+        sentence_indices, sentence_idx_file, data_dir = load_data_indices(
+            sentence_idx_file, data_file, "sentence_indices"
+        )
+
+        # If sentence indices file doesn't exists, generate and store sentence indices
+        if sentence_indices is None:
+            sentence_indices = {}
+            filenames = [data_file]
+
+            for filename in tqdm(filenames):
+                with open(filename, "rb") as f:
+                    contents = f.read()
+                    newline_indices = find_newlines(contents)
+
+                if os.path.isdir(data_dir):
+                    # Only keep the parts of the filepath that are invariant to
+                    # the dataset's location on disk
+                    filename = os.path.basename(filename)
+
+                # In python, arrays are much more space-efficient than lists
+                sentence_indices[filename] = array.array("I", newline_indices)
+
+            # Save sentence indices so we don't have to do this again
+            with open(sentence_idx_file, "wb") as f:
+                pickle.dump(sentence_indices, f)
+
+        corpus_size = 0
+        empty_files = []
+
+        # Find total number of newlines across entire corpus and remove files
+        # without any newlines
+        for filename in sentence_indices:
+            if len(sentence_indices[filename]) <= 1:
+                empty_files.append(filename)
+            else:
+                corpus_size += len(sentence_indices[filename])
+
+        for filename in empty_files:
+            del sentence_indices[filename]
+
+        self.corpus_size = corpus_size
+        self.dataset = data_dir
+        self.filenames = list(sentence_indices.keys())
+        self.mask_probability = mask_prob
+        self.max_seq_length = max_seq_length
+        self.sentence_indices = sentence_indices
+        self.vocab_size = self.tokenizer.vocab_size
+        self.short_seq_prob = short_seq_prob
+        self.seq_a_ratio = seq_a_ratio
+
+    def __len__(self):
+        return self.corpus_size
+
+    def __getitem__(self, idx: int, min_doc_length: Optional[int] = 16):
+        # Each sequence has three special tokens, as follows:
+        # tokenizer.cls_token <document a> tokenizer.sep_token <document b> tokenizer.eos_token
+        num_special_tokens = 3
+
+        max_num_tokens = self.max_seq_length - num_special_tokens
+        target_seq_length = max_num_tokens
+        if random.random() < self.short_seq_prob:
+            # TODO: maybe introduce an argument to control this.
+            target_seq_length = random.randint(2, max_num_tokens)
+
+        # prefer the seq_a to be slightly longer than seq_b, 0.6 by default
+        target_seq_length_a = int(round(target_seq_length * self.seq_a_ratio))
+        target_seq_length_b = target_seq_length - target_seq_length_a
+
+        def get_document(filepath, offset):
+            # Retrieve a specific line from a file and return as a document
+            if os.path.isdir(self.dataset):
+                filepath = os.path.join(self.dataset, filepath)
+
+            with open(filepath, "rb") as f:
+                f.seek(offset)
+                doc_text = f.readline()[:-1].decode("utf-8", errors="ignore")
+                document = self.tokenizer.text_to_ids(doc_text)
+
+            return document
+
+        def match_target_seq_length(
+            document: str, target_seq_length: int, filename: str, line_idx: int, sentence_indices: Dict[str, dict]
+        ):
+            # If document is shorter than target sequence length,
+            # append the next line or take a random line as replacement.
+            num_lines = len(sentence_indices[filename])
+
+            while len(document) < target_seq_length:
+                if line_idx < (num_lines - 1):
+                    # append the next line
+                    line_idx += 1
+                else:
+                    # current line is the last line, take a random one
+                    line_idx = random.randrange(num_lines)
+                    document = []
+
+                offset = sentence_indices[filename][line_idx]
+                document += get_document(filename, offset)
+
+            return document, line_idx
+
+        # Take sequence A from a random file and a random line
+        a_filename = random.choice(self.filenames)
+        a_line_idx = random.randrange(len(self.sentence_indices[a_filename]))
+        a_line_offset = self.sentence_indices[a_filename][a_line_idx]
+        a_document = get_document(a_filename, a_line_offset)
+        a_document, a_line_idx = match_target_seq_length(
+            a_document, target_seq_length_a, a_filename, a_line_idx, self.sentence_indices
+        )
+
+        is_last_line = a_line_idx >= (len(self.sentence_indices[a_filename]) - 1)
+        # About 50% of the time, B is a random sentence from the corpus
+        take_random_b = (random.random() < 0.5) or is_last_line
+
+        if take_random_b:
+            # This should rarely go for more than one iteration for large
+            # corpora. However, just to be careful, we try to make sure that
+            # the random document is not the same as the document
+            # we're processing.
+            for _ in range(10):
+                b_filename = random.choice(self.filenames)
+                b_line_idx = random.choice(range(len(self.sentence_indices[b_filename])))
+                if b_filename != a_filename:
+                    break
+                else:
+                    # Take another line from the same file
+                    b_line_pos = self.sentence_indices[b_filename][b_line_idx]
+                    a_line_pos = self.sentence_indices[a_filename][a_line_idx]
+                    # TODO unclear about the following check
+                    if abs(b_line_pos - a_line_pos) > max_num_tokens:
+                        break
+                    else:
+                        pass
+        else:
+            b_filename = a_filename
+            b_line_idx = a_line_idx + 1
+
+        is_next = int(not take_random_b)
+        b_line_pos = self.sentence_indices[b_filename][b_line_idx]
+        b_document = get_document(b_filename, b_line_pos)
+        b_document, b_line_idx = match_target_seq_length(
+            b_document, target_seq_length_b, b_filename, b_line_idx, self.sentence_indices
+        )
+
+        def truncate_seq_pair(a, b, max_num_tokens):
+            # Truncates a pair of sequences to a maximum sequence length
+            while (len(a) + len(b)) > max_num_tokens:
+                # Truncate the longer sequence
+                if len(a) > len(b):
+                    trunc_document = a
+                else:
+                    trunc_document = b
+
+                if len(trunc_document) <= 1:
+                    raise ValueError(
+                        "Input text corpora probably too small. "
+                        "Failed to truncate sequence pair to "
+                        "maximum sequence legnth."
+                    )
+
+                # Randomly truncate from the front or the back
+                if random.random() < 0.5:
+                    del trunc_document[0]
+                else:
+                    trunc_document.pop()
+
+        truncate_seq_pair(a_document, b_document, max_num_tokens)
+
+        output_ids = (
+            [self.tokenizer.cls_id] + a_document + [self.tokenizer.sep_id] + b_document + [self.tokenizer.eos_id]
+        )
+
+        input_ids, output_mask = self.mask_ids(output_ids)
+
+        input_mask = np.zeros(self.max_seq_length, dtype=np.long)
+        input_mask[: len(input_ids)] = 1
+
+        input_type_ids = np.zeros(self.max_seq_length, dtype=np.int)
+        input_type_ids[len(a_document) + 2 : len(output_ids) + 1] = 1
+
+        padding_length = max(0, self.max_seq_length - len(input_ids))
+        if padding_length > 0:
+            input_ids.extend([self.tokenizer.pad_id] * padding_length)
+            output_ids.extend([self.tokenizer.pad_id] * padding_length)
+            output_mask.extend([0] * padding_length)
+
+        # TODO: wrap the return value with () for consistent style.
+        return (
+            np.array(input_ids),
+            input_type_ids,
+            np.array(input_mask, dtype=np.long),
+            np.array(output_ids),
+            np.array(output_mask, dtype=np.float32),
+            is_next,
+        )
+
+    def mask_ids(self, ids: List[int]):
+        """
+        Args:
+          ids: list of token ids representing a chunk of text
+        Returns:
+          masked_ids: list of input tokens with some of the entries masked
+            according to the following protocol from the original BERT paper:
+            each token is masked with a probability of 15% and is replaced with
+            1) the [MASK] token 80% of the time,
+            2) random token 10% of the time,
+            3) the same token 10% of the time.
+          output_mask: list of binary variables which indicate what tokens has
+            been masked (to calculate the loss function for these tokens only)
+        """
+
+        # Whole-word masking by default, as it gives better performance.
+        cand_indexes = [[ids[0]]]
+        for tid in ids[1:]:
+            token = self.tokenizer.ids_to_tokens([tid])[0]
+            is_suffix = token.startswith('\u2581')
+            if is_suffix:
+                # group together with its previous token to form a whole-word
+                cand_indexes[-1].append(tid)
+            else:
+                cand_indexes.append([tid])
+
+        masked_ids, output_mask = [], []
+        mask_id = self.tokenizer.token_to_id("[MASK]")
+
+        for word_ids in cand_indexes:
+            is_special = (word_ids[0] == self.tokenizer.cls_id) or (word_ids[0] == self.tokenizer.sep_id)
+            if is_special or (random.random() > self.mask_probability):
+                output_mask.extend([0] * len(word_ids))
+                masked_ids.extend(word_ids)
+            else:
+                output_mask.extend([1] * len(word_ids))
+                p = random.random()
+                # for 80%, replace with mask
+                if p < 0.8:
+                    masked_ids.extend([mask_id] * len(word_ids))
+                # for 10%, replace by a random token
+                elif p < 0.9:
+                    for _ in word_ids:
+                        # randomly select a valid word
+                        random_word = random.randrange(self.vocab_size)
+                        while random_word in (self.tokenizer.cls_id, self.tokenizer.sep_id):
+                            random_word = random.randrange(self.vocab_size)
+                        masked_ids.append(random_word)
+                # for 10%, use same token
+                else:
+                    masked_ids.extend(word_ids)
+
+        return masked_ids, output_mask
+
+
+class BertPretrainingPreprocessedDataset(Dataset):
+    """
+    Dataset for already preprocessed data.
+    """
+
+    def __init__(self, input_file: str, max_predictions_per_seq: int):
+        """
+        Args:
+            input_file: data file in hdf5 format with preprocessed data in array format
+            max_predictions_per_seq: maximum number of masked tokens per sequence. Need to be consistent with data in input file.
+        """
+        self.input_file = input_file
+        self.max_predictions_per_seq = max_predictions_per_seq
+        f = load_h5(input_file)
+        keys = [
+            'input_ids',
+            'input_mask',
+            'segment_ids',
+            'masked_lm_positions',
+            'masked_lm_ids',
+            'next_sentence_labels',
+        ]
+        self.inputs = [np.asarray(f[key][:]) for key in keys]
+        f.close()
+
+    def __len__(self):
+        'Denotes the total number of samples'
+        return len(self.inputs[0])
+
+    def __getitem__(self, index: int):
+        [input_ids, input_mask, segment_ids, masked_lm_positions, masked_lm_ids, next_sentence_labels] = [
+            input[index].astype(np.int64) for input in self.inputs
+        ]
+
+        output_mask = np.zeros_like(input_ids)
+        output_ids = input_ids.copy()
+
+        index = self.max_predictions_per_seq
+        padded_mask_indices = (masked_lm_positions == 0).nonzero()
+        if len(padded_mask_indices[0]) != 0:
+            index = padded_mask_indices[0][0]
+
+        output_mask[masked_lm_positions[:index]] = 1.0
+        output_ids[masked_lm_positions[:index]] = masked_lm_ids[:index]
+
+        # input_mask = np.asarray(input_mask, dtype=np.float32)
+        # output_mask = np.asarray(output_mask, dtype=np.float32)
+        return (input_ids, segment_ids, input_mask, output_ids, output_mask, next_sentence_labels)
+
+
+class BertPretrainingPreprocessedDataloader(DataLoader):
+    """
+    Dataloader for already preprocessed data in hdf5 files that is already in the format expected by BERT model.
+    """
+
+    def __init__(self, data_files: List[str], max_predictions_per_seq: int, batch_size: int, seed: Optional[int] = 42):
+        """
+        Args:
+            data_files: list of data files in hdf5 format with preprocessed data in array format
+            max_predictions_per_seq: maximum number of masked tokens per sequence. Need to be consistent with data in input file.
+            batch_size: batch size per gpu per forward pass
+            seed: seed to ensure each gpu process opens the same data file in each iteration
+        """
+        super().__init__(None, batch_size=batch_size)
+        self.random = random.Random(seed)
+        self.data_files = data_files
+        self.max_predictions_per_seq = max_predictions_per_seq
+
+    # def __len__(self):
+    #     return sum([len(load_h5(data_file)['input_ids']) for data_file in self.data_files])//(self.batch_size)
+
+    def __iter__(self):
+        self.random.shuffle(self.data_files)
+        for data_file in self.data_files:
+            train_data = BertPretrainingPreprocessedDataset(
+                input_file=data_file, max_predictions_per_seq=self.max_predictions_per_seq
+            )
+            train_sampler = DistributedSampler(train_data)
+            # print("---")
+            # print(os.getpid(), train_sampler.rank, train_sampler.num_replicas, train_sampler.num_samples)
+            # print("---")
+            train_dataloader = DataLoader(
+                dataset=train_data, sampler=train_sampler, batch_size=self.batch_size, shuffle=False,
+            )
+            for x in train_dataloader:
+                yield x

tts/nemo/collections/nlp/data/language_modeling/sentence_dataset.py

@@ -0,0 +1,273 @@
+# 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.
+
+"""Pytorch Dataset with sentences packed into batches by length."""
+import io
+import json
+import logging
+import pickle
+from collections import OrderedDict
+from typing import Any
+
+import braceexpand
+import numpy as np
+import webdataset as wd
+from torch.utils.data import IterableDataset
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import dataset_to_ids
+from nemo.core import Dataset
+
+__all__ = ['SentenceDataset', 'TarredSentenceDataset']
+
+
+class SentenceDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer: Any,
+        dataset: Any,
+        tokens_in_batch: int = 1024,
+        clean: bool = False,
+        cache_ids: bool = False,
+        max_seq_length: int = 512,
+        min_seq_length: int = 1,
+    ):
+
+        self.tokenizer = tokenizer
+        self.tokens_in_batch = tokens_in_batch
+
+        ids = dataset_to_ids(dataset, tokenizer, cache_ids=cache_ids)
+        if clean:
+            ids = self.clean(ids, max_tokens=max_seq_length, min_tokens=min_seq_length)
+        self.batch_sent_ids, self.batch_elem_lengths = self.pack_data_into_batches(ids)
+        self.batches = self.pad_batches(ids)
+
+    def __len__(self):
+        return len(self.batches)
+
+    def __getitem__(self, idx):
+        ids = self.batches[idx]
+        mask = (ids != self.tokenizer.pad_id).astype(np.int32)
+        return ids, mask
+
+    def pad_batches(self, ids):
+        """
+        Augments source and target ids in the batches with padding symbol
+        to make the lengths of all sentences in the batches equal.
+        """
+
+        batches = []
+        for batch_elem_len, batch_sent_ids in zip(self.batch_elem_lengths, self.batch_sent_ids):
+            batch = self.tokenizer.pad_id * np.ones((len(batch_sent_ids), batch_elem_len), dtype=np.int)
+            for i, sentence_idx in enumerate(batch_sent_ids):
+                batch[i][: len(ids[sentence_idx])] = ids[sentence_idx]
+            batches.append(batch)
+        return batches
+
+    def pack_data_into_batches(self, ids):
+        """
+        Takes two lists of source and target sentences, sorts them, and packs
+        into batches to minimize the use of padding tokens. Returns a list of
+        batches where each batch contains indices of sentences included into it
+        """
+
+        # create buckets sorted by the number of src tokens
+        # each bucket is also sorted by the number of tgt tokens
+        buckets = {}
+        for i, line_ids in enumerate(ids):
+            len_ = len(line_ids)
+            if len_ not in buckets:
+                buckets[len_] = [i]
+            else:
+                buckets[len_].append(i)
+
+        for b_idx in buckets:
+            buckets[b_idx] = sorted(buckets[b_idx])
+
+        buckets = OrderedDict(sorted(buckets.items()))
+
+        batches = []
+        batch_elem_lengths = []
+        curr_batch = []
+        len_of_longest_sent = 0
+        for sent_len, bucket in buckets.items():
+            for sent_i in bucket:
+                if sent_len * (len(curr_batch) + 1) > self.tokens_in_batch:
+                    if not curr_batch:
+                        raise ValueError(
+                            f"The limitation on number of tokens in batch {self.tokens_in_batch} is too strong."
+                            f"Several sentences contain {sent_len} tokens."
+                        )
+                    batches.append(curr_batch)
+                    batch_elem_lengths.append(sent_len)
+                    curr_batch = []
+                curr_batch.append(sent_i)
+            len_of_longest_sent = sent_len
+        if curr_batch:
+            batches.append(curr_batch)
+            batch_elem_lengths.append(len_of_longest_sent)
+        return batches, batch_elem_lengths
+
+    def clean(self, ids, max_tokens=None, min_tokens=None):
+        """
+        Cleans source and target sentences to get rid of noisy data.
+        Specifically, a pair of sentences is removed if
+          -- either source or target is longer than *max_tokens*
+          -- either source or target is shorter than *min_tokens*
+          -- absolute difference between source and target is larger than
+             *max_tokens_diff*
+          -- one sentence is *max_tokens_ratio* times longer than the other
+        """
+
+        ids_ = []
+        for i in range(len(ids)):
+            len_ = len(ids[i])
+            if (max_tokens is not None and len_ > max_tokens) or (min_tokens is not None and len_ < min_tokens):
+                continue
+            ids_.append(ids[i])
+        return ids_
+
+
+class TarredSentenceDataset(IterableDataset):
+    """
+    A similar Dataset to the SentenceDataset, but which loads tarred tokenized pickle files.
+    Accepts a single JSON metadata file containing the total number of batches
+    as well as the path(s) to the tarball(s) containing the wav files. 
+    Valid formats for the text_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar.gz', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
+    Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
+    This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
+    Supported opening braces - { <=> (, [, < and the special tag _OP_.
+    Supported closing braces - } <=> ), ], > and the special tag _CL_.
+    For SLURM based tasks, we suggest the use of the special tags for ease of use.
+    See the WebDataset documentation for more information about accepted data and input formats.
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+    Additionally, please note that the len() of this DataLayer is assumed to be the number of tokens
+    of the text data. An incorrect manifest length may lead to some DataLoader issues down the line.
+    Args:
+        text_tar_filepaths: Either a list of tokenized text tarball filepaths, or a
+            string (can be brace-expandable).
+        metadata_path (str): Path to the metadata manifest.
+        encoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
+        decoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
+        reverse_lang_direction (bool): When True, swaps the source and target directions when returning minibatches.
+    """
+
+    def __init__(
+        self,
+        text_tar_filepaths: str,
+        metadata_path: str,
+        tokenizer: str,
+        shuffle_n: int = 1,
+        shard_strategy: str = "scatter",
+        global_rank: int = 0,
+        world_size: int = 0,
+    ):
+        super(TarredSentenceDataset, self).__init__()
+
+        self.tokenizer = tokenizer
+        self.pad_id = tokenizer.pad_id
+
+        valid_shard_strategies = ['scatter', 'replicate']
+        if shard_strategy not in valid_shard_strategies:
+            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+
+        self.metadata = metadata
+
+        if isinstance(text_tar_filepaths, str):
+            # Replace '(', '[', '<' and '_OP_' with '{'
+            brace_keys_open = ['(', '[', '<', '_OP_']
+            for bkey in brace_keys_open:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "{")
+
+            # Replace ')', ']', '>' and '_CL_' with '}'
+            brace_keys_close = [')', ']', '>', '_CL_']
+            for bkey in brace_keys_close:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "}")
+
+        if isinstance(text_tar_filepaths, str):
+            # Brace expand
+            text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
+
+        if shard_strategy == 'scatter':
+            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+            if len(text_tar_filepaths) % world_size != 0:
+                logging.warning(
+                    f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
+                    f"by number of distributed workers ({world_size})."
+                )
+            begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
+            end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
+            logging.info('Begin Index : %d' % (begin_idx))
+            logging.info('End Index : %d' % (end_idx))
+            text_tar_filepaths = text_tar_filepaths[begin_idx:end_idx]
+            logging.info(
+                "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
+            )
+
+        elif shard_strategy == 'replicate':
+            logging.info("All tarred dataset shards will be replicated across all nodes.")
+
+        else:
+            raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
+
+        self.tarpath = text_tar_filepaths
+
+        # Put together WebDataset
+        self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
+
+        if shuffle_n > 0:
+            self._dataset = self._dataset.shuffle(shuffle_n)
+        else:
+            logging.info("WebDataset will not shuffle files within the tar files.")
+
+        self._dataset = self._dataset.rename(pkl='pkl', key='__key__').to_tuple('pkl', 'key').map(f=self._build_sample)
+
+    def _build_sample(self, fname):
+        # Load file
+        pkl_file, _ = fname
+        pkl_file = io.BytesIO(pkl_file)
+        data = pickle.load(pkl_file)  # loads np.int64 vector
+        pkl_file.close()
+        ids = data["src"]
+        mask = (ids != self.pad_id).astype(np.int32)
+        return ids, mask
+
+    def __iter__(self):
+        return self._dataset.__iter__()
+
+    def __len__(self):
+        return self.metadata['num_batches']

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

@@ -0,0 +1,18 @@
+# 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.data.machine_translation.machine_translation_dataset import (
+    TarredTranslationDataset,
+    TranslationDataset,
+)

tts/nemo/collections/nlp/data/machine_translation/machine_translation_dataset.py

@@ -0,0 +1,445 @@
+# 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.
+
+"""Pytorch Dataset for training Neural Machine Translation."""
+
+import io
+import json
+import pickle
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import braceexpand
+import numpy as np
+import webdataset as wd
+from torch.utils.data import IterableDataset
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import dataset_to_ids
+from nemo.core import Dataset
+from nemo.utils import logging
+
+__all__ = ['TranslationDataset', 'TarredTranslationDataset']
+
+
+@dataclass
+class TranslationDataConfig:
+    src_file_name: Optional[Any] = None  # Any = str or List[str]
+    tgt_file_name: Optional[Any] = None  # Any = str or List[str]
+    use_tarred_dataset: bool = False
+    tar_files: Optional[Any] = None  # Any = str or List[str]
+    metadata_file: Optional[Any] = None  # Any = str or List[str]
+    lines_per_dataset_fragment: Optional[int] = 1000000
+    num_batches_per_tarfile: Optional[int] = 1000
+    shard_strategy: Optional[str] = 'scatter'
+    tokens_in_batch: int = 512
+    clean: bool = False
+    max_seq_length: int = 512
+    min_seq_length: int = 1
+    cache_ids: bool = False
+    cache_data_per_node: bool = False
+    use_cache: bool = False
+    shuffle: bool = False
+    num_samples: int = -1
+    drop_last: bool = False
+    pin_memory: bool = False
+    num_workers: int = 8
+    reverse_lang_direction: bool = False
+    load_from_tarred_dataset: bool = False
+    metadata_path: Optional[str] = None
+    tar_shuffle_n: int = 100
+    n_preproc_jobs: int = -2
+    tar_file_prefix: str = 'parallel'
+    concat_sampling_technique: Optional[str] = 'temperature'
+    concat_sampling_temperature: Optional[int] = 5
+    concat_sampling_probabilities: Optional[List[float]] = None
+
+
+class TranslationDataset(Dataset):
+    def __init__(
+        self,
+        dataset_src: str,
+        dataset_tgt: str,
+        tokens_in_batch: int = 1024,
+        clean: bool = False,
+        max_seq_length: int = 512,
+        min_seq_length: int = 1,
+        max_seq_length_diff: int = 512,
+        max_seq_length_ratio: int = 512,
+        cache_ids: bool = False,
+        cache_data_per_node: bool = False,
+        use_cache: bool = False,
+        reverse_lang_direction: bool = False,
+        prepend_id: int = None,
+    ):
+        self.dataset_src = dataset_src
+        self.dataset_tgt = dataset_tgt
+        self.tokens_in_batch = tokens_in_batch
+        self.cache_ids = cache_ids
+        self.use_cache = use_cache
+        self.clean = clean
+        self.cache_data_per_node = cache_data_per_node
+        self.max_seq_length = max_seq_length
+        self.min_seq_length = min_seq_length
+        self.max_seq_length_diff = max_seq_length_diff
+        self.max_seq_length_ratio = max_seq_length_ratio
+        self.reverse_lang_direction = reverse_lang_direction
+        self.prepend_id = prepend_id
+
+        # deprecation warnings for cache_ids, use_cache, and cache_data_per_node
+        if self.cache_ids is True or self.use_cache is True or self.cache_data_per_node is True:
+            logging.warning(
+                'Deprecation warning. self.cache_ids, self.use_cache, and self.cache_data_per_node will be removed. Data caching to be done with tarred datasets moving forward.'
+            )
+
+    def batchify(self, tokenizer_src, tokenizer_tgt):
+        src_ids = dataset_to_ids(
+            self.dataset_src,
+            tokenizer_src,
+            cache_ids=self.cache_ids,
+            cache_data_per_node=self.cache_data_per_node,
+            use_cache=self.use_cache,
+        )
+        tgt_ids = dataset_to_ids(
+            self.dataset_tgt,
+            tokenizer_tgt,
+            cache_ids=self.cache_ids,
+            cache_data_per_node=self.cache_data_per_node,
+            use_cache=self.use_cache,
+        )
+        if self.clean:
+            src_ids, tgt_ids = self.clean_src_and_target(
+                src_ids,
+                tgt_ids,
+                max_tokens=self.max_seq_length,
+                min_tokens=self.min_seq_length,
+                max_tokens_diff=self.max_seq_length_diff,
+                max_tokens_ratio=self.max_seq_length_ratio,
+            )
+        self.src_pad_id = tokenizer_src.pad_id
+        self.tgt_pad_id = tokenizer_tgt.pad_id
+
+        self.batch_indices = self.pack_data_into_batches(src_ids, tgt_ids)
+        self.batches = self.pad_batches(src_ids, tgt_ids, self.batch_indices)
+
+    def __len__(self):
+        return len(self.batches)
+
+    def __getitem__(self, idx):
+        src_ids = self.batches[idx]["src"]
+        tgt = self.batches[idx]["tgt"]
+        if self.reverse_lang_direction:
+            src_ids, tgt = tgt, src_ids
+        labels = tgt[:, 1:]
+        tgt_ids = tgt[:, :-1]
+        if self.prepend_id:
+            src_ids = np.insert(src_ids, 0, self.prepend_id, axis=-1)
+        src_mask = (src_ids != self.src_pad_id).astype(np.int32)
+        tgt_mask = (tgt_ids != self.tgt_pad_id).astype(np.int32)
+        return src_ids, src_mask, tgt_ids, tgt_mask, labels
+
+    def pad_batches(self, src_ids, tgt_ids, batch_indices):
+        """
+        Augments source and target ids in the batches with padding symbol
+        to make the lengths of all sentences in the batches equal.
+        """
+
+        batches = {}
+        for batch_idx, b in enumerate(batch_indices):
+            src_len = max([len(src_ids[i]) for i in b])
+            tgt_len = max([len(tgt_ids[i]) for i in b])
+            src_ids_ = self.src_pad_id * np.ones((len(b), src_len), dtype=np.int)
+            tgt_ids_ = self.tgt_pad_id * np.ones((len(b), tgt_len), dtype=np.int)
+            for i, sentence_idx in enumerate(b):
+                src_ids_[i][: len(src_ids[sentence_idx])] = src_ids[sentence_idx]
+                tgt_ids_[i][: len(tgt_ids[sentence_idx])] = tgt_ids[sentence_idx]
+            batches[batch_idx] = {"src": src_ids_, "tgt": tgt_ids_}
+        return batches
+
+    def pack_data_into_batches(self, src_ids, tgt_ids):
+        """
+        Takes two lists of source and target sentences, sorts them, and packs
+        into batches to minimize the use of padding tokens. Returns a list of
+        batches where each batch contains indices of sentences included into it
+        """
+
+        # create buckets sorted by the number of src tokens
+        # each bucket is also sorted by the number of tgt tokens
+        buckets = {}
+        for i, src_id in enumerate(src_ids):
+            src_len, tgt_len = len(src_id), len(tgt_ids[i])
+            if src_len not in buckets:
+                buckets[src_len] = [(tgt_len, i)]
+            else:
+                buckets[src_len].append((tgt_len, i))
+
+        for b_idx in buckets:
+            buckets[b_idx] = sorted(buckets[b_idx])
+
+        buckets = OrderedDict(sorted(buckets.items()))
+        indices = list(buckets.keys())
+
+        batches = [[]]
+        num_batches = 0
+        batch_size = 0
+        i = 0
+        src_len = 0
+        tgt_len = 0
+
+        while i < len(buckets):
+            while buckets[indices[i]]:
+
+                i_src = max(src_len, indices[i])
+                i_tgt = max(tgt_len, buckets[indices[i]][0][0])
+
+                try:
+                    ip1_src = max(src_len, indices[i + 1])
+                    ip1_tgt = max(tgt_len, buckets[indices[i + 1]][0][0])
+                except IndexError:
+                    ip1_src = i_src + 1
+                    ip1_tgt = i_tgt + 1
+
+                if i_src + i_tgt <= ip1_src + ip1_tgt:
+                    src_len = i_src
+                    tgt_len = i_tgt
+                    _, idx = buckets[indices[i]].pop(0)
+                else:
+                    src_len = ip1_src
+                    tgt_len = ip1_tgt
+                    _, idx = buckets[indices[i + 1]].pop(0)
+
+                batches[num_batches].append(idx)
+                batch_size += 1
+
+                if batch_size * (src_len + tgt_len) > self.tokens_in_batch:
+
+                    num_examples_to_split = len(batches[num_batches])
+                    batches_to_evict = 8 * ((num_examples_to_split - 1) // 8)
+
+                    if batches_to_evict == 0:
+                        batches_to_evict = num_examples_to_split
+
+                    batches.append(batches[num_batches][batches_to_evict:])
+                    batches[num_batches] = batches[num_batches][:batches_to_evict]
+                    batch_size = num_examples_to_split - batches_to_evict
+
+                    num_batches += 1
+                    if batch_size > 0:
+                        src_len = max([len(src_ids[j]) for j in batches[num_batches]])
+                        tgt_len = max([len(tgt_ids[j]) for j in batches[num_batches]])
+                    else:
+                        src_len = 0
+                        tgt_len = 0
+                    break
+
+            if not buckets[indices[i]]:
+                i = i + 1
+
+        if not batches[-1]:
+            batches.pop(-1)
+
+        return batches
+
+    def clean_src_and_target(
+        self,
+        src_ids,
+        tgt_ids,
+        max_tokens=None,
+        min_tokens=None,
+        max_tokens_diff=None,
+        max_tokens_ratio=None,
+        filter_equal_src_and_dest=False,
+    ):
+        """
+        Cleans source and target sentences to get rid of noisy data.
+        Specifically, a pair of sentences is removed if
+          -- either source or target is longer than *max_tokens*
+          -- either source or target is shorter than *min_tokens*
+          -- absolute difference between source and target is larger than
+             *max_tokens_diff*
+          -- one sentence is *max_tokens_ratio* times longer than the other
+        """
+
+        if len(src_ids) != len(tgt_ids):
+            raise ValueError("Source and target corpora have different lengths!")
+        src_ids_, tgt_ids_ = [], []
+        for i in range(len(src_ids)):
+            src_len, tgt_len = len(src_ids[i]), len(tgt_ids[i])
+            if (
+                (max_tokens is not None and (src_len > max_tokens or tgt_len > max_tokens))
+                or (min_tokens is not None and (src_len < min_tokens or tgt_len < min_tokens))
+                or (filter_equal_src_and_dest and src_ids[i] == tgt_ids[i])
+                or (max_tokens_diff is not None and np.abs(src_len - tgt_len) > max_tokens_diff)
+            ):
+                continue
+            if max_tokens_ratio is not None:
+                ratio = max(src_len - 2, 1) / max(tgt_len - 2, 1)
+                if ratio > max_tokens_ratio or ratio < (1 / max_tokens_ratio):
+                    continue
+            src_ids_.append(src_ids[i])
+            tgt_ids_.append(tgt_ids[i])
+        return src_ids_, tgt_ids_
+
+
+class TarredTranslationDataset(IterableDataset):
+    """
+    A similar Dataset to the TranslationDataset, but which loads tarred tokenized pickle files.
+    Accepts a single JSON metadata file containing the total number of batches
+    as well as the path(s) to the tarball(s) containing the pickled parallel dataset batch files.
+    Valid formats for the text_tar_filepaths argument include:
+    (1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar', or
+    (2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
+    Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference.
+    This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements.
+    Supported opening braces - { <=> (, [, < and the special tag _OP_.
+    Supported closing braces - } <=> ), ], > and the special tag _CL_.
+    For SLURM based tasks, we suggest the use of the special tags for ease of use.
+    See the WebDataset documentation for more information about accepted data and input formats.
+    If using multiple processes the number of shards should be divisible by the number of workers to ensure an
+    even split among workers. If it is not divisible, logging will give a warning but training will proceed.
+    Additionally, please note that the len() of this DataLayer is assumed to be the number of tokens
+    of the text data. An incorrect manifest length may lead to some DataLoader issues down the line.
+    Args:
+        text_tar_filepaths: Either a list of tokenized text tarball filepaths, or a
+            string (can be brace-expandable).
+        metadata_path (str): Path to the metadata manifest.
+        encoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
+        decoder_tokenizer: Autokenizer wrapped BPE tokenizer model, such as YTTM
+        shuffle_n (int): How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+            Defaults to 0.
+        shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp.
+            -   `scatter`: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   `replicate`: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
+        reverse_lang_direction (bool): When True, swaps the source and target directions when returning minibatches.
+        prepend_id (int): Prepends the specificed token id to the start of every source sentence. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        text_tar_filepaths: str,
+        metadata_path: str,
+        encoder_tokenizer: str,
+        decoder_tokenizer: str,
+        shuffle_n: int = 1,
+        shard_strategy: str = "scatter",
+        global_rank: int = 0,
+        world_size: int = 1,
+        reverse_lang_direction: bool = False,
+        prepend_id: int = None,
+    ):
+        super(TarredTranslationDataset, self).__init__()
+
+        self.encoder_tokenizer = encoder_tokenizer
+        self.decoder_tokenizer = decoder_tokenizer
+        self.reverse_lang_direction = reverse_lang_direction
+        self.src_pad_id = encoder_tokenizer.pad_id
+        self.tgt_pad_id = decoder_tokenizer.pad_id
+        self.prepend_id = prepend_id
+
+        valid_shard_strategies = ['scatter', 'replicate']
+        if shard_strategy not in valid_shard_strategies:
+            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+
+        self.metadata = metadata
+
+        if isinstance(text_tar_filepaths, str):
+            # Replace '(', '[', '<' and '_OP_' with '{'
+            brace_keys_open = ['(', '[', '<', '_OP_']
+            for bkey in brace_keys_open:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "{")
+
+            # Replace ')', ']', '>' and '_CL_' with '}'
+            brace_keys_close = [')', ']', '>', '_CL_']
+            for bkey in brace_keys_close:
+                if bkey in text_tar_filepaths:
+                    text_tar_filepaths = text_tar_filepaths.replace(bkey, "}")
+
+        if isinstance(text_tar_filepaths, str):
+            # Brace expand
+            text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
+
+        if shard_strategy == 'scatter':
+            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+            if len(text_tar_filepaths) % world_size != 0:
+                logging.warning(
+                    f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
+                    f"by number of distributed workers ({world_size})."
+                )
+            begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
+            end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
+            logging.info('Begin Index : %d' % (begin_idx))
+            logging.info('End Index : %d' % (end_idx))
+            text_tar_filepaths = text_tar_filepaths[begin_idx:end_idx]
+            logging.info(
+                "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
+            )
+            self.length = self.metadata['num_batches'] // world_size
+
+        elif shard_strategy == 'replicate':
+            logging.info("All tarred dataset shards will be replicated across all nodes.")
+            self.length = self.metadata['num_batches']
+
+        else:
+            raise ValueError(f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}")
+
+        self.tarpath = text_tar_filepaths
+
+        # Put together WebDataset
+        self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
+
+        if shuffle_n > 0:
+            self._dataset = self._dataset.shuffle(shuffle_n)
+        else:
+            logging.info("WebDataset will not shuffle files within the tar files.")
+
+        self._dataset = self._dataset.rename(pkl='pkl', key='__key__').to_tuple('pkl', 'key').map(f=self._build_sample)
+
+    def _build_sample(self, fname):
+        # Load file
+        pkl_file, _ = fname
+        pkl_file = io.BytesIO(pkl_file)
+        data = pickle.load(pkl_file)  # loads np.int64 vector
+        pkl_file.close()
+        src_ids = data["src"]
+        tgt = data["tgt"]
+        if self.reverse_lang_direction:
+            src_ids, tgt = tgt, src_ids
+        labels = tgt[:, 1:]
+        tgt_ids = tgt[:, :-1]
+        if self.prepend_id:
+            src_ids = np.insert(src_ids, 0, self.prepend_id, axis=-1)
+        src_mask = (src_ids != self.src_pad_id).astype(np.int32)
+        tgt_mask = (tgt_ids != self.tgt_pad_id).astype(np.int32)
+        return src_ids, src_mask, tgt_ids, tgt_mask, labels
+
+    def __iter__(self):
+        return self._dataset.__iter__()
+
+    def __len__(self):
+        return self.length

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

@@ -0,0 +1,18 @@
+# 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.data.neural_machine_translation.neural_machine_translation_dataset import (
+    NeuralMachineTranslationDataset,
+)

tts/nemo/collections/nlp/data/neural_machine_translation/neural_machine_translation_dataset.py

@@ -0,0 +1,146 @@
+# 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.
+
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+from nemo.collections.common.tokenizers import AutoTokenizer
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+
+__all__ = ["NeuralMachineTranslationDataset"]
+
+
+class NeuralMachineTranslationDataset(Dataset):
+    """A dataset class that converts raw data to a dataset that can be used by NeuralMachineTranslationModel.
+
+    Args:
+        filepath: .tsv file to sequence + label.
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        encoder_tokenizer: encoder tokenizer object such as AutoTokenizer
+        decoder_tokenizer: decoder tokenizer object. If using BART or end to end model, set this to encoder_tokenizer
+        max_seq_length: max sequence length including bos and eos tokens
+        num_samples: number of samples you want to use for the dataset. If -1, use all dataset. Useful for testing.
+        convert_labels: if true, converts labels for masked lm and updates pad_id to -100
+            for hf masked loss
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            "input_ids": NeuralType(("B", "T"), ChannelType()),
+            "attention_mask": NeuralType(("B", "T"), MaskType()),
+            "decoder_input_ids": NeuralType(("B", "T"), ChannelType()),
+            "lm_labels": NeuralType(("B", "T"), ChannelType()),
+        }
+
+    def __init__(
+        self,
+        filepath: str,
+        encoder_tokenizer: AutoTokenizer,
+        decoder_tokenizer: AutoTokenizer,
+        encoder_add_special_tokens: bool,
+        decoder_add_special_tokens: bool,
+        max_seq_length: int,
+        num_samples: int = -1,
+        convert_labels: bool = False,
+    ):
+        self.filepath = filepath
+        self.encoder_tokenizer = encoder_tokenizer
+        self.decoder_tokenizer = decoder_tokenizer
+        self.encoder_add_special_tokens = encoder_add_special_tokens
+        self.decoder_add_special_tokens = decoder_add_special_tokens
+        self.max_seq_length = max_seq_length
+        self.num_samples = num_samples
+        self.convert_labels = convert_labels
+
+        if num_samples == 0:
+            raise ValueError("num_samples has to be positive.", num_samples)
+
+        if self.max_seq_length and self.max_seq_length <= 2:
+            self.max_seq_length = None
+
+        if not os.path.exists(filepath):
+            raise FileNotFoundError(
+                f"{filepath} not found. The filepath must be set in train_ds.filepath and validation_ds.filepath."
+            )
+
+        with open(filepath) as f:
+            lines = f.readlines()[1:]
+
+        if num_samples > 0:
+            lines = lines[:num_samples]
+
+        input_ids, input_masks, label_ids = [], [], []
+        for line in lines:
+            try:
+                sentence, label = line.split("\t")
+            except ValueError:
+                raise ValueError("Each line of input file should contain the format [sentence][tab][label].")
+            ids, mask = self.text_to_ids(
+                sentence, tokenizer=encoder_tokenizer, add_special_tokens=encoder_add_special_tokens
+            )
+            input_ids.append(ids)
+            input_masks.append(mask)
+            label_ids.append(
+                self.text_to_ids(label, tokenizer=decoder_tokenizer, add_special_tokens=decoder_add_special_tokens)[0]
+            )
+
+        self.input_ids = np.asarray(input_ids)
+        self.input_masks = np.asarray(input_masks)
+        self.label_ids = np.asarray(label_ids)
+
+    def text_to_ids(
+        self, text: str, tokenizer: AutoTokenizer, add_special_tokens=False
+    ) -> Tuple[List[int], List[int]]:
+        """Converts text to ids. Truncates and adds padding."""
+        text_tokens = tokenizer.text_to_ids(text.strip())
+        num_special_tokens = 2 if add_special_tokens else 0
+
+        if self.max_seq_length and self.max_seq_length > num_special_tokens:
+            text_tokens = text_tokens[: self.max_seq_length - num_special_tokens]
+        if add_special_tokens:
+            text_tokens = [tokenizer.bos_id] + text_tokens + [tokenizer.eos_id]
+        mask = [1] * len(text_tokens)
+
+        if self.max_seq_length and self.max_seq_length > num_special_tokens:
+            pad_length = self.max_seq_length - len(text_tokens)
+            text_tokens += [tokenizer.pad_id] * pad_length
+            mask += [0] * pad_length
+
+        return text_tokens, mask
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def convert_label_ids(self, label_ids: List[int]) -> Tuple[List[int], List[int]]:
+        decoder_input_ids = label_ids[:-1]
+        lm_labels = label_ids[1:].copy()
+        lm_labels[label_ids[1:] == self.decoder_tokenizer.pad_id] = -100  # for huggingface masked lm loss
+        return decoder_input_ids, lm_labels
+
+    def __getitem__(self, idx):
+        if self.convert_labels:
+            decoder_input_ids, lm_labels = self.convert_label_ids(self.label_ids[idx])
+        else:
+            decoder_input_ids = self.label_ids[idx]
+            lm_labels = self.label_ids[idx]
+        return self.input_ids[idx], self.input_masks[idx], decoder_input_ids, lm_labels

tts/nemo/collections/nlp/data/question_answering_squad/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/nlp/data/question_answering_squad/qa_dataset.py

@@ -0,0 +1,443 @@
+# 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 collections
+import json
+import os
+import pickle
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from nemo.collections.common.parts.utils import _compute_softmax
+from nemo.collections.nlp.data.question_answering_squad.qa_squad_processing import (
+    EVALUATION_MODE,
+    INFERENCE_MODE,
+    TRAINING_MODE,
+    SquadProcessor,
+    apply_no_ans_threshold,
+    convert_examples_to_features,
+    exact_match_score,
+    f1_score,
+    find_all_best_thresh,
+    get_best_indexes,
+    get_final_text,
+    make_eval_dict,
+    merge_eval,
+    normalize_answer,
+)
+from nemo.core.classes import Dataset
+from nemo.utils import logging
+
+__all__ = ['SquadDataset']
+
+
+class SquadDataset(Dataset):
+    """
+    Creates SQuAD dataset for Question Answering.
+    Args:
+        data_file (str): train.*.json eval.*.json or test.*.json.
+        tokenizer (obj): Tokenizer object, e.g. AutoTokenizer.
+        version_2_with_negative (bool): True if training should allow
+            unanswerable questions.
+        doc_stride (int): When splitting up a long document into chunks,
+            how much stride to take between chunks.
+        max_query_length (iny): All training files which have a duration less
+            than min_duration are dropped. Can't be used if the `utt2dur` file
+            does not exist. Defaults to None.
+        max_seq_length (int): All training files which have a duration more
+            than max_duration are dropped. Can't be used if the `utt2dur` file
+            does not exist. Defaults to None.
+        num_samples: number of samples you want to use for the dataset.
+            If -1, use all dataset. Useful for testing.
+        mode (str): Use TRAINING_MODE/EVALUATION_MODE/INFERENCE_MODE to define between
+            training, evaluation and inference dataset.
+        use_cache (bool): Caches preprocessed data for future usage
+    """
+
+    def __init__(
+        self,
+        data_file: str,
+        tokenizer: object,
+        doc_stride: int,
+        max_query_length: int,
+        max_seq_length: int,
+        version_2_with_negative: bool,
+        num_samples: int,
+        mode: str,
+        use_cache: bool,
+    ):
+        self.tokenizer = tokenizer
+        self.version_2_with_negative = version_2_with_negative
+        self.processor = SquadProcessor(data_file=data_file, mode=mode)
+        self.mode = mode
+        if mode not in [TRAINING_MODE, EVALUATION_MODE, INFERENCE_MODE]:
+            raise ValueError(
+                f"mode should be either {TRAINING_MODE}, {EVALUATION_MODE}, {INFERENCE_MODE} but got {mode}"
+            )
+        self.examples = self.processor.get_examples()
+
+        vocab_size = getattr(tokenizer, "vocab_size", 0)
+        cached_features_file = (
+            data_file
+            + '_cache'
+            + '_{}_{}_{}_{}_{}_{}_{}'.format(
+                mode,
+                tokenizer.name,
+                str(vocab_size),
+                str(max_seq_length),
+                str(doc_stride),
+                str(max_query_length),
+                str(num_samples),
+            )
+        )
+
+        # check number of samples. Should be either -1 not to limit or positive number
+        if num_samples == 0:
+            raise ValueError(
+                f"num_samples has to be positive or -1 (to use the entire dataset), however got {num_samples}."
+            )
+        elif num_samples > 0:
+            self.examples = self.examples[:num_samples]
+
+        if use_cache and os.path.exists(cached_features_file):
+            logging.info(f"loading from {cached_features_file}")
+            with open(cached_features_file, "rb") as reader:
+                self.features = pickle.load(reader)
+        else:
+            logging.info(f"Preprocessing data.")
+
+            self.features = convert_examples_to_features(
+                examples=self.examples,
+                tokenizer=tokenizer,
+                max_seq_length=max_seq_length,
+                doc_stride=doc_stride,
+                max_query_length=max_query_length,
+                has_groundtruth=mode != INFERENCE_MODE,
+            )
+
+            if use_cache:
+                master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+                if master_device:
+                    logging.info("  Saving train features into cached file %s", cached_features_file)
+                    with open(cached_features_file, "wb") as writer:
+                        pickle.dump(self.features, writer)
+
+    def __len__(self):
+        return len(self.features)
+
+    def __getitem__(self, idx):
+        feature = self.features[idx]
+        if self.mode == INFERENCE_MODE:
+            return (
+                np.array(feature.input_ids),
+                np.array(feature.segment_ids),
+                np.array(feature.input_mask),
+                np.array(feature.unique_id),
+            )
+        else:
+            return (
+                np.array(feature.input_ids),
+                np.array(feature.segment_ids),
+                np.array(feature.input_mask),
+                np.array(feature.unique_id),
+                np.array(feature.start_position),
+                np.array(feature.end_position),
+            )
+
+    def get_predictions(
+        self,
+        unique_ids: List[int],
+        start_logits: List[List[float]],
+        end_logits: List[List[float]],
+        n_best_size: int,
+        max_answer_length: int,
+        do_lower_case: bool,
+        version_2_with_negative: bool,
+        null_score_diff_threshold: float,
+    ):
+        example_index_to_features = collections.defaultdict(list)
+
+        unique_id_to_pos = {}
+        for index, unique_id in enumerate(unique_ids):
+            unique_id_to_pos[unique_id] = index
+
+        for feature in self.features:
+            example_index_to_features[feature.example_index].append(feature)
+
+        _PrelimPrediction = collections.namedtuple(
+            "PrelimPrediction", ["feature_index", "start_index", "end_index", "start_logit", "end_logit"]
+        )
+
+        all_predictions = collections.OrderedDict()
+        all_nbest_json = collections.OrderedDict()
+        scores_diff_json = collections.OrderedDict()
+        for (example_index, example) in enumerate(self.examples):
+
+            # finish this loop if we went through all batch examples
+            if example_index >= len(unique_ids):
+                break
+
+            features = example_index_to_features[example_index]
+
+            prelim_predictions = []
+            # keep track of the minimum score of null start+end of position 0
+            # large and positive
+            score_null = 1000000
+            # the paragraph slice with min null score
+            min_null_feature_index = 0
+            # start logit at the slice with min null score
+            null_start_logit = 0
+            # end logit at the slice with min null score
+            null_end_logit = 0
+            for (feature_index, feature) in enumerate(features):
+                pos = unique_id_to_pos[feature.unique_id]
+                start_indexes = get_best_indexes(start_logits[pos], n_best_size)
+                end_indexes = get_best_indexes(end_logits[pos], n_best_size)
+                # if we could have irrelevant answers,
+                # get the min score of irrelevant
+                if version_2_with_negative:
+                    feature_null_score = start_logits[pos][0] + end_logits[pos][0]
+                    if feature_null_score < score_null:
+                        score_null = feature_null_score
+                        min_null_feature_index = feature_index
+                        null_start_logit = start_logits[pos][0]
+                        null_end_logit = end_logits[pos][0]
+                for start_index in start_indexes:
+                    for end_index in end_indexes:
+                        # We could hypothetically create invalid predictions,
+                        # e.g., predict that the start of the span is in the
+                        # question. We throw out all invalid predictions.
+                        if start_index >= len(feature.tokens):
+                            continue
+                        if end_index >= len(feature.tokens):
+                            continue
+                        if start_index not in feature.token_to_orig_map:
+                            continue
+                        if end_index not in feature.token_to_orig_map:
+                            continue
+                        if not feature.token_is_max_context.get(start_index, False):
+                            continue
+                        if end_index < start_index:
+                            continue
+                        length = end_index - start_index + 1
+                        if length > max_answer_length:
+                            continue
+                        prelim_predictions.append(
+                            _PrelimPrediction(
+                                feature_index=feature_index,
+                                start_index=start_index,
+                                end_index=end_index,
+                                start_logit=start_logits[pos][start_index],
+                                end_logit=end_logits[pos][end_index],
+                            )
+                        )
+
+            if version_2_with_negative:
+                prelim_predictions.append(
+                    _PrelimPrediction(
+                        feature_index=min_null_feature_index,
+                        start_index=0,
+                        end_index=0,
+                        start_logit=null_start_logit,
+                        end_logit=null_end_logit,
+                    )
+                )
+            prelim_predictions = sorted(prelim_predictions, key=lambda x: (x.start_logit + x.end_logit), reverse=True)
+
+            _NbestPrediction = collections.namedtuple("NbestPrediction", ["text", "start_logit", "end_logit"])
+
+            seen_predictions = {}
+            nbest = []
+            for pred in prelim_predictions:
+                if len(nbest) >= n_best_size:
+                    break
+                feature = features[pred.feature_index]
+                if pred.start_index > 0:  # this is a non-null prediction
+                    tok_tokens = feature.tokens[pred.start_index : (pred.end_index + 1)]
+                    orig_doc_start = feature.token_to_orig_map[pred.start_index]
+                    orig_doc_end = feature.token_to_orig_map[pred.end_index]
+                    orig_tokens = example.doc_tokens[orig_doc_start : (orig_doc_end + 1)]
+                    tok_text = " ".join(tok_tokens)
+
+                    # De-tokenize WordPieces that have been split off.
+                    tok_text = tok_text.replace(" ##", "")
+                    tok_text = tok_text.replace("##", "")
+
+                    # Clean whitespace
+                    tok_text = tok_text.strip()
+                    tok_text = " ".join(tok_text.split())
+                    orig_text = " ".join(orig_tokens)
+
+                    final_text = get_final_text(tok_text, orig_text, do_lower_case)
+                    if final_text in seen_predictions:
+                        continue
+
+                    seen_predictions[final_text] = True
+                else:
+                    final_text = ""
+                    seen_predictions[final_text] = True
+
+                nbest.append(_NbestPrediction(text=final_text, start_logit=pred.start_logit, end_logit=pred.end_logit))
+
+            # if we didn't include the empty option in the n-best, include it
+            if version_2_with_negative:
+                if "" not in seen_predictions:
+                    nbest.append(_NbestPrediction(text="", start_logit=null_start_logit, end_logit=null_end_logit))
+
+                # In very rare edge cases we could only
+                # have single null pred. We just create a nonce prediction
+                # in this case to avoid failure.
+                if len(nbest) == 1:
+                    nbest.insert(0, _NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0))
+
+            # In very rare edge cases we could have no valid predictions. So we
+            # just create a nonce prediction in this case to avoid failure.
+            if not nbest:
+                nbest.append(_NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0))
+
+            assert len(nbest) >= 1
+
+            total_scores = []
+            best_non_null_entry = None
+            for entry in nbest:
+                total_scores.append(entry.start_logit + entry.end_logit)
+                if not best_non_null_entry:
+                    if entry.text:
+                        best_non_null_entry = entry
+
+            probs = _compute_softmax(total_scores)
+
+            nbest_json = []
+            for (i, entry) in enumerate(nbest):
+                output = collections.OrderedDict()
+                output["question"] = example.question_text
+                output["text"] = entry.text
+                output["probability"] = probs[i]
+                output["start_logit"] = (
+                    entry.start_logit if isinstance(entry.start_logit, float) else list(entry.start_logit)
+                )
+                output["end_logit"] = entry.end_logit if isinstance(entry.end_logit, float) else list(entry.end_logit)
+                nbest_json.append(output)
+
+            assert len(nbest_json) >= 1
+            if not version_2_with_negative:
+                all_predictions[example.qas_id] = nbest_json[0]["text"]
+            else:
+                # predict "" iff the null score -
+                # the score of best non-null > threshold
+                score_diff = score_null - best_non_null_entry.start_logit - (best_non_null_entry.end_logit)
+                scores_diff_json[example.qas_id] = score_diff
+                if score_diff > null_score_diff_threshold:
+                    all_predictions[example.qas_id] = ""
+                else:
+                    all_predictions[example.qas_id] = best_non_null_entry.text
+            all_nbest_json[example.qas_id] = nbest_json
+
+        return all_predictions, all_nbest_json, scores_diff_json
+
+    def evaluate_predictions(
+        self,
+        all_predictions: Dict[str, str],
+        no_answer_probs: Optional[float] = None,
+        no_answer_probability_threshold: float = 1.0,
+    ):
+        qas_id_to_has_answer = {
+            example.qas_id: bool(example.answers) for example in self.examples[: len(all_predictions)]
+        }
+        has_answer_qids = [qas_id for qas_id, has_answer in qas_id_to_has_answer.items() if has_answer]
+        no_answer_qids = [qas_id for qas_id, has_answer in qas_id_to_has_answer.items() if not has_answer]
+        if no_answer_probs is None:
+            no_answer_probs = {k: 0.0 for k in all_predictions}
+
+        exact, f1 = self.get_raw_scores(all_predictions)
+
+        exact_threshold = apply_no_ans_threshold(
+            exact, no_answer_probs, qas_id_to_has_answer, no_answer_probability_threshold
+        )
+        f1_threshold = apply_no_ans_threshold(
+            f1, no_answer_probs, qas_id_to_has_answer, no_answer_probability_threshold
+        )
+
+        evaluation = make_eval_dict(exact_threshold, f1_threshold)
+
+        if has_answer_qids:
+            has_ans_eval = make_eval_dict(exact_threshold, f1_threshold, qid_list=has_answer_qids)
+            merge_eval(evaluation, has_ans_eval, "HasAns")
+
+        if no_answer_qids:
+            no_ans_eval = make_eval_dict(exact_threshold, f1_threshold, qid_list=no_answer_qids)
+            merge_eval(evaluation, no_ans_eval, "NoAns")
+
+        if no_answer_probs:
+            find_all_best_thresh(evaluation, all_predictions, exact, f1, no_answer_probs, qas_id_to_has_answer)
+
+        return evaluation["best_exact"], evaluation["best_f1"]
+
+    def get_raw_scores(self, preds: Dict[str, str]):
+        """
+        Computes the exact and f1 scores from the examples
+        and the model predictions
+        """
+        exact_scores = {}
+        f1_scores = {}
+
+        for example in self.examples:
+            qas_id = example.qas_id
+            gold_answers = [answer["text"] for answer in example.answers if normalize_answer(answer["text"])]
+
+            if not gold_answers:
+                # For unanswerable questions,
+                # only correct answer is empty string
+                gold_answers = [""]
+
+            if qas_id not in preds:
+                logging.warning("Missing prediction for %s" % qas_id)
+                continue
+
+            prediction = preds[qas_id]
+            exact_scores[qas_id] = max(exact_match_score(a, prediction) for a in gold_answers)
+            f1_scores[qas_id] = max(f1_score(a, prediction) for a in gold_answers)
+
+        return exact_scores, f1_scores
+
+    def evaluate(
+        self,
+        unique_ids: List[str],
+        start_logits: List[List[float]],
+        end_logits: List[List[float]],
+        n_best_size: int,
+        max_answer_length: int,
+        do_lower_case: bool,
+        version_2_with_negative: bool,
+        null_score_diff_threshold: float,
+    ):
+        (all_predictions, all_nbest_json, scores_diff_json) = self.get_predictions(
+            unique_ids,
+            start_logits,
+            end_logits,
+            n_best_size,
+            max_answer_length,
+            do_lower_case,
+            version_2_with_negative,
+            null_score_diff_threshold,
+        )
+
+        exact_match, f1 = self.evaluate_predictions(all_predictions)
+
+        return exact_match, f1, all_predictions, all_nbest_json

tts/nemo/collections/nlp/data/question_answering_squad/qa_squad_processing.py

@@ -0,0 +1,644 @@
+# 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 collections
+import json
+from typing import Dict, List, Optional
+
+from tqdm import tqdm
+from transformers.models.bert.tokenization_bert import BasicTokenizer
+
+from nemo.collections.nlp.data.data_utils import DataProcessor, is_whitespace, normalize_answer
+from nemo.utils import logging
+
+"""
+Utility functions for Question Answering NLP tasks
+Some parts of this code were adapted from the HuggingFace library at
+https://github.com/huggingface/transformers
+"""
+
+TRAINING_MODE = "train"
+EVALUATION_MODE = "eval"
+INFERENCE_MODE = "infer"
+
+
+def _get_tokens(s):
+    """get normalized tokens"""
+    if not s:
+        return []
+    return normalize_answer(s).split()
+
+
+def get_best_indexes(logits, n_best_size):
+    """Get the n-best logits from a list."""
+    index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)
+
+    best_indexes = []
+    for i in range(len(index_and_score)):
+        if i >= n_best_size:
+            break
+        best_indexes.append(index_and_score[i][0])
+    return best_indexes
+
+
+def get_final_text(pred_text: str, orig_text: str, do_lower_case: bool, verbose_logging: bool = False):
+    """Project the tokenized prediction back to the original text.
+    When we created the data, we kept track of the alignment between original
+    (whitespace tokenized) tokens and our WordPiece tokenized tokens. So
+    now `orig_text` contains the span of our original text corresponding to
+    the span that we predicted.
+
+    However, `orig_text` may contain extra characters that we don't want in
+    our prediction.
+
+    For example, let's say:
+      pred_text = steve smith
+      orig_text = Steve Smith's
+
+    We don't want to return `orig_text` because it contains the extra "'s".
+
+    We don't want to return `pred_text` because it's already been normalized
+    (the SQuAD eval script also does punctuation stripping/lower casing but
+    our tokenizer does additional normalization like stripping accent
+    characters).
+
+    What we really want to return is "Steve Smith".
+
+    Therefore, we have to apply a semi-complicated alignment heuristic
+    between `pred_text` and `orig_text` to get a character-to-character
+    alignment. This can fail in certain cases in which case we just return
+    `orig_text`."""
+
+    def _strip_spaces(text):
+        ns_chars = []
+        ns_to_s_map = collections.OrderedDict()
+        for (i, c) in enumerate(text):
+            if c == " ":
+                continue
+            ns_to_s_map[len(ns_chars)] = i
+            ns_chars.append(c)
+        ns_text = "".join(ns_chars)
+        return ns_text, ns_to_s_map
+
+    # We first tokenize `orig_text`, strip whitespace from the result
+    # and `pred_text`, and check if they are the same length. If they are
+    # NOT the same length, the heuristic has failed. If they are the same
+    # length, we assume the characters are one-to-one aligned.
+    tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
+
+    tok_text = " ".join(tokenizer.tokenize(orig_text))
+
+    start_position = tok_text.find(pred_text)
+    if start_position == -1:
+        if verbose_logging:
+            logging.warning("Unable to find text: '%s' in '%s'" % (pred_text, orig_text))
+        return orig_text
+    end_position = start_position + len(pred_text) - 1
+
+    (orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text)
+    (tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text)
+
+    if len(orig_ns_text) != len(tok_ns_text):
+        if verbose_logging:
+            logging.warning(
+                "Length not equal after stripping spaces: '%s' vs '%s'", orig_ns_text, tok_ns_text,
+            )
+        return orig_text
+
+    # We then project the characters in `pred_text` back to `orig_text` using
+    # the character-to-character alignment.
+    tok_s_to_ns_map = {}
+    for (i, tok_index) in tok_ns_to_s_map.items():
+        tok_s_to_ns_map[tok_index] = i
+
+    orig_start_position = None
+    if start_position in tok_s_to_ns_map:
+        ns_start_position = tok_s_to_ns_map[start_position]
+        if ns_start_position in orig_ns_to_s_map:
+            orig_start_position = orig_ns_to_s_map[ns_start_position]
+
+    if orig_start_position is None:
+        if verbose_logging:
+            logging.warning("Couldn't map start position")
+        return orig_text
+
+    orig_end_position = None
+    if end_position in tok_s_to_ns_map:
+        ns_end_position = tok_s_to_ns_map[end_position]
+        if ns_end_position in orig_ns_to_s_map:
+            orig_end_position = orig_ns_to_s_map[ns_end_position]
+
+    if orig_end_position is None:
+        if verbose_logging:
+            logging.warning("Couldn't map end position")
+        return orig_text
+
+    output_text = orig_text[orig_start_position : (orig_end_position + 1)]
+    return output_text
+
+
+def f1_score(prediction, ground_truth):
+    """computes f1 score between prediction and ground truth"""
+    prediction_tokens = _get_tokens(prediction)
+    ground_truth_tokens = _get_tokens(ground_truth)
+    common = collections.Counter(prediction_tokens) & collections.Counter(ground_truth_tokens)
+    num_same = sum(common.values())
+    if len(ground_truth_tokens) == 0 or len(prediction_tokens) == 0:
+        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
+        return int(ground_truth_tokens == prediction_tokens)
+    if num_same == 0:
+        return 0
+    precision = 1.0 * num_same / len(prediction_tokens)
+    recall = 1.0 * num_same / len(ground_truth_tokens)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1
+
+
+def exact_match_score(prediction, ground_truth):
+    """computes exact match between prediction and ground truth"""
+    return int(normalize_answer(prediction) == normalize_answer(ground_truth))
+
+
+def apply_no_ans_threshold(scores, na_probs, qid_to_has_ans, na_prob_thresh):
+    """Applies no answer threshhold"""
+    new_scores = {}
+    for qid, s in scores.items():
+        pred_na = na_probs[qid] > na_prob_thresh
+        if pred_na:
+            new_scores[qid] = float(not qid_to_has_ans[qid])
+        else:
+            new_scores[qid] = s
+    return new_scores
+
+
+def make_eval_dict(exact_scores, f1_scores, qid_list=None):
+    """returns dictionary with formatted evaluation scores"""
+    if not qid_list:
+        total = len(exact_scores)
+        return collections.OrderedDict(
+            [
+                ("exact", 100.0 * sum(exact_scores.values()) / total),
+                ("f1", 100.0 * sum(f1_scores.values()) / total),
+                ("total", total),
+            ]
+        )
+    else:
+        total = len(qid_list)
+        return collections.OrderedDict(
+            [
+                ("exact", 100.0 * sum(exact_scores[k] for k in qid_list) / total),
+                ("f1", 100.0 * sum(f1_scores[k] for k in qid_list) / total),
+                ("total", total),
+            ]
+        )
+
+
+def merge_eval(main_eval, new_eval, prefix):
+    """merges 2 evaluation dictionaries into the first one by adding prefix as key for name collision handling"""
+    for k in new_eval:
+        main_eval["%s_%s" % (prefix, k)] = new_eval[k]
+
+
+def find_all_best_thresh(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans):
+    """
+    Find best threshholds to maximize all evaluation metrics.
+    """
+    best_exact, exact_thresh = _find_best_thresh(preds, exact_raw, na_probs, qid_to_has_ans)
+    best_f1, f1_thresh = _find_best_thresh(preds, f1_raw, na_probs, qid_to_has_ans)
+
+    main_eval["best_exact"] = best_exact
+    main_eval["best_exact_thresh"] = exact_thresh
+    main_eval["best_f1"] = best_f1
+    main_eval["best_f1_thresh"] = f1_thresh
+
+
+def _find_best_thresh(preds, scores, na_probs, qid_to_has_ans):
+    """
+    Find best threshhold to maximize evaluation metric
+    """
+    num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k])
+    cur_score = num_no_ans
+    best_score = cur_score
+    best_thresh = 0.0
+    qid_list = sorted(na_probs, key=lambda k: na_probs[k])
+    for _, qid in enumerate(qid_list):
+        if qid not in scores:
+            continue
+        if qid_to_has_ans[qid]:
+            diff = scores[qid]
+        else:
+            if preds[qid]:
+                diff = -1
+            else:
+                diff = 0
+        cur_score += diff
+        if cur_score > best_score:
+            best_score = cur_score
+            best_thresh = na_probs[qid]
+    return 100.0 * best_score / len(scores), best_thresh
+
+
+class SquadProcessor(DataProcessor):
+    """
+    Processor for the SQuAD data set.
+    used by the version 1.1 and version 2.0 of SQuAD, respectively.
+
+    Args:
+        data_file: data file path
+        mode: TRAINING_MODE/EVALUATION_MODE/INFERENCE_MODE for creating training/evaluation/inference dataset
+    """
+
+    def __init__(self, data_file: str, mode: str):
+        self.data_file = data_file
+        self.mode = mode
+
+    def get_examples(self):
+        if self.data_file is None:
+            raise ValueError(f"{self.mode} data file is None.")
+
+        with open(self.data_file, "r", encoding="utf-8") as reader:
+            input_data = json.load(reader)["data"]
+        return self._create_examples(input_data, set_type=self.mode)
+
+    def _create_examples(self, input_data, set_type):
+        examples = []
+        for entry in tqdm(input_data):
+            title = entry["title"]
+            for paragraph in entry["paragraphs"]:
+                context_text = paragraph["context"]
+                for qa in paragraph["qas"]:
+                    qas_id = qa["id"]
+                    question_text = qa["question"]
+                    start_position_character = None
+                    answer_text = None
+                    answers = []
+                    if "is_impossible" in qa:
+                        is_impossible = qa["is_impossible"]
+                    else:
+                        is_impossible = False
+
+                    if not is_impossible:
+                        if set_type in [TRAINING_MODE, EVALUATION_MODE]:
+                            answer = qa["answers"][0]
+                            answer_text = answer["text"]
+                            start_position_character = answer["answer_start"]
+                        if set_type == EVALUATION_MODE:
+                            answers = qa["answers"]
+
+                    example = SquadExample(
+                        qas_id=qas_id,
+                        question_text=question_text,
+                        context_text=context_text,
+                        answer_text=answer_text,
+                        start_position_character=start_position_character,
+                        title=title,
+                        is_impossible=is_impossible,
+                        answers=answers,
+                    )
+
+                    examples.append(example)
+        return examples
+
+
+class SquadExample(object):
+    """
+    A single training/test example for the Squad dataset, as loaded from disk.
+    Args:
+        qas_id: The example's unique identifier
+        question_text: The question string
+        context_text: The context string
+        answer_text: The answer string
+        start_position_character: The character position of the start of
+            the answer, 0 indexed
+        title: The title of the example
+        answers: None by default, this is used during evaluation.
+            Holds answers as well as their start positions.
+        is_impossible: False by default, set to True if the example has
+            no possible answer.
+    """
+
+    def __init__(
+        self,
+        qas_id: str,
+        question_text: str,
+        context_text: str,
+        answer_text: str,
+        start_position_character: int,
+        title: str,
+        answers: List[str] = [],
+        is_impossible: bool = False,
+    ):
+        self.qas_id = qas_id
+        self.question_text = question_text
+        self.context_text = context_text
+        self.answer_text = answer_text
+        self.title = title
+        self.is_impossible = is_impossible
+        self.answers = answers
+
+        self.start_position, self.end_position = 0, 0
+
+        doc_tokens = []
+        char_to_word_offset = []
+        prev_is_whitespace = True
+
+        # Split on whitespace so that different tokens
+        # may be attributed to their original position.
+        # ex: context_text = ["hi yo"]
+        #     char_to_word_offset = [0, 0, 0, 1, 1]
+        #     doc_tokens = ["hi", "yo"]
+        for c in self.context_text:
+            if is_whitespace(c):
+                prev_is_whitespace = True
+            else:
+                if prev_is_whitespace:
+                    doc_tokens.append(c)
+                else:
+                    doc_tokens[-1] += c
+                prev_is_whitespace = False
+            char_to_word_offset.append(len(doc_tokens) - 1)
+
+        self.doc_tokens = doc_tokens
+        self.char_to_word_offset = char_to_word_offset
+
+        # Start end end positions only has a value during evaluation.
+        if start_position_character is not None and not is_impossible:
+            # start_position is index of word, end_position inclusive
+            self.start_position = char_to_word_offset[start_position_character]
+            self.end_position = char_to_word_offset[
+                min(start_position_character + len(answer_text) - 1, len(char_to_word_offset) - 1)
+            ]
+
+
+def convert_examples_to_features(
+    examples: List[object],
+    tokenizer: object,
+    max_seq_length: int,
+    doc_stride: int,
+    max_query_length: int,
+    has_groundtruth: bool,
+):
+    """Loads a data file into a list of `InputBatch`s."""
+
+    unique_id = 1000000000
+
+    features = []
+    for (example_index, example) in enumerate(examples):
+        query_tokens = tokenizer.text_to_tokens(example.question_text)
+
+        if len(query_tokens) > max_query_length:
+            query_tokens = query_tokens[0:max_query_length]
+
+        # context: index of token -> index of word
+        tok_to_orig_index = []
+        # context: index of word -> index of first token in token list
+        orig_to_tok_index = []
+        # context without white spaces after tokenization
+        all_doc_tokens = []
+        # doc tokens is word separated context
+        for (i, token) in enumerate(example.doc_tokens):
+            orig_to_tok_index.append(len(all_doc_tokens))
+            sub_tokens = tokenizer.text_to_tokens(token)
+            for sub_token in sub_tokens:
+                tok_to_orig_index.append(i)
+                all_doc_tokens.append(sub_token)
+
+        # idx of query token start and end in context
+        tok_start_position = None
+        tok_end_position = None
+        if has_groundtruth and example.is_impossible:
+            tok_start_position = -1
+            tok_end_position = -1
+        if has_groundtruth and not example.is_impossible:
+            tok_start_position = orig_to_tok_index[example.start_position]
+            if example.end_position < len(example.doc_tokens) - 1:
+                tok_end_position = orig_to_tok_index[example.end_position + 1] - 1
+            else:
+                tok_end_position = len(all_doc_tokens) - 1
+
+            (tok_start_position, tok_end_position) = _improve_answer_span(
+                all_doc_tokens, tok_start_position, tok_end_position, tokenizer, example.answer_text
+            )
+        # The -3 accounts for tokenizer.cls_token, tokenizer.sep_token and tokenizer.sep_token
+        # doc_spans contains all possible contexts options of given length
+        max_tokens_for_doc = max_seq_length - len(query_tokens) - 3
+        _DocSpan = collections.namedtuple("DocSpan", ["start", "length"])
+        doc_spans = []
+        start_offset = 0
+        while start_offset < len(all_doc_tokens):
+            length = len(all_doc_tokens) - start_offset
+            if length > max_tokens_for_doc:
+                length = max_tokens_for_doc
+            doc_spans.append(_DocSpan(start=start_offset, length=length))
+            if start_offset + length == len(all_doc_tokens):
+                break
+            start_offset += min(length, doc_stride)
+
+        for (doc_span_index, doc_span) in enumerate(doc_spans):
+            tokens = []
+            # maps context tokens idx in final input -> word idx in context
+            token_to_orig_map = {}
+            token_is_max_context = {}
+            segment_ids = []
+            tokens.append(tokenizer.cls_token)
+            segment_ids.append(0)
+            for token in query_tokens:
+                tokens.append(token)
+                segment_ids.append(0)
+            tokens.append(tokenizer.sep_token)
+            segment_ids.append(0)
+
+            for i in range(doc_span.length):
+                split_token_index = doc_span.start + i
+                token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
+
+                is_max_context = _check_is_max_context(doc_spans, doc_span_index, split_token_index)
+                token_is_max_context[len(tokens)] = is_max_context
+                tokens.append(all_doc_tokens[split_token_index])
+                segment_ids.append(1)
+            tokens.append(tokenizer.sep_token)
+            segment_ids.append(1)
+
+            input_ids = tokenizer.tokens_to_ids(tokens)
+
+            # The mask has 1 for real tokens and 0 for padding tokens.
+            # Only real tokens are attended to.
+            input_mask = [1] * len(input_ids)
+
+            # Zero-pad up to the sequence length.
+            while len(input_ids) < max_seq_length:
+                input_ids.append(tokenizer.pad_id)
+                input_mask.append(0)
+                segment_ids.append(0)
+
+            assert len(input_ids) == max_seq_length
+            assert len(input_mask) == max_seq_length
+            assert len(segment_ids) == max_seq_length
+
+            # calculate start and end position in final array
+            # of tokens in answer if no answer,
+            # 0 for both pointing to tokenizer.cls_token
+            start_position = None
+            end_position = None
+            if has_groundtruth and not example.is_impossible:
+                doc_start = doc_span.start
+                doc_end = doc_span.start + doc_span.length - 1
+                out_of_span = False
+                if not (tok_start_position >= doc_start and tok_end_position <= doc_end):
+                    out_of_span = True
+                if out_of_span:
+                    start_position = 0
+                    end_position = 0
+                else:
+                    doc_offset = len(query_tokens) + 2
+                    start_position = tok_start_position - doc_start + doc_offset
+                    end_position = tok_end_position - doc_start + doc_offset
+            if has_groundtruth and example.is_impossible:
+                # if our document chunk does not contain
+                # an annotation we throw it out, since there is nothing
+                # to predict.
+                start_position = 0
+                end_position = 0
+
+            if example_index < 1:
+                logging.info("*** Example ***")
+                logging.info("unique_id: %s" % (unique_id))
+                logging.info("example_index: %s" % (example_index))
+                logging.info("doc_span_index: %s" % (doc_span_index))
+                logging.info("tokens: %s" % " ".join(tokens))
+                logging.info(
+                    "token_to_orig_map: %s" % " ".join(["%d:%d" % (x, y) for (x, y) in token_to_orig_map.items()])
+                )
+                logging.info(
+                    "token_is_max_context: %s"
+                    % " ".join(["%d:%s" % (x, y) for (x, y) in token_is_max_context.items()])
+                )
+                logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
+                logging.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
+                logging.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
+                if has_groundtruth and example.is_impossible:
+                    logging.info("impossible example")
+                if has_groundtruth and not example.is_impossible:
+                    answer_text = " ".join(tokens[start_position : (end_position + 1)])
+                    logging.info("start_position: %d" % (start_position))
+                    logging.info("end_position: %d" % (end_position))
+                    logging.info("answer: %s" % (answer_text))
+            if example_index % 100 == 0:
+                logging.info(f"Finished processing: {example_index}")
+            features.append(
+                InputFeatures(
+                    unique_id=unique_id,
+                    example_index=example_index,
+                    doc_span_index=doc_span_index,
+                    tokens=tokens,
+                    token_to_orig_map=token_to_orig_map,
+                    token_is_max_context=token_is_max_context,
+                    input_ids=input_ids,
+                    input_mask=input_mask,
+                    segment_ids=segment_ids,
+                    start_position=start_position,
+                    end_position=end_position,
+                    is_impossible=example.is_impossible,
+                )
+            )
+            unique_id += 1
+
+    return features
+
+
+def _improve_answer_span(
+    doc_tokens: List[str], input_start: int, input_end: int, tokenizer: object, orig_answer_text: str
+):
+    """Returns tokenized answer spans that
+    better match the annotated answer."""
+    tok_answer_text = " ".join(tokenizer.text_to_tokens(orig_answer_text))
+
+    for new_start in range(input_start, input_end + 1):
+        for new_end in range(input_end, new_start - 1, -1):
+            text_span = " ".join(doc_tokens[new_start : (new_end + 1)])
+            if text_span == tok_answer_text:
+                return (new_start, new_end)
+
+    return (input_start, input_end)
+
+
+def _check_is_max_context(doc_spans, cur_span_index, position):
+    """Check if this is the 'max context' doc span for the token.
+    Because of the sliding window approach taken to scoring documents,
+    a single token can appear in multiple documents.
+    Example:
+        Doc: the man went to the store and bought a gallon of milk
+        Span A: the man went to the
+        Span B: to the store and bought
+        Span C: and bought a gallon of
+        ...
+    Now the word 'bought' will have two scores from spans B and C. We only
+    want to consider the score with "maximum context", which we define as
+    the *minimum* of its left and right context (the *sum* of left and
+    right context will always be the same, of course).
+    In the example the maximum context for 'bought' would be span C since
+    it has 1 left context and 3 right context, while span B has 4 left context
+    and 0 right context.
+    Code adapted from the code by the Google AI and HuggingFace.
+    """
+    best_score = None
+    best_span_index = None
+    for (span_index, doc_span) in enumerate(doc_spans):
+        end = doc_span.start + doc_span.length - 1
+        if position < doc_span.start:
+            continue
+        if position > end:
+            continue
+        num_left_context = position - doc_span.start
+        num_right_context = end - position
+        score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
+        if best_score is None or score > best_score:
+            best_score = score
+            best_span_index = span_index
+
+    return cur_span_index == best_span_index
+
+
+class InputFeatures(object):
+    """A single set of features of data."""
+
+    def __init__(
+        self,
+        unique_id: int,
+        example_index: int,
+        doc_span_index: int,
+        tokens: List[str],
+        token_to_orig_map: Dict[int, int],
+        token_is_max_context: Dict[int, bool],
+        input_ids: List[int],
+        input_mask: List[int],
+        segment_ids: List[int],
+        start_position: Optional[int] = None,
+        end_position: Optional[int] = None,
+        is_impossible: Optional[int] = None,
+    ):
+        self.unique_id = unique_id
+        self.example_index = example_index
+        self.doc_span_index = doc_span_index
+        self.tokens = tokens
+        self.token_to_orig_map = token_to_orig_map
+        self.token_is_max_context = token_is_max_context
+        self.input_ids = input_ids
+        self.input_mask = input_mask
+        self.segment_ids = segment_ids
+        self.start_position = start_position
+        self.end_position = end_position
+        self.is_impossible = is_impossible

tts/nemo/collections/nlp/data/text_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.data.text_classification.text_classification_dataset import (
+    TextClassificationDataset,
+    calc_class_weights,
+)

tts/nemo/collections/nlp/data/text_classification/text_classification_dataset.py

@@ -0,0 +1,297 @@
+# 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
+import pickle
+import random
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.data_utils.data_preprocessing import (
+    fill_class_weights,
+    get_freq_weights,
+    get_label_stats,
+    get_stats,
+)
+from nemo.collections.nlp.parts.utils_funcs import list2str
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
+from nemo.utils import logging
+from nemo.utils.env_var_parsing import get_envint
+
+__all__ = ['TextClassificationDataset', 'calc_class_weights']
+
+
+class TextClassificationDataset(Dataset):
+    """A dataset class that converts from raw data to
+    a dataset that can be used by DataLayerNM.
+    Args:
+        input_file: file to sequence + label.
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        tokenizer: tokenizer object such as AutoTokenizer
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        num_samples: number of samples you want to use for the dataset.
+            If -1, use all dataset. Useful for testing.
+        shuffle: Shuffles the dataset after loading.
+        use_cache: Enables caching to use pickle format to store and read data from
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'label': NeuralType(('B',), LabelsType()),
+        }
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        input_file: str = None,
+        queries: List[str] = None,
+        max_seq_length: int = -1,
+        num_samples: int = -1,
+        shuffle: bool = False,
+        use_cache: bool = False,
+    ):
+        if not input_file and not queries:
+            raise ValueError("Either input_file or queries should be passed to the text classification dataset.")
+
+        if input_file and not os.path.exists(input_file):
+            raise FileNotFoundError(
+                f'Data file `{input_file}` not found! Each line of the data file should contain text sequences, where '
+                f'words are separated with spaces and the label separated by [TAB] following this format: '
+                f'[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]'
+            )
+
+        self.input_file = input_file
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.num_samples = num_samples
+        self.shuffle = shuffle
+        self.use_cache = use_cache
+        self.vocab_size = self.tokenizer.vocab_size
+        self.pad_id = tokenizer.pad_id
+
+        self.features = None
+        labels, all_sents = [], []
+        if input_file:
+            data_dir, filename = os.path.split(input_file)
+            vocab_size = getattr(tokenizer, "vocab_size", 0)
+            tokenizer_name = tokenizer.name
+            cached_features_file = os.path.join(
+                data_dir,
+                f"cached_{filename}_{tokenizer_name}_{max_seq_length}_{vocab_size}_{num_samples}_{self.pad_id}_{shuffle}.pkl",
+            )
+
+            if get_envint("LOCAL_RANK", 0) == 0:
+                if use_cache and os.path.exists(cached_features_file):
+                    logging.warning(
+                        f"Processing of {input_file} is skipped as caching is enabled and a cache file "
+                        f"{cached_features_file} already exists."
+                    )
+                    logging.warning(
+                        f"You may need to delete the cache file if any of the processing parameters (eg. tokenizer) or "
+                        f"the data are updated."
+                    )
+                else:
+                    with open(input_file, "r") as f:
+                        lines = f.readlines()
+                        logging.info(f'Read {len(lines)} examples from {input_file}.')
+                        if num_samples > 0:
+                            lines = lines[:num_samples]
+                            logging.warning(
+                                f"Parameter 'num_samples' is set, so just the first {len(lines)} examples are kept."
+                            )
+
+                        if shuffle:
+                            random.shuffle(lines)
+
+                        for index, line in enumerate(lines):
+                            if index % 20000 == 0:
+                                logging.debug(f"Processing line {index}/{len(lines)}")
+                            line_splited = line.strip().split()
+                            try:
+                                label = int(line_splited[-1])
+                            except ValueError:
+                                logging.debug(f"Skipping line {line}")
+                                continue
+                            labels.append(label)
+                            sent_words = line_splited[:-1]
+                            all_sents.append(sent_words)
+                    verbose = True
+
+                    self.features = self.get_features(
+                        all_sents=all_sents,
+                        tokenizer=tokenizer,
+                        max_seq_length=max_seq_length,
+                        labels=labels,
+                        verbose=verbose,
+                    )
+                    with open(cached_features_file, 'wb') as out_file:
+                        pickle.dump(self.features, out_file, protocol=pickle.HIGHEST_PROTOCOL)
+        else:
+            for query in queries:
+                all_sents.append(query.strip().split())
+            labels = [-1] * len(all_sents)
+            verbose = False
+            self.features = self.get_features(
+                all_sents=all_sents, tokenizer=tokenizer, max_seq_length=max_seq_length, labels=labels, verbose=verbose
+            )
+
+        # wait until the master process writes to the processed data files
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        if input_file:
+            with open(cached_features_file, "rb") as input_file:
+                self.features = pickle.load(input_file)
+
+    def __len__(self):
+        return len(self.features)
+
+    def __getitem__(self, idx):
+        return self.features[idx]
+
+    def _collate_fn(self, batch):
+        """collate batch of input_ids, segment_ids, input_mask, and label
+        Args:
+            batch:  A list of tuples of (input_ids, segment_ids, input_mask, label).
+        """
+        max_length = 0
+        for input_ids, segment_ids, input_mask, label in batch:
+            if len(input_ids) > max_length:
+                max_length = len(input_ids)
+
+        padded_input_ids = []
+        padded_segment_ids = []
+        padded_input_mask = []
+        labels = []
+        for input_ids, segment_ids, input_mask, label in batch:
+            if len(input_ids) < max_length:
+                pad_width = max_length - len(input_ids)
+                padded_input_ids.append(np.pad(input_ids, pad_width=[0, pad_width], constant_values=self.pad_id))
+                padded_segment_ids.append(np.pad(segment_ids, pad_width=[0, pad_width], constant_values=self.pad_id))
+                padded_input_mask.append(np.pad(input_mask, pad_width=[0, pad_width], constant_values=self.pad_id))
+            else:
+                padded_input_ids.append(input_ids)
+                padded_segment_ids.append(segment_ids)
+                padded_input_mask.append(input_mask)
+            labels.append(label)
+
+        return (
+            torch.LongTensor(padded_input_ids),
+            torch.LongTensor(padded_segment_ids),
+            torch.LongTensor(padded_input_mask),
+            torch.LongTensor(labels),
+        )
+
+    @staticmethod
+    def get_features(all_sents, tokenizer, max_seq_length, labels=None, verbose=True):
+        """Encode a list of sentences into a list of tuples of (input_ids, segment_ids, input_mask, label)."""
+        features = []
+        sent_lengths = []
+        too_long_count = 0
+        for sent_id, sent in enumerate(all_sents):
+            if sent_id % 1000 == 0:
+                logging.debug(f"Encoding sentence {sent_id}/{len(all_sents)}")
+            sent_subtokens = [tokenizer.cls_token]
+            for word in sent:
+                word_tokens = tokenizer.text_to_tokens(word)
+                sent_subtokens.extend(word_tokens)
+
+            if max_seq_length > 0 and len(sent_subtokens) + 1 > max_seq_length:
+                sent_subtokens = sent_subtokens[:max_seq_length]
+                too_long_count += 1
+
+            sent_subtokens.append(tokenizer.sep_token)
+            sent_lengths.append(len(sent_subtokens))
+
+            input_ids = [tokenizer.tokens_to_ids(t) for t in sent_subtokens]
+
+            # The mask has 1 for real tokens and 0 for padding tokens.
+            # Only real tokens are attended to.
+            input_mask = [1] * len(input_ids)
+            segment_ids = [0] * len(input_ids)
+
+            if verbose and sent_id < 2:
+                logging.info("*** Example ***")
+                logging.info(f"example {sent_id}: {sent}")
+                logging.info("subtokens: %s" % " ".join(sent_subtokens))
+                logging.info("input_ids: %s" % list2str(input_ids))
+                logging.info("segment_ids: %s" % list2str(segment_ids))
+                logging.info("input_mask: %s" % list2str(input_mask))
+                logging.info("label: %s" % labels[sent_id] if labels else "**Not Provided**")
+
+            label = labels[sent_id] if labels else -1
+            features.append([np.asarray(input_ids), np.asarray(segment_ids), np.asarray(input_mask), label])
+
+        if max_seq_length > -1 and too_long_count > 0:
+            logging.warning(
+                f'Found {too_long_count} out of {len(all_sents)} sentences with more than {max_seq_length} subtokens. '
+                f'Truncated long sentences from the end.'
+            )
+        if verbose:
+            get_stats(sent_lengths)
+        return features
+
+
+def calc_class_weights(file_path: str, num_classes: int):
+    """
+    iterates over a data file and calculate the weights of each class to be used for class_balancing
+    Args:
+        file_path: path to the data file
+        num_classes: number of classes in the dataset
+    """
+
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"Could not find data file {file_path} to calculate the class weights!")
+
+    with open(file_path, 'r') as f:
+        input_lines = f.readlines()
+
+    labels = []
+    for input_line in input_lines:
+        parts = input_line.strip().split()
+        try:
+            label = int(parts[-1])
+        except ValueError:
+            raise ValueError(
+                f'No numerical labels found for {file_path}. Labels should be integers and separated by [TAB] at the end of each line.'
+            )
+        labels.append(label)
+
+    logging.info(f'Calculating stats of {file_path}...')
+    total_sents, sent_label_freq, max_id = get_label_stats(labels, f'{file_path}_sentence_stats.tsv', verbose=False)
+    if max_id >= num_classes:
+        raise ValueError(f'Found an invalid label in {file_path}! Labels should be from [0, num_classes-1].')
+
+    class_weights_dict = get_freq_weights(sent_label_freq)
+
+    logging.info(f'Total Sentence: {total_sents}')
+    logging.info(f'Sentence class frequencies: {sent_label_freq}')
+
+    logging.info(f'Class Weights: {class_weights_dict}')
+    class_weights = fill_class_weights(weights=class_weights_dict, max_id=num_classes - 1)
+
+    return class_weights

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

@@ -0,0 +1,13 @@
+# 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.

tts/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py

@@ -0,0 +1,466 @@
+# 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.
+
+__all__ = ['BertPunctuationCapitalizationDataset', 'BertPunctuationCapitalizationInferDataset']
+
+import itertools
+import os
+import pickle
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.data_utils.data_preprocessing import get_label_stats, get_stats
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
+from nemo.utils import logging
+
+
+def get_features(
+    queries: List[str],
+    max_seq_length: int,
+    tokenizer: TokenizerSpec,
+    punct_label_ids: dict = None,
+    capit_label_ids: dict = None,
+    pad_label: str = 'O',
+    punct_labels_lines=None,
+    capit_labels_lines=None,
+    ignore_extra_tokens=False,
+    ignore_start_end: Optional[bool] = False,
+):
+    """
+    Processes the data and returns features.
+
+    Args:
+        queries: text sequences
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer: such as AutoTokenizer
+        pad_label: pad value use for labels. By default, it's the neutral label.
+        punct_label_ids: dict to map punctuation labels to label ids.
+            Starts with pad_label->0 and then increases in alphabetical order.
+            Required for training and evaluation, not needed for inference.
+        capit_label_ids: dict to map labels to label ids. Starts
+            with pad_label->0 and then increases in alphabetical order.
+            Required for training and evaluation, not needed for inference.
+        punct_labels: list of labels for every word in a sequence (str)
+        capit_labels: list of labels for every word in a sequence (str)
+        ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
+        ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
+
+    Returns:
+        all_input_ids: input ids for all tokens
+        all_segment_ids: token type ids
+        all_input_mask: attention mask to use for BERT model
+        all_subtokens_mask: masks out all subwords besides the first one
+        all_loss_mask: loss mask to mask out tokens during training
+        punct_all_labels: all labels for punctuation task (ints)
+        capit_all_labels: all labels for capitalization task (ints)
+        punct_label_ids: label (str) to id (int) map for punctuation task
+        capit_label_ids: label (str) to id (int) map for capitalization task
+    """
+    all_subtokens = []
+    all_loss_mask = []
+    all_subtokens_mask = []
+    all_segment_ids = []
+    all_input_ids = []
+    all_input_mask = []
+    sent_lengths = []
+    punct_all_labels = []
+    capit_all_labels = []
+    with_label = False
+
+    if punct_labels_lines and capit_labels_lines:
+        with_label = True
+
+    for i, query in enumerate(queries):
+        words = query.strip().split()
+
+        # add bos token
+        subtokens = [tokenizer.cls_token]
+        loss_mask = [1 - ignore_start_end]
+        subtokens_mask = [0]
+        if with_label:
+            pad_id = punct_label_ids[pad_label]
+            punct_labels = [pad_id]
+            punct_query_labels = [punct_label_ids[lab] for lab in punct_labels_lines[i]]
+
+            capit_labels = [pad_id]
+            capit_query_labels = [capit_label_ids[lab] for lab in capit_labels_lines[i]]
+
+        for j, word in enumerate(words):
+            word_tokens = tokenizer.text_to_tokens(word)
+            subtokens.extend(word_tokens)
+
+            loss_mask.append(1)
+            loss_mask.extend([int(not ignore_extra_tokens)] * (len(word_tokens) - 1))
+
+            subtokens_mask.append(1)
+            subtokens_mask.extend([0] * (len(word_tokens) - 1))
+
+            if with_label:
+                punct_labels.extend([punct_query_labels[j]] * len(word_tokens))
+                capit_labels.extend([capit_query_labels[j]] * len(word_tokens))
+
+        # add eos token
+        subtokens.append(tokenizer.sep_token)
+        loss_mask.append(1 - ignore_start_end)
+        subtokens_mask.append(0)
+        sent_lengths.append(len(subtokens))
+        all_subtokens.append(subtokens)
+        all_loss_mask.append(loss_mask)
+        all_subtokens_mask.append(subtokens_mask)
+        all_input_mask.append([1] * len(subtokens))
+
+        if with_label:
+            punct_labels.append(pad_id)
+            punct_all_labels.append(punct_labels)
+            capit_labels.append(pad_id)
+            capit_all_labels.append(capit_labels)
+
+    max_seq_length = min(max_seq_length, max(sent_lengths))
+    logging.info(f'Max length: {max_seq_length}')
+    get_stats(sent_lengths)
+    too_long_count = 0
+
+    for i, subtokens in enumerate(all_subtokens):
+        if len(subtokens) > max_seq_length:
+            subtokens = [tokenizer.cls_token] + subtokens[-max_seq_length + 1 :]
+            all_input_mask[i] = [1] + all_input_mask[i][-max_seq_length + 1 :]
+            all_loss_mask[i] = [int(not ignore_start_end)] + all_loss_mask[i][-max_seq_length + 1 :]
+            all_subtokens_mask[i] = [0] + all_subtokens_mask[i][-max_seq_length + 1 :]
+
+            if with_label:
+                punct_all_labels[i] = [pad_id] + punct_all_labels[i][-max_seq_length + 1 :]
+                capit_all_labels[i] = [pad_id] + capit_all_labels[i][-max_seq_length + 1 :]
+            too_long_count += 1
+
+        all_input_ids.append(tokenizer.tokens_to_ids(subtokens))
+
+        if len(subtokens) < max_seq_length:
+            extra = max_seq_length - len(subtokens)
+            all_input_ids[i] = all_input_ids[i] + [0] * extra
+            all_loss_mask[i] = all_loss_mask[i] + [0] * extra
+            all_subtokens_mask[i] = all_subtokens_mask[i] + [0] * extra
+            all_input_mask[i] = all_input_mask[i] + [0] * extra
+
+            if with_label:
+                punct_all_labels[i] = punct_all_labels[i] + [pad_id] * extra
+                capit_all_labels[i] = capit_all_labels[i] + [pad_id] * extra
+
+        all_segment_ids.append([0] * max_seq_length)
+
+    logging.info(f'{too_long_count} are longer than {max_seq_length}')
+
+    for i in range(min(len(all_input_ids), 5)):
+        logging.info("*** Example ***")
+        logging.info("i: %s" % (i))
+        logging.info("subtokens: %s" % " ".join(list(map(str, all_subtokens[i]))))
+        logging.info("loss_mask: %s" % " ".join(list(map(str, all_loss_mask[i]))))
+        logging.info("input_mask: %s" % " ".join(list(map(str, all_input_mask[i]))))
+        logging.info("subtokens_mask: %s" % " ".join(list(map(str, all_subtokens_mask[i]))))
+        if with_label:
+            logging.info("punct_labels: %s" % " ".join(list(map(str, punct_all_labels[i]))))
+            logging.info("capit_labels: %s" % " ".join(list(map(str, capit_all_labels[i]))))
+
+    return (
+        all_input_ids,
+        all_segment_ids,
+        all_input_mask,
+        all_subtokens_mask,
+        all_loss_mask,
+        punct_all_labels,
+        capit_all_labels,
+        punct_label_ids,
+        capit_label_ids,
+    )
+
+
+class BertPunctuationCapitalizationDataset(Dataset):
+    """
+    Creates dataset to use during training for punctuaion and capitalization tasks with a pretrained model.
+    For dataset to use during inference without labels, see BertPunctuationCapitalizationInferDataset.
+
+    Args:
+        text_file: file to sequences, each line should a sentence, no header.
+        label_file: file to labels, each line corresponds to word labels for a sentence in the text_file. No header.
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer: such as AutoTokenizer
+        num_samples: number of samples you want to use for the dataset.
+            If -1, use all dataset. Useful for testing.
+        pad_label: pad value use for labels.
+            by default, it's the neutral label.
+        punct_label_ids and capit_label_ids (dict):
+            dict to map labels to label ids.
+            Starts with pad_label->0 and then increases in alphabetical order
+            For dev set use label_ids generated during training to support
+            cases when not all labels are present in the dev set.
+            For training set label_ids should be None or loaded from cache
+        ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
+        ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
+        use_cache: whether to use processed data cache or not
+        get_label_frequencies: whether to generate label frequencies
+        punct_label_ids_file and capit_label_ids_file: name of the files to save in .nemo
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports. """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+            'loss_mask': NeuralType(('B', 'T'), MaskType()),
+            'punct_labels': NeuralType(('B', 'T'), LabelsType()),
+            'capit_labels': NeuralType(('B', 'T'), LabelsType()),
+        }
+
+    def __init__(
+        self,
+        text_file: str,
+        label_file: str,
+        max_seq_length: int,
+        tokenizer: TokenizerSpec,
+        num_samples: int = -1,
+        pad_label: str = 'O',
+        punct_label_ids: Dict[str, int] = None,
+        capit_label_ids: Dict[str, int] = None,
+        ignore_extra_tokens: bool = False,
+        ignore_start_end: bool = False,
+        use_cache: bool = True,
+        get_label_frequencies: bool = False,
+        punct_label_ids_file: str = 'punct_label_ids.csv',
+        capit_label_ids_file: str = 'capit_label_ids.csv',
+    ):
+        """ Initializes BertPunctuationCapitalizationDataset. """
+
+        if not (os.path.exists(text_file) and os.path.exists(label_file)):
+            raise FileNotFoundError(
+                f'{text_file} or {label_file} not found. The data should be splitted 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).'
+            )
+
+        # Cache features
+        data_dir = os.path.dirname(text_file)
+        filename = os.path.basename(text_file)
+
+        if not filename.endswith('.txt'):
+            raise ValueError("{text_file} should have extension .txt")
+
+        filename = filename[:-4]
+        vocab_size = getattr(tokenizer, "vocab_size", 0)
+        features_pkl = os.path.join(
+            data_dir,
+            "cached_{}_{}_{}_{}_{}".format(
+                filename, tokenizer.name, str(max_seq_length), str(vocab_size), str(num_samples)
+            ),
+        )
+
+        self.punct_label_ids_file = os.path.join(data_dir, punct_label_ids_file)
+        self.capit_label_ids_file = os.path.join(data_dir, capit_label_ids_file)
+
+        master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+        cache_files_exist = (
+            os.path.exists(features_pkl)
+            and os.path.exists(self.punct_label_ids_file)
+            and os.path.exists(self.capit_label_ids_file)
+        )
+        features = None
+        if master_device and not (cache_files_exist and use_cache):
+            if num_samples == 0:
+                raise ValueError("num_samples has to be positive", num_samples)
+            logging.info(f'Processing {text_file}')
+            with open(text_file, 'r') as f:
+                text_lines = f.readlines()
+
+            # Collect all possible labels
+            punct_unique_labels = set()
+            capit_unique_labels = set()
+            punct_labels_lines = []
+            capit_labels_lines = []
+            with open(label_file, 'r') as f:
+                for line in f:
+                    line = line.strip().split()
+
+                    # extract punctuation and capitalization labels
+                    punct_line, capit_line = zip(*line)
+                    punct_labels_lines.append(punct_line)
+                    capit_labels_lines.append(capit_line)
+
+                    punct_unique_labels.update(punct_line)
+                    capit_unique_labels.update(capit_line)
+
+            if len(punct_labels_lines) != len(text_lines):
+                raise ValueError("Labels file should contain labels for every word")
+
+            dataset = list(zip(text_lines, punct_labels_lines, capit_labels_lines))
+
+            if num_samples > 0:
+                dataset = dataset[:num_samples]
+
+            dataset = list(zip(*dataset))
+            text_lines = dataset[0]
+            punct_labels_lines = dataset[1]
+            capit_labels_lines = dataset[2]
+
+            # for dev/test sets use label mapping from training set
+            if punct_label_ids:
+                if len(punct_label_ids) != len(punct_unique_labels):
+                    logging.info(
+                        'Not all labels from the specified'
+                        + 'label_ids dictionary are present in the'
+                        + 'current dataset. Using the provided'
+                        + 'label_ids dictionary.'
+                    )
+                else:
+                    logging.info('Using the provided label_ids dictionary.')
+            else:
+                logging.info(
+                    'Creating a new label to label_id dictionary.'
+                    + ' It\'s recommended to use label_ids generated'
+                    + ' during training for dev/test sets to avoid'
+                    + ' errors if some labels are not'
+                    + ' present in the dev/test sets.'
+                    + ' For training set label_ids should be None.'
+                )
+
+                def create_label_ids(unique_labels, pad_label=pad_label):
+                    label_ids = {pad_label: 0}
+                    if pad_label in unique_labels:
+                        unique_labels.remove(pad_label)
+                    for label in sorted(unique_labels):
+                        label_ids[label] = len(label_ids)
+                    return label_ids
+
+                punct_label_ids = create_label_ids(punct_unique_labels)
+                capit_label_ids = create_label_ids(capit_unique_labels)
+
+            self._save_label_ids(punct_label_ids, self.punct_label_ids_file)
+            self._save_label_ids(capit_label_ids, self.capit_label_ids_file)
+
+            features = get_features(
+                text_lines,
+                max_seq_length,
+                tokenizer,
+                pad_label=pad_label,
+                punct_labels_lines=punct_labels_lines,
+                capit_labels_lines=capit_labels_lines,
+                punct_label_ids=punct_label_ids,
+                capit_label_ids=capit_label_ids,
+                ignore_extra_tokens=ignore_extra_tokens,
+                ignore_start_end=ignore_start_end,
+            )
+
+            pickle.dump(features, open(features_pkl, "wb"))
+            logging.info(f'Features saved to {features_pkl}')
+
+        # wait until the master process writes to the processed data files
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        if features is None:
+            features = pickle.load(open(features_pkl, 'rb'))
+            logging.info(f'Features restored from {features_pkl}')
+
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_subtokens_mask = features[3]
+        self.all_loss_mask = features[4]
+        self.punct_all_labels = features[5]
+        self.capit_all_labels = features[6]
+        self.punct_label_ids = features[7]
+        self.capit_label_ids = features[8]
+
+        if get_label_frequencies:
+            self.punct_label_frequencies = self._calculate_label_frequencies(self.punct_all_labels, data_dir, 'punct')
+            self.capit_label_frequencies = self._calculate_label_frequencies(self.capit_all_labels, data_dir, 'capit')
+
+    def _calculate_label_frequencies(self, all_labels: List[int], data_dir: str, name: str) -> Dict[str, float]:
+        """ Calculates labels frequencies """
+        merged_labels = itertools.chain.from_iterable(all_labels)
+        logging.info('Three most popular labels')
+        _, label_frequencies, _ = get_label_stats(merged_labels, data_dir + '/label_count_' + name + '.tsv')
+        return label_frequencies
+
+    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 __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.long),
+            np.array(self.all_subtokens_mask[idx]),
+            np.array(self.all_loss_mask[idx]),
+            np.array(self.punct_all_labels[idx]),
+            np.array(self.capit_all_labels[idx]),
+        )
+
+
+class BertPunctuationCapitalizationInferDataset(Dataset):
+    """
+    Creates dataset to use during inference for punctuation and capitalization tasks with a pretrained model.
+    For dataset to use during training with labels, see BertPunctuationCapitalizationDataset.
+
+    Args:
+        queries file to sequences, each line should a sentence, no header.
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer: such as AutoTokenizer
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+        }
+
+    def __init__(self, queries: List[str], max_seq_length: int, tokenizer: TokenizerSpec):
+        """ Initializes BertPunctuationCapitalizationInferDataset. """
+        features = get_features(queries=queries, max_seq_length=max_seq_length, tokenizer=tokenizer)
+
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_subtokens_mask = features[3]
+
+    def __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.float32),
+            np.array(self.all_subtokens_mask[idx]),
+        )

tts/nemo/collections/nlp/data/token_classification/token_classification_dataset.py

@@ -0,0 +1,338 @@
+# 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.
+
+"""
+Utility functions for Token Classification NLP tasks
+Some parts of this code were adapted from the HuggingFace library at
+https://github.com/huggingface/pytorch-pretrained-BERT
+"""
+
+import os
+import pickle
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.collections.nlp.data.data_utils.data_preprocessing import get_stats
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import ChannelType, LabelsType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['BertTokenClassificationDataset', 'BertTokenClassificationInferDataset']
+
+
+def get_features(
+    queries: List[str],
+    tokenizer: TokenizerSpec,
+    max_seq_length: int = -1,
+    label_ids: dict = None,
+    pad_label: str = 'O',
+    raw_labels: List[str] = None,
+    ignore_extra_tokens: bool = False,
+    ignore_start_end: bool = False,
+):
+    """
+    Processes the data and returns features.
+    Args:
+        queries: text sequences
+        tokenizer: such as AutoTokenizer
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP], when -1 - use the max len from the data
+        pad_label: pad value use for labels. By default, it's the neutral label.
+        raw_labels: list of labels for every word in a sequence
+        label_ids: dict to map labels to label ids.
+            Starts with pad_label->0 and then increases in alphabetical order.
+            Required for training and evaluation, not needed for inference.
+        ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
+        ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
+    """
+    all_subtokens = []
+    all_loss_mask = []
+    all_subtokens_mask = []
+    all_segment_ids = []
+    all_input_ids = []
+    all_input_mask = []
+    sent_lengths = []
+    all_labels = []
+    with_label = False
+
+    if raw_labels is not None:
+        with_label = True
+
+    for i, query in enumerate(queries):
+        words = query.strip().split()
+
+        # add bos token
+        subtokens = [tokenizer.cls_token]
+        loss_mask = [1 - ignore_start_end]
+        subtokens_mask = [0]
+        if with_label:
+            pad_id = label_ids[pad_label]
+            labels = [pad_id]
+            query_labels = [label_ids[lab] for lab in raw_labels[i]]
+
+        for j, word in enumerate(words):
+            word_tokens = tokenizer.text_to_tokens(word)
+            subtokens.extend(word_tokens)
+
+            loss_mask.append(1)
+            loss_mask.extend([int(not ignore_extra_tokens)] * (len(word_tokens) - 1))
+
+            subtokens_mask.append(1)
+            subtokens_mask.extend([0] * (len(word_tokens) - 1))
+
+            if with_label:
+                labels.extend([query_labels[j]] * len(word_tokens))
+        # add eos token
+        subtokens.append(tokenizer.sep_token)
+        loss_mask.append(1 - ignore_start_end)
+        subtokens_mask.append(0)
+        sent_lengths.append(len(subtokens))
+        all_subtokens.append(subtokens)
+        all_loss_mask.append(loss_mask)
+        all_subtokens_mask.append(subtokens_mask)
+        all_input_mask.append([1] * len(subtokens))
+
+        if with_label:
+            labels.append(pad_id)
+            all_labels.append(labels)
+
+    max_seq_length_data = max(sent_lengths)
+    max_seq_length = min(max_seq_length, max_seq_length_data) if max_seq_length > 0 else max_seq_length_data
+    logging.info(f'Setting Max Seq length to: {max_seq_length}')
+    get_stats(sent_lengths)
+    too_long_count = 0
+
+    for i, subtokens in enumerate(all_subtokens):
+        if len(subtokens) > max_seq_length:
+            subtokens = [tokenizer.cls_token] + subtokens[-max_seq_length + 1 :]
+            all_input_mask[i] = [1] + all_input_mask[i][-max_seq_length + 1 :]
+            all_loss_mask[i] = [int(not ignore_start_end)] + all_loss_mask[i][-max_seq_length + 1 :]
+            all_subtokens_mask[i] = [0] + all_subtokens_mask[i][-max_seq_length + 1 :]
+
+            if with_label:
+                all_labels[i] = [pad_id] + all_labels[i][-max_seq_length + 1 :]
+            too_long_count += 1
+
+        all_input_ids.append(tokenizer.tokens_to_ids(subtokens))
+
+        if len(subtokens) < max_seq_length:
+            extra = max_seq_length - len(subtokens)
+            all_input_ids[i] = all_input_ids[i] + [0] * extra
+            all_loss_mask[i] = all_loss_mask[i] + [0] * extra
+            all_subtokens_mask[i] = all_subtokens_mask[i] + [0] * extra
+            all_input_mask[i] = all_input_mask[i] + [0] * extra
+
+            if with_label:
+                all_labels[i] = all_labels[i] + [pad_id] * extra
+
+        all_segment_ids.append([0] * max_seq_length)
+
+    logging.warning(f'{too_long_count} are longer than {max_seq_length}')
+
+    for i in range(min(len(all_input_ids), 1)):
+        logging.info("*** Example ***")
+        logging.info("i: %s", i)
+        logging.info("subtokens: %s", " ".join(list(map(str, all_subtokens[i]))))
+        logging.info("loss_mask: %s", " ".join(list(map(str, all_loss_mask[i]))))
+        logging.info("input_mask: %s", " ".join(list(map(str, all_input_mask[i]))))
+        logging.info("subtokens_mask: %s", " ".join(list(map(str, all_subtokens_mask[i]))))
+        if with_label:
+            logging.info("labels: %s", " ".join(list(map(str, all_labels[i]))))
+    return (all_input_ids, all_segment_ids, all_input_mask, all_subtokens_mask, all_loss_mask, all_labels)
+
+
+class BertTokenClassificationDataset(Dataset):
+    """
+    Creates dataset to use during training for token classification tasks with a pretrained model.
+
+    Converts from raw data to an instance that can be used by Dataloader.
+    For dataset to use during inference without labels, see BertTokenClassificationInferDataset.
+
+    Args:
+        text_file: file to sequences, each line should a sentence, no header.
+        label_file: file to labels, each line corresponds to word labels for a sentence in the text_file. No header.
+        max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+        tokenizer: such as AutoTokenizer
+        num_samples: number of samples you want to use for the dataset.
+            If -1, use all dataset. Useful for testing.
+        pad_label: pad value use for labels. By default, it's the neutral label.
+        label_ids: label_ids (dict): dict to map labels to label ids.
+            Starts with pad_label->0 and then increases in alphabetical order
+            For dev set use label_ids generated during training to support
+            cases when not all labels are present in the dev set.
+            For training set label_ids should be None.
+        ignore_extra_tokens: whether to ignore extra tokens in the loss_mask
+        ignore_start_end: whether to ignore bos and eos tokens in the loss_mask
+        use_cache: whether to use processed data cache or not
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+            'loss_mask': NeuralType(('B', 'T'), MaskType()),
+            'labels': NeuralType(('B', 'T'), LabelsType()),
+        }
+
+    def __init__(
+        self,
+        text_file: str,
+        label_file: str,
+        max_seq_length: int,
+        tokenizer: TokenizerSpec,
+        num_samples: int = -1,
+        pad_label: str = 'O',
+        label_ids: Dict[str, int] = None,
+        ignore_extra_tokens: bool = False,
+        ignore_start_end: bool = False,
+        use_cache: bool = True,
+    ):
+        """ Initializes BertTokenClassificationDataset. """
+
+        data_dir = os.path.dirname(text_file)
+        filename = os.path.basename(text_file)
+
+        if not filename.endswith('.txt'):
+            raise ValueError("{text_file} should have extension .txt")
+
+        vocab_size = getattr(tokenizer, "vocab_size", 0)
+        features_pkl = os.path.join(
+            data_dir,
+            "cached_{}_{}_{}_{}_{}".format(
+                filename, tokenizer.name, str(max_seq_length), str(vocab_size), str(num_samples)
+            ),
+        )
+
+        master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+        features = None
+        if master_device and (not use_cache or not os.path.exists(features_pkl)):
+            if num_samples == 0:
+                raise ValueError("num_samples has to be positive", num_samples)
+
+            with open(text_file, 'r') as f:
+                text_lines = f.readlines()
+
+            labels_lines = []
+            with open(label_file, 'r') as f:
+                for line in f:
+                    line = line.strip().split()
+                    labels_lines.append(line)
+
+            if len(labels_lines) != len(text_lines):
+                raise ValueError("Labels file should contain labels for every word")
+
+            if num_samples > 0:
+                dataset = list(zip(text_lines, labels_lines))
+                dataset = dataset[:num_samples]
+
+                dataset = list(zip(*dataset))
+                text_lines = dataset[0]
+                labels_lines = dataset[1]
+
+            features = get_features(
+                queries=text_lines,
+                max_seq_length=max_seq_length,
+                tokenizer=tokenizer,
+                pad_label=pad_label,
+                raw_labels=labels_lines,
+                label_ids=label_ids,
+                ignore_extra_tokens=ignore_extra_tokens,
+                ignore_start_end=ignore_start_end,
+            )
+
+            pickle.dump(features, open(features_pkl, "wb"))
+            logging.info(f'features saved to {features_pkl}')
+
+        # wait until the master process writes to the processed data files
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        if features is None:
+            features = pickle.load(open(features_pkl, 'rb'))
+            logging.info(f'features restored from {features_pkl}')
+
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_subtokens_mask = features[3]
+        self.all_loss_mask = features[4]
+        self.all_labels = features[5]
+
+    def __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.long),
+            np.array(self.all_subtokens_mask[idx]),
+            np.array(self.all_loss_mask[idx]),
+            np.array(self.all_labels[idx]),
+        )
+
+
+class BertTokenClassificationInferDataset(Dataset):
+    """
+    Creates dataset to use during inference for token classification tasks with a pretrained model.
+    For dataset to use during training with labels, see BertTokenClassificationDataset.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports.
+               """
+        return {
+            'input_ids': NeuralType(('B', 'T'), ChannelType()),
+            'segment_ids': NeuralType(('B', 'T'), ChannelType()),
+            'input_mask': NeuralType(('B', 'T'), MaskType()),
+            'subtokens_mask': NeuralType(('B', 'T'), MaskType()),
+        }
+
+    def __init__(
+        self, queries: List[str], max_seq_length: int, tokenizer: TokenizerSpec,
+    ):
+        """
+        Initializes BertTokenClassificationInferDataset
+        Args:
+            queries: text sequences
+            max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
+            tokenizer: such as AutoTokenizer
+        """
+        features = get_features(queries=queries, max_seq_length=max_seq_length, tokenizer=tokenizer)
+
+        self.all_input_ids = features[0]
+        self.all_segment_ids = features[1]
+        self.all_input_mask = features[2]
+        self.all_subtokens_mask = features[3]
+
+    def __len__(self):
+        return len(self.all_input_ids)
+
+    def __getitem__(self, idx):
+        return (
+            np.array(self.all_input_ids[idx]),
+            np.array(self.all_segment_ids[idx]),
+            np.array(self.all_input_mask[idx], dtype=np.long),
+            np.array(self.all_subtokens_mask[idx]),
+        )

tts/nemo/collections/nlp/data/token_classification/token_classification_utils.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pickle
+from typing import Dict
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import (
+    fill_class_weights,
+    get_freq_weights,
+    get_label_stats,
+)
+from nemo.utils import logging
+
+__all__ = ['get_label_ids']
+
+
+def get_label_ids(
+    label_file: str,
+    is_training: bool = False,
+    pad_label: str = 'O',
+    label_ids_dict: Dict[str, int] = None,
+    get_weights: bool = True,
+    class_labels_file_artifact='label_ids.csv',
+):
+    """
+    Generates str to int labels mapping for training data or checks correctness of the label_ids_dict
+    file for non-training files or if label_ids_dict is specified
+
+    Args:
+        label_file: the path of the label file to process
+        is_training: indicates whether the label_file is used for training
+        pad_label: token used for padding
+        label_ids_dict: str label name to int ids mapping. Required for non-training data.
+            If specified, the check that all labels from label_file are present in label_ids_dict will be performed.
+            For training data, if label_ids_dict is None, a new mapping will be generated from label_file.
+        get_weights: set to True to calculate class weights, required for Weighted Loss.
+        class_labels_file_artifact: name of the file to save in .nemo
+    """
+    if not os.path.exists(label_file):
+        raise ValueError(f'File {label_file} was not found.')
+
+    logging.info(f'Processing {label_file}')
+    if not is_training and label_ids_dict is None:
+        raise ValueError(
+            f'For non training data, label_ids_dict created during preprocessing of the training data '
+            f'should be provided'
+        )
+
+    # collect all labels from the label_file
+    data_dir = os.path.dirname(label_file)
+    unique_labels = set(pad_label)
+    all_labels = []
+    with open(label_file, 'r') as f:
+        for line in f:
+            line = line.strip().split()
+            all_labels.extend(line)
+            unique_labels.update(line)
+
+    # check that all labels from label_file are present in the specified label_ids_dict
+    # or generate label_ids_dict from data (for training only)
+    if label_ids_dict:
+        logging.info(f'Using provided labels mapping {label_ids_dict}')
+        for name in unique_labels:
+            if name not in label_ids_dict:
+                raise ValueError(f'{name} class from {label_file} not found in the provided mapping: {label_ids_dict}')
+    else:
+        label_ids_dict = {pad_label: 0}
+        if pad_label in unique_labels:
+            unique_labels.remove(pad_label)
+        for label in sorted(unique_labels):
+            label_ids_dict[label] = len(label_ids_dict)
+
+    label_ids_filename = os.path.join(data_dir, class_labels_file_artifact)
+    if is_training:
+        with open(label_ids_filename, 'w') as f:
+            labels, _ = zip(*sorted(label_ids_dict.items(), key=lambda x: x[1]))
+            f.write('\n'.join(labels))
+        logging.info(f'Labels mapping {label_ids_dict} saved to : {label_ids_filename}')
+
+    # calculate label statistics
+    base_name = os.path.splitext(os.path.basename(label_file))[0]
+    stats_file = os.path.join(data_dir, f'{base_name}_label_stats.tsv')
+    if os.path.exists(stats_file) and not is_training and not get_weights:
+        logging.info(f'{stats_file} found, skipping stats calculation.')
+    else:
+        all_labels = [label_ids_dict[label] for label in all_labels]
+        logging.info(f'Three most popular labels in {label_file}:')
+        total_labels, label_frequencies, max_id = get_label_stats(all_labels, stats_file)
+        logging.info(f'Total labels: {total_labels}. Label frequencies - {label_frequencies}')
+
+    if get_weights:
+        class_weights_pkl = os.path.join(data_dir, f'{base_name}_weights.p')
+        if os.path.exists(class_weights_pkl):
+            class_weights = pickle.load(open(class_weights_pkl, 'rb'))
+            logging.info(f'Class weights restored from {class_weights_pkl}')
+        else:
+            class_weights_dict = get_freq_weights(label_frequencies)
+            logging.info(f'Class Weights: {class_weights_dict}')
+            class_weights = fill_class_weights(class_weights_dict, max_id)
+
+            pickle.dump(class_weights, open(class_weights_pkl, "wb"))
+            logging.info(f'Class weights saved to {class_weights_pkl}')
+    else:
+        class_weights = None
+
+    return label_ids_dict, label_ids_filename, class_weights

tts/nemo/collections/nlp/losses/__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.losses.sgd_loss import SGDDialogueStateLoss

tts/nemo/collections/nlp/losses/sgd_loss.py

@@ -0,0 +1,218 @@
+# 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 torch
+
+from nemo.core.classes import Loss, Typing, typecheck
+from nemo.core.neural_types import ChannelType, LabelsType, LogitsType, LossType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['SGDDialogueStateLoss']
+
+
+class SGDDialogueStateLoss(Loss, Typing):
+    """
+    Neural module which implements loss for SGD model.
+    """
+
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+            logit_intent_status: Output of SGD model
+            intent_status: intent label
+            logit_req_slot_status: Output of SGD model
+            requested_slot_status: Takes value 1 if the corresponding slot is requested, 0 otherwise
+            logit_cat_slot_status: Output of SGD model
+            categorical_slot_status: The status of each categorical slot in the service
+            logit_cat_slot_value_status: Output of SGD model
+            categorical_slot_value_status: Takes value 1 if the corresponding slot value is correct, 0 otherwise
+            logit_noncat_slot_status: Output of SGD model
+            noncategorical_slot_status: The status of each noncategorical slot in the service\
+            logit_spans: Output of SGD model
+            noncategorical_slot_value_start: The index of the starting subword corresponding to the slot span for a non-categorical slot value
+            noncategorical_slot_value_end: The index of the ending (inclusive) subword corresponding to the slot span for a non-categorical slot value
+            task_mask: Mask contains 1 if its the current task, 0 otherwise
+        """
+        return {
+            "logit_intent_status": NeuralType(('B', 'T'), LogitsType()),
+            "intent_status": NeuralType(('B'), LabelsType()),
+            "logit_req_slot_status": NeuralType(('B', 'T'), LogitsType()),
+            "requested_slot_status": NeuralType(('B'), LabelsType()),
+            "logit_cat_slot_status": NeuralType(('B', 'T'), LogitsType()),
+            "categorical_slot_status": NeuralType(('B'), LabelsType()),
+            "logit_cat_slot_value_status": NeuralType(('B', 'T'), LogitsType()),
+            "categorical_slot_value_status": NeuralType(('B'), LabelsType()),
+            "logit_noncat_slot_status": NeuralType(('B', 'T'), LogitsType()),
+            "noncategorical_slot_status": NeuralType(('B'), LabelsType()),
+            "logit_spans": NeuralType(('B', 'T', 'D'), LogitsType()),
+            "noncategorical_slot_value_start": NeuralType(('B'), LabelsType()),
+            "noncategorical_slot_value_end": NeuralType(('B'), LabelsType()),
+            "task_mask": NeuralType(('B', 'T'), ChannelType()),
+        }
+
+    @property
+    def output_types(self):
+        """
+        Returns definitions of module output ports.
+        loss:
+            NeuralType(None)
+        """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    def __init__(self, reduction: str = 'mean'):
+        """
+        Args:
+            reduction: specifies the reduction to apply to the final loss, choose 'mean' or 'sum'
+        """
+        super().__init__()
+
+        if reduction not in ['mean', 'sum']:
+            logging.warning(f'{reduction} reduction is not supported. Setting reduction to "mean"')
+            reduction = 'mean'
+
+        self.reduction = reduction
+        self._cross_entropy = torch.nn.CrossEntropyLoss(reduction=self.reduction)
+        self._cross_entropy_bin = torch.nn.BCEWithLogitsLoss(reduction=self.reduction)
+
+    def _helper(self, logits, labels, loss_mask=None):
+        """
+        flattens logits and labels according loss mask
+        Args:
+            logits: logits
+            labels: labels
+            loss_mask: loss mask
+        Returns:
+            logits_flatten: flattened logits where loss mask is true
+            labels_flatten: flattened labels where loss mask is true
+        """
+        logits_flatten = torch.flatten(logits, start_dim=0, end_dim=-2)
+        labels_flatten = torch.flatten(labels, start_dim=0, end_dim=-1)
+
+        if loss_mask is not None:
+            if loss_mask.dtype is not torch.bool:
+                loss_mask = loss_mask > 0.5
+            loss_mask_flatten = torch.flatten(loss_mask, start_dim=0, end_dim=-1)
+            logits_flatten = logits_flatten[loss_mask_flatten]
+            labels_flatten = labels_flatten[loss_mask_flatten]
+
+        return logits_flatten, labels_flatten
+
+    @typecheck()
+    def forward(
+        self,
+        logit_intent_status,
+        intent_status,
+        logit_req_slot_status,
+        requested_slot_status,
+        logit_cat_slot_status,
+        categorical_slot_status,
+        logit_cat_slot_value_status,
+        categorical_slot_value_status,
+        logit_noncat_slot_status,
+        noncategorical_slot_status,
+        logit_spans,
+        noncategorical_slot_value_start,
+        noncategorical_slot_value_end,
+        task_mask,
+    ):
+        # Intent loss
+
+        old_logit_intent_status = logit_intent_status
+        logit_intent_status, intent_status = self._helper(logit_intent_status, intent_status, task_mask[:, 0])
+        if len(intent_status) == 0:
+            intent_loss = torch.clamp(torch.max(old_logit_intent_status.view(-1)), 0, 0)
+        else:
+            intent_loss = self._cross_entropy_bin(logit_intent_status.squeeze(dim=-1), intent_status)
+
+        old_logit_req_slot_status = logit_req_slot_status
+        logit_req_slot_status, requested_slot_status = self._helper(
+            logit_req_slot_status, requested_slot_status, task_mask[:, 1]
+        )
+        if len(requested_slot_status) == 0:
+            requested_slot_loss = torch.clamp(torch.max(old_logit_req_slot_status.view(-1)), 0, 0)
+        else:
+            requested_slot_loss = self._cross_entropy_bin(logit_req_slot_status.squeeze(dim=-1), requested_slot_status)
+
+        old_logit_cat_slot_status = logit_cat_slot_status
+        logit_cat_slot_status, categorical_slot_status = self._helper(
+            logit_cat_slot_status, categorical_slot_status, task_mask[:, 2]
+        )
+        if len(categorical_slot_status) == 0:
+            cat_slot_status_loss = torch.clamp(torch.max(old_logit_cat_slot_status.view(-1)), 0, 0)
+        else:
+            cat_slot_status_loss = self._cross_entropy(logit_cat_slot_status, categorical_slot_status,)
+        old_logit_cat_slot_value_status = logit_cat_slot_value_status
+        logit_cat_slot_value_status, categorical_slot_value_status = self._helper(
+            logit_cat_slot_value_status, categorical_slot_value_status, task_mask[:, 3]
+        )
+        if len(categorical_slot_value_status) == 0:
+            cat_slot_value_status_loss = torch.clamp(torch.max(old_logit_cat_slot_value_status.view(-1)), 0, 0)
+        else:
+            cat_slot_value_status_loss = self._cross_entropy_bin(
+                logit_cat_slot_value_status.squeeze(dim=-1), categorical_slot_value_status
+            )
+
+        old_logit_noncat_slot_status = logit_noncat_slot_status
+        logit_noncat_slot_status, noncategorical_slot_status = self._helper(
+            logit_noncat_slot_status, noncategorical_slot_status, task_mask[:, 4]
+        )
+        if len(noncategorical_slot_status) == 0:
+            noncat_slot_status_loss = torch.clamp(torch.max(old_logit_noncat_slot_status.view(-1)), 0, 0)
+        else:
+            noncat_slot_status_loss = self._cross_entropy(logit_noncat_slot_status, noncategorical_slot_status,)
+
+        logit_noncat_slot_start, logit_noncat_slot_end = torch.unbind(logit_spans, dim=-1)
+
+        _, max_num_tokens = logit_noncat_slot_start.size()
+        old_logit_noncat_slot_start = logit_noncat_slot_start
+        logit_noncat_slot_start, noncategorical_slot_value_start = self._helper(
+            logit_noncat_slot_start, noncategorical_slot_value_start, task_mask[:, 5]
+        )
+        if len(noncategorical_slot_value_start) == 0:
+            span_start_loss = torch.clamp(torch.max(old_logit_noncat_slot_start.view(-1)), 0, 0)
+        else:
+            span_start_loss = self._cross_entropy(logit_noncat_slot_start, noncategorical_slot_value_start)
+
+        old_logit_noncat_slot_end = logit_noncat_slot_end
+        logit_noncat_slot_end, noncategorical_slot_value_end = self._helper(
+            logit_noncat_slot_end, noncategorical_slot_value_end, task_mask[:, 5]
+        )
+        if len(noncategorical_slot_value_end) == 0:
+            span_end_loss = torch.clamp(torch.max(old_logit_noncat_slot_end.view(-1)), 0, 0)
+        else:
+            span_end_loss = self._cross_entropy(logit_noncat_slot_end, noncategorical_slot_value_end)
+
+        losses = {
+            "intent_loss": intent_loss,
+            "requested_slot_loss": requested_slot_loss,
+            "cat_slot_status_loss": cat_slot_status_loss,
+            "cat_slot_value_status_loss": cat_slot_value_status_loss,
+            "noncat_slot_status_loss": noncat_slot_status_loss,
+            "span_start_loss": span_start_loss,
+            "span_end_loss": span_end_loss,
+        }
+
+        total_loss = sum(losses.values())
+        if self.reduction == 'mean':
+            total_loss = total_loss / len(losses)
+        else:
+            batch_size = logit_intent_status.shape[0]
+            total_loss = total_loss / batch_size
+        return total_loss

tts/nemo/collections/nlp/metrics/__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.metrics.classification_report import ClassificationReport
+from nemo.collections.nlp.metrics.sequence_perplexity import SequencePerplexity

tts/nemo/collections/nlp/metrics/classification_report.py

@@ -0,0 +1,173 @@
+# 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 Any, Dict, Optional
+
+import torch
+from torchmetrics import Metric
+from torchmetrics.utilities.data import METRIC_EPS
+
+__all__ = ['ClassificationReport']
+
+
+class ClassificationReport(Metric):
+    """
+    This metric computes the number of True Positive, False Negative, and False Positive examples per class.
+    When doing distributed training/evaluation the result of res=ClassificationReport(predictions, labels) calls
+    will be all-reduced between all workers using SUM operations.
+
+    If used with PytorchLightning LightningModule, include TPs, FNs, and FPs inside validation_step results.
+    Then aggregate them at the end of validation epoch to correctly compute validation precision, recall, f1
+    using get_precision_recall_f1().
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            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):
+            ...
+            # 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)
+
+    Args:
+        num_classes: number of classes in the dataset
+        label_ids (optional): label name to label id mapping
+        mode: how to compute the average
+        dist_sync_on_step: sync across ddp
+        process_group: which processes to sync across
+    Return:
+        aggregated precision, recall, f1, report
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        label_ids: Dict[str, int] = None,
+        mode: str = 'macro',
+        dist_sync_on_step: bool = False,
+        process_group: Optional[Any] = None,
+    ):
+        super().__init__(dist_sync_on_step=dist_sync_on_step, process_group=process_group)
+        self.num_classes = num_classes
+        if label_ids:
+            self.ids_to_labels = {v: k for k, v in label_ids.items()}
+        else:
+            self.ids_to_labels = None
+        self.mode = mode
+
+        self.add_state("tp", default=torch.zeros(num_classes), dist_reduce_fx='sum', persistent=False)
+        self.add_state("fn", default=torch.zeros(num_classes), dist_reduce_fx='sum', persistent=False)
+        self.add_state("fp", default=torch.zeros(num_classes), dist_reduce_fx='sum', persistent=False)
+        self.add_state(
+            "num_examples_per_class", default=torch.zeros(num_classes), dist_reduce_fx='sum', persistent=False
+        )
+
+    def update(self, predictions: torch.Tensor, labels: torch.Tensor):
+        TP = []
+        FN = []
+        FP = []
+        for label_id in range(self.num_classes):
+            current_label = labels == label_id
+            label_predicted = predictions == label_id
+
+            TP.append((label_predicted == current_label)[label_predicted].sum())
+            FP.append((label_predicted != current_label)[label_predicted].sum())
+            FN.append((label_predicted != current_label)[current_label].sum())
+
+        tp = torch.tensor(TP).to(predictions.device)
+        fn = torch.tensor(FN).to(predictions.device)
+        fp = torch.tensor(FP).to(predictions.device)
+        num_examples_per_class = tp + fn
+
+        self.tp += tp
+        self.fn += fn
+        self.fp += fp
+        self.num_examples_per_class += num_examples_per_class
+
+    def compute(self):
+        """
+        Aggregates and then calculates logs classification report similar to sklearn.metrics.classification_report.
+        Typically used during epoch_end.
+        Return:
+            aggregated precision, recall, f1, report
+        """
+        total_examples = torch.sum(self.num_examples_per_class)
+        num_non_empty_classes = torch.nonzero(self.num_examples_per_class).size(0)
+
+        precision = torch.true_divide(self.tp * 100, (self.tp + self.fp + METRIC_EPS))
+        recall = torch.true_divide(self.tp * 100, (self.tp + self.fn + METRIC_EPS))
+        f1 = torch.true_divide(2 * precision * recall, (precision + recall + METRIC_EPS))
+
+        report = '\n{:50s}   {:10s}   {:10s}   {:10s}   {:10s}'.format('label', 'precision', 'recall', 'f1', 'support')
+        for i in range(len(self.tp)):
+            label = f'label_id: {i}'
+            if self.ids_to_labels and i in self.ids_to_labels:
+                label = f'{self.ids_to_labels[i]} ({label})'
+
+            report += '\n{:50s}   {:8.2f}   {:8.2f}   {:8.2f}   {:8.0f}'.format(
+                label, precision[i], recall[i], f1[i], self.num_examples_per_class[i]
+            )
+
+        micro_precision = torch.true_divide(torch.sum(self.tp) * 100, torch.sum(self.tp + self.fp) + METRIC_EPS)
+        micro_recall = torch.true_divide(torch.sum(self.tp) * 100, torch.sum(self.tp + self.fn) + METRIC_EPS)
+        micro_f1 = torch.true_divide(2 * micro_precision * micro_recall, (micro_precision + micro_recall + METRIC_EPS))
+
+        macro_precision = torch.sum(precision) / num_non_empty_classes
+        macro_recall = torch.sum(recall) / num_non_empty_classes
+        macro_f1 = torch.sum(f1) / num_non_empty_classes
+        weighted_precision = torch.sum(precision * self.num_examples_per_class) / total_examples
+        weighted_recall = torch.sum(recall * self.num_examples_per_class) / total_examples
+        weighted_f1 = torch.sum(f1 * self.num_examples_per_class) / total_examples
+
+        report += "\n-------------------"
+
+        report += '\n{:50s}   {:8.2f}   {:8.2f}   {:8.2f}   {:8.0f}'.format(
+            'micro avg', micro_precision, micro_recall, micro_f1, total_examples
+        )
+
+        report += '\n{:50s}   {:8.2f}   {:8.2f}   {:8.2f}   {:8.0f}'.format(
+            'macro avg', macro_precision, macro_recall, macro_f1, total_examples
+        )
+        report += (
+            '\n{:50s}   {:8.2f}   {:8.2f}   {:8.2f}   {:8.0f}'.format(
+                'weighted avg', weighted_precision, weighted_recall, weighted_f1, total_examples
+            )
+            + '\n'
+        )
+
+        self.total_examples = total_examples
+
+        if self.mode == 'macro':
+            return macro_precision, macro_recall, macro_f1, report
+        elif self.mode == 'weighted':
+            return weighted_precision, weighted_recall, weighted_f1, report
+        elif self.mode == 'micro':
+            return micro_precision, micro_recall, micro_f1, report
+        elif self.mode == 'all':
+            return precision, recall, f1, report
+        else:
+            raise ValueError(
+                f'{self.mode} mode is not supported. Choose "macro" to get aggregated numbers \
+            or "all" to get values for each class.'
+            )

tts/nemo/collections/nlp/metrics/sequence_perplexity.py

@@ -0,0 +1,78 @@
+# 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 torch
+from torchmetrics import Metric
+
+__all__ = ['SequencePerplexity']
+
+
+class SequencePerplexity(Metric):
+    """
+    This class computes mean perplexity across the batches of sequences.
+
+    You have to provide ``log_probs`` (float tensor of shape [batch_size x seq_length x vocab_size]) and
+    ``labels`` (int tensor of shape [batch_size x seq_length] with values from the range [0, vocab_size-1])
+    to the :meth:`update` method. If some of the sequences are shorter than seq_length, you can also provide
+    an optional argument ``mask`` (bool tensor of shape [batch_size x seq_length]) which masks out tokens
+    not participating in perplexity computation.
+
+    See :doc:`PyTorch Lightning Metrics<pytorch-lightning:metrics>` for the metric usage instructions.
+
+    Args:
+        compute_on_step:
+            Forward only calls ``update()`` and returns ``None`` if this is set to ``False``. default: ``True``
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()`` before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: ``None`` (which selects the entire
+                world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When ``None``, DDP will be used
+                to perform the allgather.
+    """
+
+    def __init__(self, compute_on_step=True, dist_sync_on_step=False, process_group=None, dist_sync_fn=None):
+        super().__init__(
+            compute_on_step=compute_on_step,
+            dist_sync_on_step=dist_sync_on_step,
+            process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+        )
+
+        # Total sum of exponentiated average negative log likelihoods
+        self.add_state('perplexities_sum', default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx='sum')
+        # Total number of sequences in all batches
+        self.add_state('num_sequences', default=torch.tensor(0, dtype=torch.int64), dist_reduce_fx='sum')
+
+    def update(self, log_probs: torch.Tensor, labels: torch.Tensor, mask=None):
+
+        if mask is None:
+            mask = torch.ones_like(labels)
+        if mask.dtype is not log_probs.dtype:
+            mask = mask.to(log_probs.dtype)
+
+        target_log_probs = log_probs.gather(2, labels.unsqueeze(2)).squeeze(2)
+        avg_neg_ll = -(target_log_probs * mask).sum(dim=-1) / mask.sum(dim=-1)
+        ppl = avg_neg_ll.exp()
+        self.num_sequences += ppl.numel()
+        self.perplexities_sum += ppl.sum()
+
+    def compute(self):
+        """
+        Returns perplexity across all workers and resets to 0 :attr:`perplexities_sum` and :attr:`num_sequences`.
+        """
+        if self.num_sequences.eq(0):
+            return None
+        return self.perplexities_sum / self.num_sequences

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

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

@@ -0,0 +1,33 @@
+# 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.modules.common import (
+    AlbertEncoder,
+    BertEncoder,
+    BertModule,
+    DistilBertEncoder,
+    MegatronBertEncoder,
+    RobertaEncoder,
+    SequenceClassifier,
+    SequenceRegression,
+    SequenceTokenClassifier,
+    get_lm_model,
+    get_megatron_lm_models_list,
+    get_pretrained_lm_models_list,
+    get_tokenizer,
+    get_tokenizer_list,
+)
+from nemo.collections.nlp.modules.dialogue_state_tracking.sgd_decoder import SGDDecoder
+from nemo.collections.nlp.modules.dialogue_state_tracking.sgd_encoder import SGDEncoder

tts/nemo/collections/nlp/modules/common/__init__.py

@@ -0,0 +1,34 @@
+# 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 nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.collections.nlp.modules.common.huggingface import (
+    AlbertEncoder,
+    BertEncoder,
+    DistilBertEncoder,
+    RobertaEncoder,
+)
+from nemo.collections.nlp.modules.common.lm_utils import (
+    get_lm_model,
+    get_megatron_lm_models_list,
+    get_pretrained_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.megatron import MegatronBertEncoder
+from nemo.collections.nlp.modules.common.sequence_classifier import SequenceClassifier
+from nemo.collections.nlp.modules.common.sequence_regression import SequenceRegression
+from nemo.collections.nlp.modules.common.sequence_token_classifier import SequenceTokenClassifier
+from nemo.collections.nlp.modules.common.token_classifier import BertPretrainingTokenClassifier, TokenClassifier
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer, get_tokenizer_list

tts/nemo/collections/nlp/modules/common/bert_module.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+from typing import Dict, Optional
+
+import torch
+
+from nemo.core.classes import NeuralModule
+from nemo.core.classes.exportable import Exportable
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['BertModule']
+
+
+class BertModule(NeuralModule, Exportable):
+    @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 {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """Restores module/model with weights"""
+        pass
+
+    @classmethod
+    def save_to(self, save_path: str):
+        """Saves module/model with weights"""
+        pass
+
+    def restore_weights(self, restore_path: str):
+        """Restores module/model's weights"""
+        logging.info(f"Restoring weights from {restore_path}")
+
+        if not os.path.exists(restore_path):
+            logging.warning(f'Path {restore_path} not found')
+            return
+
+        pretrained_dict = torch.load(restore_path)
+
+        # backward compatibility with NeMo0.11
+        if "state_dict" in pretrained_dict.keys():
+            pretrained_dict = pretrained_dict["state_dict"]
+
+        # remove prefix from pretrained dict
+        m = re.match("^bert.*?\.", list(pretrained_dict.keys())[0])
+        if m:
+            prefix = m.group(0)
+            pretrained_dict = {k[len(prefix) :]: v for k, v in pretrained_dict.items()}
+        model_dict = self.state_dict()
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+
+        # starting with transformers 3.1.0, embeddings.position_ids is added to the model's state dict and could be
+        # missing in checkpoints trained with older transformers version
+        if 'embeddings.position_ids' in model_dict and 'embeddings.position_ids' not in pretrained_dict:
+            pretrained_dict['embeddings.position_ids'] = model_dict['embeddings.position_ids']
+
+        assert len(pretrained_dict) == len(model_dict)
+        model_dict.update(pretrained_dict)
+        self.load_state_dict(model_dict)
+        logging.info(f"Weights for {type(self).__name__} restored from {restore_path}")
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        sample = next(self.parameters())
+        input_ids = torch.randint(low=0, high=2048, size=(2, 16), device=sample.device)
+        attention_mask = torch.randint(low=0, high=1, size=(2, 16), device=sample.device)
+        return tuple([input_ids, attention_mask, attention_mask])

tts/nemo/collections/nlp/modules/common/classifier.py

@@ -0,0 +1,87 @@
+# 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 torch import nn as nn
+
+from nemo.collections.common.parts import transformer_weights_init
+from nemo.core.classes import Exportable, NeuralModule, typecheck
+from nemo.core.neural_types import ChannelType, NeuralType
+
+__all__ = ['Classifier']
+
+
+class Classifier(NeuralModule, Exportable):
+    """
+    A baseclass for modules to perform various classification tasks.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module input ports.
+        We implement it here since all NLP classifiers have the same inputs
+        """
+        return {"hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    def __init__(self, hidden_size: int, dropout: float = 0.0,) -> None:
+        """
+        Initializes the Classifier base module.
+        Args:
+            hidden_size: the size of the hidden dimension
+            dropout: dropout to apply to the input hidden states
+        """
+        super().__init__()
+        self._hidden_size = hidden_size
+        self.dropout = nn.Dropout(dropout)
+
+    def post_init(self, use_transformer_init: bool):
+        """
+        Common post-processing to be called at the end of concrete Classifiers init methods
+        Args:
+          use_transformer_init : whether or not to apply transformer_weights_init
+        """
+        if use_transformer_init:
+            self.apply(lambda module: transformer_weights_init(module, xavier=False))
+
+    def input_example(self):
+        """
+        Generates input examples for tracing etc.
+        Returns:
+            A tuple of input examples.
+        """
+        bs = 8
+        seq = 64
+        sample = next(self.parameters())
+        input_example = torch.randn(bs, seq, self._hidden_size).to(sample.device).to(sample.dtype)
+        return tuple([input_example])
+
+    def save_to(self, save_path: str):
+        """
+        Saves the module to the specified path.
+        Args:
+            save_path: Path to where to save the module.
+        """
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """
+        Restores the module from the specified path.
+        Args:
+            restore_path: Path to restore the module from.
+        """
+        pass

tts/nemo/collections/nlp/modules/common/decoder_module.py

@@ -0,0 +1,58 @@
+# 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 abc import ABC
+from typing import Any, Dict, Optional
+
+from nemo.core.classes import NeuralModule
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+
+__all__ = ['DecoderModule']
+
+
+class DecoderModule(NeuralModule, ABC):
+    """ Base class for decoder neural module to be used in NLP models. """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "decoder_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "encoder_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
+            "encoder_embeddings": NeuralType(('B', 'T', 'D'), ChannelType(), optional=True),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        raise NotImplementedError
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        raise NotImplementedError
+
+    @property
+    def embedding(self) -> Optional[Any]:
+        raise NotImplementedError
+
+    @property
+    def decoder(self) -> Optional[Any]:
+        raise NotImplementedError
+
+    @property
+    def max_sequence_length(self) -> Optional[int]:
+        raise NotImplementedError

tts/nemo/collections/nlp/modules/common/encoder_module.py

@@ -0,0 +1,41 @@
+# 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 abc import ABC
+from typing import Dict, Optional
+
+from nemo.core.classes import NeuralModule
+from nemo.core.neural_types import ChannelType, MaskType, NeuralType
+from nemo.utils import logging
+
+__all__ = ['EncoderModule']
+
+
+class EncoderModule(NeuralModule, ABC):
+    """ Base class for encoder neural module to be used in NLP models. """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "input_ids": NeuralType(('B', 'T'), ChannelType()),
+            "encoder_mask": NeuralType(('B', 'T'), MaskType()),
+        }
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"last_hidden_states": NeuralType(('B', 'T', 'D'), ChannelType())}
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        raise NotImplementedError

tts/nemo/collections/nlp/modules/common/huggingface/__init__.py

@@ -0,0 +1,22 @@
+# 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.modules.common.huggingface.albert import AlbertEncoder
+from nemo.collections.nlp.modules.common.huggingface.bert import BertEncoder
+from nemo.collections.nlp.modules.common.huggingface.distilbert import DistilBertEncoder
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import (
+    get_huggingface_lm_model,
+    get_huggingface_pretrained_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.huggingface.roberta import RobertaEncoder

tts/nemo/collections/nlp/modules/common/huggingface/albert.py

@@ -0,0 +1,33 @@
+# 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 transformers import AlbertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['AlbertEncoder']
+
+
+class AlbertEncoder(AlbertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/bert.py

@@ -0,0 +1,33 @@
+# 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 transformers import BertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['BertEncoder']
+
+
+class BertEncoder(BertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/distilbert.py

@@ -0,0 +1,34 @@
+# Copyright 2020 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 transformers import DistilBertModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['DistilBertEncoder']
+
+
+class DistilBertEncoder(DistilBertModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        # distilBert does not use token_type_ids as the most of the other Bert models
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask)[0]
+        return res

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_decoder.py

@@ -0,0 +1,79 @@
+# 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 Optional
+
+from hydra.utils import instantiate
+from transformers import AutoConfig, AutoModel
+
+from nemo.collections.nlp.modules.common.decoder_module import DecoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.utils import logging
+
+
+class HuggingFaceDecoderModule(DecoderModule):
+    """Gets HuggingFace based model to be used as an Decoder in NeMo NLP.
+    Use the model_name arg to get a named model architecture. 
+    Available model names can be found with get_huggingface_pretrained_lm_models_list() or
+    by going to https://huggingface.co/models.
+    Use the pretrained arg to get the named model architecture with or without pretrained weights.
+
+    If model_name is None, then we can pass in a custom configuration via the config_dict.
+    For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+        config_dict={
+            '_target_': 'transformers.BertConfig',
+            'hidden_size': 1536
+        } 
+
+
+    Args:
+        model_name (Optional[str]): Named model architecture from HuggingFace. Defaults to None.
+        pretrained (bool): Use True to get pretrained weights. 
+                                    False will use the same architecture but with randomly initialized weights.
+                                    Defaults to False.
+        config_dict (Optional[dict], optional): Use for custom configuration of the HuggingFace model. Defaults to None.
+        checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        pretrained: bool = False,
+        config_dict: Optional[dict] = None,
+        checkpoint_file: Optional[str] = None,
+    ):
+        super().__init__()
+        model = None
+        if model_name is not None:
+            if model_name in get_huggingface_pretrained_lm_models_list():
+                if pretrained:
+                    model = AutoModel.from_pretrained(model_name)
+                else:
+                    cfg = AutoConfig.from_pretrained(model_name)
+                    model = AutoModel.from_config(cfg)
+            else:
+                logging.error(f'{model_name} not found in list of HuggingFace pretrained models')
+        else:
+            cfg = instantiate(config_dict)
+            model = AutoModel.from_config(cfg)
+        self._hidden_size = model.config.hidden_size
+        self._vocab_size = model.config.vocab_size
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        return self._hidden_size
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        return self._vocab_size

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_encoder.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from hydra.utils import instantiate
+from transformers import AutoConfig, AutoModel
+
+from nemo.collections.nlp.modules.common.encoder_module import EncoderModule
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.core.classes.common import typecheck
+from nemo.utils import logging
+
+
+class HuggingFaceEncoderModule(EncoderModule):
+    """ Class for using HuggingFace encoders in NeMo NLP."""
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        pretrained: bool = False,
+        config_dict: Optional[dict] = None,
+        checkpoint_file: Optional[str] = None,
+    ):
+        """Gets HuggingFace based model to be used as an Encoder in NeMo NLP.
+        Use the model_name arg to get a named model architecture. 
+        Available model names can be found with get_huggingface_pretrained_lm_models_list() or
+        by going to https://huggingface.co/models.
+        Use the pretrained arg to get the named model architecture with or without pretrained weights.
+
+        If model_name is None, then we can pass in a custom configuration via the config_dict.
+        For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+            config_dict={
+                '_target_': 'transformers.BertConfig',
+                'hidden_size': 1536
+            } 
+
+
+        Args:
+            model_name (Optional[str]): Named model architecture from HuggingFace. Defaults to None.
+            pretrained (bool): Use True to get pretrained weights. 
+                                        False will use the same architecture but with randomly initialized weights.
+                                        Defaults to False.
+            config_dict (Optional[dict], optional): Use for custom configuration of the HuggingFace model. Defaults to None.
+            checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+        """
+        super().__init__()
+
+        if checkpoint_file:
+            raise NotImplementedError('Restoring from checkpoint file not implemented yet.')
+
+        model = None
+        if model_name is not None:
+            if model_name in get_huggingface_pretrained_lm_models_list(include_external=True):
+                if pretrained:
+                    config_dict.pop('vocab_size')
+                    if config_dict:
+                        raise ValueError(
+                            f'When using pretrained model, config_dict should be None or empty. Got: {config_dict}'
+                        )
+                    model = AutoModel.from_pretrained(model_name)
+                else:
+                    cfg = AutoConfig.from_pretrained(model_name)
+                    model = AutoModel.from_config(cfg)
+            else:
+                logging.error(f'{model_name} not found in list of HuggingFace pretrained models')
+        else:
+            if pretrained:
+                raise ValueError(f'If not using model_name, then pretrained should be False. Got: {pretrained}.')
+            cfg = instantiate(config_dict)
+            model = AutoModel.from_config(cfg)
+        self._hidden_size = model.config.hidden_size
+        self._vocab_size = model.config.vocab_size
+
+        self._encoder = model
+
+    @typecheck()
+    def forward(self, input_ids, encoder_mask):
+        encoder_hidden_states = self._encoder.forward(input_ids=input_ids, attention_mask=encoder_mask)[0]
+        return encoder_hidden_states
+
+    @property
+    def hidden_size(self) -> Optional[int]:
+        return self._hidden_size
+
+    @property
+    def vocab_size(self) -> Optional[int]:
+        return self._vocab_size

tts/nemo/collections/nlp/modules/common/huggingface/huggingface_utils.py

@@ -0,0 +1,126 @@
+# 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 List, Optional
+
+from transformers import (
+    ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    ALL_PRETRAINED_CONFIG_ARCHIVE_MAP,
+    BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+    AlbertConfig,
+    AutoModel,
+    BertConfig,
+    DistilBertConfig,
+    RobertaConfig,
+)
+
+from nemo.collections.nlp.modules.common.huggingface.albert import AlbertEncoder
+from nemo.collections.nlp.modules.common.huggingface.bert import BertEncoder
+from nemo.collections.nlp.modules.common.huggingface.distilbert import DistilBertEncoder
+from nemo.collections.nlp.modules.common.huggingface.roberta import RobertaEncoder
+from nemo.utils import logging
+
+__all__ = ["get_huggingface_lm_model", "get_huggingface_pretrained_lm_models_list"]
+
+
+HUGGINGFACE_MODELS = {
+    "BertModel": {
+        "default": "bert-base-uncased",
+        "class": BertEncoder,
+        "config": BertConfig,
+        "pretrained_model_list": BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "DistilBertModel": {
+        "default": "distilbert-base-uncased",
+        "class": DistilBertEncoder,
+        "config": DistilBertConfig,
+        "pretrained_model_list": DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "RobertaModel": {
+        "default": "roberta-base",
+        "class": RobertaEncoder,
+        "config": RobertaConfig,
+        "pretrained_model_list": ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+    "AlbertModel": {
+        "default": "albert-base-v2",
+        "class": AlbertEncoder,
+        "config": AlbertConfig,
+        "pretrained_model_list": ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+    },
+}
+
+
+def get_huggingface_lm_model(
+    pretrained_model_name: str, config_dict: Optional[dict] = None, config_file: Optional[str] = None,
+):
+    """
+    Returns lm model instantiated with Huggingface
+
+    Args:
+        pretrained_mode_name: specify this to instantiate pretrained model from Huggingface,
+            e.g. bert-base-cased. For entire list, see get_huggingface_pretrained_lm_models_list().
+        config_dict: model configuration dictionary used to instantiate Huggingface model from scratch
+        config_file: path to model configuration file used to instantiate Huggingface model from scratch
+
+    Returns:
+        BertModule
+    """
+
+    try:
+        automodel = AutoModel.from_pretrained(pretrained_model_name)
+    except Exception as e:
+        raise ValueError(f"{pretrained_model_name} is not supported by HuggingFace. {e}")
+
+    model_type = type(automodel).__name__
+    if model_type in HUGGINGFACE_MODELS:
+        model_class = HUGGINGFACE_MODELS[model_type]["class"]
+        if config_file:
+            if not os.path.exists(config_file):
+                logging.warning(
+                    f"Config file was not found at {config_file}. Will attempt to use config_dict or pretrained_model_name."
+                )
+            else:
+                config_class = HUGGINGFACE_MODELS[model_type]["config"]
+                return model_class(config_class.from_json_file(config_file))
+        if config_dict:
+            config_class = HUGGINGFACE_MODELS[model_type]["config"]
+            return model_class(config=config_class(**config_dict))
+        else:
+            return model_class.from_pretrained(pretrained_model_name)
+    else:
+        raise ValueError(f"Use HuffingFace API directly in NeMo for {pretrained_model_name}")
+
+
+def get_huggingface_pretrained_lm_models_list(include_external: bool = False,) -> List[str]:
+    """
+    Returns the list of pretrained HuggingFace language models
+    
+    Args:
+        include_external if true includes all HuggingFace model names, not only those supported language models in NeMo.
+    
+    Returns the list of HuggingFace models
+    """
+
+    huggingface_models = []
+    if include_external:
+        huggingface_models = list(ALL_PRETRAINED_CONFIG_ARCHIVE_MAP.keys())
+    else:
+        for model in HUGGINGFACE_MODELS:
+            model_names = HUGGINGFACE_MODELS[model]["pretrained_model_list"]
+            huggingface_models.extend(model_names)
+    return huggingface_models

tts/nemo/collections/nlp/modules/common/huggingface/roberta.py

@@ -0,0 +1,34 @@
+# 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 transformers import RobertaModel
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+
+__all__ = ['RobertaEncoder']
+
+
+class RobertaEncoder(RobertaModel, BertModule):
+    """
+    Wraps around the Huggingface transformers implementation repository for easy use within NeMo.
+    """
+
+    @typecheck()
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        res = super().forward(input_ids=input_ids, attention_mask=attention_mask)[0]
+        return res

tts/nemo/collections/nlp/modules/common/lm_utils.py

@@ -0,0 +1,179 @@
+# 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 dataclasses import dataclass
+from typing import List, Optional, Union
+
+from attr import asdict
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+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.huggingface.huggingface_utils import (
+    get_huggingface_lm_model,
+    get_huggingface_pretrained_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import (
+    get_megatron_lm_model,
+    get_megatron_lm_models_list,
+)
+from nemo.collections.nlp.modules.common.transformer.transformer import NeMoTransformerConfig
+from nemo.collections.nlp.modules.common.transformer.transformer_utils import (
+    get_huggingface_transformer,
+    get_nemo_transformer,
+)
+from nemo.utils import logging
+
+__all__ = ['get_pretrained_lm_models_list', 'get_lm_model']
+
+
+def get_pretrained_lm_models_list(include_external: bool = False) -> List[str]:
+    """
+    Returns the list of supported pretrained model names
+
+    Args:
+        include_external if true includes all HuggingFace model names, not only those supported language models in NeMo.
+
+    """
+    return get_megatron_lm_models_list() + get_huggingface_pretrained_lm_models_list(include_external=include_external)
+
+
+def get_lm_model(
+    pretrained_model_name: str,
+    config_dict: Optional[dict] = None,
+    config_file: Optional[str] = None,
+    checkpoint_file: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+) -> BertModule:
+    """
+    Helper function to instantiate a language model encoder, either from scratch or a pretrained model.
+    If only pretrained_model_name are passed, a pretrained model is returned.
+    If a configuration is passed, whether as a file or dictionary, the model is initialized with random weights.
+
+    Args:
+        pretrained_model_name: pretrained model name, for example, bert-base-uncased or megatron-bert-cased.
+            See get_pretrained_lm_models_list() for full list.
+        config_dict: path to the model configuration dictionary
+        config_file: path to the model configuration file
+        checkpoint_file: path to the pretrained model checkpoint
+        vocab_file: path to vocab_file to be used with Megatron-LM
+
+    Returns:
+        Pretrained BertModule
+    """
+
+    # check valid model type
+    if not pretrained_model_name or pretrained_model_name not in get_pretrained_lm_models_list(include_external=False):
+        logging.warning(
+            f'{pretrained_model_name} is not in get_pretrained_lm_models_list(include_external=False), '
+            f'will be using AutoModel from HuggingFace.'
+        )
+
+    # warning when user passes both configuration dict and file
+    if config_dict and config_file:
+        logging.warning(
+            f"Both config_dict and config_file were found, defaulting to use config_file: {config_file} will be used."
+        )
+
+    if "megatron" in pretrained_model_name:
+        model, checkpoint_file = get_megatron_lm_model(
+            config_dict=config_dict,
+            config_file=config_file,
+            pretrained_model_name=pretrained_model_name,
+            checkpoint_file=checkpoint_file,
+            vocab_file=vocab_file,
+        )
+    else:
+        model = get_huggingface_lm_model(
+            config_dict=config_dict, config_file=config_file, pretrained_model_name=pretrained_model_name,
+        )
+
+    if checkpoint_file and os.path.exists(checkpoint_file):
+        model.restore_weights(restore_path=checkpoint_file)
+
+    return model
+
+
+# @dataclass
+# class TransformerConfig:
+#     library: str = 'nemo'
+#     model_name: Optional[str] = None
+#     pretrained: bool = False
+#     config_dict: Optional[dict] = None
+#     checkpoint_file: Optional[str] = None
+#     encoder: bool = True
+
+
+def get_transformer(
+    library: str = 'nemo',
+    model_name: Optional[str] = None,
+    pretrained: bool = False,
+    config_dict: Optional[dict] = None,
+    checkpoint_file: Optional[str] = None,
+    encoder: bool = True,
+    pre_ln_final_layer_norm=True,
+) -> Union[EncoderModule, DecoderModule]:
+    """Gets Transformer based model to be used as an Encoder or Decoder in NeMo NLP.
+       First choose the library to get the transformer from. This can be huggingface,
+       megatron, or nemo. Use the model_name arg to get a named model architecture
+       and use the pretrained arg to get the named model architecture with pretrained weights.
+
+       If model_name is None, then we can pass in a custom configuration via the config_dict.
+       For example, to instantiate a HuggingFace BERT model with custom configuration we would do:
+       encoder = get_transformer(library='huggingface',
+                                 config_dict={
+                                     '_target_': 'transformers.BertConfig',
+                                     'hidden_size': 1536
+                                 }) 
+
+
+    Args:
+        library (str, optional): Can be 'nemo', 'huggingface', or 'megatron'. Defaults to 'nemo'.
+        model_name (Optional[str], optional): Named model architecture from the chosen library. Defaults to None.
+        pretrained (bool, optional): Use True to get pretrained weights. 
+                                     False will use the same architecture but with randomly initialized weights.
+                                     Defaults to False.
+        config_dict (Optional[dict], optional): Use for custom configuration of transformer. Defaults to None.
+        checkpoint_file (Optional[str], optional): Provide weights for the transformer from a local checkpoint. Defaults to None.
+        encoder (bool, optional): True returns an EncoderModule, False returns a DecoderModule. Defaults to True.
+
+    Returns:
+        Union[EncoderModule, DecoderModule]: Ensures that Encoder/Decoder will work in EncDecNLPModel
+    """
+
+    model = None
+
+    if library == 'nemo':
+        if isinstance(config_dict, NeMoTransformerConfig):
+            config_dict = asdict(config_dict)
+        model = get_nemo_transformer(
+            model_name=model_name,
+            pretrained=pretrained,
+            config_dict=config_dict,
+            encoder=encoder,
+            pre_ln_final_layer_norm=pre_ln_final_layer_norm,
+        )
+
+        if checkpoint_file is not None:
+            if os.path.isfile(checkpoint_file):
+                raise ValueError(f'Loading transformer weights from checkpoint file has not been implemented yet.')
+
+    elif library == 'huggingface':
+        model = get_huggingface_transformer(
+            model_name=model_name, pretrained=pretrained, config_dict=config_dict, encoder=encoder
+        )
+
+    return model

tts/nemo/collections/nlp/modules/common/megatron/__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.modules.common.megatron.megatron_bert import MegatronBertEncoder
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import (
+    get_megatron_checkpoint,
+    get_megatron_lm_models_list,
+)

tts/nemo/collections/nlp/modules/common/megatron/megatron_bert.py

@@ -0,0 +1,216 @@
+# 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
+
+import torch
+from megatron import get_args, initialize_megatron
+from megatron.checkpointing import set_checkpoint_version
+from megatron.model import get_language_model
+from megatron.model.bert_model import bert_attention_mask_func, bert_extended_attention_mask, bert_position_ids
+from megatron.mpu import (
+    get_model_parallel_group,
+    model_parallel_is_initialized,
+    set_pipeline_model_parallel_rank,
+    set_pipeline_model_parallel_world_size,
+)
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.modules.common.bert_module import BertModule
+from nemo.core.classes import typecheck
+from nemo.utils import logging
+from nemo.utils.app_state import AppState
+
+__all__ = ['MegatronBertEncoder']
+
+
+def complete_lazy_init(self):
+    # finish megatron-lm initialization
+    if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
+        self._lazy_init_fn()
+        self._lazy_init_fn = None
+
+
+class MegatronBertEncoder(BertModule):
+    """
+    MegatronBERT wraps around the Megatron Language model
+    from https://github.com/NVIDIA/Megatron-LM
+
+    Args:
+        config_file (str): path to model configuration file.
+        vocab_file (str): path to vocabulary file.
+        tokenizer_type (str): tokenizer type, currently only 'BertWordPieceLowerCase' supported.
+    """
+
+    def __init__(self, model_name, config, vocab_file, model_parallel_size=None, model_parallel_rank=None):
+
+        super().__init__()
+
+        self._model_parallel_size = model_parallel_size
+        self._model_parallel_rank = model_parallel_rank
+        self._restore_path = None
+        self._app_state = None
+        self._model_name = model_name
+
+        if not os.path.exists(vocab_file):
+            raise ValueError(f'Vocab file not found at {vocab_file}')
+
+        # convert config to dictionary
+        if isinstance(config, DictConfig):
+            config = OmegaConf.to_container(config)
+        config["vocab_file"] = vocab_file
+        config['tokenizer_type'] = 'BertWordPieceLowerCase'
+        config['lazy_mpu_init'] = True
+        config['onnx_safe'] = True
+
+        # if 'model_parallel_size' in config:
+        if self._model_parallel_size is not None:
+            app_state = AppState()
+            self._app_state = app_state
+
+            # must be set for model parallel megatron-lm
+            os.environ["WORLD_SIZE"] = str(app_state.world_size)
+            os.environ["RANK"] = str(self._model_parallel_rank)
+
+            extra_args_provider = self._update_megatron_args(tensor_model_parallel_size=self._model_parallel_size)
+
+        else:
+            extra_args_provider = self._update_megatron_args()
+
+        # configure globals for megatron
+        set_pipeline_model_parallel_rank(0)  # pipeline model parallelism not implemented in NeMo
+        set_pipeline_model_parallel_world_size(1)  # pipeline model parallelism not implemented in NeMo
+
+        # Initialize part of Megatron global state that is needed for its constructor.
+        # We set 'lazy_mpu_init' flag on to make Megatron do only the initialization that does not depend
+        # on ddp be initialized yet (and we don't want Megatron to initialize DDP itself either)
+        # and to return a hook for us to call after PTL has torch.distributed initialized.
+        # (or if no PTL in case of inference - then we'll initialize torch.distributed)
+        # We call and clear this hook on first call to forward()
+        self._lazy_init_fn = initialize_megatron(
+            extra_args_provider=extra_args_provider, args_defaults=config, ignore_unknown_args=True
+        )
+
+        # read Megatron arguments back
+        args = get_args()
+        logging.info(f'Megatron-lm argparse args: {args}')
+
+        self.language_model, self._language_model_key = get_language_model(
+            attention_mask_func=bert_attention_mask_func, num_tokentypes=2, add_pooler=False
+        )
+
+        self.config = OmegaConf.create(config)
+        # key used for checkpoints
+        self._hidden_size = self.language_model.hidden_size
+
+    def _update_megatron_args(
+        self,
+        micro_batch_size=1,
+        tensor_model_parallel_size=1,
+        scaled_masked_softmax_fusion=False,
+        bias_gelu_fusion=False,
+        bias_dropout_fusion=False,
+    ):
+        def extra_args_provider(parser):
+            parser.set_defaults(micro_batch_size=micro_batch_size)
+            parser.set_defaults(tensor_model_parallel_size=tensor_model_parallel_size)
+            parser.set_defaults(scaled_masked_softmax_fusion=scaled_masked_softmax_fusion)
+            parser.set_defaults(bias_gelu_fusion=bias_gelu_fusion)
+            parser.set_defaults(bias_dropout_fusion=bias_dropout_fusion)
+
+            return parser
+
+        return extra_args_provider
+
+    def complete_lazy_init(self):
+        # finish megatron-lm initialization
+        if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
+            self._lazy_init_fn()
+            self._lazy_init_fn = None
+
+    @property
+    def hidden_size(self):
+        """
+        Property returning hidden size.
+
+        Returns:
+            Hidden size.
+        """
+        return self._hidden_size
+
+    @typecheck()
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        app_state = AppState()
+        if app_state.model_parallel_size is None:
+            self.complete_lazy_init()
+
+        extended_attention_mask = bert_extended_attention_mask(attention_mask)
+        position_ids = bert_position_ids(input_ids)
+
+        sequence_output = self.language_model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=extended_attention_mask,
+            tokentype_ids=token_type_ids,
+        )
+        return sequence_output
+
+    def _load_checkpoint(self, filename):
+        """Helper function for loading megatron checkpoints.
+
+        Args:
+            filename (str): Path to megatron checkpoint.
+        """
+        state_dict = torch.load(filename, map_location='cpu')
+        if 'checkpoint_version' in state_dict:
+            if state_dict['checkpoint_version'] is not None:
+                set_checkpoint_version(state_dict['checkpoint_version'])
+                logging.info(
+                    f"Megatron-lm checkpoint version found. Setting checkpoint_version to {state_dict['checkpoint_version']}."
+                )
+        else:
+            logging.warning('Megatron-lm checkpoint version not found. Setting checkpoint_version to 0.')
+            set_checkpoint_version(0)
+        # to load from Megatron pretrained checkpoint
+        if 'model' in state_dict:
+            self.language_model.load_state_dict(state_dict['model'][self._language_model_key])
+        else:
+            self.load_state_dict(state_dict)
+
+        logging.info(f"Checkpoint loaded from from {filename}")
+
+    def restore_weights(self, restore_path: str):
+        """Restores module/model's weights.
+           For model parallel checkpoints the directory structure
+           should be restore_path/mp_rank_0X/model_optim_rng.pt
+
+        Args:
+            restore_path (str): restore_path should a file or a directory if using model parallel
+        """
+        self._restore_path = restore_path
+
+        if os.path.isfile(restore_path):
+            self._load_checkpoint(restore_path)
+        elif os.path.isdir(restore_path):
+            # need model parallel groups to restore model parallel checkpoints
+            if model_parallel_is_initialized():
+                model_parallel_rank = torch.distributed.get_rank(group=get_model_parallel_group())
+                mp_restore_path = f'{restore_path}/mp_rank_{model_parallel_rank:02d}/model_optim_rng.pt'
+                self._load_checkpoint(mp_restore_path)
+            else:
+                logging.info(f'torch.distributed not initialized yet. Will not restore model parallel checkpoint')
+        else:
+            logging.error(f'restore_path: {restore_path} must be a file or directory.')

tts/nemo/collections/nlp/modules/common/megatron/megatron_utils.py

@@ -0,0 +1,307 @@
+# Copyright 2020 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 glob
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import wget
+from torch.hub import _get_torch_home
+
+from nemo.collections.nlp.modules.common.megatron.megatron_bert import MegatronBertEncoder
+from nemo.utils import AppState, logging
+
+__all__ = [
+    "get_megatron_lm_model",
+    "get_megatron_lm_models_list",
+    "get_megatron_checkpoint",
+    "is_lower_cased_megatron",
+    "get_megatron_tokenizer",
+]
+
+
+torch_home = _get_torch_home()
+
+if not isinstance(torch_home, str):
+    logging.info("Torch home not found, caching megatron in cwd")
+    torch_home = os.getcwd()
+
+MEGATRON_CACHE = os.path.join(torch_home, "megatron")
+
+
+CONFIGS = {"345m": {"hidden_size": 1024, "num_attention_heads": 16, "num_layers": 24, "max_position_embeddings": 512}}
+
+MEGATRON_CONFIG_MAP = {
+    "megatron-bert-345m-uncased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.0/files/release/mp_rank_00/model_optim_rng.pt",
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "megatron-bert-345m-cased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/files/release/mp_rank_00/model_optim_rng.pt",
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+    "megatron-bert-uncased": {
+        "config": None,
+        "checkpoint": None,
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "megatron-bert-cased": {
+        "config": None,
+        "checkpoint": None,
+        "vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+    "biomegatron-bert-345m-uncased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/MegatronBERT.pt",
+        "vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/vocab.txt",
+        "do_lower_case": True,
+        "tokenizer_name": "bert-large-uncased",
+    },
+    "biomegatron-bert-345m-cased": {
+        "config": CONFIGS["345m"],
+        "checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/MegatronBERT.pt",
+        "vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/vocab.txt",
+        "do_lower_case": False,
+        "tokenizer_name": "bert-large-cased",
+    },
+}
+
+
+def get_megatron_lm_model(
+    pretrained_model_name: str,
+    config_dict: Optional[dict] = None,
+    config_file: Optional[str] = None,
+    checkpoint_file: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+) -> Tuple[MegatronBertEncoder, str]:
+    """
+    Returns MegatronBertEncoder and a default or user specified path to the checkpoint file
+
+    Args:
+        pretrained_mode_name: model name from MEGATRON_CONFIG_MAP
+            for example: megatron-bert-cased
+        config_dict: model configuration parameters
+        config_file: path to model configuration file. Takes precedence over config_dict if both supplied.
+        checkpoint_file: path to checkpoint file or directory if using model parallel.
+        vocab_file: path to vocab file
+
+    Returns:
+        model: MegatronBertEncoder
+        checkpoint_file: path to checkpoint file or directory
+    """
+    config = None
+    # get default config and checkpoint
+    if config_file:
+        with open(config_file) as f:
+            config = json.load(f)
+            # replace dashes with underscores in config keys
+            fixed_config = {}
+            for key in config.keys():
+                fixed_key = key.replace("-", "_")
+                if fixed_key == 'max_seq_length':
+                    fixed_key = 'max_position_embeddings'
+                fixed_config[fixed_key] = config[key]
+            # 'vocab_size" no longer used.
+            if 'vocab_size' in fixed_config:
+                fixed_config.pop('vocab_size')
+            config = fixed_config
+    elif config_dict:
+        config = config_dict
+    elif pretrained_model_name in get_megatron_lm_models_list():
+        config = get_megatron_config(pretrained_model_name)
+    else:
+        raise ValueError(f"{pretrained_model_name} is not supported")
+
+    if config is None:
+        raise ValueError(f"config_file or config_dict is required for {pretrained_model_name}")
+
+    if not checkpoint_file:
+        checkpoint_file = get_megatron_checkpoint(pretrained_model_name)
+
+    if not vocab_file:
+        vocab_file = get_megatron_vocab_file(pretrained_model_name)
+
+    app_state = AppState()
+    if app_state.model_parallel_size is not None and app_state.model_parallel_rank is not None:
+        # model parallel already known from .nemo restore
+        model_parallel_size = app_state.model_parallel_size
+        model_parallel_rank = app_state.model_parallel_rank
+    elif os.path.isdir(checkpoint_file):
+        # starting training from megatron-lm checkpoint
+        mp_ranks = glob.glob(os.path.join(checkpoint_file, 'mp_rank*'))
+        model_parallel_size = len(mp_ranks)
+        app_state.model_parallel_size = model_parallel_size
+        logging.info(
+            (
+                f'restore_path: {checkpoint_file} is a directory. '
+                f'Assuming megatron model parallelism with '
+                f'model_parallel_size: {model_parallel_size}'
+            )
+        )
+        # try to get local rank from global
+        local_rank = None
+        try:
+            local_rank = int(os.environ['LOCAL_RANK'])
+        except:
+            logging.info('Global variable LOCAL_RANK not yet specified')
+        if local_rank is not None:
+            app_state.local_rank = local_rank
+        else:
+            # if local is None then we are on the main process
+            local_rank = 0
+        model_parallel_rank = compute_model_parallel_rank(local_rank, model_parallel_size)
+        app_state.model_parallel_rank = model_parallel_rank
+    else:
+        model_parallel_size = None
+        model_parallel_rank = None
+
+    model = MegatronBertEncoder(
+        model_name=pretrained_model_name,
+        config=config,
+        vocab_file=vocab_file,
+        model_parallel_size=model_parallel_size,
+        model_parallel_rank=model_parallel_rank,
+    )
+
+    return model, checkpoint_file
+
+
+def compute_model_parallel_rank(local_rank, model_parallel_size):
+    return local_rank % model_parallel_size
+
+
+def get_megatron_lm_models_list() -> List[str]:
+    """
+    Returns the list of supported Megatron-LM models
+    """
+    return list(MEGATRON_CONFIG_MAP.keys())
+
+
+def get_megatron_config(pretrained_model_name: str) -> Dict[str, int]:
+    """
+    Returns Megatron-LM model config file
+
+    Args:
+        pretrained_model_name (str): pretrained model name
+
+    Returns:
+        config (dict): contains model configuration: number of hidden layers, number of attention heads, etc
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["config"]
+
+
+def _check_megatron_name(pretrained_model_name: str) -> None:
+    megatron_model_list = get_megatron_lm_models_list()
+    if pretrained_model_name not in megatron_model_list:
+        raise ValueError(f'For Megatron-LM models, choose from the following list: {megatron_model_list}')
+
+
+def get_megatron_vocab_file(pretrained_model_name: str) -> str:
+    """
+    Gets vocabulary file from cache or downloads it
+
+    Args:
+        pretrained_model_name: pretrained model name
+
+    Returns:
+        path: path to the vocab file
+    """
+    _check_megatron_name(pretrained_model_name)
+    url = MEGATRON_CONFIG_MAP[pretrained_model_name]["vocab"]
+
+    path = os.path.join(MEGATRON_CACHE, pretrained_model_name + "_vocab")
+    path = _download(path, url)
+    return path
+
+
+def get_megatron_checkpoint(pretrained_model_name: str) -> str:
+    """
+    Gets checkpoint file from cache or downloads it
+    Args:
+        pretrained_model_name: pretrained model name
+    Returns:
+        path: path to model checkpoint
+    """
+    _check_megatron_name(pretrained_model_name)
+    url = MEGATRON_CONFIG_MAP[pretrained_model_name]["checkpoint"]
+    path = os.path.join(MEGATRON_CACHE, pretrained_model_name)
+    return _download(path, url)
+
+
+def _download(path: str, url: str):
+    """
+    Gets a file from cache or downloads it
+
+    Args:
+        path: path to the file in cache
+        url: url to the file
+    Returns:
+        path: path to the file in cache
+    """
+    if url is None:
+        return None
+
+    if not os.path.exists(path):
+        master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+        if not os.path.exists(path):
+            if master_device:
+                os.makedirs(MEGATRON_CACHE, exist_ok=True)
+                logging.info(f"Downloading from {url}")
+                wget.download(url, path)
+            # wait until the master process downloads the file and writes it to the cache dir
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+    return path
+
+
+def is_lower_cased_megatron(pretrained_model_name):
+    """
+    Returns if the megatron is cased or uncased
+
+    Args:
+        pretrained_model_name (str): pretrained model name
+    Returns:
+        do_lower_cased (bool): whether the model uses lower cased data
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["do_lower_case"]
+
+
+def get_megatron_tokenizer(pretrained_model_name: str):
+    """
+    Takes a pretrained_model_name for megatron such as "megatron-bert-cased" and returns the according 
+    tokenizer name for tokenizer instantiating.
+
+    Args:
+        pretrained_model_name: pretrained_model_name for megatron such as "megatron-bert-cased"
+    Returns: 
+        tokenizer name for tokenizer instantiating
+    """
+    _check_megatron_name(pretrained_model_name)
+    return MEGATRON_CONFIG_MAP[pretrained_model_name]["tokenizer_name"]

tts/nemo/collections/nlp/modules/common/sequence_classifier.py

@@ -0,0 +1,72 @@
+# 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
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, LogprobsType, NeuralType
+
+__all__ = ['SequenceClassifier']
+
+
+class SequenceClassifier(Classifier):
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'D'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'D'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+        idx_conditioned_on: int = 0,
+    ):
+        """
+        Initializes the SequenceClassifier module.
+        Args:
+            hidden_size: the hidden size of the mlp head on the top of the encoder
+            num_classes: number of the classes to predict
+            num_layers: number of the linear layers of the mlp head on the top of the encoder
+            activation: type of activations between layers of the mlp head
+            log_softmax: applies the log softmax on the output
+            dropout: the dropout used for the mlp head
+            use_transformer_init: initializes the weights with the same approach used in Transformer
+            idx_conditioned_on: index of the token to use as the sequence representation for the classification task, default is the first token
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.log_softmax = log_softmax
+        self._idx_conditioned_on = idx_conditioned_on
+        self.mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        hidden_states = self.dropout(hidden_states)
+        logits = self.mlp(hidden_states[:, self._idx_conditioned_on])
+        return logits

tts/nemo/collections/nlp/modules/common/sequence_regression.py

@@ -0,0 +1,69 @@
+# 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
+
+from torch import Tensor
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import NeuralType, RegressionValuesType
+
+__all__ = ['SequenceRegression']
+
+
+class SequenceRegression(Classifier):
+    """
+    Args:
+        hidden_size: the hidden size of the mlp head on the top of the encoder
+        num_layers: number of the linear layers of the mlp head on the top of the encoder
+        activation: type of activations between layers of the mlp head
+        dropout: the dropout used for the mlp head
+        use_transformer_init: initializes the weights with the same approach used in Transformer
+        idx_conditioned_on: index of the token to use as the sequence representation for the classification task,
+            default is the first token
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {"preds": NeuralType(tuple('B'), RegressionValuesType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+        idx_conditioned_on: int = 0,
+    ):
+        """ Initializes the SequenceRegression module. """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self._idx_conditioned_on = idx_conditioned_on
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes=1, num_layers=num_layers, activation=activation, log_softmax=False,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        """ Forward pass through the module.
+
+        Args:
+            hidden_states: hidden states for each token in a sequence, for example, BERT module output
+        """
+        hidden_states = self.dropout(hidden_states)
+        preds = self.mlp(hidden_states[:, self._idx_conditioned_on])
+        return preds.view(-1)

tts/nemo/collections/nlp/modules/common/sequence_token_classifier.py

@@ -0,0 +1,80 @@
+# 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
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, NeuralType
+
+__all__ = ['SequenceTokenClassifier']
+
+
+class SequenceTokenClassifier(Classifier):
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            "intent_logits": NeuralType(('B', 'D'), LogitsType()),
+            "slot_logits": NeuralType(('B', 'T', 'D'), LogitsType()),
+        }
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_intents: int,
+        num_slots: int,
+        num_layers: int = 2,
+        activation: str = 'relu',
+        log_softmax: bool = False,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ):
+        """
+        Initializes the SequenceTokenClassifier module, could be used for tasks that train sequence and
+        token classifiers jointly, for example, for intent detection and slot tagging task.
+        Args:
+            hidden_size: hidden size of the mlp head on the top of the encoder
+            num_intents: number of the intents to predict
+            num_slots: number of the slots to predict
+            num_layers: number of the linear layers of the mlp head on the top of the encoder
+            activation: type of activations between layers of the mlp head
+            log_softmax: applies the log softmax on the output
+            dropout: the dropout used for the mlp head
+            use_transformer_init: initializes the weights with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.intent_mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_intents,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.slot_mlp = MultiLayerPerceptron(
+            hidden_size=hidden_size,
+            num_classes=num_slots,
+            num_layers=num_layers,
+            activation=activation,
+            log_softmax=log_softmax,
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        hidden_states = self.dropout(hidden_states)
+        # intent is classified by first hidden position
+        intent_logits = self.intent_mlp(hidden_states[:, 0])
+        slot_logits = self.slot_mlp(hidden_states)
+        return intent_logits, slot_logits

tts/nemo/collections/nlp/modules/common/token_classifier.py

@@ -0,0 +1,165 @@
+# 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 Dict, Optional
+
+from omegaconf.omegaconf import MISSING
+from torch import nn as nn
+
+from nemo.collections.common.parts import MultiLayerPerceptron
+from nemo.collections.nlp.modules.common.classifier import Classifier
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import LogitsType, LogprobsType, NeuralType
+
+__all__ = ['BertPretrainingTokenClassifier', 'TokenClassifier']
+
+ACT2FN = {"gelu": nn.functional.gelu, "relu": nn.functional.relu}
+
+
+@dataclass
+class TokenClassifierConfig:
+    num_layers: int = 1
+    activation: str = 'relu'
+    log_softmax: bool = True
+    dropout: float = 0.0
+    use_transformer_init: bool = True
+
+
+class TokenClassifier(Classifier):
+    """
+    A module to perform token level classification tasks such as Named entity recognition.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module output ports.
+        """
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'T', 'C'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'T', 'C'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 1,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ) -> None:
+
+        """
+        Initializes the Token Classifier module.
+
+        Args:
+            hidden_size: the size of the hidden dimension
+            num_classes: number of classes
+            num_layers: number of fully connected layers in the multilayer perceptron (MLP)
+            activation: activation to usee between fully connected layers in the MLP
+            log_softmax: whether to apply softmax to the output of the MLP
+            dropout: dropout to apply to the input hidden states
+            use_transformer_init: whether to initialize the weights of the classifier head with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+        self.log_softmax = log_softmax
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes, num_layers=num_layers, activation=activation, log_softmax=log_softmax
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        """
+        Performs the forward step of the module.
+        Args:
+            hidden_states: batch of hidden states (for example, from the BERT encoder module)
+                [BATCH_SIZE x SEQ_LENGTH x HIDDEN_SIZE]
+        Returns: logits value for each class [BATCH_SIZE x SEQ_LENGTH x NUM_CLASSES]
+        """
+        hidden_states = self.dropout(hidden_states)
+        logits = self.mlp(hidden_states)
+        return logits
+
+
+class BertPretrainingTokenClassifier(Classifier):
+    """
+    A module to perform token level classification tasks for Bert pretraining.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns definitions of module output ports.
+        """
+        if not self.log_softmax:
+            return {"logits": NeuralType(('B', 'T', 'C'), LogitsType())}
+        else:
+            return {"log_probs": NeuralType(('B', 'T', 'C'), LogprobsType())}
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_classes: int,
+        num_layers: int = 1,
+        activation: str = 'relu',
+        log_softmax: bool = True,
+        dropout: float = 0.0,
+        use_transformer_init: bool = True,
+    ) -> None:
+
+        """
+        Initializes the Token Classifier module.
+
+        Args:
+            hidden_size: the size of the hidden dimension
+            num_classes: number of classes
+            num_layers: number of fully connected layers in the multilayer perceptron (MLP)
+            activation: activation to usee between fully connected layers in the MLP
+            log_softmax: whether to apply softmax to the output of the MLP
+            dropout: dropout to apply to the input hidden states
+            use_transformer_init: whether to initialize the weights of the classifier head with the same approach used in Transformer
+        """
+        super().__init__(hidden_size=hidden_size, dropout=dropout)
+
+        self.log_softmax = log_softmax
+
+        if activation not in ACT2FN:
+            raise ValueError(f'activation "{activation}" not found')
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.act = ACT2FN[activation]
+        self.norm = nn.LayerNorm(hidden_size, eps=1e-12)
+        self.mlp = MultiLayerPerceptron(
+            hidden_size, num_classes, num_layers=num_layers, activation=activation, log_softmax=log_softmax
+        )
+        self.post_init(use_transformer_init=use_transformer_init)
+
+    @typecheck()
+    def forward(self, hidden_states):
+        """
+        Performs the forward step of the module.
+        Args:
+            hidden_states: batch of hidden states (for example, from the BERT encoder module)
+                [BATCH_SIZE x SEQ_LENGTH x HIDDEN_SIZE]
+        Returns: logits value for each class [BATCH_SIZE x SEQ_LENGTH x NUM_CLASSES]
+        """
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.act(hidden_states)
+        transform = self.norm(hidden_states)
+        logits = self.mlp(transform)
+        return logits

tts/nemo/collections/nlp/modules/common/tokenizer_utils.py

@@ -0,0 +1,144 @@
+# 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.path
+from dataclasses import MISSING, dataclass
+from os import path
+from typing import Dict, List, Optional
+
+import nemo
+from nemo.collections.common.tokenizers.char_tokenizer import CharTokenizer
+from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+from nemo.collections.common.tokenizers.word_tokenizer import WordTokenizer
+from nemo.collections.common.tokenizers.youtokentome_tokenizer import YouTokenToMeTokenizer
+from nemo.collections.nlp.modules.common.huggingface.huggingface_utils import get_huggingface_pretrained_lm_models_list
+from nemo.collections.nlp.modules.common.lm_utils import get_pretrained_lm_models_list
+from nemo.collections.nlp.modules.common.megatron.megatron_utils import get_megatron_tokenizer
+from nemo.utils import logging
+
+__all__ = ['get_tokenizer', 'get_tokenizer_list']
+
+
+def get_tokenizer_list() -> List[str]:
+    """
+    Returns all all supported tokenizer names
+    """
+    s = set(get_pretrained_lm_models_list())
+    s.update(set(get_huggingface_pretrained_lm_models_list(include_external=True)))
+    return ["sentencepiece", "char", "word"] + list(s)
+
+
+@dataclass
+class TokenizerConfig:
+    library: str = MISSING
+    tokenizer_model: Optional[str] = None
+    vocab_size: Optional[int] = None
+    vocab_file: Optional[str] = None
+    special_tokens: Optional[Dict[str, str]] = None
+    bpe_dropout: Optional[float] = 0.0
+    coverage: Optional[float] = 0.999
+    training_sample_size: Optional[int] = None
+
+
+def get_tokenizer(
+    tokenizer_name: str,
+    tokenizer_model: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+    special_tokens: Optional[Dict[str, str]] = None,
+    use_fast: Optional[bool] = False,
+    bpe_dropout: Optional[float] = 0.0,
+):
+    """
+    Args:
+        tokenizer_name: sentencepiece or pretrained model from the hugging face list,
+            for example: bert-base-cased
+            To see the list of all HuggingFace pretrained models, use: nemo_nlp.modules.common.get_huggingface_pretrained_lm_models_list()
+        tokenizer_model: tokenizer model file of sentencepiece or youtokentome
+        special_tokens: dict of special tokens
+        vocab_file: path to vocab file
+        use_fast: (only for HuggingFace AutoTokenizer) set to True to use fast HuggingFace tokenizer
+        bpe_dropout: (only supported by YTTM tokenizer) BPE dropout tries to corrupt the standard segmentation procedure of BPE to help
+            model better learn word compositionality and become robust to segmentation errors. 
+            It has emperically been shown to improve inference time BLEU scores.
+    """
+    if special_tokens is None:
+        special_tokens_dict = {}
+    else:
+        special_tokens_dict = special_tokens
+
+    if 'megatron' in tokenizer_name:
+        if vocab_file is None:
+            vocab_file = nemo.collections.nlp.modules.common.megatron.megatron_utils.get_megatron_vocab_file(
+                tokenizer_name
+            )
+        tokenizer_name = get_megatron_tokenizer(tokenizer_name)
+
+    if tokenizer_name == 'sentencepiece':
+        return nemo.collections.common.tokenizers.sentencepiece_tokenizer.SentencePieceTokenizer(
+            model_path=tokenizer_model, special_tokens=special_tokens, legacy=True
+        )
+    elif tokenizer_name == 'yttm':
+        return YouTokenToMeTokenizer(model_path=tokenizer_model, bpe_dropout=bpe_dropout)
+    elif tokenizer_name == 'word':
+        return WordTokenizer(vocab_file=vocab_file, **special_tokens_dict)
+    elif tokenizer_name == 'char':
+        return CharTokenizer(vocab_file=vocab_file, **special_tokens_dict)
+
+    return AutoTokenizer(
+        pretrained_model_name=tokenizer_name, vocab_file=vocab_file, **special_tokens_dict, use_fast=use_fast
+    )
+
+
+def get_nmt_tokenizer(
+    library: str = 'yttm',
+    model_name: Optional[str] = None,
+    tokenizer_model: Optional[str] = None,
+    vocab_file: Optional[str] = None,
+    special_tokens: Optional[Dict[str, str]] = None,
+    use_fast: Optional[bool] = False,
+    bpe_dropout: Optional[float] = 0.0,
+):
+    """
+    Args:
+        model_name: if using a pretrained model from NeMo or HuggingFace
+        tokenizer_model: tokenizer model file of sentencepiece or youtokentome
+        special_tokens: dict of special tokens
+        vocab_file: path to vocab file
+        use_fast: (only for HuggingFace AutoTokenizer) set to True to use fast HuggingFace tokenizer
+        bpe_dropout: (only supported by YTTM tokenizer) BPE dropout tries to corrupt the standard segmentation procedure of BPE to help
+            model better learn word compositionality and become robust to segmentation errors. 
+            It has emperically been shown to improve inference time BLEU scores.        
+    """
+    if special_tokens is None:
+        special_tokens_dict = {}
+    else:
+        special_tokens_dict = special_tokens
+
+    if library == 'yttm':
+        logging.info(f'Getting YouTokenToMeTokenizer with model: {tokenizer_model}.')
+        return YouTokenToMeTokenizer(model_path=tokenizer_model, bpe_dropout=bpe_dropout)
+    elif library == 'huggingface':
+        logging.info(f'Getting HuggingFace AutoTokenizer with pretrained_model_name: {model_name}')
+        return AutoTokenizer(
+            pretrained_model_name=model_name, vocab_file=vocab_file, **special_tokens_dict, use_fast=use_fast
+        )
+    elif library == 'sentencepiece':
+        logging.info(f'Getting SentencePiece with model: {model_name}')
+        return nemo.collections.common.tokenizers.sentencepiece_tokenizer.SentencePieceTokenizer(
+            model_path=tokenizer_model, special_tokens=special_tokens_dict
+        )
+    else:
+        raise NotImplementedError(
+            'Currently we only support "yttm", "huggingface", and "sentencepiece" tokenizer library.'
+        )