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"""16This script is designed to extract features from different layers of a pretrained SSL model.17The extracted features will be in *.npy format, and in the shape of [L, D, T], where L is the 18number of layers, D is the feature dimension, and T is the time dimension.19 20Example usage:21 22python extract_features.py \23    --model_path="nvidia/ssl_en_nest_large_v1.0" \24    --input=<path to input manifest, or a dir containing audios, or path to audio> \25    --output=<output directory to store features and manifest> \26    --layers="all" \27    --batch_size=8 \28    --workers=8 \29    --max_cache=1000 # save features every 1000 samples to avoid OOM in system memory30"""31 32 33import argparse34import os35import tempfile36from pathlib import Path37from typing import List38 39import lightning.pytorch as pl40import numpy as np41import torch42from tqdm import tqdm43 44from nemo.collections.asr.data.audio_to_text_dataset import get_char_dataset45from nemo.collections.asr.models import EncDecDenoiseMaskedTokenPredModel46from nemo.collections.asr.modules import ConformerMultiLayerFeatureExtractor47from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest48from nemo.collections.common.data.utils import move_data_to_device49from nemo.collections.common.parts.preprocessing.manifest import get_full_path50from nemo.core.classes.common import typecheck51from nemo.utils import logging52 53typecheck.set_typecheck_enabled(enabled=False)54 55parser = argparse.ArgumentParser(description="Extract audio features using an SSL model")56parser.add_argument(57    "--model_path",58    type=str,59    required=True,60    help="Path to the .nemo model file or a pretrained model name from the NGC/HF model hub",61)62parser.add_argument(63    "-i",64    "--input",65    type=str,66    required=True,67    help="Path to the input audio file, or list of files, directory or jsonl manifest",68)69parser.add_argument(70    "-o", "--output", type=str, required=True, help="Path to the output directory that contains .npy file"71)72parser.add_argument(73    "-l",74    "--layers",75    type=str,76    default="all",77    help="Layers to extract features from, use 'all' to extract from all layer, 'last' for last layer, "78    "or comma-separated indices of the target layers (e.g. '0,1,2')",79)80parser.add_argument("-b", "--batch_size", type=int, default=8, help="Batch size for feature extraction")81parser.add_argument("-w", "--workers", type=int, default=8, help="Number of workers for feature extraction")82parser.add_argument("-d", "--device", type=str, default="cuda", help="Device to use for feature extraction")83parser.add_argument("-t", "--type", type=str, default="wav", help="audio file type, only needed for directory input")84parser.add_argument("--use_amp", action="store_true", help="Use automatic mixed precision")85parser.add_argument(86    "--amp_dtype",87    type=str,88    default="float16",89    choices=["float16", "bfloat16"],90    help="Data type for automatic mixed precision",91)92parser.add_argument("-mc", "--max_cache", type=int, default=-1, help="Max cache size before saving features")93args = parser.parse_args()94 95 96def get_input_manifest(input: str) -> List[dict]:97    """98    Build manifest from input path or directory99    """100    if input.endswith(".json") or input.endswith(".jsonl") and os.path.isfile(input):101        logging.info(f"Reading manifest from: {input}")102        manifest = [103            {"audio_filepath": str(get_full_path(item["audio_filepath"], input)), "duration": None, "text": "-"}104            for item in read_manifest(input)105        ]106    elif os.path.isdir(input):107        logging.info(f"Creating manifest from directory: {input}")108        manifest = [109            {"audio_filepath": str(p), "duration": None, "text": "-"} for p in Path(input).rglob(f"*.{args.type}")110        ]111        logging.info(f"Found {len(manifest)} items of {args.type} files")112    elif os.path.isfile(input):113        logging.info(f"Reading single file: {input}")114        manifest = [{"audio_filepath": Path(input).absolute.as_posix(), "duration": None, "text": "-"}]115    else:116        raise ValueError(f"Invalid input: {input}")117    return manifest118 119 120def load_model(model_path):121    """122    Load SSL model from local or pretrained123    """124    if model_path.endswith(".nemo") and os.path.isfile(model_path):125        logging.info(f"Loading model from local: {model_path}")126        model = EncDecDenoiseMaskedTokenPredModel.restore_from(model_path)127    else:128        logging.info(f"Loading model from pretrained: {model_path}")129        model = EncDecDenoiseMaskedTokenPredModel.from_pretrained(model_name=model_path)130    return model131 132 133class FeatureExtractor(pl.LightningModule):134    """135    Wrapper class for extracting features from SSL model136    """137 138    def __init__(self, ssl_model: EncDecDenoiseMaskedTokenPredModel, layer: str = "all"):139        super().__init__()140        self.preprocessor = ssl_model.preprocessor141        self.encoder = ssl_model.encoder142        self.layer_idx_list = None143        self.sample_rate = ssl_model.cfg.sample_rate144        if layer == "all":145            self.layer_idx_list = None146        elif layer == "last":147            self.layer_idx_list = [len(self.encoder.layers) - 1]148        else:149            try:150                self.layer_idx_list = [int(l) for l in layer.split(",")]151            except Exception as e:152                raise ValueError(f"Invalid layer argument: {layer}. Error: {e}")153        self.feature_extractor = ConformerMultiLayerFeatureExtractor(154            self.encoder, aggregator=None, layer_idx_list=self.layer_idx_list155        )156 157    def forward(158        self,159        input_signal=None,160        input_signal_length=None,161        processed_signal=None,162        processed_signal_length=None,163    ):164        """165        Forward pass to extract features, same input interface as EncDecDenoiseMaskedTokenPredModel.forward166        """167        has_input_signal = input_signal is not None and input_signal_length is not None168        has_processed_signal = processed_signal is not None and processed_signal_length is not None169        if (has_input_signal ^ has_processed_signal) == False:170            raise ValueError(171                f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "172                " with ``processed_signal`` and ``processed_signal_len`` arguments."173            )174        if not has_processed_signal:175            processed_signal, processed_signal_length = self.preprocessor(176                input_signal=input_signal,177                length=input_signal_length,178            )179        encoded, encoded_len = self.feature_extractor(audio_signal=processed_signal, length=processed_signal_length)180        return encoded, encoded_len181 182 183def maybe_save_features(output_dir, results, max_cache, manifest):184    """185    Check if the cache is full and save features to disk186    """187    if len(results) == 0 or max_cache < 0 or len(results) < max_cache:188        return189    os.makedirs(output_dir, exist_ok=True)190    logging.info(f"Saving {len(results)} features to {output_dir}")191 192    for sample_id, audio_file, features_np in tqdm(results, desc="Saving features", total=len(results)):193        filename = str(audio_file).replace("/", "_").replace(".", "_")194        if len(filename) > 256:195            filename = filename[-256:]196        output_path = os.path.join(output_dir, f"{filename}.npy")197        np.save(output_path, features_np)198        manifest[sample_id]["feature_path"] = output_path199 200    logging.info(f"Saved {len(results)} features to {output_dir}")201    results.clear()202 203 204def extract_features(args):205    """206    Main function to extract and save features from SSL model207    """208 209    logging.info(f"Extracting features using params: {vars(args)}")210 211    # Load model212    model = load_model(args.model_path)213    feature_extractor = FeatureExtractor(model, args.layers)214    device = torch.device(args.device)215    feature_extractor.to(device)216 217    # Load data218    logging.info(f"Building dataset from input: {args.input}")219    tmp_manifest = tempfile.NamedTemporaryFile(mode="w", delete=False)220    manifest = get_input_manifest(args.input)221    write_manifest(tmp_manifest.name, manifest)222    total_num_samples = len(manifest)223 224    # Build dataloader225    config = {226        "manifest_filepath": tmp_manifest.name,227        "sample_rate": feature_extractor.sample_rate,228        "return_sample_id": True,229    }230    dataset = get_char_dataset(config)231    logging.info(f"Built dataset with {len(dataset)} samples")232    dataloader = torch.utils.data.DataLoader(233        dataset=dataset,234        collate_fn=dataset.collate_fn,235        batch_size=args.batch_size,236        shuffle=False,237        num_workers=args.workers,238        pin_memory=True,239        drop_last=False,240    )241 242    # Extract features243    indices = set()244    results = []245    amp_dtype = torch.float16 if args.amp_dtype == "float16" else torch.bfloat16246    logging.info(f"Extracting features using AMP: {args.use_amp}, dtype: {amp_dtype}")247    with torch.amp.autocast('cuda' if torch.cuda.is_available() else 'cpu', dtype=amp_dtype, enabled=args.use_amp):248        with torch.inference_mode():249            for batch in tqdm(dataloader, desc="Extracting features"):250                batch = move_data_to_device(batch, device)251                audio_signal, audio_signal_len, _, _, sample_id = batch252                features, features_len = feature_extractor(253                    input_signal=audio_signal, input_signal_length=audio_signal_len254                )255                batch_size = features[0].size(0)256                num_layers = len(features)257                for i in range(batch_size):258                    sid_i = sample_id[i]259                    if sid_i in indices:260                        logging.warning(f"Skipping duplicated sample_id: {sample_id}")261                        continue262 263                    feat_i_len = features_len[0][i]264                    feat_i = []265                    for j in range(num_layers):266                        feat_i.append(features[j][i][:, :feat_i_len])267 268                    feat_i_np = torch.stack(feat_i, dim=0).cpu().numpy()269 270                    indices.add(sid_i)271                    results.append((sid_i, manifest[sid_i]['audio_filepath'], feat_i_np))272 273                maybe_save_features(args.output, results, args.max_cache, manifest)274 275    maybe_save_features(args.output, results, 0, manifest)276 277    output_manifest = Path(args.output) / "features.json"278    write_manifest(output_manifest, manifest)279    os.remove(tmp_manifest.name)280    logging.info(f"Extracted features from {total_num_samples} samples to {args.output}")281    logging.info(f"Manifest saved to: {output_manifest}")282 283 284if __name__ == "__main__":285    extract_features(args)286