# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. 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. import io import os from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional import librosa import soundfile as sf import torch.utils.data from nemo.collections.asr.data.audio_to_text import expand_sharded_filepaths from nemo.collections.asr.parts.preprocessing.segment import available_formats as valid_sf_formats from nemo.collections.asr.parts.utils.manifest_utils import read_manifest from nemo.collections.tts.parts.preprocessing.feature_processors import FeatureProcessor from nemo.collections.tts.parts.utils.tts_dataset_utils import ( filter_dataset_by_duration, get_weighted_sampler, load_audio, sample_audio, stack_tensors, ) from nemo.core.classes import Dataset, IterableDataset from nemo.utils import logging from nemo.utils import webdataset as wds from nemo.utils.decorators import experimental from nemo.utils.distributed import webdataset_split_by_workers VALID_FILE_FORMATS = ';'.join(['wav', 'mp3', 'flac', 'opus'] + [fmt.lower() for fmt in valid_sf_formats.keys()]) @dataclass class DatasetMeta: manifest_path: Path audio_dir: Path sample_weight: float = 1.0 audio_tar_filepaths: Optional[List[str]] = None @dataclass class DatasetSample: dataset_name: str manifest_entry: dict audio_dir: Path def audio_collate_fn(batch: List[dict]): dataset_name_list = [] audio_filepath_list = [] audio_list = [] audio_len_list = [] for example in batch: dataset_name_list.append(example["dataset_name"]) audio_filepath_list.append(example["audio_filepath"]) audio_list.append(example["audio"]) audio_len_list.append(example["audio_len"]) batch_audio_len = torch.IntTensor(audio_len_list) audio_max_len = int(batch_audio_len.max().item()) batch_audio = stack_tensors(audio_list, max_lens=[audio_max_len]) batch_dict = { "dataset_names": dataset_name_list, "audio_filepaths": audio_filepath_list, "audio": batch_audio, "audio_lens": batch_audio_len, } return batch_dict def preprocess_manifest( dataset_name: str, dataset: DatasetMeta, min_duration: float, max_duration: float, ): entries = read_manifest(dataset.manifest_path) filtered_entries, total_hours, filtered_hours = filter_dataset_by_duration( entries=entries, min_duration=min_duration, max_duration=max_duration ) logging.info(dataset_name) logging.info(f"Original # of files: {len(entries)}") logging.info(f"Filtered # of files: {len(filtered_entries)}") logging.info(f"Original duration: {total_hours:.2f} hours") logging.info(f"Filtered duration: {filtered_hours:.2f} hours") samples = [] sample_weights = [] for entry in filtered_entries: sample = DatasetSample(dataset_name=dataset_name, manifest_entry=entry, audio_dir=Path(dataset.audio_dir)) samples.append(sample) sample_weights.append(dataset.sample_weight) return samples, sample_weights @experimental class VocoderDataset(Dataset): """ Class for processing and loading Vocoder training examples. Args: dataset_meta: Dict of dataset names (string) to dataset metadata. sample_rate: Sample rate to load audio as. If the audio is stored at a different sample rate, then it will be resampled. n_samples: Optional int, if provided then n_samples samples will be randomly sampled from the full audio file. weighted_sampling_steps_per_epoch: Optional int, If provided, then data will be sampled (with replacement) based on the sample weights provided in the dataset metadata. If None, then sample weights will be ignored. feature_processors: Optional, list of feature processors to run on training examples. min_duration: Optional float, if provided audio files in the training manifest shorter than 'min_duration' will be ignored. max_duration: Optional float, if provided audio files in the training manifest longer than 'max_duration' will be ignored. trunc_duration: Optional int, if provided audio will be truncated to at most 'trunc_duration' seconds. volume_norm: Whether to apply volume normalization to loaded audio. """ def __init__( self, dataset_meta: Dict, sample_rate: int, n_samples: Optional[int] = None, weighted_sampling_steps_per_epoch: Optional[int] = None, feature_processors: Optional[Dict[str, FeatureProcessor]] = None, min_duration: Optional[float] = None, max_duration: Optional[float] = None, trunc_duration: Optional[float] = None, volume_norm: bool = False, ): super().__init__() self.sample_rate = sample_rate self.n_samples = n_samples self.trunc_duration = trunc_duration self.volume_norm = volume_norm self.weighted_sampling_steps_per_epoch = weighted_sampling_steps_per_epoch self.load_precomputed_mel = False if feature_processors: logging.info(f"Found feature processors {feature_processors.keys()}") self.feature_processors = list(feature_processors.values()) else: self.feature_processors = [] self.data_samples = [] self.sample_weights = [] for dataset_name, dataset_info in dataset_meta.items(): dataset = DatasetMeta(**dataset_info) samples, weights = preprocess_manifest( dataset_name=dataset_name, dataset=dataset, min_duration=min_duration, max_duration=max_duration, ) self.data_samples += samples self.sample_weights += weights def get_sampler(self, batch_size: int, world_size: int) -> Optional[torch.utils.data.Sampler]: if not self.weighted_sampling_steps_per_epoch: return None sampler = get_weighted_sampler( sample_weights=self.sample_weights, batch_size=batch_size, world_size=world_size, num_steps=self.weighted_sampling_steps_per_epoch, ) return sampler def __len__(self): return len(self.data_samples) def __getitem__(self, index): data = self.data_samples[index] if self.n_samples: audio_array, _, audio_filepath_rel = sample_audio( manifest_entry=data.manifest_entry, audio_dir=data.audio_dir, sample_rate=self.sample_rate, n_samples=self.n_samples, volume_norm=self.volume_norm, ) else: audio_array, _, audio_filepath_rel = load_audio( manifest_entry=data.manifest_entry, audio_dir=data.audio_dir, sample_rate=self.sample_rate, max_duration=self.trunc_duration, volume_norm=self.volume_norm, ) audio = torch.tensor(audio_array, dtype=torch.float32) audio_len = audio.shape[0] example = { "dataset_name": data.dataset_name, "audio_filepath": audio_filepath_rel, "audio": audio, "audio_len": audio_len, } for processor in self.feature_processors: processor.process(example) return example def collate_fn(self, batch): return audio_collate_fn(batch) class TarredVocoderDataset(IterableDataset): """ A similar Dataset to the VocoderDataset, but loads tarred audio files. Accepts a single comma-separated JSON manifest file (in the same style as for the VocoderDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, and duration of audio. Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...]. See the WebDataset documentation for more information about accepted data and input formats. If using multiple processes the number of shards should be divisible by the number of workers to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven! Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line. Args: dataset_meta: Dict of dataset names (string) to dataset metadata. audio_tar_filepaths: Either a list of audio tarball filepaths, or a string (can be brace-expandable). sample_rate: Sample rate to load audio as. If the audio is stored at a different sample rate, then it will be resampled. n_samples: Optional int, if provided then n_samples samples will be randomly sampled from the full audio file. shuffle_n (int): How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0. min_duration: Optional float, if provided audio files in the training manifest shorter than 'min_duration' will be ignored. max_duration: Optional float, if provided audio files in the training manifest longer than 'max_duration' will be ignored. trunc_duration: Optional int, if provided audio will be truncated to at most 'trunc_duration' seconds. feature_processors: Optional, list of feature processors to run on training examples. shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp. - `scatter`: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime. - `replicate`: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of `shuffle_n`. .. warning:: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets. global_rank (int): Worker rank, used for partitioning shards. Defaults to 0. world_size (int): Total number of processes, used for partitioning shards. Defaults to 0. """ def __init__( self, dataset_meta: Dict, sample_rate: int, n_samples: Optional[int] = None, shuffle_n: int = 0, min_duration: float = 0.1, max_duration: Optional[float] = None, trunc_duration: Optional[float] = None, feature_processors: Optional[Dict[str, FeatureProcessor]] = None, shard_strategy: str = "scatter", global_rank: int = 0, world_size: int = 2, **kwargs, ): super().__init__() if len(kwargs) > 0: logging.warning( f"Arguments {kwargs.keys()} does not support for TarredVocoderDataset, they will be ignored." ) self.sample_rate = sample_rate self.n_samples = n_samples if trunc_duration: self.trunc_samples = int(trunc_duration * self.sample_rate) else: self.trunc_samples = None if feature_processors: logging.info(f"Found feature processors {feature_processors.keys()}") self.feature_processors = list(feature_processors.values()) else: self.feature_processors = [] self.data_samples = [] self.audio_tar_filepaths = [] for dataset_name, dataset_info in dataset_meta.items(): audio_tar_filepaths = dataset_info.audio_tar_filepaths self.audio_tar_filepaths += [audio_tar_filepaths] dataset = DatasetMeta(**dataset_info) samples, _ = preprocess_manifest( dataset_name=dataset_name, dataset=dataset, min_duration=min_duration, max_duration=max_duration, ) self.data_samples += samples self.file_id_to_sample_map = {} for sample in self.data_samples: file_id = os.path.splitext(os.path.basename(sample.manifest_entry["audio_filepath"]))[0] if file_id not in self.file_id_to_sample_map: self.file_id_to_sample_map[file_id] = sample else: raise ValueError( f"Duplicate file_id {file_id} found in manifest {sample.manifest_entry['audio_filepath']}" ) logging.info(f"world size: {world_size}") audio_tar_filepaths = expand_sharded_filepaths( sharded_filepaths=audio_tar_filepaths, global_rank=global_rank, world_size=world_size, shard_strategy=shard_strategy, ) self._dataset = wds.DataPipeline( wds.SimpleShardList(urls=audio_tar_filepaths), webdataset_split_by_workers, wds.shuffle(shuffle_n), wds.tarfile_to_samples(), wds.rename(audio=VALID_FILE_FORMATS, key='__key__'), wds.to_tuple('audio', 'key'), self._filter, wds.map(self._build_sample), ) def _filter(self, iterator): class FilteredIterator: def __init__(self, file_id_to_sample_map): self.iterator = iterator self.file_id_to_sample_map = file_id_to_sample_map def __iter__(self): return self def __next__(self): while True: audio_bytes, audio_filename = next(self.iterator) file_id = os.path.splitext(os.path.basename(audio_filename))[0] if file_id in self.file_id_to_sample_map: return audio_bytes, audio_filename return FilteredIterator(self.file_id_to_sample_map) def _build_sample(self, tup): audio_bytes, audio_filename = tup file_id = os.path.splitext(os.path.basename(audio_filename))[0] data = self.file_id_to_sample_map[file_id] audio_array, sr = sf.read(file=io.BytesIO(audio_bytes), dtype='float32') if sr != self.sample_rate: logging.warning( f"Sample rate of {sr} does not match target sample rate of {self.sample_rate}. Resampling audio." ) audio_array = librosa.core.resample(audio_array, orig_sr=sr, target_sr=self.sample_rate) audio_array = torch.from_numpy(audio_array) if self.n_samples: len_audio = audio_array.shape[0] if len_audio > self.n_samples: start = torch.randint(0, len_audio - self.n_samples, (1,)) audio_array = audio_array[start : start + self.n_samples] else: audio_array = audio_array[: self.n_samples] if self.trunc_samples: audio_array = audio_array[: self.trunc_samples] audio_len = torch.tensor(audio_array.shape[0]) example = { "dataset_name": data.dataset_name, "audio_filepath": audio_filename, "audio": audio_array, "audio_len": audio_len, } for processor in self.feature_processors: processor.process(example) return example def get_sampler(self, batch_size: int, world_size: int): """ Currently sampler is not supported for tarred dataset. """ return None def collate_fn(self, batch): return audio_collate_fn(batch) def __iter__(self): return self._dataset.__iter__() def __len__(self): return len(self.file_id_to_sample_map)