Introduce code from Rust monorepo. This is a mess

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

@@ -0,0 +1,133 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Punctuation and capitalization model with pretrained BERT-like models
+
+pretrained_model: null # pretrained Punctuation and Capitalization model from list_available_models(), for example:
+# punctuation_en_bert or punctuation_en_distilbert
+# or your_model.nemo
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 3
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: false  # Provided by exp_manager
+  logger: false  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: Punctuation_and_Capitalization  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manger to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a model checkpoint callback
+
+model:
+  punct_label_ids: null # will be filled during training
+  capit_label_ids: null # will be filled during training
+  class_labels:
+    punct_labels_file: punct_label_ids.csv
+    capit_labels_file: capit_label_ids.csv
+  dataset:
+    data_dir: ??? # /path/to/training/data
+    # parameters for dataset preprocessing
+    max_seq_length: 128
+    pad_label: 'O'
+    ignore_extra_tokens: false
+    ignore_start_end: false
+    use_cache: true
+    # shared among dataloaders
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    text_file: text_train.txt
+    labels_file: labels_train.txt
+    shuffle: true
+    num_samples: -1
+    batch_size: 64
+
+  validation_ds:
+    # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for
+    # evaluation is needed, specify ds_item, otherwise by default data_dir is used
+    ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets)
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  test_ds:
+    # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for
+    # evaluation is needed, specify ds_item, otherwise by default data_dir is used
+    ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets)
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+  punct_head:
+    punct_num_fc_layers: 1
+    fc_dropout: 0.1
+    activation: 'relu'
+    use_transformer_init: True
+
+  capit_head:
+    capit_num_fc_layers: 1
+    fc_dropout: 0.1
+    activation: 'relu'
+    use_transformer_init: true
+
+  optim:
+    name: adam
+    lr: 1e-4
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

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

@@ -0,0 +1,119 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Token Classification tasks (for example, Named Entity Recognition) with pretrained BERT-like models
+
+pretrained_model: null # pretrained TokenClassification model from list_available_models() or path to a .nemo file,
+# for example: ner_en_bert or your_model.nemo
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU
+  num_nodes: 1
+  max_epochs: 5
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  gradient_clip_val: 0.0
+  amp_level: O0 # O1/O2 for mixed precision
+  precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
+  accelerator: ddp
+  checkpoint_callback: False  # Provided by exp_manager
+  logger: False  # Provided by exp_manager
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 1.0  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+
+exp_manager:
+  exp_dir: null  # exp_dir for your experiment, if None, defaults to "./nemo_experiments"
+  name: token_classification_model  # The name of your model
+  create_tensorboard_logger: true  # Whether you want exp_manager to create a tb logger
+  create_checkpoint_callback: true  # Whether you want exp_manager to create a model checkpoint callback
+
+model:
+  label_ids: null # will be filled during training
+  class_labels:
+    class_labels_file: label_ids.csv # will be generated during training and saved in .nemo file
+  dataset:
+    data_dir: ??? # /path/to/data
+    class_balancing: null # choose from [null, weighted_loss]. Weighted_loss enables the weighted class balancing of the loss, may be used for handling unbalanced classes
+    max_seq_length: 128
+    pad_label: 'O'
+    ignore_extra_tokens: false
+    ignore_start_end: false
+    use_cache: true
+    # shared among dataloaders
+    num_workers:  2
+    pin_memory: false
+    drop_last: false
+
+  train_ds:
+    text_file: text_train.txt
+    labels_file: labels_train.txt
+    shuffle: true
+    num_samples: -1
+    batch_size: 64
+
+  validation_ds:
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  test_ds:
+    text_file: text_dev.txt
+    labels_file: labels_dev.txt
+    shuffle: false
+    num_samples: -1
+    batch_size: 64
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # only used if tokenizer is sentencepiece
+    special_tokens: null
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    lm_checkpoint: null
+    config_file: null # json file, precedence over config
+    config: null
+
+
+  head:
+    num_fc_layers: 2
+    fc_dropout: 0.5
+    activation: 'relu'
+    use_transformer_init: True
+
+  optim:
+    name: adam
+    lr: 5e-5
+    weight_decay: 0.00
+
+    sched:
+      name: WarmupAnnealing
+      # Scheduler params
+      warmup_steps: null
+      warmup_ratio: 0.1
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

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

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

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

