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 is DeepSeek-V4, encode with encode_messages and split each assistant turn into prompt and completion; reject encoded text left after the final assistant turn.
- 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-mask checks (packed and unpacked), DeepSeek-V4 detection including unk_token_id fallbacks, and model_name forwarding in hf_data_processing.preprocessing_pipeline.

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
   )
 
 
@@ -572,6 +572,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,84 @@ 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 whether the model or tokenizer is DeepSeek-V4."""
+  for name in (model_name, getattr(tokenizer_model, "model_name", None), getattr(tokenizer_model, "name_or_path", None)):
+    if name:
+      name_lower = str(name).lower()
+      if any(x in name_lower for x in ("deepseek4", "deepseek-v4", "deepseek_v4")):
+        return True
+  try:
+    if hasattr(tokenizer_model, "get_vocab"):
+      vocab = tokenizer_model.get_vocab()
+      if "<｜User｜>" in vocab and "<｜Assistant｜>" in vocab and "｜DSML｜" in vocab:
+        return True
+    elif hasattr(tokenizer_model, "convert_tokens_to_ids"):
+      unk_id = getattr(tokenizer_model, "unk_token_id", None)
+      dsml_id = tokenizer_model.convert_tokens_to_ids("｜DSML｜")
+      user_id = tokenizer_model.convert_tokens_to_ids("<｜User｜>")
+      if dsml_id is not None and dsml_id != unk_id and user_id is not None and user_id != unk_id:
+        return True
+  except Exception:  # pylint: disable=broad-exception-caught
+    pass
+  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:
+    completion = match.group(2)
+    prompt = text[cursor : match.start(2)]
+    messages.extend([prompt, completion])
+    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 +398,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 = []