nvidia-nemo--speech/tests/collections/tts/test_torch_tts.py

# Copyright (c) 2020, 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 json
import os
from pathlib import Path

import pytest
import torch

from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import EnglishPhonemesTokenizer
from nemo.collections.tts.data.dataset import TTSDataset
from nemo.collections.tts.g2p.models.en_us_arpabet import EnglishG2p
from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir


class TestTTSDataset:
    @pytest.mark.unit
    @pytest.mark.run_only_on('CPU')
    def test_dataset(self, test_data_dir):
        manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
        sup_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/sup')

        dataset = TTSDataset(
            manifest_filepath=manifest_path,
            sample_rate=22050,
            sup_data_types=["pitch"],
            sup_data_path=sup_path,
            text_tokenizer=EnglishPhonemesTokenizer(
                punct=True,
                stresses=True,
                chars=True,
                space=' ',
                apostrophe=True,
                pad_with_space=True,
                g2p=EnglishG2p(),
            ),
        )

        dataloader = torch.utils.data.DataLoader(dataset, 2, collate_fn=dataset._collate_fn)
        data, _, _, _, _, _ = next(iter(dataloader))

    @pytest.mark.unit
    @pytest.mark.run_only_on('CPU')
    def test_raise_exception_on_not_supported_sup_data_types(self, test_data_dir):
        manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
        sup_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/sup')
        with pytest.raises(NotImplementedError):
            dataset = TTSDataset(
                manifest_filepath=manifest_path,
                sample_rate=22050,
                sup_data_types=["not_supported_sup_data_type"],
                sup_data_path=sup_path,
                text_tokenizer=EnglishPhonemesTokenizer(
                    punct=True,
                    stresses=True,
                    chars=True,
                    space=' ',
                    apostrophe=True,
                    pad_with_space=True,
                    g2p=EnglishG2p(),
                ),
            )

    @pytest.mark.unit
    @pytest.mark.run_only_on('CPU')
    def test_raise_exception_on_not_supported_window(self, test_data_dir):
        manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
        sup_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/sup')
        with pytest.raises(NotImplementedError):
            dataset = TTSDataset(
                manifest_filepath=manifest_path,
                sample_rate=22050,
                sup_data_types=["pitch"],
                sup_data_path=sup_path,
                window="not_supported_window",
                text_tokenizer=EnglishPhonemesTokenizer(
                    punct=True,
                    stresses=True,
                    chars=True,
                    space=' ',
                    apostrophe=True,
                    pad_with_space=True,
                    g2p=EnglishG2p(),
                ),
            )

    @pytest.mark.unit
    @pytest.mark.run_only_on('CPU')
    @pytest.mark.parametrize("sup_data_type", ["voiced_mask", "p_voiced"])
    def test_raise_exception_on_missing_pitch_sup_data_type_if_use_voiced(self, test_data_dir, sup_data_type):
        manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
        sup_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/sup')
        with pytest.raises(ValueError):
            dataset = TTSDataset(
                manifest_filepath=manifest_path,
                sample_rate=22050,
                sup_data_types=[sup_data_type],
                sup_data_path=sup_path,
                window="hann",
                text_tokenizer=EnglishPhonemesTokenizer(
                    punct=True,
                    stresses=True,
                    chars=True,
                    space=' ',
                    apostrophe=True,
                    pad_with_space=True,
                    g2p=EnglishG2p(),
                ),
            )

    @pytest.mark.unit
    @pytest.mark.run_only_on('CPU')
    @pytest.mark.parametrize(
        "sup_data_types, output_indices",
        [
            (["p_voiced", "pitch", "voiced_mask"], [-4, -3, -1]),
            (["voiced_mask", "pitch"], [-3, -2]),
            (["pitch", "p_voiced"], [-3, -1]),
            (["pitch"], [-2]),
        ],
    )
    def test_save_voiced_items_if_pt_file_not_exist(self, test_data_dir, sup_data_types, output_indices, tmp_path):
        manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
        sup_path = tmp_path / "sup_data"
        print(f"sup_path={sup_path}")
        dataset = TTSDataset(
            manifest_filepath=manifest_path,
            sample_rate=22050,
            sup_data_types=sup_data_types,
            sup_data_path=sup_path,
            text_tokenizer=EnglishPhonemesTokenizer(
                punct=True,
                stresses=True,
                chars=True,
                space=' ',
                apostrophe=True,
                pad_with_space=True,
                g2p=EnglishG2p(),
            ),
        )

        # load manifest
        audio_filepaths = []
        with open(manifest_path, 'r', encoding="utf-8") as fjson:
            for line in fjson:
                audio_filepaths.append(json.loads(line)["audio_filepath"])
        base_data_dir = get_base_dir(audio_filepaths)

        dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=1, collate_fn=dataset._collate_fn)
        for batch, audio_filepath in zip(dataloader, audio_filepaths):
            rel_audio_path = Path(audio_filepath).relative_to(base_data_dir).with_suffix("")
            rel_audio_path_as_text_id = str(rel_audio_path).replace("/", "_")

            for sup_data_type, output_index in zip(sup_data_types, output_indices):
                sup_data = batch[output_index]
                sup_data = sup_data.squeeze(0)
                assert sup_data is not None
                assert torch.equal(sup_data, torch.load(f"{sup_path}/{sup_data_type}/{rel_audio_path_as_text_id}.pt"))

                if sup_data_type == "pitch":
                    pitch_lengths = batch[output_index + 1]
                    pitch_lengths = pitch_lengths.squeeze(0)
                    assert pitch_lengths is not None

            # test pitch, voiced_mask, and p_voiced do not have the same values.
            if len(sup_data_types) == 3:
                x = torch.load(f"{sup_path}/{sup_data_types[0]}/{rel_audio_path_as_text_id}.pt")
                y = torch.load(f"{sup_path}/{sup_data_types[1]}/{rel_audio_path_as_text_id}.pt")
                z = torch.load(f"{sup_path}/{sup_data_types[2]}/{rel_audio_path_as_text_id}.pt")
                assert not torch.equal(x, y)
                assert not torch.equal(x, z)