diff --git a/tensorrt_llm/evaluate/json_mode_eval.py b/tensorrt_llm/evaluate/json_mode_eval.py index ce35d474b333..1b4b4892cddc 100644 --- a/tensorrt_llm/evaluate/json_mode_eval.py +++ b/tensorrt_llm/evaluate/json_mode_eval.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,7 +14,7 @@ # limitations under the License. import json import os -from typing import Iterable, List, Optional +from typing import Any, Iterable, List, Optional, Union import click import datasets @@ -27,6 +27,12 @@ from ..sampling_params import GuidedDecodingParams, SamplingParams from .interface import Evaluator +# A Harmony (gpt-oss) generation prompt stops at the assistant role: the model +# writes the channel header itself, and its answer is the message body of the +# final channel. +_HARMONY_ASSISTANT_START = "<|start|>assistant" +_HARMONY_FINAL_MESSAGE_START = "<|channel|>final<|message|>" + class JsonModeEval(Evaluator): @@ -55,6 +61,16 @@ def __init__(self, else: self.num_samples = min(num_samples, self.data.num_rows) + def do_apply_chat_template(self, llm: Any, + prompt: Union[str, List[dict]]) -> str: + prompt = super().do_apply_chat_template(llm, prompt) + # The JSON guide binds the completion from its first token. Open the + # Harmony final channel in the prompt so that the guide constrains the + # answer instead of the message header. + if prompt.endswith(_HARMONY_ASSISTANT_START): + prompt += _HARMONY_FINAL_MESSAGE_START + return prompt + def generate_samples(self) -> Iterable[tuple]: for i, sample in enumerate(self.data): if i >= self.num_samples: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 308de1728474..2e217737c496 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -174,6 +174,7 @@ l0_cpu: - unittest/docs/test_llmapi_reference_docs.py - unittest/dynamo - unittest/evaluate/test_image_generation_eval_cli.py + - unittest/evaluate/test_json_mode_eval.py - unittest/executor/test_base_worker.py ISOLATION - unittest/executor/test_fatal_error_health_check.py - unittest/executor/test_ipc.py diff --git a/tests/unittest/evaluate/test_json_mode_eval.py b/tests/unittest/evaluate/test_json_mode_eval.py new file mode 100644 index 000000000000..2b60db58d7bc --- /dev/null +++ b/tests/unittest/evaluate/test_json_mode_eval.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 +from types import SimpleNamespace + +import datasets +import pytest + +from tensorrt_llm.evaluate.json_mode_eval import JsonModeEval + +pytestmark = pytest.mark.cpu_only + +_SCHEMA = json.dumps( + {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} +) + + +class _ChatTemplateTokenizer: + """Renders messages and then appends a fixed generation prompt.""" + + def __init__(self, generation_prompt: str) -> None: + self.generation_prompt = generation_prompt + + def apply_chat_template(self, messages, tokenize, add_generation_prompt, **kwargs) -> str: + assert not tokenize and add_generation_prompt + rendered = "".join(f"<{m['role']}>{m['content']}" for m in messages) + return rendered + self.generation_prompt + + +@pytest.fixture +def evaluator(monkeypatch: pytest.MonkeyPatch) -> JsonModeEval: + sample = { + "prompt": [ + {"role": "system", "content": "Answer in JSON."}, + {"role": "user", "content": "The name is Ada."}, + ], + "completion": '{"name": "Ada"}', + "schema": _SCHEMA, + } + monkeypatch.setattr( + datasets, "load_dataset", lambda *args, **kwargs: datasets.Dataset.from_list([sample]) + ) + return JsonModeEval(dataset_path="unused") + + +def _render_first_prompt(evaluator: JsonModeEval, generation_prompt: str) -> str: + llm = SimpleNamespace(tokenizer=_ChatTemplateTokenizer(generation_prompt)) + prompt, sampling_args, _, _ = next(iter(evaluator.generate_samples())) + assert sampling_args["guided_decoding"].json is not None + return evaluator.do_apply_chat_template(llm, prompt) + + +def test_harmony_prompt_opens_the_final_channel(evaluator: JsonModeEval) -> None: + rendered = _render_first_prompt(evaluator, "<|start|>assistant") + + assert rendered.endswith("<|start|>assistant<|channel|>final<|message|>") + + +@pytest.mark.parametrize( + "generation_prompt", + ["<|im_start|>assistant\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"], +) +def test_other_generation_prompts_are_unchanged( + evaluator: JsonModeEval, generation_prompt: str +) -> None: + rendered = _render_first_prompt(evaluator, generation_prompt) + + assert rendered.endswith(generation_prompt) + assert "<|channel|>" not in rendered