@@ -0,0 +1,124 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from nemo.utils import logging
+
+
+def __convert_data(in_file: str, out_text_f: str, out_labels_f: str, max_length: int):
+    """
+    Convert data from the IOB format to NeMo accepted format described below.
+    in_file should be in the IOB format, see example here:
+    https://www.clips.uantwerpen.be/conll2003/ner/.
+
+    Args:
+        in_file: input file name
+        out_text_f: output file with text
+        out_labels_f: output file with labels
+        max_length: use -1 to leave the examples' length as is, otherwise long examples will be split into multiple
+            examples
+    After the conversion, the dataset is split into 2 files: text.txt
+    and labels.txt.
+    Each line of the text.txt file contains text sequences, where words
+    are separated with spaces. The labels.txt file contains corresponding
+    labels for each word in text.txt, the labels are separated with spaces.
+    Each line of the files should follow the format:
+    [WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and
+    [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).
+
+    """
+    in_file = open(in_file, 'r')
+
+    if max_length == -1:
+        with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
+            for line in in_file:
+                if line == '\n':
+                    out_text.write(line)
+                    out_labels.write(line)
+                else:
+                    line = line.split()
+                    out_text.write(line[0] + ' ')
+                    out_labels.write(line[-1] + ' ')
+
+    else:
+        words = []
+        labels = []
+        with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels:
+            lines = in_file.readlines()
+            for line_id, line in enumerate(lines):
+                logging.info(f"{line_id} {len(lines)}")
+                contends = line.strip()
+                if len(contends) == 0:
+                    assert len(words) == len(labels)
+                    if len(words) > max_length:
+                        # split if the sentence is longer than max_length
+                        while len(words) > max_length:
+                            tmplabel = labels[:max_length]
+                            for iidx in range(len(tmplabel)):
+                                if tmplabel.pop() == 'O':
+                                    break
+                            l = ' '.join([label for label in labels[: len(tmplabel) + 1] if len(label) > 0])
+                            w = ' '.join([word for word in words[: len(tmplabel) + 1] if len(word) > 0])
+
+                            out_text.write(w + "\n")
+                            out_labels.write(l + "\n")
+                            words = words[len(tmplabel) + 1 :]
+                            labels = labels[len(tmplabel) + 1 :]
+
+                    if len(words) == 0:
+                        continue
+                    l = ' '.join([label for label in labels if len(label) > 0])
+                    w = ' '.join([word for word in words if len(word) > 0])
+
+                    out_text.write(w + "\n")
+                    out_labels.write(l + "\n")
+                    words = []
+                    labels = []
+                    continue
+
+                word = line.strip().split()[0]
+                label = line.strip().split()[-1]
+                words.append(word)
+                labels.append(label)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Convert data from IOB format to the format compatible with \
+        nlp/examples/token_classification/scripts/token_classification_train.py and \
+        token_classification_evaluate.py'
+    )
+    parser.add_argument("--data_file", required=True, type=str, help='path to a file in IOB format')
+    parser.add_argument(
+        "--max_length",
+        default=-1,
+        type=int,
+        help='use -1 to leave the examples\'s length as is, '
+        'otherwise long examples will be split into multiple examples',
+    )
+    args = parser.parse_args()
+
+    data_dir, basename = os.path.split(args.data_file)
+    prefix = os.path.splitext(basename)[0]
+    if not os.path.exists(args.data_file):
+        raise FileNotFoundError(f"{args.data_file} not found")
+
+    logging.info(f'Processing {args.data_file}')
+    out_text = os.path.join(data_dir, 'text_' + prefix + '.txt')
+    out_labels = os.path.join(data_dir, 'labels_' + prefix + '.txt')
+
+    __convert_data(args.data_file, out_text, out_labels, args.max_length)
+    logging.info(f'Processing of the {args.data_file} is complete')

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

