def __init__(
self,
manifest_filepath: str,
parser: Union[str, Callable],
sample_rate: int,
int_values: bool = False,
augmentor: 'nemo.collections.asr.parts.perturb.AudioAugmentor' = None,
max_duration: Optional[int] = None,
min_duration: Optional[int] = None,
max_utts: int = 0,
trim: bool = False,
bos_id: Optional[int] = None,
eos_id: Optional[int] = None,
pad_id: int = 0,
):
self.parser = parser
self.collection = collections.ASRAudioText(
manifests_files=manifest_filepath.split(','),
parser=parser,
min_duration=min_duration,
max_duration=max_duration,
max_number=max_utts,
)
self.featurizer = WaveformFeaturizer(sample_rate=sample_rate,
int_values=int_values,
augmentor=augmentor)
self.trim = trim
self.eos_id = eos_id
self.bos_id = bos_id
self.pad_id = pad_id
def _setup_dataloader_from_config(self, config: Optional[Dict]):
if config.get('manifest_filepath') is None:
return
if 'augmentor' in config:
augmentor = process_augmentations(config['augmentor'])
else:
augmentor = None
featurizer = WaveformFeaturizer(sample_rate=config['sample_rate'],
int_values=config.get(
'int_values', False),
augmentor=augmentor)
dataset = AudioLabelDataset(
manifest_filepath=config['manifest_filepath'],
labels=config['labels'],
featurizer=featurizer,
max_duration=config.get('max_duration', None),
min_duration=config.get('min_duration', None),
trim=config.get('trim_silence', True),
load_audio=config.get('load_audio', True),
)
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=config['batch_size'],
collate_fn=dataset.collate_fn,
drop_last=config.get('drop_last', False),
shuffle=config['shuffle'],
num_workers=config.get('num_workers', 0),
pin_memory=config.get('pin_memory', False),
)
def __init__(self,
*,
manifest_filepath: str,
labels: List[str],
batch_size: int,
sample_rate: int = 16000,
int_values: bool = False,
num_workers: int = 0,
shuffle: bool = True,
min_duration: Optional[float] = 0.1,
max_duration: Optional[float] = None,
trim_silence: bool = False,
drop_last: bool = False,
load_audio: bool = True,
augmentor: Optional[Union[AudioAugmentor,
Dict[str, Dict[str, Any]]]] = None,
num_classes: int = 35,
class_dists=None,
class_probs=None,
probs_num=0):
super(BalancedAudioToSpeechLabelDataLayer, self).__init__()
self._manifest_filepath = manifest_filepath
self._labels = labels
self._sample_rate = sample_rate
if augmentor is not None:
augmentor = self._process_augmentations(augmentor)
self._featurizer = WaveformFeaturizer(sample_rate=sample_rate,
int_values=int_values,
augmentor=augmentor)
dataset_params = {
'manifest_filepath': manifest_filepath,
'labels': labels,
'featurizer': self._featurizer,
'max_duration': max_duration,
'min_duration': min_duration,
'trim': trim_silence,
'load_audio': load_audio,
}
self._dataset = AudioLabelDataset(**dataset_params)
labels = []
for sample in self._dataset.collection:
labels.append(self._dataset.label2id[sample.label])
self._dataloader = torch.utils.data.DataLoader(
dataset=self._dataset,
batch_sampler=BalancedBatchSampler(labels,
n_classes=num_classes,
n_samples=batch_size //
num_classes,
class_dists=class_dists,
class_probs=class_probs,
probs_num=probs_num),
# TODO replace with kwargs
collate_fn=partial(seq_collate_fn, token_pad_value=0),
num_workers=num_workers,
)
def _setup_dataloader_from_config(self, config: Optional[Dict]):
if config.get('manifest_filepath') is None:
return
if 'augmentor' in config:
augmentor = process_augmentations(config['augmentor'])
else:
augmentor = None
featurizer = WaveformFeaturizer(sample_rate=config['sample_rate'],
int_values=config.get(
'int_values', False),
augmentor=augmentor)
if 'vad_stream' in config and config['vad_stream']:
print("Perform streaming frame-level VAD")
dataset = AudioToSpeechLabelDataSet(
manifest_filepath=config['manifest_filepath'],
labels=config['labels'],
featurizer=featurizer,
max_duration=config.get('max_duration', None),
min_duration=config.get('min_duration', None),
trim=config.get('trim_silence', True),
load_audio=config.get('load_audio', True),
time_length=config.get('time_length', 0.31),
shift_length=config.get('shift_length', 0.01),
)
batch_size = 1
collate_func = dataset.vad_frame_seq_collate_fn
else:
dataset = AudioLabelDataset(
manifest_filepath=config['manifest_filepath'],
labels=config['labels'],
featurizer=featurizer,
max_duration=config.get('max_duration', None),
min_duration=config.get('min_duration', None),
trim=config.get('trim_silence', True),
load_audio=config.get('load_audio', True),
)
batch_size = config['batch_size']
collate_func = dataset.collate_fn
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
collate_fn=collate_func,
drop_last=config.get('drop_last', False),
shuffle=config['shuffle'],
num_workers=config.get('num_workers', 0),
pin_memory=config.get('pin_memory', False),
)
def __init__(self,
audio_batch: List[Union[str, BytesIO]],
sample_rate: int,
int_values: bool,
trim=False) -> None:
"""Dataset reader for AudioInferDataLayer.
Args:
audio_batch: Batch to be read. Elements could be either paths to audio files or Binary I/O objects.
sample_rate: Audio files sample rate.
int_values: If true, load samples as 32-bit integers.
trim: Trim leading and trailing silence from an audio signal if True.
"""
self.audio_batch = audio_batch
self.featurizer = WaveformFeaturizer(sample_rate=sample_rate,
int_values=int_values)
self.trim = trim
def __setup_dataloader_from_config(self, config: Optional[Dict]):
if 'augmentor' in config:
augmentor = process_augmentations(config['augmentor'])
else:
augmentor = None
featurizer = WaveformFeaturizer(sample_rate=config['sample_rate'],
int_values=config.get(
'int_values', False),
augmentor=augmentor)
self.dataset = AudioToSpeechLabelDataset(
manifest_filepath=config['manifest_filepath'],
labels=config['labels'],
featurizer=featurizer,
max_duration=config.get('max_duration', None),
min_duration=config.get('min_duration', None),
trim=False,
load_audio=config.get('load_audio', True),
time_length=config.get('time_length', 8),
shift_length=config.get('shift_length', 0.75),
)
if self.task == 'diarization':
logging.info("Setting up diarization parameters")
_collate_func = self.dataset.sliced_seq_collate_fn
batch_size = 1
shuffle = False
else:
logging.info("Setting up identification parameters")
_collate_func = self.dataset.fixed_seq_collate_fn
batch_size = config['batch_size']
shuffle = config.get('shuffle', False)
return torch.utils.data.DataLoader(
dataset=self.dataset,
batch_size=batch_size,
collate_fn=_collate_func,
drop_last=config.get('drop_last', False),
shuffle=shuffle,
num_workers=config.get('num_workers', 0),
pin_memory=config.get('pin_memory', False),
)
def test_tarred_dataset_duplicate_name(self, test_data_dir):
manifest_path = os.path.abspath(
os.path.join(
test_data_dir,
'asr/tarred_an4/tarred_duplicate_audio_manifest.json'))
# Test braceexpand loading
tarpath = os.path.abspath(
os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
featurizer = WaveformFeaturizer(sample_rate=16000,
int_values=False,
augmentor=None)
ds_braceexpand = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath,
manifest_filepath=manifest_path,
labels=self.labels,
featurizer=featurizer)
assert len(ds_braceexpand) == 6
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 6
# Test loading via list
tarpath = [
os.path.abspath(
os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar'))
for i in range(2)
]
ds_list_load = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath,
manifest_filepath=manifest_path,
labels=self.labels,
featurizer=featurizer)
count = 0
for _ in ds_list_load:
count += 1
assert count == 6
def _setup_dataloader_from_config(self, config: DictConfig):
OmegaConf.set_struct(config, False)
config.is_regression_task = self.is_regression_task
OmegaConf.set_struct(config, True)
if 'augmentor' in config:
augmentor = process_augmentations(config['augmentor'])
else:
augmentor = None
featurizer = WaveformFeaturizer(sample_rate=config['sample_rate'],
int_values=config.get(
'int_values', False),
augmentor=augmentor)
shuffle = config['shuffle']
# Instantiate tarred dataset loader or normal dataset loader
if config.get('is_tarred', False):
if ('tarred_audio_filepaths' in config
and config['tarred_audio_filepaths'] is None) or (
'manifest_filepath' in config
and config['manifest_filepath'] is None):
logging.warning(
"Could not load dataset as `manifest_filepath` is None or "
f"`tarred_audio_filepaths` is None. Provided config : {config}"
)
return None
if 'vad_stream' in config and config['vad_stream']:
logging.warning(
"VAD inference does not support tarred dataset now")
return None
shuffle_n = config.get('shuffle_n', 4 *
config['batch_size']) if shuffle else 0
dataset = audio_to_label_dataset.get_tarred_classification_label_dataset(
featurizer=featurizer,
config=OmegaConf.to_container(config),
shuffle_n=shuffle_n,
global_rank=self.global_rank,
world_size=self.world_size,
)
shuffle = False
batch_size = config['batch_size']
collate_func = dataset.collate_fn
else:
if 'manifest_filepath' in config and config[
'manifest_filepath'] is None:
logging.warning(
f"Could not load dataset as `manifest_filepath` is None. Provided config : {config}"
)
return None
if 'vad_stream' in config and config['vad_stream']:
logging.info("Perform streaming frame-level VAD")
dataset = audio_to_label_dataset.get_speech_label_dataset(
featurizer=featurizer,
config=OmegaConf.to_container(config))
batch_size = 1
collate_func = dataset.vad_frame_seq_collate_fn
else:
dataset = audio_to_label_dataset.get_classification_label_dataset(
featurizer=featurizer,
config=OmegaConf.to_container(config))
batch_size = config['batch_size']
collate_func = dataset.collate_fn
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
collate_fn=collate_func,
drop_last=config.get('drop_last', False),
shuffle=shuffle,
num_workers=config.get('num_workers', 0),
pin_memory=config.get('pin_memory', False),
)
def __init__(
self,
audio_tar_filepaths: Union[str, List[str]],
manifest_filepath: str,
parser: Callable,
sample_rate: int,
int_values: bool = False,
augmentor: Optional[
'nemo.collections.asr.parts.perturb.AudioAugmentor'] = None,
shuffle_n: int = 0,
min_duration: Optional[float] = None,
max_duration: Optional[float] = None,
max_utts: int = 0,
trim: bool = False,
bos_id: Optional[int] = None,
eos_id: Optional[int] = None,
add_misc: bool = False,
pad_id: int = 0,
shard_strategy: str = "scatter",
global_rank: int = 0,
world_size: int = 0,
):
self.collection = collections.ASRAudioText(
manifests_files=manifest_filepath.split(','),
parser=parser,
min_duration=min_duration,
max_duration=max_duration,
max_number=max_utts,
index_by_file_id=
True, # Must set this so the manifest lines can be indexed by file ID
)
self.featurizer = WaveformFeaturizer(sample_rate=sample_rate,
int_values=int_values,
augmentor=augmentor)
self.trim = trim
self.eos_id = eos_id
self.bos_id = bos_id
self.pad_id = pad_id
self._add_misc = add_misc
valid_shard_strategies = ['scatter', 'replicate']
if shard_strategy not in valid_shard_strategies:
raise ValueError(
f"`shard_strategy` must be one of {valid_shard_strategies}")
if isinstance(audio_tar_filepaths, str):
# Replace '(' and '[' with '{'
brace_keys_open = ['(', '[', '<', '_OP_']
for bkey in brace_keys_open:
if bkey in audio_tar_filepaths:
audio_tar_filepaths = audio_tar_filepaths.replace(
bkey, "{")
# Replace ')' and ']' with '}'
brace_keys_close = [')', ']', '>', '_CL_']
for bkey in brace_keys_close:
if bkey in audio_tar_filepaths:
audio_tar_filepaths = audio_tar_filepaths.replace(
bkey, "}")
# Check for distributed and partition shards accordingly
if world_size > 1:
if isinstance(audio_tar_filepaths, str):
# Brace expand
audio_tar_filepaths = list(
braceexpand.braceexpand(audio_tar_filepaths))
if shard_strategy == 'scatter':
logging.info(
"All tarred dataset shards will be scattered evenly across all nodes."
)
if len(audio_tar_filepaths) % world_size != 0:
logging.warning(
f"Number of shards in tarred dataset ({len(audio_tar_filepaths)}) is not divisible "
f"by number of distributed workers ({world_size}).")
begin_idx = (len(audio_tar_filepaths) //
world_size) * global_rank
end_idx = begin_idx + (len(audio_tar_filepaths) // world_size)
audio_tar_filepaths = audio_tar_filepaths[begin_idx:end_idx]
logging.info(
"Partitioning tarred dataset: process (%d) taking shards [%d, %d)",
global_rank, begin_idx, end_idx)
elif shard_strategy == 'replicate':
logging.info(
"All tarred dataset shards will be replicated across all nodes."
)
else:
raise ValueError(
f"Invalid shard strategy ! Allowed values are : {valid_shard_strategies}"
)
# Put together WebDataset
self._dataset = (
wd.Dataset(audio_tar_filepaths).shuffle(shuffle_n).rename(
audio='wav', key='__key__').to_tuple('audio', 'key').pipe(
self._filter).map(f=self._build_sample))
def __init__(
self,
manifest_filepath: str,
mappings_filepath: str,
sample_rate: int,
max_duration: Optional[float] = None,
min_duration: Optional[float] = None,
ignore_file: Optional[str] = None,
trim: bool = False,
load_supplementary_values=True, # Set to False for validation
):
"""
Dataset that loads audio, phonemes and their durations, pitches per frame, and energies per frame
for FastSpeech 2 from paths described in a JSON manifest (see the AudioDataset documentation for details
on the manifest format), as well as a mappings file for word to phones and phones to indices.
The text in the manifest is ignored; instead, the phoneme indices for prediction come from the
duration files.
For each sample, paths for duration, energy, and pitch files are inferred from the manifest's audio
filepaths by replacing '/wavs' with '/phoneme_durations', '/pitches', and '/energies', and swapping out
the file extension to '.pt', '.npy', and '.npy' respectively.
For example, given manifest audio path `/data/LJSpeech/wavs/LJ001-0001.wav`, the inferred duration and
phonemes file path would be `/data/LJSpeech/phoneme_durations/LJ001-0001.pt`.
Note that validation datasets only need the audio files and phoneme & duration files, set
`load_supplementary_values` to False for validation sets.
Args:
manifest_filepath (str): Path to the JSON manifest file that lists audio files.
mappings_filepath (str): Path to a JSON mappings file that contains mappings "word2phones" and
"phone2idx". The latter is used to determine the padding index.
sample_rate (int): Target sample rate of the audio.
max_duration (float): If audio exceeds this length in seconds, it is filtered from the dataset.
Defaults to None, which does not filter any audio.
min_duration (float): If audio is shorter than this length in seconds, it is filtered from the dataset.
Defaults to None, which does not filter any audio.
ignore_file (str): Optional pickled file which contains a list of files to ignore (e.g. files that
contain OOV words).
Defaults to None.
trim (bool): Whether to use librosa.effects.trim on the audio clip.
Defaults to False.
load_supplementary_values (bool): Whether or not to load pitch and energy files. Set this to False for
validation datasets.
Defaults to True.
"""
super().__init__()
# Retrieve mappings from file
with open(mappings_filepath, 'r') as f:
mappings = json.load(f)
self.word2phones = mappings['word2phones']
self.phone2idx = mappings['phone2idx']
# Load data from manifests
audio_files = []
total_duration = 0
if isinstance(manifest_filepath, str):
manifest_filepath = [manifest_filepath]
for manifest_file in manifest_filepath:
with open(expanduser(manifest_file), 'r') as f:
logging.info(f"Loading dataset from {manifest_file}.")
for line in f:
item = json.loads(line)
audio_files.append({
"audio_filepath": item["audio_filepath"],
"duration": item["duration"]
})
total_duration += item["duration"]
total_dataset_len = len(audio_files)
logging.info(
f"Loaded dataset with {total_dataset_len} files totalling {total_duration/3600:.2f} hours."
)
self.data = []
if load_supplementary_values:
dataitem = py_collections.namedtuple(
typename='AudioTextEntity',
field_names='audio_file duration text_tokens pitches energies')
else:
dataitem = py_collections.namedtuple(
typename='AudioTextEntity',
field_names='audio_file duration text_tokens')
if ignore_file:
logging.info(f"using {ignore_file} to prune dataset.")
with open(ignore_file, "rb") as f:
wavs_to_ignore = set(pickle.load(f))
pruned_duration = 0
pruned_items = 0
for item in audio_files:
audio_path = item['audio_filepath']
LJ_id = os.path.splitext(os.path.basename(audio_path))[0]
# Prune data according to min/max_duration & the ignore file
if (min_duration and item["duration"] < min_duration) or (
max_duration and item["duration"] > max_duration):
pruned_duration += item["duration"]
pruned_items += 1
continue
if ignore_file and (LJ_id in wavs_to_ignore):
pruned_items += 1
pruned_duration += item["duration"]
wavs_to_ignore.remove(LJ_id)
continue
# Else not pruned, load additional info
# Phoneme durations and text token indices from durations file
dur_path = audio_path.replace('/wavs/',
'/phoneme_durations/').replace(
'.wav', '.pt')
duration_info = torch.load(dur_path)
durs = duration_info['token_duration']
text_tokens = duration_info['text_encoded']
if load_supplementary_values:
# Load pitch file (F0s)
pitch_path = audio_path.replace('/wavs/', '/pitches/').replace(
'.wav', '.npy')
pitches = torch.from_numpy(
np.load(pitch_path).astype(dtype='float32'))
# Load energy file (L2-norm of the amplitude of each STFT frame of an utterance)
energies_path = audio_path.replace('/wavs/',
'/energies/').replace(
'.wav', '.npy')
energies = torch.from_numpy(np.load(energies_path))
self.data.append(
dataitem(
audio_file=item['audio_filepath'],
duration=durs,
pitches=torch.clamp(pitches, min=1e-5),
energies=energies,
text_tokens=text_tokens,
))
else:
self.data.append(
dataitem(
audio_file=item['audio_filepath'],
duration=durs,
text_tokens=text_tokens,
))
logging.info(
f"Pruned {pruned_items} files and {pruned_duration/3600:.2f} hours."
)
logging.info(
f"Final dataset contains {len(self.data)} files and {(total_duration-pruned_duration)/3600:.2f} hours."
)
self.featurizer = WaveformFeaturizer(sample_rate=sample_rate)
self.trim = trim
self.load_supplementary_values = load_supplementary_values
def main(cfg):
logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
enrollment_manifest = cfg.data.enrollment_manifest
test_manifest = cfg.data.test_manifest
out_manifest = cfg.data.out_manifest
sample_rate = cfg.data.sample_rate
backend = cfg.backend.backend_model.lower()
if backend == 'cosine_similarity':
model_path = cfg.backend.cosine_similarity.model_path
batch_size = cfg.backend.cosine_similarity.batch_size
if model_path.endswith('.nemo'):
speaker_model = EncDecSpeakerLabelModel.restore_from(model_path)
else:
speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_path)
enroll_embs, _, enroll_truelabels, enroll_id2label = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
enrollment_manifest,
batch_size,
sample_rate,
device=device,
)
test_embs, _, _, _ = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
test_manifest,
batch_size,
sample_rate,
device=device,
)
# length normalize
enroll_embs = enroll_embs / (np.linalg.norm(
enroll_embs, ord=2, axis=-1, keepdims=True))
test_embs = test_embs / (np.linalg.norm(
test_embs, ord=2, axis=-1, keepdims=True))
# reference embedding
reference_embs = []
keyslist = list(enroll_id2label.keys())
for label_id in keyslist:
indices = np.where(enroll_truelabels == label_id)
embedding = (enroll_embs[indices].sum(
axis=0).squeeze()) / len(indices)
reference_embs.append(embedding)
reference_embs = np.asarray(reference_embs)
scores = np.matmul(test_embs, reference_embs.T)
matched_labels = scores.argmax(axis=-1)
elif backend == 'neural_classifier':
model_path = cfg.backend.neural_classifier.model_path
batch_size = cfg.backend.neural_classifier.batch_size
if model_path.endswith('.nemo'):
speaker_model = EncDecSpeakerLabelModel.restore_from(model_path)
else:
speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_path)
featurizer = WaveformFeaturizer(sample_rate=sample_rate)
dataset = AudioToSpeechLabelDataset(
manifest_filepath=enrollment_manifest,
labels=None,
featurizer=featurizer)
enroll_id2label = dataset.id2label
if speaker_model.decoder.final.out_features != len(enroll_id2label):
raise ValueError(
"number of labels mis match. Make sure you trained or finetuned neural classifier with labels from enrollement manifest_filepath"
)
_, test_logits, _, _ = EncDecSpeakerLabelModel.get_batch_embeddings(
speaker_model,
test_manifest,
batch_size,
sample_rate,
device=device,
)
matched_labels = test_logits.argmax(axis=-1)
with open(test_manifest, 'rb') as f1, open(out_manifest,
'w',
encoding='utf-8') as f2:
lines = f1.readlines()
for idx, line in enumerate(lines):
line = line.strip()
item = json.loads(line)
item['infer'] = enroll_id2label[matched_labels[idx]]
json.dump(item, f2)
f2.write('\n')
logging.info(
"Inference labels have been written to {} manifest file".format(
out_manifest))