Fix chat template prompt/completion formatting and SFT masking for Gemma 4 reasoning

[csgoogle]

Description

Fixes conversational round formatting in apply_chat_template and SFT prompt masking for reasoning models (such as Gemma 4).

Problem & Evidence on Unmodified Code

Previously in apply_chat_template, the prompt segment was generated eagerly when iterating over the user message:

prompt_in_chat_template = tokenizer_model.apply_chat_template(
    round_msgs, add_generation_prompt=True, tokenize=False
)

For Gemma 4, add_generation_prompt=True emits <bos><|turn>user\n...<turn|>\n<|turn>model\n<|channel>thought\n<channel|>.

When running unit tests against the unmodified code, 5 out of 8 tests failed:

1. Assistant turn WITH reasoning (test_single_turn_with_thinking - FAIL on unmodified code):
   The completion returned by _get_completion_in_chat_template started after <|channel>thought\n with {reasoning}\n<channel|>{content}<turn|>\n. Concatenating prompt + completion produced:

     <|turn>model
     <|channel>thought
   - <channel|>We will reformulate this using a single vector reduction pass.
   + We will reformulate this using a single vector reduction pass.
     <channel|>```python

   The thought channel was prematurely closed by <channel|> before the reasoning trace ever began, followed by duplicate closing tags.

2. Assistant turn WITHOUT reasoning (test_single_turn_without_thinking - FAIL on unmodified code):
   Because the prompt was generated before inspecting the assistant message, the empty thought channel remained in the prompt:

     <|turn>model
   - <|channel>thought
   - <channel|>The capital of France is Paris.<turn|>
   + The capital of France is Paris.<turn|>

   A phantom <|channel>thought\n<channel|> was injected into normal conversational turns where ground truth had no thought channel.

---

Solution

We extracted turn splitting into a clean, reusable helper: _split_turn_into_prompt_and_completion(tokenizer_model, round_msgs):

def _split_turn_into_prompt_and_completion(tokenizer_model, round_msgs):
  """Splits a conversational round (system + user + assistant) into prompt and completion formatted strings."""
  full_tokens = extract_token_ids(
      tokenizer_model.apply_chat_template(round_msgs, add_generation_prompt=False, tokenize=True)
  )
  prompt_tokens = extract_token_ids(
      tokenizer_model.apply_chat_template(round_msgs[:-1], add_generation_prompt=True, tokenize=True)
  )

  # Find the longest common prefix where prompt ends and completion begins
  common_len = 0
  for fid, pid in zip(full_tokens, prompt_tokens):
    if fid == pid:
      common_len += 1
    else:
      break

  if common_len == 0:
    raise ValueError(
        "Chat template generation prompt mismatch: no common prefix tokens found.\n"
        f"Full conversation tokens: {full_tokens} ('{tokenizer_model.decode(full_tokens)}')\n"
        f"Generation prompt tokens: {prompt_tokens} ('{tokenizer_model.decode(prompt_tokens)}')\n"
        "Cannot determine completion boundary."
    )

  prompt_str = tokenizer_model.decode(full_tokens[:common_len], skip_special_tokens=False)
  completion_str = tokenizer_model.decode(full_tokens[common_len:], skip_special_tokens=False)
  return prompt_str, completion_str

And simplified apply_chat_template:

def apply_chat_template(example, tokenizer_model, data_column_name):
  messages = []
  is_prompt = []
  round_msgs = []
  try:
    for idx, message in enumerate(example[data_column_name]):
      if message["role"] == "system":
        if idx != 0:
          raise ValueError(f"System message found at index {idx}. System messages must be at index 0.")
        round_msgs.append(message)
      elif message["role"] == "user":
        round_msgs.append(message)
      elif message["role"] == "assistant":
        round_msgs.append(message)
        prompt_str, completion_str = _split_turn_into_prompt_and_completion(tokenizer_model, round_msgs)
        messages.extend([prompt_str, completion_str])
        is_prompt.extend([True, False])
        round_msgs.clear()
  except ValueError as e:
    max_logging.log(f"Unable to apply chat template: {e}")
    raise e
  example["is_prompt"] = is_prompt
  example[data_column_name] = messages
  return example

Guarantees

• Exact Token Invariant: Slicing full_tokens into full_tokens[:common_len] and full_tokens[common_len:] guarantees prompt_str + completion_str == tokenizer.apply_chat_template(round_msgs, add_generation_prompt=False, tokenize=False).
• Model Agnostic: Automatically works for thinking models (Gemma 4, Qwen 3, DeepSeek R1) and standard models (Llama 2/3) without requiring any model-specific branches or hardcoded strings.

Tests

• python3 -m unittest -v tests/unit/chat_template_sft_test.py (8/8 tests pass)
• python3 -m unittest -v tests.post_training.unit.sft_data_processing_test.SFTChatTemplateLogicTest (4/4 tests pass)

[gemini-code-assist]

Code Review

This pull request refactors the chat template formatting in input_pipeline_utils.py to determine prompt and completion boundaries using token-level common prefix matching, which correctly handles Gemma 4 thinking/reasoning channels. It also adds comprehensive unit tests to verify SFT prompt masking and thinking channel boundaries. The reviewer noted that directly loading the gated google/gemma-4-31b-it model from Hugging Face in unit tests will fail in unauthenticated CI/CD environments, and suggested using a local path or environment variable fallback.

[SurbhiJainUSC]

nit: let's keep it as full_id and prompt_id