diff --git a/nemo/collections/common/prompts/canary2.py b/nemo/collections/common/prompts/canary2.py index c6ba2e2bc847..bef6dcfadc71 100644 --- a/nemo/collections/common/prompts/canary2.py +++ b/nemo/collections/common/prompts/canary2.py @@ -216,7 +216,7 @@ def canary2(cut: Cut, prompt: Canary2PromptFormatter) -> dict[str, torch.Tensor] eos = prompt.tokenizer.eos else: # SPE eos = prompt.tokenizer.token_to_id(CANARY_EOS) - assert eos > -1, "Invalid tokenizer: tokenizer.token_to_id('{CANARY_EOS}') returned {eos}" + assert eos > -1, f"Invalid tokenizer: tokenizer.token_to_id('{CANARY_EOS}') returned {eos}" assert ( ans["answer_ids"][-1].item() == eos ), f"Expected the last token in answer_ids to be EOS, but we got {ans['answer_ids']}" diff --git a/tests/collections/common/test_canary2_prompt.py b/tests/collections/common/test_canary2_prompt.py new file mode 100644 index 000000000000..ae3f789b95ac --- /dev/null +++ b/tests/collections/common/test_canary2_prompt.py @@ -0,0 +1,39 @@ +# 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. + +""" +Tests for Canary 2.0 prompt formatting. +""" + +from unittest.mock import MagicMock + +import pytest +from lhotse.testing.dummies import dummy_cut + +from nemo.collections.common.prompts.canary2 import canary2 + + +def test_canary2_assert_message_includes_eos_id(): + """When the tokenizer returns -1 for CANARY_EOS, the assertion should report the actual id.""" + from lhotse import SupervisionSegment + + cut = dummy_cut(0, duration=1.0, supervisions=[SupervisionSegment("", "", 0, 1.0, text="hello")]) + cut.custom = {"source_lang": "en", "target_lang": "en"} + + prompt = MagicMock() + prompt.encode_dialog.return_value = {"answer_ids": MagicMock()} + prompt.tokenizer.token_to_id.return_value = -1 + + with pytest.raises(AssertionError, match="token_to_id.*returned -1"): + canary2(cut, prompt)