fix(prompts): supervise all SpeechLLM assistant turns

@@ -127,6 +127,10 @@ SALMDataset Structure ^^^^^^^^^^^^^^^^^^^^^  Data used for SALM can be either regular speech-to-text data (in any NeMo or Lhotse format), or a dataset of multi-turn conversions.++With ``prompt_format: nemotron-nano-v3`` or ``prompt_format: nemotron3p5``,+training loss is computed over responses from all assistant turns.+ For the most part, please refer to :doc:`the ASR datasets documentation <../asr/datasets>` for details on data formats and multimodal dataloading.  When using speech-to-text data, you'll need read it with a special ``lhotse_as_conversation`` data reader

@@ -58,6 +58,8 @@ class NemotronNanoV3PromptFormatter(PromptFormatter):     def encode_dialog(self, turns: list[dict], enable_thinking: bool = True) -> dict[str, torch.Tensor]:         """Encode a dialog for Nemotron Nano v3 with <think> reasoning support. +        Training loss is computed over responses from all assistant turns.+         Args:             turns: List of turns with "role" and "slots"/"content" keys.             enable_thinking: If True, inference prefix ends with ``<think>\\n``;@@ -116,15 +118,15 @@ def encode_dialog(self, turns: list[dict], enable_thinking: bool = True) -> dict         # 6) Tokenize all turns.         turn_tokens = []         turn_token_counts = []-        turn_mask_values = []+        loss_mask = []          if self.INSERT_BOS:             turn_tokens.append(self.tokenizer.bos)             turn_token_counts.append(1)-            turn_mask_values.append(False)+            loss_mask.append(False)          is_inference = turns[-1]["role"] != self.OUTPUT_ROLE-        for idx, turn in enumerate(turns):+        for turn in turns:             role = turn["role"]             expected_slots = self.get_slots(role)             slot_values = turn.get("slots", {})@@ -134,8 +136,10 @@ def encode_dialog(self, turns: list[dict], enable_thinking: bool = True) -> dict             tokens = self.encode_turn(template, expected_slots, slot_values)             turn_tokens.extend(tokens)             turn_token_counts.append(len(tokens))-            # Loss mask only on the last assistant turn.-            turn_mask_values.append(role == self.OUTPUT_ROLE and idx == len(turns) - 1)+            if not is_inference and role == self.OUTPUT_ROLE:+                loss_mask.extend(self._assistant_loss_mask(tokens, template, expected_slots, slot_values))+            else:+                loss_mask.extend([False] * len(tokens))          # 7) Append inference prefix with thinking toggle.         if is_inference and self.INFERENCE_PREFIX is not None:@@ -146,31 +150,49 @@ def encode_dialog(self, turns: list[dict], enable_thinking: bool = True) -> dict             inference_tokens = self._apply_tokenizer(inference_prefix)             turn_tokens.extend(inference_tokens)             turn_token_counts.append(len(inference_tokens))-            turn_mask_values.append(False)+            loss_mask.extend([False] * len(inference_tokens))          # Insert EOS only when the last turn comes from the OUTPUT_ROLE.         if self.INSERT_EOS and not is_inference:             turn_tokens.append(self.tokenizer.eos)             turn_token_counts[-1] += 1-            turn_mask_values.append(True)+            loss_mask.append(True)          ans = {"input_ids": torch.tensor(turn_tokens, dtype=torch.long)}-        if turn_mask_values[-1]:+        if not is_inference:             ans["context_ids"] = ans["input_ids"][: -turn_token_counts[-1]]             ans["answer_ids"] = ans["input_ids"][-turn_token_counts[-1] :]-            ans["mask"] = torch.tensor(-                [-                    turn_mask_values[turn_idx]-                    for turn_idx, turn_len in enumerate(turn_token_counts)-                    for _ in range(turn_len)-                ],-                dtype=torch.bool,-            )+            ans["mask"] = torch.tensor(loss_mask, dtype=torch.bool)         else:             ans["context_ids"] = ans["input_ids"]          return ans +    def _assistant_loss_mask(+        self, tokens: list[int], template: str, expected_slots: dict, slot_values: dict+    ) -> list[bool]:+        # The caller provides the assistant header at generation time. Match only+        # leading thinking prefills present in this response; never mask reasoning+        # or a closing </think> that the model must generate after reasoning.+        message = slot_values["message"]+        prefix = template.split("|message|", 1)[0]+        for thinking_prefix in ("<think></think>", "<think>\n", "<think>"):+            if message.startswith(thinking_prefix):+                prefix += thinking_prefix+                break+        prefix_tokens = self._apply_tokenizer(prefix, lang=slot_values.get(self.PROMPT_LANGUAGE_SLOT))+        start = _common_prefix_length(tokens, prefix_tokens)++        # Keep the EOT target but exclude formatting after it. Tokenize complete+        # turns as before: concatenating independently encoded prefix/body pieces+        # would change BPE segmentation. If a boundary merges tokens, supervise+        # the ambiguous token(s) rather than dropping any response/EOT targets.+        through_eot = self.encode_turn(template.removesuffix("\n"), expected_slots, slot_values)+        end = _common_prefix_length(tokens, through_eot)+        if end < len(through_eot):+            end = len(tokens)+        return [False] * start + [True] * (end - start) + [False] * (len(tokens) - end)+  @registered_prompt_format_fn(Cut, NemotronNanoV3PromptFormatter) def nemotron_nano_v3(cut: Cut, prompt: NemotronNanoV3PromptFormatter):@@ -194,3 +216,10 @@ def nemotron_nano_v3(cut: Cut, prompt: NemotronNanoV3PromptFormatter):         turns.append({"role": "assistant", "content": answer})      return prompt.encode_dialog(turns)+++def _common_prefix_length(tokens: list[int], prefix: list[int]) -> int:+    for i, (token, prefix_token) in enumerate(zip(tokens, prefix)):+        if token != prefix_token:+            return i+    return min(len(tokens), len(prefix))

@@ -34,7 +34,7 @@ def test_nemotron3p5_training_basic(bpe_tokenizer_with_think):         == "<|im_start|>system\n<|im_end|>\n <|im_start|>user\nTEST<|im_end|>\n "         "<|im_start|>assistant\n<think></think>TEST<|im_end|>\n"     )-    assert ans["mask"].tolist() == [False] * len(ans["context_ids"]) + [True] * len(ans["answer_ids"])+    assert bpe_tokenizer_with_think.ids_to_text(ans["input_ids"][ans["mask"]].tolist()) == "TEST<|im_end|>"   def test_nemotron3p5_inference_generation_prompt(bpe_tokenizer_with_think):

@@ -16,6 +16,7 @@  import pytest +from nemo.collections.common.prompts.nemotron3p5 import Nemotron3p5PromptFormatter from nemo.collections.common.prompts.nemotron_nano_v3 import NemotronNanoV3PromptFormatter  # ──────────────────────────────────────────────────────────────────────@@ -124,6 +125,136 @@ def test_nemotron_nano_v3_training_multiturn_past_asst_no_think(bpe_tokenizer_wi     # fmt: on  +@pytest.mark.parametrize("formatter_cls", [NemotronNanoV3PromptFormatter, Nemotron3p5PromptFormatter])+@pytest.mark.parametrize("insert_bos,insert_eos", [(False, False), (True, False), (False, True), (True, True)])+@pytest.mark.parametrize("historical_answer", ["TEST", ""])+def test_nemotron_training_supervises_all_assistant_turns(+    bpe_tokenizer_with_think, formatter_cls, insert_bos, insert_eos, historical_answer+):+    tokenizer = bpe_tokenizer_with_think+    formatter = formatter_cls(tokenizer)+    formatter.INSERT_BOS = insert_bos+    formatter.INSERT_EOS = insert_eos+    ans = formatter.encode_dialog(+        [+            {"role": "user", "content": "TEST"},+            {"role": "assistant", "content": f"<think>SYSTEM</think>{historical_answer}"},+            {"role": "tool", "content": "TEST"},+            {"role": "user", "content": "SYSTEM"},+            {"role": "assistant", "content": "TEST"},+        ]+    )+    # Tokenize each rendered turn independently, including the normalized history.+    rendered_turns = [+        "<|im_start|>system\n<|im_end|>\n",+        "<|im_start|>user\nTEST<|im_end|>\n",+        f"<|im_start|>assistant\n<think></think>{historical_answer}<|im_end|>\n",+        "<|im_start|>tool\nTEST<|im_end|>\n",+        "<|im_start|>user\nSYSTEM<|im_end|>\n",+        "<|im_start|>assistant\n<think></think>TEST<|im_end|>\n",+    ]+    chunks = [tokenizer.text_to_ids(turn) for turn in rendered_turns]+    expected_ids = [tokenizer.bos] if insert_bos else []+    expected_mask = [False] if insert_bos else []+    for chunk, supervised in zip(chunks, [False, False, True, False, False, True]):+        expected_ids.extend(chunk)+        if supervised:+            # This fixture has atomic tags/newlines: only response text + EOT carry loss.+            prefix_len = len(tokenizer.text_to_ids("<|im_start|>assistant\n<think></think>"))+            expected_mask.extend([False] * prefix_len + [True] * (len(chunk) - prefix_len - 1) + [False])+        else:+            expected_mask.extend([False] * len(chunk))+    if insert_eos:+        expected_ids.append(tokenizer.eos)+        expected_mask.append(True)++    assert ans["input_ids"].tolist() == expected_ids+    assert ans["mask"].tolist() == expected_mask+    # Generation references still describe the final turn, not all supervised turns.+    final_answer = chunks[-1] + ([tokenizer.eos] if insert_eos else [])+    assert ans["answer_ids"].tolist() == final_answer+    assert ans["context_ids"].tolist() == expected_ids[: -len(final_answer)]+++@pytest.mark.parametrize("formatter_cls", [NemotronNanoV3PromptFormatter, Nemotron3p5PromptFormatter])+@pytest.mark.parametrize(+    "message,supervised_text",+    [+        ("TEST", "TEST<|im_end|>"),+        ("", "<|im_end|>"),+        ("<think></think>TEST", "TEST<|im_end|>"),+        ("<think>\nSYSTEM</think>TEST", "SYSTEM</think>TEST<|im_end|>"),+        ("<think>SYSTEM</think>TEST", "SYSTEM</think>TEST<|im_end|>"),+        ("<think>\n</think>TEST", "</think>TEST<|im_end|>"),+        ("<think>TEST", "TEST<|im_end|>"),+    ],+)+def test_nemotron_masks_prefill_but_supervises_generated_content(+    bpe_tokenizer_with_think, formatter_cls, message, supervised_text+):+    formatter = formatter_cls(bpe_tokenizer_with_think)+    ans = formatter.encode_dialog([{"role": "assistant", "content": message}])+    supervised_ids = ans["input_ids"][ans["mask"]].tolist()+    assert bpe_tokenizer_with_think.ids_to_text(supervised_ids) == supervised_text+    assert not ans["mask"][: len(ans["context_ids"])].any()+++@pytest.mark.parametrize("formatter_cls", [NemotronNanoV3PromptFormatter, Nemotron3p5PromptFormatter])+@pytest.mark.parametrize(+    "message,supervised_text",+    [("TEST", "TEST"), ("", ""), ("<think>\nSYSTEM</think>TEST", "SYSTEM</think>TEST")],+)+def test_nemotron_supervises_assistant_before_tool(bpe_tokenizer_with_think, formatter_cls, message, supervised_text):+    formatter = formatter_cls(bpe_tokenizer_with_think)+    ans = formatter.encode_dialog(+        [+            {"role": "user", "content": "TEST"},+            {"role": "assistant", "content": message},+            {"role": "tool", "content": "SYSTEM"},+            {"role": "assistant", "content": "TEST"},+        ]+    )+    assert bpe_tokenizer_with_think.ids_to_text(ans["input_ids"][ans["mask"]].tolist()) == (+        supervised_text + "<|im_end|>TEST<|im_end|>"+    )+++@pytest.mark.parametrize("formatter_cls", [NemotronNanoV3PromptFormatter, Nemotron3p5PromptFormatter])+def test_nemotron_mask_keeps_tokens_straddling_prefill_boundary(bpe_tokenizer_with_think, formatter_cls, monkeypatch):+    formatter = formatter_cls(bpe_tokenizer_with_think)+    # A token may merge the final prefill character with the first answer character.+    # Likewise, a token may merge the end marker with the formatting newline.+    encodings = {+        "<|im_start|>system\n<|im_end|>\n": [1, 2],+        "<|im_start|>assistant\n<think></think>": [3, 4, 5],+        "<|im_start|>assistant\n<think></think>TEST<|im_end|>": [3, 4, 6, 7],+        "<|im_start|>assistant\n<think></think>TEST<|im_end|>\n": [3, 4, 6, 8],+    }+    monkeypatch.setattr(formatter, "_apply_tokenizer", lambda text, **kwargs: encodings[text])+    ans = formatter.encode_dialog([{"role": "assistant", "content": "TEST"}])+    assert ans["input_ids"].tolist() == [1, 2, 3, 4, 6, 8]+    assert ans["mask"].tolist() == [False, False, False, False, True, True]+++@pytest.mark.parametrize("formatter_cls", [NemotronNanoV3PromptFormatter, Nemotron3p5PromptFormatter])+@pytest.mark.parametrize("enable_thinking", [False, True])+def test_nemotron_multiturn_inference_has_no_training_mask(bpe_tokenizer_with_think, formatter_cls, enable_thinking):+    formatter = formatter_cls(bpe_tokenizer_with_think)+    ans = formatter.encode_dialog(+        [+            {"role": "user", "content": "TEST"},+            {"role": "assistant", "content": "TEST"},+            {"role": "user", "content": "TEST"},+        ],+        enable_thinking=enable_thinking,+    )+    assert set(ans) == {"input_ids", "context_ids"}+    assert ans["context_ids"].tolist() == ans["input_ids"].tolist()+    suffix = "<think>\n" if enable_thinking else "<think></think>"+    prefix_ids = bpe_tokenizer_with_think.text_to_ids(f"<|im_start|>assistant\n{suffix}")+    assert ans["input_ids"][-len(prefix_ids) :].tolist() == prefix_ids++ def test_nemotron_nano_v3_history_thinking_truncation(bpe_tokenizer_with_think):     """Multi-turn: earlier assistant thinking replaced with <think></think>."""     formatter = NemotronNanoV3PromptFormatter(bpe_tokenizer_with_think)