Train an RNN language model

If you have enough text data, you can train a neural network language model (NNLM) to improve the WER of your E2E ASR system. This tutorial shows you how to train an RNNLM from scratch.

Hint

For how to use an NNLM during decoding, please refer to the following tutorials: Shallow fusion for Transducer, LODR for RNN Transducer, LM rescoring for Transducer

Note

This tutorial is based on the LibriSpeech recipe. Please check it out for the necessary python scripts for this tutorial. We use the LibriSpeech LM-corpus as the LM training set for illustration purpose. You can also collect your own data. The data format is quite simple: each line should contain a complete sentence, and words should be separated by space.

First, let’s download the training data for the RNNLM. This can be done via the following command:

$ wget https://www.openslr.org/resources/11/librispeech-lm-norm.txt.gz
$ gzip -d librispeech-lm-norm.txt.gz

As we are training a BPE-level RNNLM, we need to tokenize the training text, which requires a BPE tokenizer. This can be achieved by executing the following command:

$ # if you don't have the BPE
$ GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-2023-05-15
$ cd icefall-asr-librispeech-zipformer-2023-05-15/data/lang_bpe_500
$ git lfs pull --include bpe.model
$ cd ../../..

$ ./local/prepare_lm_training_data.py \
    --bpe-model icefall-asr-librispeech-zipformer-2023-05-15/data/lang_bpe_500/bpe.model \
    --lm-data librispeech-lm-norm.txt \
    --lm-archive data/lang_bpe_500/lm_data.pt

Now, you should have a file name lm_data.pt file store under the directory data/lang_bpe_500. This is the packed training data for the RNNLM. We then sort the training data according to its sentence length.

$ # This could take a while (~ 20 minutes), feel free to grab a cup of coffee :)
$ ./local/sort_lm_training_data.py \
    --in-lm-data data/lang_bpe_500/lm_data.pt \
    --out-lm-data data/lang_bpe_500/sorted_lm_data.pt \
    --out-statistics data/lang_bpe_500/lm_data_stats.txt

The aforementioned steps can be repeated to create a a validation set for you RNNLM. Let’s say you have a validation set in valid.txt, you can just set --lm-data valid.txt and --lm-archive data/lang_bpe_500/lm-data-valid.pt when calling ./local/prepare_lm_training_data.py.

After completing the previous steps, the training and testing sets for training RNNLM are ready. The next step is to train the RNNLM model. The training command is as follows:

$ # assume you are in the icefall root directory
$ cd rnn_lm
$ ln -s ../../egs/librispeech/ASR/data .
$ cd ..
$ ./rnn_lm/train.py \
    --world-size 4 \
    --exp-dir ./rnn_lm/exp \
    --start-epoch 0 \
    --num-epochs 10 \
    --use-fp16 0 \
    --tie-weights 1 \
    --embedding-dim 2048 \
    --hidden-dim 2048 \
    --num-layers 3 \
    --batch-size 300 \
    --lm-data rnn_lm/data/lang_bpe_500/sorted_lm_data.pt \
    --lm-data-valid rnn_lm/data/lang_bpe_500/sorted_lm_data.pt

Note

You can adjust the RNNLM hyper parameters to control the size of the RNNLM, such as embedding dimension and hidden state dimension. For more details, please run ./rnn_lm/train.py --help.

Note

The training of RNNLM can take a long time (usually a couple of days).