# Copyright (c) 2023, 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. import os import numpy as np import pytest import torch from omegaconf import DictConfig from nemo.collections.asr.parts.utils.data_simulation_utils import ( DataAnnotator, SpeechSampler, add_silence_to_alignments, binary_search_alignments, get_cleaned_base_path, get_split_points_in_alignments, normalize_audio, read_noise_manifest, ) from nemo.collections.asr.parts.utils.manifest_utils import get_ctm_line @pytest.fixture() def annotator(): cfg = get_data_simulation_configs() return DataAnnotator(cfg) @pytest.fixture() def sampler(): cfg = get_data_simulation_configs() sampler = SpeechSampler(cfg) # Must get session-wise randomized silence/overlap mean sampler.get_session_overlap_mean() sampler.get_session_silence_mean() return sampler def get_data_simulation_configs(): config_dict = { 'data_simulator': { 'manifest_filepath': '???', 'sr': 16000, 'random_seed': 42, 'multiprocessing_chunksize': 10000, 'session_config': {'num_speakers': 4, 'num_sessions': 60, 'session_length': 600}, 'session_params': { 'max_audio_read_sec': 20, 'sentence_length_params': [0.4, 0.05], 'dominance_var': 0.11, 'min_dominance': 0.05, 'turn_prob': 0.875, 'min_turn_prob': 0.5, 'mean_silence': 0.15, 'mean_silence_var': 0.01, 'per_silence_var': 900, 'per_silence_min': 0.0, 'per_silence_max': -1, 'mean_overlap': 0.1, 'mean_overlap_var': 0.01, 'per_overlap_var': 900, 'per_overlap_min': 0.0, 'per_overlap_max': -1, 'start_window': True, 'window_type': 'hamming', 'window_size': 0.05, 'start_buffer': 0.1, 'split_buffer': 0.1, 'release_buffer': 0.1, 'normalize': True, 'normalization_type': 'equal', 'normalization_var': 0.1, 'min_volume': 0.75, 'max_volume': 1.25, 'end_buffer': 0.5, }, 'outputs': { 'output_dir': '???', 'output_filename': 'multispeaker_session', 'overwrite_output': True, 'output_precision': 3, }, 'background_noise': { 'add_bg': False, 'background_manifest': None, 'num_noise_files': 10, 'snr': 60, 'snr_min': None, }, 'segment_augmentor': { 'add_seg_aug': False, 'augmentor': { 'gain': {'prob': 0.5, 'min_gain_dbfs': -10.0, 'max_gain_dbfs': 10.0}, }, }, 'session_augmentor': { 'add_sess_aug': False, 'augmentor': { 'white_noise': {'prob': 1.0, 'min_level': -90, 'max_level': -46}, }, }, 'speaker_enforcement': {'enforce_num_speakers': True, 'enforce_time': [0.25, 0.75]}, 'segment_manifest': {'window': 0.5, 'shift': 0.25, 'step_count': 50, 'deci': 3}, } } return DictConfig(config_dict) def generate_words_and_alignments(sample_index): if sample_index == 0: words = ['', 'hello', 'world'] alignments = [0.5, 1.0, 1.5] elif sample_index == 1: words = ["", "stephanos", "dedalos", ""] alignments = [0.51, 1.31, 2.04, 2.215] elif sample_index == 2: words = ['', 'hello', 'world', '', 'welcome', 'to', 'nemo', ''] alignments = [0.5, 1.0, 1.5, 1.7, 1.8, 2.2, 2.7, 2.8] else: raise ValueError(f"sample_index {sample_index} not supported") speaker_id = 'speaker_0' return words, alignments, speaker_id class TestGetCtmLine: @pytest.mark.unit @pytest.mark.parametrize("conf", [0, 1]) def test_wrong_type_conf_values(self, conf): # Test with wrong integer confidence values with pytest.raises(ValueError): result = get_ctm_line( source="test_source", channel=1, start_time=0.123, duration=0.456, token="word", conf=conf, type_of_token="lex", speaker="speaker1", ) expected = f"test_source 1 0.12 0.46 word {conf} lex speaker1\n" assert result == expected, f"Failed on valid conf value {conf}" @pytest.mark.unit @pytest.mark.parametrize("conf", [0.0, 0.5, 1.0, 0.01, 0.99]) def test_valid_conf_values(self, conf): # Test with valid confidence values output_precision = 2 result = get_ctm_line( source="test_source", channel=1, start_time=0.123, duration=0.456, token="word", conf=conf, type_of_token="lex", speaker="speaker1", output_precision=output_precision, ) expected = "test_source 1 0.12 0.46 word" + f" {conf:.{output_precision}f} lex speaker1\n" assert result == expected, f"Failed on valid conf value {conf}" @pytest.mark.unit @pytest.mark.parametrize("conf", [-0.1, 1.1, 2, -1, 100, -100]) def test_invalid_conf_ranges(self, conf): # Test with invalid confidence values with pytest.raises(ValueError): get_ctm_line( source="test_source", channel=1, start_time=0.123, duration=0.456, token="word", conf=conf, type_of_token="lex", speaker="speaker1", ) @pytest.mark.unit @pytest.mark.parametrize( "start_time, duration, output_precision", [(0.123, 0.456, 2), (1.0, 2.0, 1), (0.0, 0.0, 2), (0.01, 0.99, 3), (1.23, 4.56, 2)], ) def test_valid_start_time_duration_with_precision(self, start_time, duration, output_precision): # Test with valid beginning time, duration values and output precision confidence = 0.5 result = get_ctm_line( source="test_source", channel=1, start_time=start_time, duration=duration, token="word", conf=confidence, type_of_token="lex", speaker="speaker1", output_precision=output_precision, ) expected_start_time = ( f"{start_time:.{output_precision}f}" # Adjusted to match the output format with precision ) expected_duration = f"{duration:.{output_precision}f}" # Adjusted to match the output format with precision expected_confidence = ( f"{confidence:.{output_precision}f}" # Adjusted to match the output format with precision ) expected = f"test_source 1 {expected_start_time} {expected_duration} word {expected_confidence} lex speaker1\n" assert ( result == expected ), f"Failed on valid start_time {start_time}, duration {duration} with precision {output_precision}" @pytest.mark.unit def test_valid_input(self): # Test with completely valid inputs result = get_ctm_line( source="test_source", channel=1, start_time=0.123, duration=0.456, token="word", conf=0.789, type_of_token="lex", speaker="speaker1", ) expected = "test_source 1 0.12 0.46 word 0.79 lex speaker1\n" assert result == expected, "Failed on valid input" @pytest.mark.unit @pytest.mark.parametrize( "start_time, duration", [ ("not a float", 1.0), (1.0, "not a float"), (1, 2.0), # Integers should be converted to float (2.0, 3), # Same as above ], ) def test_invalid_types_for_time_duration(self, start_time, duration): # Test with invalid types for start_time and duration with pytest.raises(ValueError): get_ctm_line( source="test_source", channel=1, start_time=start_time, duration=duration, token="word", conf=0.5, type_of_token="lex", speaker="speaker1", ) @pytest.mark.unit @pytest.mark.parametrize("conf", [-0.1, 1.1, "not a float"]) def test_invalid_conf_values(self, conf): # Test with invalid values for conf with pytest.raises(ValueError): get_ctm_line( source="test_source", channel=1, start_time=0.123, duration=0.456, token="word", conf=conf, type_of_token="lex", speaker="speaker1", ) @pytest.mark.unit def test_default_values(self): # Test with missing optional parameters result = get_ctm_line( source="test_source", channel=None, start_time=0.123, duration=0.456, token="word", conf=None, type_of_token=None, speaker=None, ) expected = "test_source 1 0.12 0.46 word NA unknown NA\n" assert result == expected, "Failed on default values" class TestDataSimulatorUtils: # TODO: add tests for all util functions @pytest.mark.parametrize("max_audio_read_sec", [2.5, 3.5, 4.5]) @pytest.mark.parametrize("min_alignment_count", [2, 3, 4]) def test_binary_search_alignments(self, max_audio_read_sec, min_alignment_count): inds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] alignments = [0.5, 11.0, 11.5, 12.0, 13.0, 14.0, 14.5, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 30, 40.0] offset_max = binary_search_alignments(inds, max_audio_read_sec, min_alignment_count, alignments) assert max_audio_read_sec <= alignments[-1 * min_alignment_count] - alignments[inds[offset_max]] @pytest.mark.parametrize("sample_len", [100, 16000]) @pytest.mark.parametrize("gain", [0.1, 0.5, 1.0, 2.0, 5.0]) def test_normalize_audio(self, sample_len, gain): array_raw = np.random.randn(sample_len) array_input = torch.from_numpy(gain * array_raw / np.max(np.abs(array_raw))) norm_array = normalize_audio(array_input) assert torch.max(torch.abs(norm_array)) == 1.0 assert torch.min(torch.abs(norm_array)) < 1.0 @pytest.mark.parametrize("output_dir", [os.path.join(os.getcwd(), "test_dir")]) def test_get_cleaned_base_path(self, output_dir): result_path = get_cleaned_base_path(output_dir, overwrite_output=True) assert os.path.exists(result_path) and not os.path.isfile(result_path) result_path = get_cleaned_base_path(output_dir, overwrite_output=False) assert os.path.exists(result_path) and not os.path.isfile(result_path) os.rmdir(result_path) assert not os.path.exists(result_path) @pytest.mark.parametrize( "words, alignments, answers", [ (['', 'hello', 'world'], [0.5, 1.0, 1.5], [[0, 16000.0]]), ( ['', 'hello', 'world', '', 'welcome', 'to', 'nemo', ''], [0.27, 1.0, 1.7, 2.7, 2.8, 3.2, 3.7, 3.9], [[0, (1.7 + 0.5) * 16000], [(2.7 - 0.5) * 16000, (3.9 - 0.27) * 16000]], ), ], ) @pytest.mark.parametrize("sr", [16000]) @pytest.mark.parametrize("split_buffer", [0.5]) @pytest.mark.parametrize("new_start", [0.0]) def test_get_split_points_in_alignments(self, words, alignments, sr, new_start, split_buffer, answers): sentence_audio_len = sr * (alignments[-1] - alignments[0]) splits = get_split_points_in_alignments(words, alignments, split_buffer, sr, sentence_audio_len, new_start) assert len(splits) == len(answers) for k, interval in enumerate(splits): assert abs(answers[k][0] - interval[0]) < 1e-4 assert abs(answers[k][1] - interval[1]) < 1e-4 @pytest.mark.parametrize( "alignments, words", [(['hello', 'world'], [1.0, 1.5]), (['', 'hello', 'world'], [0.0, 1.0, 1.5])] ) def test_add_silence_to_alignments(self, alignments, words): """ Test add_silence_to_alignments function. """ audio_manifest = { 'audio_filepath': 'test.wav', 'alignments': alignments, 'words': words, } audio_manifest = add_silence_to_alignments(audio_manifest) if words[0] == '': assert audio_manifest['alignments'] == [0.0] + alignments assert audio_manifest['words'] == [''] + words else: assert audio_manifest['alignments'] == alignments assert audio_manifest['words'] == words class TestDataAnnotator: def test_init(self, annotator): assert isinstance(annotator, DataAnnotator) def test_create_new_rttm_entry(self, annotator): words, alignments, speaker_id = generate_words_and_alignments(sample_index=0) start, end = alignments[0], alignments[-1] rttm_list = annotator.create_new_rttm_entry( words=words, alignments=alignments, start=start, end=end, speaker_id=speaker_id ) assert rttm_list[0] == f"{start} {end} {speaker_id}" def test_create_new_json_entry(self, annotator): words, alignments, speaker_id = generate_words_and_alignments(sample_index=0) start, end = alignments[0], alignments[-1] test_wav_filename = '/path/to/test_wav_filename.wav' test_rttm_filename = '/path/to/test_rttm_filename.rttm' test_ctm_filename = '/path/to/test_ctm_filename.ctm' text = " ".join(words) one_line_json_dict = annotator.create_new_json_entry( text=text, wav_filename=test_wav_filename, start=start, length=end - start, speaker_id=speaker_id, rttm_filepath=test_rttm_filename, ctm_filepath=test_ctm_filename, ) start = round(float(start), annotator._params.data_simulator.outputs.output_precision) length = round(float(end - start), annotator._params.data_simulator.outputs.output_precision) meta = { "audio_filepath": test_wav_filename, "offset": start, "duration": length, "label": speaker_id, "text": text, "num_speakers": annotator._params.data_simulator.session_config.num_speakers, "rttm_filepath": test_rttm_filename, "ctm_filepath": test_ctm_filename, "uem_filepath": None, } assert one_line_json_dict == meta def test_create_new_ctm_entry(self, annotator): words, alignments, speaker_id = generate_words_and_alignments(sample_index=0) session_name = 'test_session' ctm_list = annotator.create_new_ctm_entry( words=words, alignments=alignments, session_name=session_name, speaker_id=speaker_id, start=alignments[0] ) assert ctm_list[0] == ( alignments[1], get_ctm_line( source=session_name, channel="1", start_time=alignments[1], duration=float(alignments[1] - alignments[0]), token=words[1], conf=None, type_of_token='lex', speaker=speaker_id, ), ) assert ctm_list[1] == ( alignments[2], get_ctm_line( source=session_name, channel="1", start_time=alignments[2], duration=float(alignments[2] - alignments[1]), token=words[2], conf=None, type_of_token='lex', speaker=speaker_id, ), ) class TestSpeechSampler: def test_init(self, sampler): assert isinstance(sampler, SpeechSampler) def test_init_overlap_params(self, sampler): sampler._init_overlap_params() assert sampler.per_silence_min_len is not None assert sampler.per_silence_max_len is not None assert type(sampler.per_silence_min_len) == int assert type(sampler.per_silence_max_len) == int def test_init_silence_params(self, sampler): sampler._init_overlap_params() assert sampler.per_overlap_min_len is not None assert sampler.per_overlap_max_len is not None assert type(sampler.per_overlap_min_len) == int assert type(sampler.per_overlap_max_len) == int @pytest.mark.parametrize("mean", [0.1, 0.2, 0.3]) @pytest.mark.parametrize("var", [0.05, 0.07]) def test_get_session_silence_mean_pass(self, sampler, mean, var): sampler.mean_silence = mean sampler.mean_silence_var = var sampled_silence_mean = sampler.get_session_silence_mean() assert 0 <= sampled_silence_mean <= 1 @pytest.mark.parametrize("mean", [0.5]) @pytest.mark.parametrize("var", [0.5, 0.6]) def test_get_session_silence_mean_fail(self, sampler, mean, var): """ This test should raise `ValueError` because `mean_silence_var` should be less than `mean_silence * (1 - mean_silence)`. """ sampler.mean_silence = mean sampler.mean_silence_var = var with pytest.raises(ValueError) as execinfo: sampler.get_session_silence_mean() assert "ValueError" in str(execinfo) and "mean_silence_var" in str(execinfo) @pytest.mark.parametrize("mean", [0.1, 0.2, 0.3]) @pytest.mark.parametrize("var", [0.05, 0.07]) def test_get_session_overlap_mean_pass(self, sampler, mean, var): sampler.mean_overlap = mean sampler.mean_overlap_var = var sampled_overlap_mean = sampler.get_session_overlap_mean() assert 0 <= sampled_overlap_mean <= 1 @pytest.mark.parametrize("mean", [0.4, 0.5]) @pytest.mark.parametrize("var", [0.3, 0.8]) def test_get_session_overlap_mean_fail(self, sampler, mean, var): """ This test should raise `ValueError` because `mean_overlap_var` should be less than `mean_overlap * (1 - mean_overlap)`. """ sampler.mean_overlap = mean sampler.mean_overlap_var = var sampler._params = DictConfig(sampler._params) with pytest.raises(ValueError) as execinfo: sampler.get_session_overlap_mean() assert "ValueError" in str(execinfo) and "mean_overlap_var" in str(execinfo) @pytest.mark.parametrize("non_silence_len_samples", [16000, 32000]) @pytest.mark.parametrize("running_overlap_len_samples", [8000, 12000]) def test_sample_from_overlap_model(self, sampler, non_silence_len_samples, running_overlap_len_samples): sampler.get_session_overlap_mean() sampler.running_overlap_len_samples = running_overlap_len_samples overlap_amount = sampler.sample_from_overlap_model(non_silence_len_samples=non_silence_len_samples) assert type(overlap_amount) == int assert 0 <= overlap_amount @pytest.mark.parametrize("running_len_samples", [8000, 16000]) @pytest.mark.parametrize("running_overlap_len_samples", [8000, 12000]) def test_sample_from_silence_model(self, sampler, running_len_samples, running_overlap_len_samples): sampler.get_session_silence_mean() self.running_overlap_len_samples = running_overlap_len_samples silence_amount = sampler.sample_from_silence_model(running_len_samples=running_len_samples) assert type(silence_amount) == int assert 0 <= silence_amount @pytest.mark.with_downloads() @pytest.mark.parametrize("num_noise_files", [1, 2, 4]) def test_sample_noise_manifest(self, sampler, num_noise_files, test_data_dir): sampler.num_noise_files = num_noise_files manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/an4_val.json')) noise_manifest = read_noise_manifest(add_bg=True, background_manifest=manifest_path) sampled_noise_manifests = sampler.sample_noise_manifest(noise_manifest=noise_manifest) assert len(sampled_noise_manifests) == num_noise_files @pytest.mark.parametrize("running_speech_len_samples", [32000, 64000]) @pytest.mark.parametrize("running_overlap_len_samples", [16000, 32000]) @pytest.mark.parametrize("running_len_samples", [64000, 96000]) @pytest.mark.parametrize("non_silence_len_samples", [16000, 32000]) def test_silence_vs_overlap_selector( self, sampler, running_overlap_len_samples, running_speech_len_samples, running_len_samples, non_silence_len_samples, ): sampler.running_overlap_len_samples = running_overlap_len_samples sampler.running_speech_len_samples = running_speech_len_samples add_overlap = sampler.silence_vs_overlap_selector( running_len_samples=running_len_samples, non_silence_len_samples=non_silence_len_samples ) assert type(add_overlap) == bool