RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15# Copyright 2017 Johns Hopkins University (Shinji Watanabe)16#17# Licensed under the Apache License, Version 2.0 (the "License");18# you may not use this file except in compliance with the License.19# You may obtain a copy of the License at20#21#     http://www.apache.org/licenses/LICENSE-2.022#23# Unless required by applicable law or agreed to in writing, software24# distributed under the License is distributed on an "AS IS" BASIS,25# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.26# See the License for the specific language governing permissions and27# limitations under the License.28 29"""30# Changes to script31Change the script to import the NeMo model class you would like to load a checkpoint for,32then update the model constructor to use this model class. This can be found by the line:33<<< Change model class here ! >>>34By default, this script imports and creates the `EncDecCTCModelBPE` class but it can be35changed to any NeMo Model.36# Run the script37## Saving a .nemo model file (loaded with ModelPT.restore_from(...))38HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \39    --config-path="<path to config directory>" \40    --config-name="<config name>" \41    name=<name of the averaged checkpoint> \42    +checkpoint_dir=<OPTIONAL: directory of checkpoint> \43    +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\"44## Saving an averaged pytorch checkpoint (loaded with torch.load(...))45HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \46    --config-path="<path to config directory>" \47    --config-name="<config name>" \48    name=<name of the averaged checkpoint> \49     +checkpoint_dir=<OPTIONAL: directory of checkpoint> \50    +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\" \51    +save_ckpt_only=true52"""53 54import os55 56import lightning.pytorch as pl57import torch58from omegaconf import OmegaConf, open_dict59 60# Change this import to the model you would like to average61from nemo.collections.asr.models import EncDecCTCModelBPE62from nemo.core.config import hydra_runner63from nemo.utils import logging64 65 66def process_config(cfg: OmegaConf):67    """68    Process config69    """70    if 'name' not in cfg or cfg.name is None:71        raise ValueError("`cfg.name` must be provided to save a model checkpoint")72 73    if 'checkpoint_paths' not in cfg or cfg.checkpoint_paths is None:74        raise ValueError(75            "`cfg.checkpoint_paths` must be provided as a list of one or more str paths to "76            "pytorch lightning checkpoints"77        )78 79    save_ckpt_only = False80 81    with open_dict(cfg):82        name_prefix = cfg.name83        checkpoint_paths = cfg.pop('checkpoint_paths')84 85        if 'checkpoint_dir' in cfg:86            checkpoint_dir = cfg.pop('checkpoint_dir')87        else:88            checkpoint_dir = None89 90        if 'save_ckpt_only' in cfg:91            save_ckpt_only = cfg.pop('save_ckpt_only')92 93    if type(checkpoint_paths) not in (list, tuple):94        checkpoint_paths = str(checkpoint_paths).replace("[", "").replace("]", "")95        checkpoint_paths = checkpoint_paths.split(",")96        checkpoint_paths = [ckpt_path.strip() for ckpt_path in checkpoint_paths]97 98    if checkpoint_dir is not None:99        checkpoint_paths = [os.path.join(checkpoint_dir, path) for path in checkpoint_paths]100 101    return name_prefix, checkpoint_paths, save_ckpt_only102 103 104@hydra_runner(config_path=None, config_name=None)105def main(cfg):106    """107    Main function108    """109 110    logging.info("This script is deprecated and will be removed in the 25.01 release.")111 112    name_prefix, checkpoint_paths, save_ckpt_only = process_config(cfg)113 114    if not save_ckpt_only:115        trainer = pl.Trainer(**cfg.trainer)116 117        # <<< Change model class here ! >>>118        # Model architecture which will contain the averaged checkpoints119        # Change the model constructor to the one you would like (if needed)120        model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)121 122    """ < Checkpoint Averaging Logic > """123    # load state dicts124    n = len(checkpoint_paths)125    avg_state = None126 127    logging.info(f"Averaging {n} checkpoints ...")128 129    for ix, path in enumerate(checkpoint_paths):130        checkpoint = torch.load(path, map_location='cpu')131 132        if 'state_dict' in checkpoint:133            checkpoint = checkpoint['state_dict']134 135        if ix == 0:136            # Initial state137            avg_state = checkpoint138 139            logging.info(f"Initialized average state dict with checkpoint : {path}")140        else:141            # Accumulated state142            for k in avg_state:143                avg_state[k] = avg_state[k] + checkpoint[k]144 145            logging.info(f"Updated average state dict with state from checkpoint : {path}")146 147    for k in avg_state:148        if str(avg_state[k].dtype).startswith("torch.int"):149            # For int type, not averaged, but only accumulated.150            # e.g. BatchNorm.num_batches_tracked151            pass152        else:153            avg_state[k] = avg_state[k] / n154 155    # Save model156    if save_ckpt_only:157        ckpt_name = name_prefix + '-averaged.ckpt'158        torch.save(avg_state, ckpt_name)159 160        logging.info(f"Averaged pytorch checkpoint saved as : {ckpt_name}")161    else:162        # Set model state163        logging.info("Loading averaged state dict in provided model")164        model.load_state_dict(avg_state, strict=True)165 166        ckpt_name = name_prefix + '-averaged.nemo'167        model.save_to(ckpt_name)168 169        logging.info(f"Averaged model saved as : {ckpt_name}")170 171 172if __name__ == '__main__':173    main()174