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 15import difflib16import os17from typing import List18 19import nemo_run as run20from lightning.pytorch.callbacks.callback import Callback21from nemo_run.core.serialization.yaml import YamlSerializer22from nemo_run.run.torchx_backend.packaging import _serialize23 24from nemo.collections.common.tokenizers.huggingface import AutoTokenizer25from nemo.collections.llm.gpt.data.squad import SquadDataModule26from nemo.collections.llm.gpt.model import GPTModel27from nemo.collections.llm.recipes.llama3_8b import MegatronCommOverlapCallback28from nemo.lightning.base import DEFAULT_NEMO_CACHE_HOME29from nemo.utils import logging30 31DEFAULT_NEMO_HOME = os.getenv('NEMO_HOME', DEFAULT_NEMO_CACHE_HOME)32 33 34def hf_tokenizer(model_name: str) -> run.Config[AutoTokenizer]:35    """36    HuggingFace tokenizer.37 38    Args:39        model_name (str): corresponds to HuggingFace-AutoTokenizer's 'pretrained_model_name_or_path' input argument.40                For more details please refer to-41                huggingface.co/docs/transformers/v4.47.1/en/model_doc/auto#transformers.AutoTokenizer42    """43    log_msg = [44        f"`AutoTokenizer` first searches for tokenizer files locally stored in {DEFAULT_NEMO_HOME}.",45        "(from env var `NEMO_HOME`- can be changed using '-nh/--nemo_home' CLI arg).",46        "If files are missing locally, `AutoTokenizer` will try downloading from HuggingFace. In this case-",47        "make sure env vars 'HF_HUB_OFFLINE':'0' and 'HF_TOKEN':'<token_value>' are set in your sbatch script.",48        "Both of these will be set automatically if you provide '-hf/--hf_token' CLI arg.",49    ]50    logging.warning(" ".join(log_msg))51 52    return run.Config(53        AutoTokenizer,54        pretrained_model_name=model_name,55        use_fast=True,56    )57 58 59def import_ckpt_experiment(executor: run.SlurmExecutor, model: run.Config[GPTModel], source: str):60    """61    Downloads/Acceses checkpoint to be used for fine-tuning. `import_ckpt` first tries find the nemo checkpoint in62    <NEMO_HOME>/models/. For eg: for llama3 8b, the path will look like- <NEMO_HOME>/models/meta-llama/Meta-Llama-3-8B63    If missing, tries to downloads at the same location from HuggingFace and converts it nemo format.64 65    Args:66        source (str): HuggingFace URL. For eg- hf://meta-llama/Meta-Llama-3-70B67    """68    from copy import deepcopy69 70    from nemo.collections.llm import import_ckpt71 72    import_executor = deepcopy(executor)73    import_executor.ntasks_per_node = 174    import_executor.nodes = 175 76    return run.Partial(import_ckpt, model=model, source=source, overwrite=False), import_executor, "import_ckpt_exp"77 78 79def get_nemo_home(nemo_home=None):80    """81    Get NEMO_HOME path. Checks for both nemo_home argument and NEMO_HOME environment variable.82    """83    arg_nemo_set = nemo_home is True84    env_nemo_set = "NEMO_HOME" in os.environ85 86    if arg_nemo_set and env_nemo_set:87        if os.environ["NEMO_HOME"] != nemo_home:88            logging.warning(f"Using nemo_home ({nemo_home}) instead of NEMO_HOME ({os.environ['NEMO_HOME']})")89        return nemo_home90 91    if arg_nemo_set:92        return nemo_home93 94    if env_nemo_set:95        return os.environ["NEMO_HOME"]96 97    raise ValueError("Neither -nh/--nemo_home argument nor NEMO_HOME environment variable is set")98 99 100def prepare_squad_dataset(model_name: str, seq_length: int = 2048, nemo_home=None):101    """Prepare the SQuAD dataset for fine-tuning.102 103    Args:104        model_name (str): The name of the model105        seq_length (int): The sequence length to use for packing. Defaults to 2048.106        nemo_home: Optional path to NEMO home directory set via args.nemo_home107    """108    from pathlib import Path109 110    from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer111    from nemo.collections.llm.gpt.data.packed_sequence import PackedSequenceSpecs112    from nemo.collections.llm.gpt.data.squad import SquadDataModule113 114    nemo_home_path = Path(get_nemo_home(nemo_home))115    dataset_root = nemo_home_path / "datasets" / "squad"116    dataset_root.mkdir(parents=True, exist_ok=True)117 118    tokenizer = AutoTokenizer(pretrained_model_name=model_name)119 120    # Configure SquadDataModule with packing specs121    datamodule = SquadDataModule(122        dataset_root=dataset_root,123        seq_length=seq_length,124        global_batch_size=8,125        micro_batch_size=1,126        packed_sequence_specs=PackedSequenceSpecs(packed_sequence_size=seq_length),127        tokenizer=tokenizer,128        force_redownload=True,129        delete_raw=False,130        seed=1234,131    )132 133    # This will generate both JSONL and packed .bin files134    datamodule.prepare_data()135 136    # Verify the output137    packed_dir = dataset_root / "packed" / model_name.replace("/", "--")138    print(f"Packed files should be in: {packed_dir}")139    if packed_dir.exists():140        print("Files found:", list(packed_dir.glob("*")))141    else:142        raise FileNotFoundError(f"Packed dataset dir not found at {packed_dir}. Dataset download failed")143 144 145def prepare_squad_dataset_experiment(146    executor: run.SlurmExecutor, model_name: str, seq_length: int = 2048, nemo_home=None147):148    """149    Downloads and prepares the SQuAD dataset for fine-tuning.150    """151    from copy import deepcopy152 153    dataset_executor = deepcopy(executor)154    dataset_executor.ntasks_per_node = 1155    dataset_executor.nodes = 1156 157    return (158        run.Partial(159            prepare_squad_dataset,160            model_name=model_name,161            seq_length=seq_length,162            nemo_home=nemo_home,163        ),164        dataset_executor,165        "prepare_squad_dataset_exp",166    )167 168 169def isfile_train_pack_metadata(hf_model_uri: str, data_config: run.Config[SquadDataModule]) -> bool:170    """171    This method is used for fine-tuning. It checks if packed train data for a partiular172    sequence length exists locally. This is needed to set data flag (force_redownload=True)173    which avoids experiment crash in case files are missing.174    """175    datasets_dir = os.getenv("NEMO_DATASETS_CACHE", os.path.join(DEFAULT_NEMO_HOME, "datasets"))176    model_dir = hf_model_uri.replace("/", "--")177    metadata_filename = f"{data_config.seq_length}_metadata.jsonl"178 179    train_pack_metadata_filepath = os.path.join(datasets_dir, "squad", "packed", model_dir, metadata_filename)180 181    return os.path.exists(train_pack_metadata_filepath) and os.path.isfile(train_pack_metadata_filepath)182 183 184def get_comm_overlap_callback_idx(callbacks: List[Callback]) -> int | None:185    """186    nemo.lightning.Trainer has a list of callbacks defined. This method identifies index of MegatronCommOverlapCallback187    from the list defined in recipes in nemo.collections.llm.recipes. The index is needed to override ddp communication188    params189    """190    if callbacks:  # default is None in lightning191        for idx, callback in enumerate(callbacks):192            if callback.__fn_or_cls__ == MegatronCommOverlapCallback:193                return idx194    return None195 196 197def dump_config_diff_from_base_recipe(198    base_recipe: str, new_recipe: str, output_dir: str, file_name: str = "config_diff.txt"199):200    """201    Dump the config diff from the base recipe.202    """203    base_recipe_config = _serialize(base_recipe, serializer_cls=YamlSerializer)204    new_recipe_config = _serialize(new_recipe, serializer_cls=YamlSerializer)205    diff = difflib.unified_diff(206        base_recipe_config.splitlines(keepends=True),207        new_recipe_config.splitlines(keepends=True),208        fromfile="base_recipe",209        tofile="new_recipe",210        lineterm="",211    )212    diff = "".join(diff)213    print("dumping config diff to ", os.path.join(output_dir, file_name))214    with open(os.path.join(output_dir, file_name), "w") as f:215        f.write(diff)216