Add DeepSeek-V4 chat encoding to SFT preprocessing

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

benchmarks/api_server/encoding/encoding_dsv4.py → 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:
@@ -529,14 +529,17 @@ def _format_chat_template_grain(element, data_columns, tokenizer_model):
     messages = element[data_columns[0]]
 
   assert all(
-      hasattr(m, "__contains__") and "role" in m and "content" in m for m in messages
-  ), f"SFT requires a conversational format. Expected dicts with 'role' and 'content', but got: {messages}"
+      hasattr(m, "__contains__")
+      and "role" in m
+      and any(k in m for k in ("content", "tool_calls", "reasoning_content", "content_blocks"))
+      for m in messages
+  ), ("SFT requires a conversational format. Expected dicts with 'role' and content/tool_calls, " f"but got: {messages}")
 
   # Assign the standardized messages back to the primary column
   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 +561,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 +577,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,89 @@ 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.DS_TASK_SP_TOKENS['action'])})?"
+    + 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 = list(raw_messages[: last_assistant_idx + 1])
+  if example.get("tools") and not any(m.get("tools") for m in messages_to_encode):
+    target_idx = next(
+        (i for i, m in enumerate(messages_to_encode) if m.get("role") in ("system", "developer")),
+        None,
+    )
+    if target_idx is not None:
+      messages_to_encode[target_idx] = {**messages_to_encode[target_idx], "tools": example["tools"]}
+    else:
+      messages_to_encode.insert(0, {"role": "system", "content": "", "tools": example["tools"]})
+  thinking_mode = "thinking" if any(m.get("reasoning_content") for m in messages_to_encode) else "chat"
+  has_tools = bool(example.get("tools")) or any(
+      m.get("tools") or m.get("tool_calls") or m.get("role") == "tool" for m in messages_to_encode
+  )
+  text = encoding_dsv4.encode_messages(messages_to_encode, thinking_mode=thinking_mode, drop_thinking=not has_tools)
+  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 final completion boundary ({cursor}); "
+        f"the final assistant turn was not matched as a trainable completion: {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.
 
@@ -309,6 +393,8 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
       which contains the specific chat template.
     data_column_name: The name of the column in the `example` dictionary
       that contains the list of messages.
+    model_name: Optional model name used as a fallback to identify models
+      that use a custom chat encoder when `tokenizer_model.chat_template` is None.
 
   Returns:
     The modified `example` dictionary.
@@ -323,6 +409,8 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
   is_prompt = []
   round_msgs = []
   try:
+    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)
     for idx, message in enumerate(example[data_column_name]):
       if message["role"] == "system":
         if idx != 0: