artcraft/models/tts/NeMo/tutorials/nlp/Punctuation_and_Capitalization.ipynb

In [ ]:

BRANCH = 'v1.0.2'

In [ ]:

"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.

Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""
# If you're using Google Colab and not running locally, run this cell

# install NeMo
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]

In [ ]:

# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:
# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'

! pip install ipywidgets
! jupyter nbextension enable --py widgetsnbextension

# Please restart the kernel after running this cell

In [ ]:

from nemo.utils.exp_manager import exp_manager
from nemo.collections import nlp as nemo_nlp

import os
import wget 
import torch
import pytorch_lightning as pl
from omegaconf import OmegaConf

Task Description

Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words. This tutorial explains how to implement a model in NeMo that will predict punctuation and capitalization for each word in a sentence to make ASR output more readable and to boost performance of the named entity recognition, machine translation or text-to-speech models. We'll show how to train a model for this task using a pre-trained BERT model. For every word in our training dataset we’re going to predict:

• punctuation mark that should follow the word and
• whether the word should be capitalized

Dataset

This model can work with any dataset as long as it follows the format specified below. 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: [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.)

In this tutorial, we are considering only commas, periods, and question marks the rest punctuation marks were removed. To use more punctuation marks, update the dataset to include desired labels, no changes to the model needed.

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.

Download and preprocess the data¶

In this notebook we are going to use a subset of English examples from the Tatoeba collection of sentences this script will download and preprocess the Tatoeba data NeMo/examples/nlp/token_classification/get_tatoeba_data.py. Note, for further experiments with the model, set NUM_SAMPLES=-1 and consider including other datasets to improve model performance.

In [ ]:

DATA_DIR = "PATH_TO_DATA"
WORK_DIR = "PATH_TO_CHECKPOINTS_AND_LOGS"
MODEL_CONFIG = "punctuation_capitalization_config.yaml"

# model parameters
BATCH_SIZE = 128
MAX_SEQ_LENGTH = 64
LEARNING_RATE = 0.00002
NUM_SAMPLES = 10000

In [ ]:

## download get_tatoeba_data.py script to download and preprocess the Tatoeba data
os.makedirs(WORK_DIR, exist_ok=True)
if not os.path.exists(WORK_DIR + '/get_tatoeba_data.py'):
    print('Downloading get_tatoeba_data.py...')
    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/data/get_tatoeba_data.py', WORK_DIR)
else:
    print ('get_tatoeba_data.py is already exists')

In [ ]:

# download and preprocess the data
# --clean_dir flag deletes raw Tataoeba data, remove the flag to avoid multiple data downloads if you want to experiment with the dataset size
! python $WORK_DIR/get_tatoeba_data.py --data_dir $DATA_DIR --num_sample $NUM_SAMPLES --clean_dir

after execution of the above cell, your data folder will contain the following 4 files needed for training (raw Tatoeba data could be present if --clean_dir was not used):

• labels_dev.txt
• labels_train.txt
• text_dev.txt
• text_train.txt

In [ ]:

! ls -l {DATA_DIR}

In [ ]:

# let's take a look at the data 
print('Text:')
! head -n 5 {DATA_DIR}/text_train.txt

print('\nLabels:')
! head -n 5 {DATA_DIR}/labels_train.txt

Model Configuration

In the Punctuation and Capitalization Model, we are jointly training two token-level classifiers on top of the pretrained BERT model:

• one classifier to predict punctuation and
• the other one - capitalization.

The model is defined in a config file which declares multiple important sections. They are:

• model: All arguments that are related to the Model - language model, token classifiers, optimizer and schedulers, dataset and any other related information

• trainer: Any argument to be passed to PyTorch Lightning

In [ ]:

# download the model's configuration file 
config_dir = WORK_DIR + '/configs/'
os.makedirs(config_dir, exist_ok=True)
if not os.path.exists(config_dir + MODEL_CONFIG):
    print('Downloading config file...')
    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/conf/' + MODEL_CONFIG, config_dir)
else:
    print ('config file is already exists')

In [ ]:

# this line will print the entire config of the model
config_path = f'{WORK_DIR}/configs/{MODEL_CONFIG}'
print(config_path)
config = OmegaConf.load(config_path)
print(OmegaConf.to_yaml(config))

Setting up Data within the config

Among other things, the config file contains dictionaries called dataset, train_ds and validation_ds. These are configurations used to setup the Dataset and DataLoaders of the corresponding config.

If both training and evaluation files are located in the same directory, simply specify model.dataset.data_dir, like we are going to do below. However, if your evaluation files are located in a different directory, or you want to use multiple datasets for evaluation, specify paths to the directory(ies) with evaluation file(s) in the following way:

