Export inference checkpoint

To run your model on Myrtle.ai’s hardware-accelerated inference server you will need a hardware checkpoint. This is a training checkpoint with additional data such as mel-bin mean and variances as described here.

By default, train.sh will automatically create hardware checkpoints for your best and last checkpoints. If your best checkpoint is /results/RNN-T_best_checkpoint.pt, you can find the hardware checkpoint at /results/RNN-T_best_checkpoint.hw.pt.

Creating a hardware checkpoint manually

If you want more customization, you can also create a hardware checkpoint from a training checkpoint using the hardware_ckpt.py script. To create a hardware checkpoint run:

python ./caiman_asr_train/export/hardware_ckpt.py \
    --ckpt /results/RNN-T_best_checkpoint.pt \
    --config <path/to/config.yaml> \
    --output_ckpt /results/hardware_checkpoint.example.pt

where /results/RNN-T_best_checkpoint.pt is your best checkpoint.

The script should take a few seconds to run.

The generated hardware checkpoint will contain the sentencepiece model specified in the config file and the dataset mel stats.

The hardware checkpoint also records which prefixes the model accepts, so the inference server can validate a requested prefix. These are stored as a prefix_spec dict with the keys:

• valid_prefix: a regex matching exactly the prefixes seen in training, e.g. ^(?:<pnc>|<nopnc>)(?:<lang_en>|<lang_fr>|<lang_detect>)$, where no prefix is the empty string. A model trained without --conditional only accepts no prefix, ^$.
• default_prefix: the prefix the model was validated with during training (--pnc_prefix and --val_lang_prefix), e.g. <pnc><lang_en>, or None for no prefix. Training fails at startup if this doesn’t match valid_prefix.

The prefix_spec is generated automatically from the training arguments and stored in the training checkpoint. It is None, or absent, when the prefixes are unknown, e.g. for checkpoints trained before it was introduced.

The export will also find the binary n-gram generated during preprocessing, as specified by the ngram_path field in the config file, check it and print its path, as it is a separate file to keep next to the hardware checkpoint. However, this is optional, and can be skipped by passing the --skip_ngram flag:

python ./caiman_asr_train/export/hardware_ckpt.py \
    --ckpt /results/RNN-T_best_checkpoint.pt \
    --config <path/to/config.yaml> \
    --output_ckpt /results/hardware_checkpoint.example.pt
    --skip_ngram

To use an n-gram that was generated on a different dataset, use the --override_ngram_path argument:

python ./caiman_asr_train/export/hardware_ckpt.py \
    --ckpt /results/RNN-T_best_checkpoint.pt \
    --config <path/to/config.yaml> \
    --output_ckpt /results/hardware_checkpoint.example.pt \
    --override_ngram_path /path/to/ngram.binary

The n-gram file

The n-gram is not stored in the hardware checkpoint (which has only its scale factor). It can be large (many GB), and every process that loads a checkpoint would get its own copy of it, so it is a separate file that the inference server memory-maps, and any number of servers share one copy of it in memory. The export does not copy it: it prints the path of the n-gram file it used (the one found from ngram_path in the config, or --override_ngram_path), and that file is what you upload to wherever you put the hardware checkpoint.

The inference server has to be given the file with --ngram-binary-path. It memory-maps it, so any number of servers share one copy of it in memory.

The n-gram file should say which SentencePiece model it was made for: it is the KenLM binary, unchanged, followed by the hash of that model (see N-gram Language Model; scripts/generate_ngram.sh does this, and python -m caiman_asr_core.ngram_file append does it for an existing n-gram). The export refuses an n-gram file that says it was made for another SentencePiece model than the checkpoint’s. One that doesn’t say can still be used, with a warning, as nothing can check it. The inference server checks the same when it is given the file. The format is described in caiman_asr_core.ngram_file. (Loading such a file while asking KenLM for the vocabulary strings needs our KenLM fork, which stops reading them at the end of the model; plain loading, the Python module and the inference server work with stock KenLM.)