chore: import upstream snapshot with attribution

scripts/ssl/extract_features.py

@@ -0,0 +1,285 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script is designed to extract features from different layers of a pretrained SSL model.
+The extracted features will be in *.npy format, and in the shape of [L, D, T], where L is the 
+number of layers, D is the feature dimension, and T is the time dimension.
+
+Example usage:
+
+python extract_features.py \
+    --model_path="nvidia/ssl_en_nest_large_v1.0" \
+    --input=<path to input manifest, or a dir containing audios, or path to audio> \
+    --output=<output directory to store features and manifest> \
+    --layers="all" \
+    --batch_size=8 \
+    --workers=8 \
+    --max_cache=1000 # save features every 1000 samples to avoid OOM in system memory
+"""
+
+
+import argparse
+import os
+import tempfile
+from pathlib import Path
+from typing import List
+
+import lightning.pytorch as pl
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from nemo.collections.asr.data.audio_to_text_dataset import get_char_dataset
+from nemo.collections.asr.models import EncDecDenoiseMaskedTokenPredModel
+from nemo.collections.asr.modules import ConformerMultiLayerFeatureExtractor
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.parts.preprocessing.manifest import get_full_path
+from nemo.core.classes.common import typecheck
+from nemo.utils import logging
+
+typecheck.set_typecheck_enabled(enabled=False)
+
+parser = argparse.ArgumentParser(description="Extract audio features using an SSL model")
+parser.add_argument(
+    "--model_path",
+    type=str,
+    required=True,
+    help="Path to the .nemo model file or a pretrained model name from the NGC/HF model hub",
+)
+parser.add_argument(
+    "-i",
+    "--input",
+    type=str,
+    required=True,
+    help="Path to the input audio file, or list of files, directory or jsonl manifest",
+)
+parser.add_argument(
+    "-o", "--output", type=str, required=True, help="Path to the output directory that contains .npy file"
+)
+parser.add_argument(
+    "-l",
+    "--layers",
+    type=str,
+    default="all",
+    help="Layers to extract features from, use 'all' to extract from all layer, 'last' for last layer, "
+    "or comma-separated indices of the target layers (e.g. '0,1,2')",
+)
+parser.add_argument("-b", "--batch_size", type=int, default=8, help="Batch size for feature extraction")
+parser.add_argument("-w", "--workers", type=int, default=8, help="Number of workers for feature extraction")
+parser.add_argument("-d", "--device", type=str, default="cuda", help="Device to use for feature extraction")
+parser.add_argument("-t", "--type", type=str, default="wav", help="audio file type, only needed for directory input")
+parser.add_argument("--use_amp", action="store_true", help="Use automatic mixed precision")
+parser.add_argument(
+    "--amp_dtype",
+    type=str,
+    default="float16",
+    choices=["float16", "bfloat16"],
+    help="Data type for automatic mixed precision",
+)
+parser.add_argument("-mc", "--max_cache", type=int, default=-1, help="Max cache size before saving features")
+args = parser.parse_args()
+
+
+def get_input_manifest(input: str) -> List[dict]:
+    """
+    Build manifest from input path or directory
+    """
+    if input.endswith(".json") or input.endswith(".jsonl") and os.path.isfile(input):
+        logging.info(f"Reading manifest from: {input}")
+        manifest = [
+            {"audio_filepath": str(get_full_path(item["audio_filepath"], input)), "duration": None, "text": "-"}
+            for item in read_manifest(input)
+        ]
+    elif os.path.isdir(input):
+        logging.info(f"Creating manifest from directory: {input}")
+        manifest = [
+            {"audio_filepath": str(p), "duration": None, "text": "-"} for p in Path(input).rglob(f"*.{args.type}")
+        ]
+        logging.info(f"Found {len(manifest)} items of {args.type} files")
+    elif os.path.isfile(input):
+        logging.info(f"Reading single file: {input}")
+        manifest = [{"audio_filepath": Path(input).absolute.as_posix(), "duration": None, "text": "-"}]
+    else:
+        raise ValueError(f"Invalid input: {input}")
+    return manifest
+
+
+def load_model(model_path):
+    """
+    Load SSL model from local or pretrained
+    """
+    if model_path.endswith(".nemo") and os.path.isfile(model_path):
+        logging.info(f"Loading model from local: {model_path}")
+        model = EncDecDenoiseMaskedTokenPredModel.restore_from(model_path)
+    else:
+        logging.info(f"Loading model from pretrained: {model_path}")
+        model = EncDecDenoiseMaskedTokenPredModel.from_pretrained(model_name=model_path)
+    return model
+
+
+class FeatureExtractor(pl.LightningModule):
+    """
+    Wrapper class for extracting features from SSL model
+    """
+
+    def __init__(self, ssl_model: EncDecDenoiseMaskedTokenPredModel, layer: str = "all"):
+        super().__init__()
+        self.preprocessor = ssl_model.preprocessor
+        self.encoder = ssl_model.encoder
+        self.layer_idx_list = None
+        self.sample_rate = ssl_model.cfg.sample_rate
+        if layer == "all":
+            self.layer_idx_list = None
+        elif layer == "last":
+            self.layer_idx_list = [len(self.encoder.layers) - 1]
+        else:
+            try:
+                self.layer_idx_list = [int(l) for l in layer.split(",")]
+            except Exception as e:
+                raise ValueError(f"Invalid layer argument: {layer}. Error: {e}")
+        self.feature_extractor = ConformerMultiLayerFeatureExtractor(
+            self.encoder, aggregator=None, layer_idx_list=self.layer_idx_list
+        )
+
+    def forward(
+        self,
+        input_signal=None,
+        input_signal_length=None,
+        processed_signal=None,
+        processed_signal_length=None,
+    ):
+        """
+        Forward pass to extract features, same input interface as EncDecDenoiseMaskedTokenPredModel.forward
+        """
+        has_input_signal = input_signal is not None and input_signal_length is not None
+        has_processed_signal = processed_signal is not None and processed_signal_length is not None
+        if (has_input_signal ^ has_processed_signal) == False:
+            raise ValueError(
+                f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "
+                " with ``processed_signal`` and ``processed_signal_len`` arguments."
+            )
+        if not has_processed_signal:
+            processed_signal, processed_signal_length = self.preprocessor(
+                input_signal=input_signal,
+                length=input_signal_length,
+            )
+        encoded, encoded_len = self.feature_extractor(audio_signal=processed_signal, length=processed_signal_length)
+        return encoded, encoded_len
+
+
+def maybe_save_features(output_dir, results, max_cache, manifest):
+    """
+    Check if the cache is full and save features to disk
+    """
+    if len(results) == 0 or max_cache < 0 or len(results) < max_cache:
+        return
+    os.makedirs(output_dir, exist_ok=True)
+    logging.info(f"Saving {len(results)} features to {output_dir}")
+
+    for sample_id, audio_file, features_np in tqdm(results, desc="Saving features", total=len(results)):
+        filename = str(audio_file).replace("/", "_").replace(".", "_")
+        if len(filename) > 256:
+            filename = filename[-256:]
+        output_path = os.path.join(output_dir, f"{filename}.npy")
+        np.save(output_path, features_np)
+        manifest[sample_id]["feature_path"] = output_path
+
+    logging.info(f"Saved {len(results)} features to {output_dir}")
+    results.clear()
+
+
+def extract_features(args):
+    """
+    Main function to extract and save features from SSL model
+    """
+
+    logging.info(f"Extracting features using params: {vars(args)}")
+
+    # Load model
+    model = load_model(args.model_path)
+    feature_extractor = FeatureExtractor(model, args.layers)
+    device = torch.device(args.device)
+    feature_extractor.to(device)
+
+    # Load data
+    logging.info(f"Building dataset from input: {args.input}")
+    tmp_manifest = tempfile.NamedTemporaryFile(mode="w", delete=False)
+    manifest = get_input_manifest(args.input)
+    write_manifest(tmp_manifest.name, manifest)
+    total_num_samples = len(manifest)
+
+    # Build dataloader
+    config = {
+        "manifest_filepath": tmp_manifest.name,
+        "sample_rate": feature_extractor.sample_rate,
+        "return_sample_id": True,
+    }
+    dataset = get_char_dataset(config)
+    logging.info(f"Built dataset with {len(dataset)} samples")
+    dataloader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        collate_fn=dataset.collate_fn,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    # Extract features
+    indices = set()
+    results = []
+    amp_dtype = torch.float16 if args.amp_dtype == "float16" else torch.bfloat16
+    logging.info(f"Extracting features using AMP: {args.use_amp}, dtype: {amp_dtype}")
+    with torch.amp.autocast('cuda' if torch.cuda.is_available() else 'cpu', dtype=amp_dtype, enabled=args.use_amp):
+        with torch.inference_mode():
+            for batch in tqdm(dataloader, desc="Extracting features"):
+                batch = move_data_to_device(batch, device)
+                audio_signal, audio_signal_len, _, _, sample_id = batch
+                features, features_len = feature_extractor(
+                    input_signal=audio_signal, input_signal_length=audio_signal_len
+                )
+                batch_size = features[0].size(0)
+                num_layers = len(features)
+                for i in range(batch_size):
+                    sid_i = sample_id[i]
+                    if sid_i in indices:
+                        logging.warning(f"Skipping duplicated sample_id: {sample_id}")
+                        continue
+
+                    feat_i_len = features_len[0][i]
+                    feat_i = []
+                    for j in range(num_layers):
+                        feat_i.append(features[j][i][:, :feat_i_len])
+
+                    feat_i_np = torch.stack(feat_i, dim=0).cpu().numpy()
+
+                    indices.add(sid_i)
+                    results.append((sid_i, manifest[sid_i]['audio_filepath'], feat_i_np))
+
+                maybe_save_features(args.output, results, args.max_cache, manifest)
+
+    maybe_save_features(args.output, results, 0, manifest)
+
+    output_manifest = Path(args.output) / "features.json"
+    write_manifest(output_manifest, manifest)
+    os.remove(tmp_manifest.name)
+    logging.info(f"Extracted features from {total_num_samples} samples to {args.output}")
+    logging.info(f"Manifest saved to: {output_manifest}")
+
+
+if __name__ == "__main__":
+    extract_features(args)