model.validation_ds.ds_item=[PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets).

Also notice that some configs, including model.dataset.data_dir, have ??? in place of paths, this values are required to be specified by the user.

Let's now add the data directory path to the config.

In [ ]:

# in this tutorial train and dev data is located in the same folder, so it is enought to add the path of the data directory to our config
config.model.dataset.data_dir = DATA_DIR

Building the PyTorch Lightning Trainer

NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!

Let's first instantiate a Trainer object!

In [ ]:

print("Trainer config - \n")
print(OmegaConf.to_yaml(config.trainer))

In [ ]:

# lets modify some trainer configs
# checks if we have GPU available and uses it
cuda = 1 if torch.cuda.is_available() else 0
config.trainer.gpus = cuda
config.trainer.precision = 16 if torch.cuda.is_available() else 32

# For mixed precision training, use precision=16 and amp_level=O1

# Reduces maximum number of epochs to 1 for a quick training
config.trainer.max_epochs = 1

# Remove distributed training flags
config.trainer.accelerator = None

trainer = pl.Trainer(**config.trainer)

Setting up a NeMo Experiment¶

NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it!

In [ ]:

exp_dir = exp_manager(trainer, config.get("exp_manager", None))

# the exp_dir provides a path to the current experiment for easy access
exp_dir = str(exp_dir)
exp_dir

Model Training

Before initializing the model, we might want to modify some of the model configs. For example, we might want to modify the pretrained BERT model.

In [ ]:

# complete list of supported BERT-like models
nemo_nlp.modules.get_pretrained_lm_models_list()

PRETRAINED_BERT_MODEL = "bert-base-uncased"

In [ ]:

# add the specified above model parameters to the config
config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL
config.model.train_ds.batch_size = BATCH_SIZE
config.model.validation_ds.batch_size = BATCH_SIZE
config.model.optim.lr = LEARNING_RATE
config.model.train_ds.num_samples = NUM_SAMPLES
config.model.validation_ds.num_samples = NUM_SAMPLES

Now, we are ready to initialize our model. During the model initialization call, the dataset and data loaders we'll be prepared for training and evaluation. Also, the pretrained BERT model will be downloaded, note it can take up to a few minutes depending on the size of the chosen BERT model.

In [ ]:

# initialize the model
# during this stage, the dataset and data loaders we'll be prepared for training and evaluation
config.trainer.max_epochs = 3
model = nemo_nlp.models.PunctuationCapitalizationModel(cfg=config.model, trainer=trainer)

Monitoring training progress

Optionally, you can create a Tensorboard visualization to monitor training progress.

In [ ]:

try:
  from google import colab
  COLAB_ENV = True
except (ImportError, ModuleNotFoundError):
  COLAB_ENV = False

# Load the TensorBoard notebook extension
if COLAB_ENV:
  %load_ext tensorboard
  %tensorboard --logdir {exp_dir}
else:
  print("To use tensorboard, please use this notebook in a Google Colab environment.")

In [ ]:

# start the training
trainer.fit(model)

Inference

To see how the model performs, let’s run inference on a few examples.

In [ ]:

# define the list of queiries for inference
queries = [
        'we bought four shirts and one mug from the nvidia gear store in santa clara',
        'what can i do for you today',
        'how are you',
        'how is the weather in',
    ]
inference_results = model.add_punctuation_capitalization(queries)
print()

for query, result in zip(queries, inference_results):
    print(f'Query   : {query}')
    print(f'Combined: {result.strip()}\n')

Training Script

If you have NeMo installed locally, you can also train the model with nlp/token_classification/punctuation_capitalization_train.py.

To run training script, use:

python punctuation_and_capitalization_train.py model.dataset.data_dir=PATH_TO_DATA_DIR

Set NUM_SAMPLES=-1 and consider including other datasets to improve the performance of the model.

Finetuning model with your data

When we were training the model from scratch, the datasets were prepared for training during the model initialization. When we are using a pretrained Punctuation and Capitalization model, before training, we need to setup training and evaluation data.

In [ ]:

# let's reload our pretrained model
pretrained_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained('punctuation_en_distilbert')

# setup train and validation Pytorch DataLoaders
pretrained_model.update_data_dir(DATA_DIR)
pretrained_model.setup_training_data()
pretrained_model.setup_validation_data()

# and now we can create a PyTorch Lightning trainer and call `fit` again
# for this tutorial we are setting fast_dev_run to True, and the trainer will run 1 training batch and 1 validation batch
# for actual model training, disable the flag
fast_dev_run = True
trainer = pl.Trainer(gpus=1, fast_dev_run=fast_dev_run)
trainer.fit(pretrained_model)

In [ ]: