chat template overwriting, chat mode in instruction and preference recipes (#1228)

* chat template overwriting, chat mode in SFT dataset

* qwen template added

* formatting

* dpo recipe with chat mode

* extra newline fix in the template

* introducing chat template interface

* lint

* lint

* lint

* removing v0.4 compatibility

---------

Co-authored-by: Ilia Kulikov <kulikov@meta.com>

Author: uralik

src/fairseq2/assets/cards/models/qwen.yaml

@@ -64,7 +64,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_0.6b
+name: qwen3_0.6b_base
 model_family: qwen
 model_arch: qwen3_0.6b
 checkpoint: "hg://qwen/qwen3-0.6b-base"
@@ -73,7 +73,7 @@ tokenizer_family: qwen
 
 ---
 
-name: qwen3_0.6b_instruct
+name: qwen3_0.6b
 model_family: qwen
 model_arch: qwen3_0.6b
 checkpoint: "hg://qwen/qwen3-0.6b"
@@ -83,7 +83,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_1.7b
+name: qwen3_1.7b_base
 model_family: qwen
 model_arch: qwen3_1.7b
 checkpoint: "hg://qwen/qwen3-1.7b-base"
@@ -92,7 +92,7 @@ tokenizer_family: qwen
 
 ---
 
-name: qwen3_1.7b_instruct
+name: qwen3_1.7b
 model_family: qwen
 model_arch: qwen3_1.7b
 checkpoint: "hg://qwen/qwen3-1.7b"
@@ -102,7 +102,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_4b
+name: qwen3_4b_base
 model_family: qwen
 model_arch: qwen3_4b
 checkpoint: "hg://qwen/qwen3-4b-base"
@@ -111,7 +111,7 @@ tokenizer_family: qwen
 
 ---
 
-name: qwen3_4b_instruct
+name: qwen3_4b
 model_family: qwen
 model_arch: qwen3_4b
 checkpoint: "hg://qwen/qwen3-4b"
@@ -121,7 +121,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_8b
+name: qwen3_8b_base
 model_family: qwen
 model_arch: qwen3_8b
 checkpoint: "hg://qwen/qwen3-8b-base"
@@ -130,7 +130,7 @@ tokenizer_family: qwen
 
 ---
 
-name: qwen3_8b_instruct
+name: qwen3_8b
 model_family: qwen
 model_arch: qwen3_8b
 checkpoint: "hg://qwen/qwen3-8b"
@@ -140,7 +140,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_14b
+name: qwen3_14b_base
 model_family: qwen
 model_arch: qwen3_14b
 checkpoint: "hg://qwen/qwen3-14b-base"
@@ -149,7 +149,7 @@ tokenizer_family: qwen
 
 ---
 
-name: qwen3_14b_instruct
+name: qwen3_14b
 model_family: qwen
 model_arch: qwen3_14b
 checkpoint: "hg://qwen/qwen3-14b"
@@ -159,7 +159,7 @@ use_im_end: true
 
 ---
 
-name: qwen3_32b_instruct
+name: qwen3_32b
 model_family: qwen
 model_arch: qwen3_32b
 checkpoint: "hg://qwen/qwen3-32b"

src/fairseq2/data/text/tokenizers/hg.py

@@ -8,7 +8,7 @@
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, final
+from typing import TYPE_CHECKING, Any, final
 
 import torch
 from torch import Tensor
@@ -19,6 +19,9 @@
         AutoTokenizer,
         PreTrainedTokenizer,
     )
+    from transformers.tokenization_utils_base import (  # type: ignore[import-not-found, attr-defined]
+        BatchEncoding,
+    )
 except ImportError:
     _has_hg_transformers = False
 else:
@@ -44,8 +47,12 @@ def load_hg_token_model(
 
     @final
     class HuggingFaceTokenModel:
+        _tok: PreTrainedTokenizer
+
         def encode(self, text: str) -> list[int]: ...
 
+        def overwrite_chat_template(self, chat_template: str) -> None: ...
+
         def decode(
             self, token_indices: list[int], skip_special_tokens: bool = False
         ) -> str: ...
@@ -111,6 +118,9 @@ def __init__(
         def encode(self, text: str) -> list[int]:
             return self._tok.encode(text, add_special_tokens=False)
 
+        def overwrite_chat_template(self, chat_template: str) -> None:
+            self._tok.chat_template = chat_template
+
         def decode(
             self, token_indices: list[int], skip_special_tokens: bool = False
         ) -> str:
@@ -207,6 +217,11 @@ def encode_as_tokens(self, text: str) -> list[str]:
 
         return self._model.convert_ids_to_tokens(indices)
 
+    def apply_chat_template(self, chat: Any, **kwargs: Any) -> Any | BatchEncoding:
+        # pass all kwargs to HG tokenizer as is here
+        encoded_output = self._model._tok.apply_chat_template(chat, **kwargs)
+        return encoded_output
+
     @property
     @override
     def prefix_indices(self) -> Tensor | None:

src/fairseq2/data/text/tokenizers/llama.py

@@ -37,6 +37,21 @@
 )
 from fairseq2.device import Device
 
+# llama3 chat template with assistant mask support, see https://github.com/huggingface/transformers/issues/28950
+LLAMA3_CHAT_TEMPLATE = """{{- bos_token }}
+{%- for message in messages %}
+    {%- if message.role == 'user' or message.role == 'system' %}
+        {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + message['content'] | trim + '<|eot_id|>' }}
+    {%- else %}
+        {%- generation %}
+        {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' + message['content'] | trim + '<|eot_id|>' }}
+        {%- endgeneration %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}
+{%- endif %}"""
+
 
 @final
 class LLaMA3Tokenizer(TextTokenizer):
@@ -291,6 +306,7 @@ def load_llama3_hg_tokenizer(path: Path, card: AssetCard) -> TextTokenizer:
             boh_token="<|start_header_id|>",
             eoh_token="<|end_header_id|>",
         )
