conftest.py 2.8 KB
Newer Older
1
2
3
import os
import shutil

4
import pytest
5
import torch
6
import torchaudio
7
8
9


class GreedyCTCDecoder(torch.nn.Module):
10
    def __init__(self, labels, blank: int = 0):
11
        super().__init__()
12
        self.blank = blank
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
        self.labels = labels

    def forward(self, logits: torch.Tensor) -> str:
        """Given a sequence logits over labels, get the best path string

        Args:
            logits (Tensor): Logit tensors. Shape `[num_seq, num_label]`.

        Returns:
            str: The resulting transcript
        """
        best_path = torch.argmax(logits, dim=-1)  # [num_seq,]
        best_path = torch.unique_consecutive(best_path, dim=-1)
        hypothesis = []
        for i in best_path:
28
29
            if i != self.blank:
                hypothesis.append(self.labels[i])
30
        return "".join(hypothesis)
31
32
33
34
35
36
37


@pytest.fixture
def ctc_decoder():
    return GreedyCTCDecoder


moto's avatar
moto committed
38
_FILES = {
39
40
41
42
43
44
    "en": "Lab41-SRI-VOiCES-src-sp0307-ch127535-sg0042.flac",
    "de": "20090505-0900-PLENARY-16-de_20090505-21_56_00_8.flac",
    "en2": "20120613-0900-PLENARY-8-en_20120613-13_46_50_3.flac",
    "es": "20130207-0900-PLENARY-7-es_20130207-13_02_05_5.flac",
    "fr": "20121212-0900-PLENARY-5-fr_20121212-11_37_04_10.flac",
    "it": "20170516-0900-PLENARY-16-it_20170516-18_56_31_1.flac",
moto's avatar
moto committed
45
}
46
47
48
49
50
_MIXTURE_FILE = "mixture_3729-6852-0037_8463-287645-0000.wav"
_CLEAN_FILES = [
    "s1_3729-6852-0037_8463-287645-0000.wav",
    "s2_3729-6852-0037_8463-287645-0000.wav",
]
moto's avatar
moto committed
51
52


53
@pytest.fixture
54
def sample_speech(lang):
moto's avatar
moto committed
55
    if lang not in _FILES:
56
        raise NotImplementedError(f"Unexpected lang: {lang}")
moto's avatar
moto committed
57
    filename = _FILES[lang]
58
    path = torchaudio.utils.download_asset(f"test-assets/{filename}")
moto's avatar
moto committed
59
    return path
moto's avatar
moto committed
60
61


62
63
@pytest.fixture
def mixture_source():
64
    path = torchaudio.utils.download_asset(f"test-assets/{_MIXTURE_FILE}")
65
66
67
68
69
70
71
    return path


@pytest.fixture
def clean_sources():
    paths = []
    for file in _CLEAN_FILES:
72
        path = torchaudio.utils.download_asset(f"test-assets/{file}")
73
74
75
76
        paths.append(path)
    return paths


moto's avatar
moto committed
77
78
79
80
81
82
83
def pytest_addoption(parser):
    parser.addoption(
        "--use-tmp-hub-dir",
        action="store_true",
        help=(
            "When provided, tests will use temporary directory as Torch Hub directory. "
            "Downloaded models will be deleted after each test."
84
        ),
moto's avatar
moto committed
85
86
87
88
    )


@pytest.fixture(autouse=True)
89
def temp_hub_dir(tmp_path, pytestconfig):
90
    if not pytestconfig.getoption("use_tmp_hub_dir"):
moto's avatar
moto committed
91
92
93
        yield
    else:
        org_dir = torch.hub.get_dir()
94
95
        subdir = os.path.join(tmp_path, "hub")
        torch.hub.set_dir(subdir)
moto's avatar
moto committed
96
97
        yield
        torch.hub.set_dir(org_dir)
98
        shutil.rmtree(subdir, ignore_errors=True)
99
100
101
102
103
104


@pytest.fixture()
def emissions():
    path = torchaudio.utils.download_asset("test-assets/emissions-8555-28447-0012.pt")
    return torch.load(path)