N-gram Language Model

An external language model can improve ASR accuracy, especially in out-of-domain contexts with rare or specialised vocabulary. CAIMAN-ASR supports the use of KenLM n-gram language model shallow fusion corrections. We have seen a consistent WER improvement even when the N-gram is trained on ASR data transcripts. As such we automatically generate this N-gram during preprocessing and use it during validation (default for the beam decoder). See the Validation with an N-gram section for more details.

Build an N-gram Language Model

When adapting the preprocessing steps detailed here for your own dataset, you should have generated an n-gram language model trained on your transcripts. To generate an n-gram from a different dataset, see the following steps.

Preparing Data

To train an n-gram with KenLM on transcripts from ASR datasets, the data must first be prepared into the correct format - a .txt file where tokens within a sentence are space-separated and each sentence appears on a new line.

To gather the transcripts from json manifest files, run the following command inside a running container:

python caiman_asr_train/lm/prep_kenlm_data.py --data_dir /path/to/dataset/ --manifests manifest1.json manifest2.json --output_path /path/to/transcripts.txt --model_config configs/config.yaml

Training an N-gram

To train an n-gram, run the generate_ngram.sh script as follows:

./scripts/generate_ngram.sh [NGRAM_ORDER] /path/to/transcripts.txt /path/to/ngram.arpa /path/to/ngram.binary /path/to/sentencepiece.model

For example, to generate a 4-gram, set [NGRAM_ORDER] to 4 as follows:

./scripts/generate_ngram.sh 4 /path/to/transcripts.txt /path/to/ngram.arpa /path/to/ngram.binary /path/to/sentencepiece.model

The script will produce an ARPA file, which is a human-readable version of the language model, and a binary file, which allows for faster loading and is the recommended format. Binary files are the only usable format when generating hardware checkpoints, though providing an n-gram is optional. The export leaves it out of the hardware checkpoint, which has only its scale factor, and prints its path so that you can upload it with the checkpoint (see Exporting an inference checkpoint). By default, the ngram will be pruned of singleton/singular n-grams of order 3 and above.

The last argument is the SentencePiece model the transcripts were tokenized with, which is the one of the RNN-T the n-gram is for. An n-gram is only useful with the SentencePiece model it was made for, and nothing in a KenLM binary says which that is, so the script records it at the end of ngram.binary: the binary itself is unchanged (KenLM ignores what follows it), followed by the hash of the SentencePiece model. The beam decoder and the inference server check that it is the model of the checkpoint they use the n-gram with, and refuse an n-gram made for another. To add this to an n-gram that was made without it, run python -m caiman_asr_core.ngram_file append /path/to/ngram.binary /path/to/sentencepiece.model; info instead of append prints what a file says. The training image installs our fork of KenLM, which differs from upstream only in loading such a file when asked for its vocabulary strings; for everything else the extra bytes are ignored, so stock KenLM also works.

Validation with an N-gram

During validation, the n-gram language model generated during preprocessing is used, by reading from the following entries in the model configuration file:

ngram:
  ngram_path: /datasets/ngrams/NGRAM_SUBDIR
  scale_factor: 0.05

First, a binary file, named ngram.binary, in NGRAM_SUBDIR is searched for. If not found, an ARPA file - ngram.arpa - is searched for. If neither file exists, the process will crash with an error. To prevent this, use the --skip_ngram flag to disable the use of an n-gram during validation (disabled by default for the greedy decoder):

scripts/val.sh --skip_ngram

The n-gram is enabled by default for the beam decoder and disabled by default for the greedy decoder. To enable it for the greedy decoder use the --no-skip_ngram flag.

scripts/val.sh --decoder greedy --no-skip_ngram

The scale_factor adjusts the scores from the n-gram language model, and this will require tuning for your dataset. Values between 0.05 and 0.1 are empirically effective for improving WER. See the Sweep Scale Factor section below for details on running a sweep across the scale factor.

In addition, the greedy decoder also uses a top_k value, where only the top k candidates are used by the n-gram model to adjust the scores. top_k is set to 5 by default. To use a different value use the --top_k flag

scripts/val.sh --decoder greedy --no-skip_ngram

To use an n-gram that was trained on a different dataset, use the --override_ngram_path argument, which will take precedence over any n-grams in NGRAM_SUBDIR:

scripts/val.sh --override_ngram_path /path/to/ngram.binary

Sweep Scale Factor

To optimize the scale_factor for your n-gram language model, use the sweep_scale_factor.py script. This script iterates over multiple scale_factor values, performs validation, and updates your model config YAML with the best one based on WER.

Run the following command to perform a sweep:

python caiman_asr_train/lm/sweep_scale_factor.py --checkpoint /path/to/checkpoint.pt --model_config configs/config.yaml --val_manifests /path/to/manifest.json

By default, a sweep is performed across [0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.25]. To specify custom values, use the --scale_factors argument:

python caiman_asr_train/lm/sweep_scale_factor.py --scale_factors 0.1 0.2 0.3 --checkpoint /path/to/checkpoint.pt --model_config configs/config.yaml --val_manifests /path/to/manifest.json