Add DeepSeek-V4 chat encoding to SFT preprocessing

DeepSeek-V4 ships no Jinja chat_template, so apply_chat_template fails during SFT preprocessing. The official encoder was only vendored under benchmarks/api_server, which src/maxtext cannot import, and that copy had lost the blank lines in TOOLS_TEMPLATE and several docstrings, so its tool-calling prompts no longer matched the official encoder goldens.

- Move benchmarks/api_server/encoding/encoding_dsv4.py to src/maxtext/input_pipeline/encoding_dsv4.py, restore the stripped blank lines, and import it from benchmarks/api_server/{maxtext_server,server_utils}.py.
- input_pipeline_utils.apply_chat_template: when chat_template is None and the model or tokenizer name is DeepSeek-V4, encode with encode_messages and split it into prompt and completion chunks. In thinking mode, turns whose reasoning the encoder dropped stay in the prompt. Encoded text left after the final assistant turn raises ValueError.
- grain_data_processing: build the SFT tokenizer with add_bos=False and add_eos=False, as the HF and vision SFT paths do; with transformers 5.x, add_bos_token=True prepended BOS to every chunk.
- grain_data_processing and hf_data_processing: forward model_name to apply_chat_template.
- tests: golden parity with the 4 official encoder fixtures, Grain and HF SFT loss masks, name-only DeepSeek-V4 detection, thinking-mode history turns, Grain tokenizer flags, and model_name forwarding.

Author: AusarYao28

benchmarks/api_server/maxtext_server.py

@@ -59,7 +59,8 @@
     ChatMessage,
 )
 from benchmarks.api_server import server_utils
-from benchmarks.api_server.encoding import encoding_dsv32, encoding_dsv4
+from benchmarks.api_server.encoding import encoding_dsv32
+from maxtext.input_pipeline import encoding_dsv4
 
 # ----------------------------
 # Init

benchmarks/api_server/server_utils.py

@@ -32,7 +32,8 @@
 
 from benchmarks.api_server.maxtext_generator import MaxTextGenerator
 from benchmarks.api_server.server_models import LogProbsPayload
-from benchmarks.api_server.encoding import encoding_dsv32, encoding_dsv4
+from benchmarks.api_server.encoding import encoding_dsv32
+from maxtext.input_pipeline import encoding_dsv4
 
 # ----------------------------
 # Debugging

…rks/api_server/encoding/encoding_dsv4.py …/maxtext/input_pipeline/encoding_dsv4.pybenchmarks/api_server/encoding/encoding_dsv4.py renamed to src/maxtext/input_pipeline/encoding_dsv4.py benchmarks/api_server/encoding/encoding_dsv4.py renamed to src/maxtext/input_pipeline/encoding_dsv4.py

@@ -81,7 +81,9 @@
 )
 
 TOOLS_TEMPLATE = """## Tools
+
 You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
+
 <{dsml_token}tool_calls>
 <{dsml_token}invoke name="$TOOL_NAME">
 <{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
@@ -91,11 +93,17 @@
 ...
 </{dsml_token}invoke>
 </{dsml_token}tool_calls>
+
 String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
 If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
+
 Otherwise, output directly after {thinking_end_token} with tool calls or final response.
+
 ### Available Tool Schemas
+
 {tool_schemas}
+
 You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
 """
 
@@ -145,8 +153,10 @@ def tool_calls_to_openai_format(tool_calls):
 def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
   """
   Encode tool call arguments into DSML parameter format.
+
   Args:
       tool_call: Dict with "name" and "arguments" (JSON string) keys.
+
   Returns:
       DSML-formatted parameter string.
   """
@@ -173,9 +183,11 @@ def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
 def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
   """
   Decode DSML parameters back to a tool call dict.
+
   Args:
       tool_name: Name of the tool.
       tool_args: Dict mapping param_name -> (value, is_string_flag).
+
   Returns:
       Dict with "name" and "arguments" (JSON string) keys.
   """
@@ -192,8 +204,10 @@ def _decode_value(key: str, value: str, string: str):
 def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
   """
   Render tool schemas into the system prompt format.
+
   Args:
       tools: List of tool schema dicts (each with name, description, parameters).
+
   Returns:
       Formatted tools section string.
   """
@@ -244,14 +258,17 @@ def render_message(
 ) -> str:
   """
   Render a single message at the given index into its encoded string form.
+
   This is the core function that converts each message in the conversation
   into the DeepSeek-V4 format.
+
   Args:
       index: Index of the message to render.
       messages: Full list of messages in the conversation.
       thinking_mode: Either "chat" or "thinking".
       drop_thinking: Whether to drop reasoning content from earlier turns.
       reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
   Returns:
       Encoded string for this message.
   """
