[https://nvbugs/6982380][fix] Open the Harmony final channel in JSON mode eval prompts

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:

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

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