@@ -0,0 +1,99 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+The script converts raw text to the NeMo format for punctuation and capitalization task.
+
+Raw Data Format
+---------------
+
+The Punctuation and Capitalization model can work with any text dataset, although it is recommended to balance the data, especially for the punctuation task.
+Before pre-processing the data to the format expected by the model, the data should be split into train.txt and dev.txt (and optionally test.txt).
+Each line in the **train.txt/dev.txt/test.txt** should represent one or more full and/or truncated sentences.
+
+Example of the train.txt/dev.txt file:
+    When is the next flight to New York?
+    The next flight is ...
+    ....
+
+
+The `source_data_dir` structure should look like this:
+   .
+   |--sourced_data_dir
+     |-- dev.txt
+     |-- train.txt
+
+
+
+NeMo Data Format for training the model
+---------------------------------------
+
+The punctuation and capitalization model expects the data in the following format:
+
+The training and evaluation data is divided into 2 files: text.txt and labels.txt. \
+Each line of the **text.txt** file contains text sequences, where words are separated with spaces, i.e.
+
+[WORD] [SPACE] [WORD] [SPACE] [WORD], for example:
+        when is the next flight to new york
+        the next flight is ...
+        ...
+
+The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. \
+Each label in labels.txt file consists of 2 symbols:
+
+* the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed);
+* the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.)
+
+By default the following punctuation marks are considered: commas, periods, and question marks; the rest punctuation marks were removed from the data.
+This can be changed by introducing new labels in the labels.txt files
+
+Each line of the labels.txt should follow the format: [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). \
+For example, labels for the above text.txt file should be:
+
+        OU OO OO OO OO OO OU ?U
+        OU OO OO OO ...
+        ...
+
+The complete list of all possible labels for this task used in this tutorial is: OO, ,O, .O, ?O, OU, ,U, .U, ?U.
+
+Converting Raw data to NeMo format
+----------------------------------
+
+To pre-process the raw text data, stored under :code:`sourced_data_dir` (see the :ref:`raw_data_format_punct`
+section), run the following command:
+
+    python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \
+           -s <PATH_TO_THE_SOURCE_FILE>
+           -o <PATH_TO_THE_OUTPUT_DIRECTORY>
+
+"""
+
+import argparse
+import os
+
+from get_tatoeba_data import create_text_and_labels
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Prepare data for punctuation and capitalization tasks')
+    parser.add_argument("-s", "--source_file", required=True, type=str, help="Path to the source file")
+    parser.add_argument("-o", "--output_dir", required=True, type=str, help="Path to the output directory")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.source_file):
+        raise ValueError(f'{args.source_file} was not found')
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    create_text_and_labels(args.output_dir, args.source_file)
+
+    print(f'Processing of the {args.source_file} is complete')

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

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

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

@@ -0,0 +1,101 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import PunctuationCapitalizationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+"""
+This script show how to train a Punctuation and Capitalization Model.
+More details on the task and data format could be found in tutorials/nlp/Punctuation_and_Capitalization.ipynb
+
+*** Setting the configs ***
+
+The model and the PT trainer are defined in a config file which declares multiple important sections.
+The most important ones are:
+    model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer,
+            schedulers, and datasets/data loaders.
+    trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs,
+            precision level, etc.
+This script uses the `/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml` config file
+by default. You may update the config file from the file directly. 
+The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb
+
+*** Model training ***
+
+To run this script and train the model from scratch, use:
+    python punctuation_and_capitalization_train.py \
+           model.dataset.data_dir=<PATH_TO_DATA_DIR>
+
+To use one of the pretrained versions of the model and finetune it, run:
+    python punctuation_and_capitalization.py \
+    pretrained_model=punctuation_en_bert \
+    model.dataset.data_dir=<PATH_TO_DATA_DIR>
+    
+    <PATH_TO_DATA_DIR> - a directory that contains test_ds.text_file and test_ds.labels_file (see the config)
+    pretrained_model   - pretrained PunctuationCapitalization model from list_available_models() or 
+                     path to a .nemo file, for example: punctuation_en_bert or model.nemo
+
+"""
+
+
+@hydra_runner(config_path="conf", config_name="punctuation_capitalization_config")
+def main(cfg: DictConfig) -> None:
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+
+    if not cfg.pretrained_model:
+        logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+        model = PunctuationCapitalizationModel(cfg.model, trainer=trainer)
+    else:
+        if os.path.exists(cfg.pretrained_model):
+            model = PunctuationCapitalizationModel.restore_from(cfg.pretrained_model)
+        elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names():
+            model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model)
+        else:
+            raise ValueError(
+                f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}'
+            )
+
+        data_dir = cfg.model.dataset.get('data_dir', None)
+        if data_dir:
+            if not os.path.exists(data_dir):
+                raise ValueError(f'{data_dir} is not found at')
+
+            # we can also do finetuning of the pretrained model but we would need to update the data dir
+            model.update_data_dir(data_dir)
+            # setup train and validation Pytorch DataLoaders
+            model.setup_training_data()
+            model.setup_validation_data()
+            logging.info(f'Using config file of the pretrained model')
+        else:
+            raise ValueError(
+                'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \
+                with "model.dataset.data_dir" argument'
+            )
+
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    main()

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

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

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

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