Punctuation and Capitalization Model

Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words. There are two issues with non-punctuated ASR output:

  • it could be difficult to read and understand

  • models for some downstream tasks, such as named entity recognition, machine translation, or text-to-speech, are usually trained

on punctuated datasets and using raw ASR output as the input to these models could deteriorate their performance

Quick Start Guide

from nemo.collections.nlp.models import PunctuationCapitalizationModel

# to get the list of pre-trained models

# Download and load the pre-trained BERT-based model
model = PunctuationCapitalizationModel.from_pretrained("punctuation_en_bert")

# try the model on a few examples
model.add_punctuation_capitalization(['how are you', 'great how about you'])

Model Description

For each word in the input text, the Punctuation and Capitalization model:

  • predicts a punctuation mark that should follow the word (if any). By default, the model supports commas, periods, and question marks.

  • predicts if the word should be capitalized or not

In the Punctuation and Capitalization model, we are jointly training two token-level classifiers on top of a pre-trained language model, such as BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding [NLP-PUNCT1].


We recommend you try this model in a Jupyter notebook (run on Google’s Colab.): NeMo/tutorials/nlp/Punctuation_and_Capitalization.ipynb.

Connect to an instance with a GPU (Runtime -> Change runtime type -> select GPU for the hardware accelerator).

An example script on how to train the model can be found at: NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py.

An example script on how to run evaluation and inference can be found at: NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py.

The default configuration file for the model can be found at: NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml.

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 similar to the following:

  |-- dev.txt
  |-- train.txt

NeMo Data Format

The Punctuation and Capitalization model expects the data in the following format:

The training and evaluation data is divided into 2 files: - text.txt - 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)

By default, the following punctuation marks are considered: commas, periods, and question marks; the remaining 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:


The complete list of all possible labels used in this tutorial are:

  • OO

  • O

  • .O

  • ?O

  • OU

  • <blank space>

  • U

  • .U

  • ?U

Converting Raw Data to NeMo Format

To pre-process the raw text data, stored under sourced_data_dir (see the Raw Data Format section), run the following command:

python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \

Required Argument for Dataset Conversion

  • -s or --source_file: path to the raw file

  • -o or --output_dir - path to the directory to store the converted files

After the conversion, the output_dir should contain labels_*.txt and text_*.txt files. The default names for the training and evaluation in the conf/punctuation_capitalization_config.yaml are the following:

  |-- labels_dev.txt
  |-- labels_train.txt
  |-- text_dev.txt
  |-- text_train.txt

Training Punctuation and Capitalization Model

The language model is initialized with the pre-trained model from HuggingFace Transformers, unless the user provides a pre-trained checkpoint for the language model. Example of model configuration file for training the model can be found at: NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml.

The specification is roughly grouped into the following categories:

  • Parameters that describe the training process: trainer

  • Parameters that describe the datasets: model.dataset, model.train_ds, model.validation_ds

  • Parameters that describe the model: model

More details about parameters in the config file can be found below and in the model’s config file:


Data Type




Path to the pre-trained model .nemo file or pre-trained model name.



Path to the data converted to the specified above format.



Number of fully connected layers.



Activation to use between fully connected layers.



Dropout to apply to the input hidden states.



Whether to initialize the weights of the classifier head with the same approach used in Transformer.



Number of fully connected layers.



Dropout to apply to the input hidden states.



Activation function to use between fully connected layers.



Whether to initialize the weights of the classifier head with the same approach used in Transformer.



Name of the text training file located at data_dir.



Name of the labels training file located at data_dir, such as labels_train.txt.



Number of samples to use from the training dataset, -1 - to use all.



Name of the text file for evaluation, located at data_dir.



Name of the labels dev file located at data_dir, such as labels_dev.txt.



Number of samples to use from the dev set, -1 - to use all.

For more information, refer to the Model NLP section.

To train the model from scratch, run:

python examples/nlp/token_classification/punctuation_and_capitalization_train.py \
       model.dataset.data_dir=<PATH/TO/DATA_DIR> \
       trainer.gpus=[0,1] \
       optim.name=adam \