+        model.overwrite_chat_template(LLAMA3_CHAT_TEMPLATE)
     except (OSError, RuntimeError) as ex:
         raise TextTokenizerLoadError(
             card.name, f"The '{card.name}' text tokenizer model cannot be loaded. See the nested exception for details."  # fmt: skip

src/fairseq2/data/text/tokenizers/qwen.py

@@ -28,6 +28,71 @@
 )
 from fairseq2.device import Device
 
+# qwen chat template with assistant mask support. ALWAYS HAS THINKING TOKENS HERE
+QWEN_CHAT_TEMPLATE = """{%- if tools %}
+    {{- '<|im_start|>system\\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\\n\\n' }}
+    {%- endif %}
+    {{- "# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- '\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\\n</tool_call><|im_end|>\\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}
+            {%- endif %}
+        {%- endif %}
+        {{- '<|im_start|>' + message.role }}
+        {%- generation %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}
+            {%- else %}
+                {{- content }}
+            {%- endif %}
+        {%- else %}
+            {{- content }}
+        {%- endif %}
+        {{- '<|im_end|>' }}
+        {%- endgeneration %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\\n\\n</think>\\n\\n' }}
+    {%- endif %}
+{%- endif %}"""
+
 
 @final
 class QwenTokenizer(TextTokenizer):
@@ -123,6 +188,7 @@ def load_qwen_tokenizer(path: Path, card: AssetCard) -> TextTokenizer:
             boh_token=None,
             eoh_token=None,
         )
+        model.overwrite_chat_template(QWEN_CHAT_TEMPLATE)
     except (OSError, RuntimeError) as ex:
         raise TextTokenizerLoadError(
             card.name, f"The '{card.name}' text tokenizer model cannot be loaded. See the nested exception for details."  # fmt: skip

src/fairseq2/datasets/instruction.py

@@ -26,6 +26,7 @@
     read_sequence,
 )
 from fairseq2.data.text.tokenizers import TextTokenizer
+from fairseq2.data.text.tokenizers.hg import HuggingFaceTokenEncoder
 from fairseq2.datasets import (
     DataPipelineReader,
     DataReader,
@@ -56,6 +57,8 @@ class InstructionReadOptions(DataReadOptions):
     target_encode_mode: str = "prompt_response"
     """The tokenizer mode to encode the target text."""
 
+    chat_mode: bool = False
+
 
 @dataclass
 class InstructionPromptReadOptions(DataReadOptions):
@@ -230,26 +233,55 @@ def create_reader(
 
         seed += gang.rank
 
-        # Encode source and target texts.
-        source_encoder = tokenizer.create_encoder(mode=options.source_encode_mode)
-        target_encoder = tokenizer.create_encoder(mode=options.target_encode_mode)
+        if options.chat_mode is True:
+            # not passing any encoding modes here, because we use apply_chat_template here
+            encoder = tokenizer.create_encoder()
+            if not isinstance(encoder, HuggingFaceTokenEncoder):
+                raise RuntimeError(
+                    "Huggingface tokenizer must be used when chat_mode is True"
+                )
+            else:
 
-        builder.map(source_encoder, selector="src")
-        builder.map(target_encoder, selector="tgt")
+                def encoding_chat(example: dict[str, Any]) -> dict[str, Any]:
+                    id_ = example.get("id")
+                    chat = example.get("chat")
 
-        def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
-            id_ = example.get("id")
+                    encoded_output = encoder.apply_chat_template(
+                        chat,
+                        return_dict=True,
+                        return_assistant_tokens_mask=True,
+                        return_tensors="pt",
+                    )
 
-            source_indices = example["src"]
-            target_indices = example["tgt"]
+                    indices = encoded_output["input_ids"][0]
+                    target_mask = encoded_output["assistant_masks"][0].bool()
+
+                    return {"id": id_, "indices": indices, "target_mask": target_mask}
+
+                builder.map(encoding_chat)
+
+        else:
 
-            indices = torch.cat([source_indices, target_indices])
+            # Encode source and target texts.
+            source_encoder = tokenizer.create_encoder(mode=options.source_encode_mode)
+            target_encoder = tokenizer.create_encoder(mode=options.target_encode_mode)
 
-            target_mask = torch.arange(len(indices)) >= len(source_indices)
+            builder.map(source_encoder, selector="src")
+            builder.map(target_encoder, selector="tgt")
 
-            return {"id": id_, "indices": indices, "target_mask": target_mask}
+            def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
+                id_ = example.get("id")
 
-        builder.map(cat_source_and_target)
+                source_indices = example["src"]
+                target_indices = example["tgt"]
+
+                indices = torch.cat([source_indices, target_indices])
+
+                target_mask = torch.arange(len(indices)) >= len(source_indices)
+
+                return {"id": id_, "indices": indices, "target_mask": target_mask}
+
+            builder.map(cat_source_and_target)
 
         batching = options.batching
 
@@ -312,11 +344,7 @@ def to_batch(example: dict[str, Any]) -> SequenceBatch:
 
             seqs, seq_lens = indices["seqs"], indices["seq_lens"]
 
-            target_mask = example["target_mask"]["seqs"]
-
-            return SequenceBatch(
-                seqs, seq_lens, target_mask=target_mask, example=example
-            )
+            return SequenceBatch(seqs, seq_lens, example=example)
 
         pipeline = builder.map(to_batch).and_return()