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 shutil16from collections import Counter17from pathlib import Path18from typing import Dict, Optional, Union19 20import torch21 22 23def is_nemo2_checkpoint(checkpoint_path: str) -> bool:24    """25    Checks if the checkpoint is in NeMo 2.0 format.26    Args:27        checkpoint_path (str): Path to a checkpoint.28    Returns:29        bool: True if the path points to a NeMo 2.0 checkpoint; otherwise false.30    """31 32    ckpt_path = Path(checkpoint_path)33    return (ckpt_path / 'context').is_dir()34 35 36def prepare_directory_for_export(37    model_dir: Union[str, Path], delete_existing_files: bool, subdir: Optional[str] = None38) -> None:39    """40    Prepares model_dir path for the TensorRTT-LLM / vLLM export.41    Makes sure that the model_dir directory exists and is empty.42 43    Args:44        model_dir (str): Path to the target directory for the export.45        delete_existing_files (bool): Attempt to delete existing files if they exist.46        subdir (Optional[str]): Subdirectory to create inside the model_dir.47 48    Returns:49        None50    """51    model_path = Path(model_dir)52 53    if model_path.exists():54        if delete_existing_files:55            shutil.rmtree(model_path)56        elif any(model_path.iterdir()):57            raise RuntimeError(f"There are files in {model_path} folder: try setting delete_existing_files=True.")58 59    if subdir is not None:60        model_path /= subdir61    model_path.mkdir(parents=True, exist_ok=True)62 63 64def is_nemo_tarfile(path: str) -> bool:65    """66    Checks if the path exists and points to packed NeMo 1 checkpoint.67 68    Args:69        path (str): Path to possible checkpoint.70    Returns:71        bool: NeMo 1 checkpoint exists and is in '.nemo' format.72    """73    checkpoint_path = Path(path)74    return checkpoint_path.exists() and checkpoint_path.suffix == '.nemo'75 76 77# Copied from nemo.collections.nlp.parts.utils_funcs to avoid introducing extra NeMo dependencies:78def torch_dtype_from_precision(precision: Union[int, str], megatron_amp_O2: bool = True) -> torch.dtype:79    """80    Mapping from PyTorch Lighthing (PTL) precision types to corresponding PyTorch parameter data type.81 82    Args:83        precision (Union[int, str]): The PTL precision type used.84        megatron_amp_O2 (bool): A flag indicating if Megatron AMP O2 is enabled.85 86    Returns:87        torch.dtype: The corresponding PyTorch data type based on the provided precision.88    """89    if not megatron_amp_O2:90        return torch.float3291 92    if precision in ['bf16', 'bf16-mixed']:93        return torch.bfloat1694    elif precision in [16, '16', '16-mixed']:95        return torch.float1696    elif precision in [32, '32', '32-true']:97        return torch.float3298    else:99        raise ValueError(f"Could not parse the precision of '{precision}' to a valid torch.dtype")100 101 102def get_model_device_type(module: torch.nn.Module) -> str:103    """Find the device type the model is assigned to and ensure consistency."""104    # Collect device types of all parameters and buffers105    param_device_types = {param.device.type for param in module.parameters()}106    buffer_device_types = {buffer.device.type for buffer in module.buffers()}107    all_device_types = param_device_types.union(buffer_device_types)108 109    if len(all_device_types) > 1:110        raise ValueError(111            f"Model parameters and buffers are on multiple device types: {all_device_types}. "112            "Ensure all parameters and buffers are on the same device type."113        )114 115    # Return the single device type, or default to 'cpu' if no parameters or buffers116    return all_device_types.pop() if all_device_types else "cpu"117 118 119def get_example_inputs(tokenizer) -> Dict[str, torch.Tensor]:120    """Gets example data to feed to the model during ONNX export.121 122    Returns:123        Dictionary of tokenizer outputs.124    """125    example_inputs = dict(126        tokenizer(127            ["example query one", "example query two"],128            ["example passage one", "example passage two"],129            return_tensors="pt",130        )131    )132 133    return example_inputs134 135 136def validate_fp8_network(network) -> None:137    """Checks the network to ensure it's compatible with fp8 precison.138 139    Raises:140        ValueError if netowrk doesn't container Q/DQ FP8 layers141    """142 143    import tensorrt as trt144 145    quantize_dequantize_layers = []146    for layer in network:147        if layer.type in {trt.LayerType.QUANTIZE, trt.LayerType.DEQUANTIZE}:148            quantize_dequantize_layers.append(layer)149    if not quantize_dequantize_layers:150        error_msg = "No Quantize/Dequantize layers found"151        raise ValueError(error_msg)152    quantize_dequantize_layer_dtypes = Counter(layer.precision for layer in quantize_dequantize_layers)153    if trt.DataType.FP8 not in quantize_dequantize_layer_dtypes:154        error_msg = "Found Quantize/Dequantize layers. But none with FP8 precision."155        raise ValueError(error_msg)156