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
110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
# Copyright (c) 2025, 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 pytest
|
|
|
|
from nemo.collections.asr.inference.utils.state_management_utils import (
|
|
detect_overlap,
|
|
find_max_overlap,
|
|
merge_segment_tail,
|
|
merge_timesteps,
|
|
merge_word_tail,
|
|
)
|
|
from nemo.collections.asr.inference.utils.text_segment import TextSegment, Word
|
|
|
|
|
|
class TestStateManagementUtils:
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"timesteps1, timesteps2, expected_merged_timesteps",
|
|
[
|
|
([0, 1, 2, 3], [4, 5, 6, 7], [0, 1, 2, 3, 4, 5, 6, 7]),
|
|
([0, 1, 2, 3], [], [0, 1, 2, 3]),
|
|
([], [4, 5, 6, 7], [4, 5, 6, 7]),
|
|
([-1, 0, 1, 2], [0, 1, 2, 3], [0, 1, 2, 3, 4, 5, 6, 7]),
|
|
([-3, 1, 2, 3], [], [0, 4, 5, 6]),
|
|
([], [-3, 1, 2, 3], [0, 4, 5, 6]),
|
|
],
|
|
)
|
|
def test_merge_timesteps(self, timesteps1, timesteps2, expected_merged_timesteps):
|
|
merged_timesteps = merge_timesteps(timesteps1, timesteps2)
|
|
assert merged_timesteps == expected_merged_timesteps
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"state_tokens, new_tokens, limit, expected_max_overlap",
|
|
[
|
|
([0, 1, 2, 3], [2, 3, 4, 5], 4, 2),
|
|
([0, 2, 3, 4], [2, 3, 4, 5], 4, 3),
|
|
([0, 0, 0, 1], [2, 3, 4, 5], 4, 0),
|
|
],
|
|
)
|
|
def test_find_max_overlap(self, state_tokens, new_tokens, limit, expected_max_overlap):
|
|
max_overlap = find_max_overlap(state_tokens, new_tokens, limit)
|
|
assert max_overlap == expected_max_overlap
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"state_tokens, state_timesteps, new_tokens, new_timesteps, expected_overlap",
|
|
[
|
|
([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [2.0, 3.0, 4.0, 5.0], 2),
|
|
([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [1.0, 2.0, 4.0, 5.0], 2),
|
|
([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [5.0, 7.0, 8.0, 9.0], 0),
|
|
],
|
|
)
|
|
def test_detect_overlap(self, state_tokens, state_timesteps, new_tokens, new_timesteps, expected_overlap):
|
|
overlap = detect_overlap(state_tokens, state_timesteps, new_tokens, new_timesteps)
|
|
assert overlap == expected_overlap
|
|
|
|
@pytest.mark.unit
|
|
def test_merge_word_tail_without_pnc(self):
|
|
word_head = Word(text="meaning", start=0.0, end=1.0, conf=0.5)
|
|
word_tail = Word(text="ful", start=1.0, end=2.0, conf=0.6)
|
|
head, _ = merge_word_tail(word_head, word_tail, conf_aggregator=min)
|
|
|
|
assert head.text == "meaningful"
|
|
assert head.start == 0.0
|
|
assert head.end == 2.0
|
|
assert head.conf == 0.5
|
|
|
|
@pytest.mark.unit
|
|
def test_merge_word_tail_with_pnc(self):
|
|
|
|
word_head = Word(text="meaning", start=0.0, end=1.0, conf=0.5)
|
|
word_tail = Word(text="s", start=1.0, end=2.0, conf=0.6)
|
|
pnc_head = Word(text="Meaning?", start=0.0, end=1.0, conf=0.5)
|
|
new_head, new_pnc_head = merge_word_tail(word_head, word_tail, conf_aggregator=min, pnc_word_head=pnc_head)
|
|
|
|
assert new_head.text == "meanings"
|
|
assert new_head.start == 0.0
|
|
assert new_head.end == 2.0
|
|
assert new_head.conf == 0.5
|
|
assert new_pnc_head.text == "Meanings?"
|
|
assert new_pnc_head.start == 0.0
|
|
assert new_pnc_head.end == 2.0
|
|
assert new_pnc_head.conf == 0.5
|
|
|
|
@pytest.mark.unit
|
|
def test_merge_segment_tail(self):
|
|
seg1 = TextSegment(text="Good morn", start=0.0, end=1.0, conf=0.5)
|
|
seg2 = TextSegment(text="ing", start=1.0, end=2.0, conf=0.6)
|
|
merged_seg = merge_segment_tail(seg1, seg2, conf_aggregator=min)
|
|
|
|
assert merged_seg.text == "Good morning"
|
|
assert merged_seg.start == 0.0
|
|
assert merged_seg.end == 2.0
|
|
assert merged_seg.conf == 0.5
|