Implementation:Princeton nlp SimPO Apply Chat Template
| Knowledge Sources | |
|---|---|
| Domains | NLP, Data_Preprocessing |
| Last Updated | 2026-02-08 04:30 GMT |
Overview
Concrete tool for formatting preference data into model-specific prompt/chosen/rejected text strings, provided by the SimPO training script.
Description
The apply_chat_template function in run_simpo.py extends the base alignment library version with a simpo task type. When task="simpo", it splits each example into prompt (N-1 turns), chosen response (final chosen turn), and rejected response (final rejected turn), then applies the tokenizer's chat template. A key SimPO-specific behavior is stripping the BOS token from chosen and rejected text to prevent double-BOS artifacts during tokenization. It also supports swapping to a Mistral-specific chat template when the model name contains "mistral". Helper functions is_openai_format and maybe_insert_system_message validate input format and handle system message insertion.
Usage
Use after loading datasets and before passing data to SimPOTrainer. Applied via dataset.map() to transform every example in the dataset.
Code Reference
Source Location
- Repository: SimPO
- File: scripts/run_simpo.py (Lines 48-121)
- File: alignment/data.py (Lines 28-39, 111-122)
Signature
def apply_chat_template(
example: dict,
tokenizer: PreTrainedTokenizerBase,
task: Literal["sft", "generation", "rm", "simpo"],
auto_insert_empty_system_msg: bool = True,
change_template: Optional[str] = None,
) -> dict:
"""
Apply chat template to a single dataset example.
Args:
example: Dataset row with 'chosen' and 'rejected' keys
in OpenAI message format.
tokenizer: The tokenizer with a chat_template attribute.
task: Formatting mode. Use "simpo" for SimPO training.
auto_insert_empty_system_msg: Insert empty system message
if template expects one.
change_template: Set to "mistral" for Mistral models.
Returns:
Dict with 'text_prompt', 'text_chosen', 'text_rejected' keys.
"""
def is_openai_format(messages: Any) -> bool:
"""Check if messages are in OpenAI format (list of role/content dicts)."""
def maybe_insert_system_message(messages: list, tokenizer) -> None:
"""Insert empty system message if template references 'system'."""
Import
# The SimPO-specific version is defined in run_simpo.py
# For direct use:
from alignment.data import is_openai_format, maybe_insert_system_message
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| example | dict | Yes | Dataset row with "chosen" and "rejected" keys in OpenAI message format |
| tokenizer | PreTrainedTokenizerBase | Yes | Tokenizer with chat_template configured |
| task | str | Yes | Must be "simpo" for SimPO training |
| auto_insert_empty_system_msg | bool | No | Whether to insert empty system message (default: True) |
| change_template | str | No | Set to "mistral" for Mistral models |
Outputs
| Name | Type | Description |
|---|---|---|
| example["text_prompt"] | str | Formatted prompt text (all turns except final response) |
| example["text_chosen"] | str | Formatted chosen response text (BOS stripped) |
| example["text_rejected"] | str | Formatted rejected response text (BOS stripped) |
Usage Examples
Applying SimPO Chat Template
from alignment import get_tokenizer, ModelArguments, DataArguments
# Get tokenizer
tokenizer = get_tokenizer(model_args, data_args)
# Detect Mistral models for template override
change_template = "mistral" if "mistral" in model_args.model_name_or_path.lower() else None
# Apply chat template to entire dataset
raw_datasets = raw_datasets.map(
apply_chat_template,
fn_kwargs={
"tokenizer": tokenizer,
"task": "simpo",
"auto_insert_empty_system_msg": data_args.auto_insert_empty_system_msg,
"change_template": change_template,
},
num_proc=data_args.preprocessing_num_workers,
remove_columns=column_names,
desc="Formatting comparisons with prompt template",
)
# Rename columns to match SimPOTrainer expectations
for split in ["train", "test"]:
raw_datasets[split] = raw_datasets[split].rename_columns(
{"text_prompt": "prompt", "text_chosen": "chosen", "text_rejected": "rejected"}
)
Example Input/Output
# Input example (OpenAI format)
example = {
"chosen": [
{"role": "user", "content": "What is 2+2?"},
{"role": "assistant", "content": "2+2 equals 4."},
],
"rejected": [
{"role": "user", "content": "What is 2+2?"},
{"role": "assistant", "content": "The answer is 5."},
],
}
# After apply_chat_template(example, tokenizer, task="simpo"):
# example["text_prompt"] -> "<|user|>\nWhat is 2+2?<eos>\n<|assistant|>\n"
# example["text_chosen"] -> "2+2 equals 4.<eos>" (BOS stripped)
# example["text_rejected"] -> "The answer is 5.<eos>" (BOS stripped)