wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
488 lines
19 KiB
Python
488 lines
19 KiB
Python
# Copyright (c) 2022, 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 copy
|
|
import os
|
|
from functools import cached_property, lru_cache
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
from kaldialign import edit_distance
|
|
from omegaconf import DictConfig, open_dict
|
|
|
|
from nemo.collections.asr.models import ASRModel
|
|
from nemo.collections.asr.parts.mixins import mixins
|
|
from nemo.collections.asr.parts.submodules.ctc_decoding import (
|
|
CTCBPEDecoding,
|
|
CTCBPEDecodingConfig,
|
|
CTCDecoding,
|
|
CTCDecodingConfig,
|
|
)
|
|
from nemo.collections.asr.parts.submodules.ngram_lm.ngram_lm_batched import NGramGPULanguageModel
|
|
from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceConfig
|
|
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
|
|
from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
|
|
from tests.collections.asr.decoding.test_timestamps import BaseTimestampsTest
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def audio_file(test_data_dir):
|
|
return os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
|
|
|
|
|
|
CTC_MODEL = "nvidia/stt_en_conformer_ctc_small"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def kenlm_model_path(tmp_path_factory, test_data_dir):
|
|
lm_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
|
|
assert os.path.exists(lm_path), f"LM file not found: {lm_path}"
|
|
lm_nemo_path = tmp_path_factory.mktemp("lm") / f"{lm_path.name}.nemo"
|
|
NGramGPULanguageModel.from_file(lm_path, vocab_size=1024).save_to(f"{lm_nemo_path}")
|
|
return f"{lm_nemo_path}"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ctc_model():
|
|
model = ASRModel.from_pretrained(model_name=CTC_MODEL, map_location="cpu")
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def char_vocabulary():
|
|
return [' ', 'a', 'b', 'c', 'd', 'e', 'f', '.']
|
|
|
|
|
|
@pytest.fixture()
|
|
@lru_cache(maxsize=8)
|
|
def tmp_tokenizer(test_data_dir):
|
|
cfg = DictConfig({'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'})
|
|
|
|
class _TmpASRBPE(mixins.ASRBPEMixin):
|
|
def register_artifact(self, _, vocab_path):
|
|
return vocab_path
|
|
|
|
asrbpe = _TmpASRBPE()
|
|
asrbpe._setup_tokenizer(cfg)
|
|
return asrbpe.tokenizer
|
|
|
|
|
|
class TestCTCDecoding:
|
|
@pytest.mark.unit
|
|
def test_constructor(self):
|
|
cfg = CTCDecodingConfig()
|
|
vocab = char_vocabulary()
|
|
decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
|
|
assert decoding is not None
|
|
|
|
@pytest.mark.unit
|
|
def test_constructor_subword(self, tmp_tokenizer):
|
|
cfg = CTCBPEDecodingConfig()
|
|
decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
assert decoding is not None
|
|
|
|
@pytest.mark.unit
|
|
def test_char_decoding_greedy_forward(
|
|
self,
|
|
):
|
|
cfg = CTCDecodingConfig(strategy='greedy')
|
|
vocab = char_vocabulary()
|
|
decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
|
|
|
|
B, T = 4, 20
|
|
V = len(char_vocabulary()) + 1
|
|
input_signal = torch.randn(size=(B, T, V))
|
|
length = torch.randint(low=1, high=T, size=[B])
|
|
|
|
with torch.no_grad():
|
|
hypotheses = decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=False
|
|
)
|
|
texts = [hyp.text for hyp in hypotheses]
|
|
|
|
for text in texts:
|
|
assert isinstance(text, str)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('alignments', [False, True])
|
|
@pytest.mark.parametrize('timestamps', [False, True])
|
|
def test_char_decoding_greedy_forward_hypotheses(self, alignments, timestamps):
|
|
cfg = CTCDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
|
|
vocab = char_vocabulary()
|
|
decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
|
|
|
|
B, T = 4, 20
|
|
V = len(char_vocabulary()) + 1
|
|
input_signal = torch.randn(size=(B, T, V))
|
|
length = torch.randint(low=1, high=T, size=[B])
|
|
|
|
with torch.no_grad():
|
|
hyps = decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
for idx, hyp in enumerate(hyps):
|
|
assert isinstance(hyp, Hypothesis)
|
|
assert torch.is_tensor(hyp.y_sequence)
|
|
assert isinstance(hyp.text, str)
|
|
|
|
# alignments check
|
|
if alignments:
|
|
assert hyp.alignments is not None
|
|
assert isinstance(hyp.alignments, tuple)
|
|
assert len(hyp.alignments[0]) == length[idx]
|
|
assert len(hyp.alignments[1]) == length[idx]
|
|
|
|
# timestamps check
|
|
if timestamps:
|
|
BaseTimestampsTest.check_char_timestamps(hyp, decoding)
|
|
|
|
@pytest.mark.unit
|
|
def test_subword_decoding_greedy_forward(self, tmp_tokenizer):
|
|
cfg = CTCBPEDecodingConfig(strategy='greedy')
|
|
decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
|
|
B, T = 4, 20
|
|
V = decoding.tokenizer.tokenizer.vocab_size + 1
|
|
input_signal = torch.randn(size=(B, T, V))
|
|
length = torch.randint(low=1, high=T, size=[B])
|
|
|
|
with torch.no_grad():
|
|
hypotheses = decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=False
|
|
)
|
|
texts = [hyp.text for hyp in hypotheses]
|
|
|
|
for text in texts:
|
|
assert isinstance(text, str)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('alignments', [False, True])
|
|
@pytest.mark.parametrize('timestamps', [False, True])
|
|
@pytest.mark.pleasefixme
|
|
def test_subword_decoding_greedy_forward_hypotheses(self, tmp_tokenizer, alignments, timestamps):
|
|
cfg = CTCBPEDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
|
|
decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
|
|
B, T = 4, 20
|
|
V = decoding.tokenizer.tokenizer.vocab_size + 1
|
|
input_signal = torch.randn(size=(B, T, V))
|
|
length = torch.randint(low=1, high=T, size=[B])
|
|
|
|
with torch.no_grad():
|
|
hyps = decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
for idx, hyp in enumerate(hyps):
|
|
assert isinstance(hyp, Hypothesis)
|
|
assert torch.is_tensor(hyp.y_sequence)
|
|
assert isinstance(hyp.text, str)
|
|
|
|
# alignments check
|
|
if alignments:
|
|
assert hyp.alignments is not None
|
|
assert isinstance(hyp.alignments, tuple)
|
|
assert len(hyp.alignments[0]) == length[idx]
|
|
assert len(hyp.alignments[1]) == length[idx]
|
|
|
|
# timestamps check
|
|
if timestamps:
|
|
BaseTimestampsTest.check_subword_timestamps(hyp, decoding)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('alignments', [False, True])
|
|
@pytest.mark.parametrize('timestamps', [False, True])
|
|
@pytest.mark.parametrize('preserve_frame_confidence', [False, True])
|
|
@pytest.mark.parametrize('length_is_none', [False, True])
|
|
@pytest.mark.parametrize(
|
|
"logprobs_device",
|
|
[
|
|
torch.device("cpu"),
|
|
pytest.param(
|
|
torch.device("cuda"),
|
|
marks=pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason='CUDA required for test.',
|
|
),
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"length_device",
|
|
[
|
|
torch.device("cpu"),
|
|
pytest.param(
|
|
torch.device("cuda"),
|
|
marks=pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason='CUDA required for test.',
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_batched_decoding_logprobs(
|
|
self,
|
|
tmp_tokenizer,
|
|
alignments,
|
|
timestamps,
|
|
preserve_frame_confidence,
|
|
length_is_none,
|
|
logprobs_device,
|
|
length_device,
|
|
):
|
|
cfg = CTCBPEDecodingConfig(
|
|
strategy='greedy',
|
|
preserve_alignments=alignments,
|
|
compute_timestamps=timestamps,
|
|
confidence_cfg=ConfidenceConfig(preserve_frame_confidence=preserve_frame_confidence),
|
|
)
|
|
unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
|
|
cfg.strategy = 'greedy_batch'
|
|
batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
|
|
torch.manual_seed(1)
|
|
B, T = 4, 20
|
|
V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
|
|
input_signal = torch.randn(size=(B, T, V), device=logprobs_device)
|
|
# Set the blank index to a very high probability to make sure
|
|
# that we always handle at least a few blanks.
|
|
input_signal[:, 0, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
|
|
input_signal[:, 1, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
|
|
if length_is_none:
|
|
length = None
|
|
else:
|
|
length = torch.randint(low=1, high=T, size=[B], device=length_device)
|
|
|
|
with torch.inference_mode():
|
|
hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
|
|
input_signal, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
assert len(hyps) == len(batched_hyps) == B
|
|
for hyp, batched_hyp in zip(hyps, batched_hyps):
|
|
assert torch.abs(hyp.score - batched_hyp.score) <= 1e-5
|
|
assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
|
|
if timestamps:
|
|
assert hyp.timestamp == batched_hyp.timestamp
|
|
if alignments:
|
|
assert torch.all(hyp.alignments[0] == batched_hyp.alignments[0])
|
|
assert torch.all(hyp.alignments[1] == batched_hyp.alignments[1])
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('timestamps', [False, True])
|
|
@pytest.mark.parametrize('length_is_none', [False, True])
|
|
@pytest.mark.parametrize(
|
|
"labels_device",
|
|
[
|
|
torch.device("cpu"),
|
|
pytest.param(
|
|
torch.device("cuda"),
|
|
marks=pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason='CUDA required for test.',
|
|
),
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"length_device",
|
|
[
|
|
torch.device("cpu"),
|
|
pytest.param(
|
|
torch.device("cuda"),
|
|
marks=pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason='CUDA required for test.',
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_batched_decoding_labels(self, tmp_tokenizer, timestamps, length_is_none, labels_device, length_device):
|
|
cfg = CTCBPEDecodingConfig(strategy='greedy', compute_timestamps=timestamps)
|
|
unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
cfg.strategy = 'greedy_batch'
|
|
batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
|
|
|
|
torch.manual_seed(1)
|
|
B, T = 4, 20
|
|
V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
|
|
input_labels = torch.randint(V, size=(B, T), device=labels_device)
|
|
# Set some indices to blank to make sure that we always handle
|
|
# at least a few blanks.
|
|
input_labels[:, 0] = unbatched_decoding.tokenizer.tokenizer.vocab_size
|
|
input_labels[:, 1] = unbatched_decoding.tokenizer.tokenizer.vocab_size
|
|
if length_is_none:
|
|
length = None
|
|
else:
|
|
length = torch.randint(low=1, high=T, size=[B], device=length_device)
|
|
|
|
with torch.inference_mode():
|
|
hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
|
|
input_labels, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
|
|
input_labels, length, fold_consecutive=True, return_hypotheses=True
|
|
)
|
|
|
|
assert len(hyps) == len(batched_hyps) == B
|
|
for hyp, batched_hyp in zip(hyps, batched_hyps):
|
|
assert abs(hyp.score - batched_hyp.score) <= 1e-5
|
|
assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
|
|
if timestamps:
|
|
assert hyp.timestamp == batched_hyp.timestamp
|
|
|
|
|
|
class TestCTCTimestamps(BaseTimestampsTest):
|
|
"""CTC-specific timestamp tests that inherit from BaseTimestampsTest"""
|
|
|
|
@cached_property
|
|
def decoding_char(self):
|
|
cfg = CTCDecodingConfig()
|
|
vocab = char_vocabulary()
|
|
decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
|
|
return decoding
|
|
|
|
@cached_property
|
|
def decoding_subword_wpe(self):
|
|
cfg = CTCBPEDecodingConfig(compute_timestamps=True)
|
|
decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.tmp_tokenizer)
|
|
return decoding
|
|
|
|
@cached_property
|
|
def decoding_subword_bpe(self):
|
|
cfg = CTCBPEDecodingConfig(compute_timestamps=True)
|
|
decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.bpe_tokenizer)
|
|
return decoding
|
|
|
|
@pytest.mark.unit
|
|
def test_word_offsets_subword_wpe(self, tmp_tokenizer):
|
|
self.tmp_tokenizer = tmp_tokenizer
|
|
super().test_word_offsets_subword_wpe()
|
|
|
|
@pytest.mark.unit
|
|
def test_word_offsets_subword_wpe_other_delimiter(self, tmp_tokenizer):
|
|
self.tmp_tokenizer = tmp_tokenizer
|
|
super().test_word_offsets_subword_wpe_other_delimiter()
|
|
|
|
|
|
class TestCTCGreedyDecodingWithNGPU_LM:
|
|
@pytest.mark.with_downloads
|
|
@pytest.mark.unit
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
|
|
def test_ctc_decoding_gpulm(
|
|
self,
|
|
audio_file,
|
|
kenlm_model_path,
|
|
ctc_model,
|
|
):
|
|
device = torch.device("cuda")
|
|
model = ctc_model.to(device)
|
|
|
|
gt_hyp = model.transcribe([audio_file], num_workers=None)
|
|
|
|
decoding_config = copy.deepcopy(model.cfg.decoding)
|
|
with open_dict(model.decoding.cfg) as cfg:
|
|
cfg.greedy["ngram_lm_model"] = kenlm_model_path
|
|
cfg.greedy["ngram_lm_alpha"] = 0.0
|
|
model.change_decoding_strategy(cfg)
|
|
lm_hyp = model.transcribe([audio_file], num_workers=None)
|
|
|
|
assert gt_hyp[0].text == lm_hyp[0].text
|
|
assert abs(gt_hyp[0].score - lm_hyp[0].score) <= 1e-3
|
|
|
|
with open_dict(model.decoding.cfg) as cfg:
|
|
cfg.greedy["ngram_lm_model"] = kenlm_model_path
|
|
cfg.greedy["ngram_lm_alpha"] = 10.0
|
|
model.change_decoding_strategy(cfg)
|
|
lm_hyp = model.transcribe([audio_file], num_workers=None)
|
|
assert gt_hyp[0].text != lm_hyp[0].text
|
|
assert abs(gt_hyp[0].score - lm_hyp[0].score) > 1e-3
|
|
|
|
model.change_decoding_strategy(decoding_config)
|
|
|
|
|
|
class TestCTCGreedyDecodingCudaGrpahs:
|
|
"""
|
|
Tests CudaGraphs implementations from CTC models greedy decoding
|
|
"""
|
|
|
|
@pytest.mark.with_downloads
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
|
|
@pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
|
|
def test_stated_stateless(self, audio_file, kenlm_model_path, ctc_model, force_mode: str):
|
|
"""
|
|
Compares pure Pytorch and with three modes of statefull implementations for double floating point precision.
|
|
1. Pure pytorch, but statefull implementation: no_graphs
|
|
2. With CudaGrpahs: no_while_loops and full_graph.
|
|
"""
|
|
if force_mode == "full_graph":
|
|
skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
|
|
|
|
device = torch.device("cuda")
|
|
model = ctc_model.to(device)
|
|
decoding_config = copy.deepcopy(model.cfg.decoding)
|
|
|
|
with open_dict(model.decoding.cfg) as cfg:
|
|
cfg.greedy["ngram_lm_model"] = kenlm_model_path
|
|
cfg.greedy["ngram_lm_alpha"] = 0.2
|
|
cfg.greedy["allow_cuda_graphs"] = False
|
|
|
|
model.change_decoding_strategy(cfg)
|
|
|
|
actual_hypotheses = model.transcribe([audio_file], num_workers=None)
|
|
actual_transcripts = [hyp.text for hyp in actual_hypotheses]
|
|
actual_scores = [hyp.score for hyp in actual_hypotheses]
|
|
actual_timestamps = [hyp.timestamp for hyp in actual_hypotheses]
|
|
|
|
# transcribe with use implementation with cuda graphs
|
|
model.decoding.cfg["greedy"]["allow_cuda_graphs"] = True
|
|
model.change_decoding_strategy(model.decoding.cfg)
|
|
model.decoding.decoding.force_cuda_graphs_mode(mode=force_mode)
|
|
|
|
cudagraph_hypotheses = model.transcribe([audio_file], num_workers=None)
|
|
cudagraph_transcripts = [hyp.text for hyp in cudagraph_hypotheses]
|
|
cudagraph_scores = [hyp.score for hyp in cudagraph_hypotheses]
|
|
cudagraph_timestamps = [hyp.timestamp for hyp in cudagraph_hypotheses]
|
|
|
|
for batch_idx in range(len(actual_transcripts)):
|
|
assert len(actual_transcripts[batch_idx]) == len(cudagraph_transcripts[batch_idx])
|
|
assert cudagraph_scores[batch_idx] == pytest.approx(
|
|
actual_scores[batch_idx], abs=1e-2
|
|
), f"Scores mismatch for batch_idx {batch_idx}"
|
|
assert (
|
|
cudagraph_timestamps[batch_idx] == actual_timestamps[batch_idx]
|
|
), f"Timestamps mismatch for batch_idx {batch_idx}"
|
|
|
|
ref_words = actual_transcripts[batch_idx].split()
|
|
hyp_words = cudagraph_transcripts[batch_idx].split()
|
|
wer = edit_distance(ref_words, hyp_words)['total'] / max(len(ref_words), 1)
|
|
|
|
assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
|
|
|
|
for actual, fast in zip(actual_transcripts[batch_idx], cudagraph_transcripts[batch_idx]):
|
|
if actual != fast:
|
|
print("Erroneous samples in batch:", batch_idx)
|
|
print("Original transcript:", actual)
|
|
print("New transcript:", fast)
|
|
|
|
model.change_decoding_strategy(decoding_config)
|