Introduce code from Rust monorepo. This is a mess

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