@@ -408,13 +425,17 @@ def render_message(
 def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
   """
   Merge tool messages into the preceding user message using content_blocks format.
+
   DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
   are encoded as <tool_result> blocks within user messages.
+
   This function converts a standard OpenAI-format conversation (with separate
   "tool" role messages) into V4 format where tool results are merged into
   user messages.
+
   Args:
       messages: List of message dicts in OpenAI format.
+
   Returns:
       Processed message list with tool messages merged into user messages.
   """
@@ -471,8 +492,10 @@ def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict
   """
   Sort tool_result blocks within user messages by the order of tool_calls
   in the preceding assistant message.
+
   Args:
       messages: Preprocessed message list (after merge_tool_messages).
+
   Returns:
       Message list with sorted tool result blocks.
   """
@@ -519,11 +542,13 @@ def encode_messages(
 ) -> str:
   """
   Encode a list of messages into the DeepSeek-V4 prompt format.
+
   This is the main entry point for encoding conversations. It handles:
   - BOS token insertion
   - Thinking mode with optional reasoning content dropping
   - Tool message merging into user messages
   - Multi-turn conversation context
+
   Args:
       messages: List of message dicts to encode.
       thinking_mode: Either "chat" or "thinking".
@@ -532,6 +557,7 @@ def encode_messages(
                     (only keep reasoning for messages after the last user message).
       add_default_bos_token: Whether to prepend BOS token at conversation start.
       reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
   Returns:
       The encoded prompt string.
   """
@@ -578,6 +604,7 @@ def encode_messages(
 def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
   """
   Drop reasoning_content and non-essential messages before the last user message.
+
   Behavior:
   - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
   - Messages at or after the last user index are always kept.
@@ -609,6 +636,7 @@ def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, An
 def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
   """
   Read text from index until one of the stop strings is found.
+
   Returns:
       Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
   """
@@ -632,9 +660,11 @@ def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str,
 def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
   """
   Parse DSML tool calls from text starting at the given index.
+
   Args:
       index: Starting position in text.
       text: The full text to parse.
+
   Returns:
       Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
       Each tool call dict has "name" and "arguments" keys.
@@ -689,16 +719,20 @@ def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Di
 def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
   """
   Parse a model completion text into a structured assistant message.
+
   This function takes the raw text output from the model (a single assistant turn)
   and extracts:
   - reasoning_content (thinking block)
   - content (summary/response)
   - tool_calls (if any)
+
   NOTE: This function is designed to parse only correctly formatted strings and
   will raise ValueError for malformed output.
+
   Args:
       text: The raw completion text (including EOS token).
       thinking_mode: Either "chat" or "thinking".
+
   Returns:
       Dict with keys: "role", "content", "reasoning_content", "tool_calls".
       tool_calls are in OpenAI format.

src/maxtext/input_pipeline/grain_data_processing.py

@@ -509,7 +509,7 @@ def dpo_preprocessing_pipeline(
   return dataset
 
 
-def _format_chat_template_grain(element, data_columns, tokenizer_model):
+def _format_chat_template_grain(element, data_columns, tokenizer_model, model_name=None):
   """Grain-compatible mapping function to format raw columns into conversational messages."""
   # Convert raw columns to conversational messages
   if "messages" in data_columns:
@@ -536,7 +536,7 @@ def _format_chat_template_grain(element, data_columns, tokenizer_model):
   element[data_columns[0]] = messages
 
   return input_pipeline_utils.apply_chat_template(
-      element, tokenizer_model=tokenizer_model, data_column_name=data_columns[0]
+      element, tokenizer_model=tokenizer_model, data_column_name=data_columns[0], model_name=model_name
   )
 
 
@@ -558,7 +558,9 @@ def sft_preprocessing_pipeline(
   """Use grain pipeline to pre-process the dataset and return iterators for sft fine-tuning"""
   dataset = data_processing_utils.parse_and_keep_features(dataset, config, data_columns, tokenize)
 
-  tokenizer_model, pad_id = data_processing_utils.get_tokenizer_and_pad_id(config)
+  # Chunks are tokenized one by one and the chat template already emits BOS/EOS, so the tokenizer
+  # must not add them to every chunk (same as the HF SFT path and vision SFT below).
+  tokenizer_model, pad_id = data_processing_utils.get_tokenizer_and_pad_id(config, add_bos=False, add_eos=False)
   base_tokenizer_model = tokenizer_model
 
   tokenizer_model = getattr(tokenizer_model, "tokenizer", tokenizer_model)
@@ -572,6 +574,7 @@ def sft_preprocessing_pipeline(
           _format_chat_template_grain,
           data_columns=data_columns,
           tokenizer_model=tokenizer_model,
+          model_name=getattr(config, "model_name", None),
       )
   )
 

src/maxtext/input_pipeline/hf_data_processing.py

@@ -398,7 +398,11 @@ def preprocessing_pipeline(
     data_column_names = list(dataset.features.keys())
     dataset = dataset.map(
         input_pipeline_utils.apply_chat_template,
-        fn_kwargs={"tokenizer_model": tokenizer, "data_column_name": data_column_names[0]},
+        fn_kwargs={
+            "tokenizer_model": tokenizer,
+            "data_column_name": data_column_names[0],
+            "model_name": getattr(config, "model_name", None),
+        },
     )
 
   pad_id = _get_pad_id(tokenizer)

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -15,6 +15,7 @@
 """Operations used by Grain"""
 
 import dataclasses
+import re
 import warnings
 from threading import current_thread
 from typing import Any, Iterable, TYPE_CHECKING
@@ -29,6 +30,7 @@
 from grain.experimental import TFRecordIterDataset
 from maxtext.diffusion.block_diffusion import corruption as block_diffusion_corruption
 from maxtext.input_pipeline.protos import example_pb2
+from maxtext.input_pipeline import encoding_dsv4
 from maxtext.input_pipeline import tokenizer
 from maxtext.multimodal import processor as mm_processor
 from maxtext.multimodal import utils as mm_utils
@@ -298,7 +300,76 @@ def _get_completion_in_chat_template(tokenizer_model, round_msgs):
   return completion_str
 
 
-def apply_chat_template(example, tokenizer_model, data_column_name):
+_DS4_TURN_RE = re.compile(
+    re.escape(encoding_dsv4.ASSISTANT_SP_TOKEN)
+    + f"({re.escape(encoding_dsv4.thinking_start_token)}|"
+    + f"{re.escape(encoding_dsv4.thinking_end_token)})"
+    + f"(.*?{re.escape(encoding_dsv4.eos_token)})",
+    re.DOTALL,
+)
+
+
+def _is_deepseek_v4(tokenizer_model, model_name=None) -> bool:
+  """Detects DeepSeek-V4 by model or tokenizer name.
+
+  Vocabulary checks are not used: the DeepSeek-V3.2 tokenizer has the same special tokens and
+  also ships without a chat template, but uses a different encoding.
+  """
+  for name in (model_name, getattr(tokenizer_model, "model_name", None), getattr(tokenizer_model, "name_or_path", None)):
+    if name and any(x in str(name).lower() for x in ("deepseek4", "deepseek-v4", "deepseek_v4")):
+      return True
+  return False
+
+
+def _format_conversations_dsv4(example, data_column_name):
+  """Formats DeepSeek-V4 conversational data using encoding_dsv4.encode_messages."""
+  raw_messages = example[data_column_name]
+  for idx, message in enumerate(raw_messages):
+    if message.get("role") == "system" and idx != 0:
+      raise ValueError(f"System message found at index {idx}. System messages must be at index 0.")
+
+  assistants = [m for m in raw_messages if m.get("role") == "assistant"]
+  if not assistants:
+    example["is_prompt"] = []
+    example[data_column_name] = []
+    return example
+
+  last_assistant_idx = max(i for i, m in enumerate(raw_messages) if m.get("role") == "assistant")
+  messages_to_encode = raw_messages[: last_assistant_idx + 1]
+  thinking_mode = "thinking" if any(m.get("reasoning_content") for m in messages_to_encode) else "chat"
+  text = encoding_dsv4.encode_messages(messages_to_encode, thinking_mode=thinking_mode)
+  matches = list(_DS4_TURN_RE.finditer(text))
+  if len(matches) != len(assistants):
+    raise ValueError(
+        f"Mismatch between number of assistant messages ({len(assistants)}) "
+        f"and encoded turns ({len(matches)}). Encoded text: {text!r}"
+    )
+
+  messages = []
+  is_prompt = []
+  cursor = 0
+  for match in matches:
+    # In thinking mode, encode_messages drops the reasoning of assistant turns before the last user
+    # message and renders them after `</think>`. Keep those turns in the prompt so that only turns
+    # with their reasoning are trained.
+    if thinking_mode == "thinking" and match.group(1) == encoding_dsv4.thinking_end_token:
+      continue
+    messages.extend([text[cursor : match.start(2)], match.group(2)])
+    is_prompt.extend([True, False])
+    cursor = match.end(2)
+
+  if cursor != len(text):
+    raise ValueError(
+        f"Encoded text length ({len(text)}) does not match completion boundary ({cursor}). "
+        f"Trailing conversation content after the final assistant turn was dropped: {text[cursor:]!r}"
+    )
+
+  example["is_prompt"] = is_prompt
+  example[data_column_name] = messages
+  return example
+
+
+def apply_chat_template(example, tokenizer_model, data_column_name, model_name=None):
   """Formats conversational data by applying the tokenizer's chat template
   and identifying prompt/completion segments for SFT masking.
 
@@ -319,6 +390,8 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
         prompt (if applicable). `False` indicates the expected LLM
         completion, excluding the assistant's start tokens.
   """
+  if getattr(tokenizer_model, "chat_template", None) is None and _is_deepseek_v4(tokenizer_model, model_name):
+    return _format_conversations_dsv4(example, data_column_name)
   messages = []
   is_prompt = []
   round_msgs = []