The above command will start model training on GPUs 0 and 1 with Adam optimizer and learning rate of 0.0001; and the trained model is stored in the nemo_experiments/Punctuation_and_Capitalization folder.

To train from the pre-trained model, run:

python examples/nlp/token_classification/punctuation_and_capitalization_train.py \
       model.dataset.data_dir=<PATH/TO/DATA_DIR> \

Required Arguments for Training

  • model.dataset.data_dir: Path to the data_dir with the pre-processed data files.


All parameters defined in the configuration file can be changed with command arguments. For example, the sample config file mentioned above has validation_ds.batch_size set to 64. However, if you see that the GPU utilization can be optimized further by using a larger batch size, you may override to the desired value by adding the field validation_ds.batch_size=128 over the command-line. You can repeat this with any of the parameters defined in the sample configuration file.


Inference is performed by a script examples/nlp/token_classification/punctuate_capitalize_infer.py

python punctuate_capitalize_infer.py \
    --input_manifest <PATH_TO_INPUT_MANIFEST> \
    --output_manifest <PATH_TO_OUTPUT_MANIFEST> \
    --pretrained_name punctuation_en_bert \
    --max_seq_length 64 \
    --margin 16 \
    --step 8

<PATH_TO_INPUT_MANIFEST> is a path to NeMo ASR manifest with text in which you need to restore punctuation and capitalization. If manifest contains 'pred_text' key, then 'pred_text' elements will be processed. Otherwise, punctuation and capitalization will be restored in 'text' elements.

<PATH_TO_OUTPUT_MANIFEST> is a path to NeMo ASR manifest into which result will be saved. The text with restored punctuation and capitalization is saved into 'pred_text' elements if 'pred_text' key is present in input manifest. Otherwise result will be saved into 'text' elements.

Alternatively you can pass data for restoring punctuation and capitalization as plain text. See help for parameters --input_text and --output_text of the script punctuate_capitalize_infer.py.

The script punctuate_capitalize_infer.py can restore punctuation and capitalization in a text of arbitrary length. Long sequences are split into segments --max_seq_length - 2 tokens each. Each segment starts and ends with [CLS] and [SEP] tokens correspondingly. Every segment is offset to the previous one by --step tokens. For example, if every character is a token, --max_seq_length=5, --step=2, then text "hello" will be split into segments [['[CLS]', 'h', 'e', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']].

If segments overlap, then predicted probabilities for a token present in several segments are multiplied before before selecting the best candidate.

Splitting leads to pour performance of a model near edges of segments. Use parameter --margin to discard --margin probabilities predicted for --margin tokens near segment edges. For example, if every character is a token, --max_seq_length=5, --step=1, --margin=1, then text "hello" will be split into segments [['[CLS]', 'h', 'e', 'l', '[SEP]'], ['[CLS]', 'e', 'l', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']]. Before calculating final predictions, probabilities for tokens marked by asterisk are removed: [['[CLS]', 'h', 'e', 'l'*, '[SEP]'*], ['[CLS]'*, 'e'*, 'l', 'l'*, '[SEP]'*], ['[CLS]'*, 'l'*, 'l', 'o', '[SEP]']]

Model Evaluation

An example script on how to evaluate the pre-trained model, can be found at examples/nlp/token_classification/punctuation_capitalization_evaluate.py.

To start evaluation of the pre-trained model, run:

python punctuation_capitalization_evaluate.py \
       model.dataset.data_dir=<PATH/TO/DATA/DIR>  \
       pretrained_model=punctuation_en_bert \
       model.test_ds.text_file=<text_dev.txt> \

Required Arguments

  • pretrained_model: pretrained Punctuation and Capitalization model from list_available_models() or path to a .nemo

file. For example: punctuation_en_bert or your_model.nemo. - model.dataset.data_dir: path to the directory that containes model.test_ds.text_file and model.test_ds.labels_file

During evaluation of the test_ds, the script generates two classification reports: one for capitalization task and another one for punctuation task. This classification reports include the following metrics:

  • Precision

  • Recall

  • F1

More details about these metrics can be found here.



Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.