Training IndicConformer models

Setup Instructions

Create a conda environment

conda create -y -n kemo python=3.10.12
conda activate kemo

Install Pytorch

Every machine requires a different version of pytorch based on the GPU / CUDA version that the machine supports. So we recommend installing the correct version from the official repository - https://pytorch.org/get-started/locally/.

We install pytorch for CUDA 12.2 using the following command -

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

Install ffmpeg (requires sudo access)

Install ffmpeg on a linux machine using the following command

sudo apt-get update -y && sudo apt-get install -y libsndfile1 ffmpeg

Install NeMo and additional components

Since there are many installations required here, we have simplified it to a bash file

cd libraries/
bash setup_nemo.sh

Download checkpoints

1. Download parakeet models into the assets/pretrained_checkpoints folder. Please ensure to update the PROJECT_ROOT to download the model to the correct path

python utils/download_parakeet.py

2. Download tokenizers into assets/tokenizers folder. Please ensure to update the PROJECT_ROOT to download the tokenizers to the correct path

bash utils/download_tokenizers.sh

Training Instructions

All training code is present in the training/ folder. To start the training, we need to first set the PROJECT_ROOT in training/runner_hms.sh. (todo - need a cleaner way to do this in future, but for now having everything relative to the project root helps to the project to be self-contained.)

Setting up the config

The training config can be found in training/configs/config.yaml. Set the following two lines in the config file -

project_root: <root_folder>
project_name: "parakeet-training"

The project_name is used to create a wandb project.

Training the model

To use the default parameters for training a 1B parakeet model for 22 languages on the Indicvoices dataset, you can use the following parameters -

python ${PROJECT_ROOT}/training/train_hms.py \
    name=parakeet1b-hms-indicvoices \
    trainer=vanilla exp_manager=vanilla \
    model=hms \
    model/encoder=fastconformer_1B \
    model/train_ds/manifest_filepath=indicvoices \
    model/validation_ds/manifest_filepath=indicvoices \
    model/tokenizer=multisoftmax_in22 \
    model/optim=constant_lr \
    +init_from_nemo_model=${PROJECT_ROOT}/assets/pretrained_checkpoints/nvidia-parakeet-rnnt-1.1b.nemo \
    exp_manager.resume_if_exists=false

The configs directory

The configurations for the model training are stored in the directory as expanded here. The configs can be changed by updating the config files below appropriately -

 |-configs
 | |-config.yaml
 | |-exp_manager
 | | |-early_stopping.yaml
 | | |-early_stopping_callback_params
 | | | |-val_wer_0.001_patience_10.yaml
 | | | |-val_wer_0.002_patience_5.yaml
 | | |-vanilla.yaml
 | |-model
 | | |-encoder
 | | | |-conformer_100M.yaml
 | | | |-fastconformer_1B.yaml
 | | |-hms.yaml
 | | |-joint
 | | | |-language_keys
 | | | | |-in22.yaml
 | | |-optim
 | | | |-constant_lr.yaml
 | | | |-cosine_annealing.yaml
 | | |-tokenizer
 | | | |-multisoftmax_in22.yaml
 | | |-train_ds
 | | | |-manifest_filepath
 | | | | |-indicvoices.yaml
 | | |-validation_ds
 | | | |-manifest_filepath
 | | | | |-indicvoices.yaml
 | |-trainer
 | | |-grad_acc_2.yaml
 | | |-grad_acc_4.yaml
 | | |-vanilla.yaml
 |-runner_hms.sh
 |-train_hms.py