Implementation:Eric mitchell Direct preference optimization Tokenize Batch Element

Knowledge Sources	Direct Preference Optimization
Domains	Preprocessing, NLP, Text_Processing
Last Updated	2026-02-08 02:00 GMT

Overview

Concrete tool for tokenizing individual prompt-response pairs with truncation and label masking provided by the direct-preference-optimization repository.

Description

The tokenize_batch_element function tokenizes a single (prompt, chosen, rejected) triple. It handles truncation (prompt first, then response), appends EOS tokens, creates label sequences with prompt positions masked as -100, and returns a dictionary of token sequences ready for batching.

Usage

Called by get_batch_iterator for each example during data pipeline construction. In SFT mode, the sft_target is passed as both chosen and rejected arguments.

Code Reference

Source Location

Repository: direct-preference-optimization
File: preference_datasets.py
Lines: 214-277

Signature

def tokenize_batch_element(
    prompt: str,
    chosen: str,
    rejected: str,
    truncation_mode: str,
    tokenizer,
    max_length: int,
    max_prompt_length: int,
) -> Dict:
    """Tokenize a single batch element.

    Handles truncation (prompt first, then response), appends EOS,
    and creates labels with -100 for prompt tokens.
    """

Import

from preference_datasets import tokenize_batch_element

I/O Contract

Inputs

Name	Type	Required	Description
prompt	str	Yes	The prompt text (e.g., "\n\nHuman: ...\n\nAssistant:")
chosen	str	Yes	The preferred response text
rejected	str	Yes	The dispreferred response text (same as chosen in SFT mode)
truncation_mode	str	Yes	"keep_start" or "keep_end" for prompt truncation direction
tokenizer	PreTrainedTokenizer	Yes	HuggingFace tokenizer
max_length	int	Yes	Maximum combined prompt+response length
max_prompt_length	int	Yes	Maximum prompt length after truncation

Outputs

Name	Type	Description
batch	Dict	Contains: prompt (str), chosen (str), rejected (str), chosen_response_only (str), rejected_response_only (str), chosen_input_ids (List[int]), chosen_attention_mask (List[int]), chosen_labels (List[int] with -100 for prompt), rejected_input_ids, rejected_attention_mask, rejected_labels, prompt_input_ids, prompt_attention_mask

Usage Examples

Tokenizing a Preference Pair

from preference_datasets import tokenize_batch_element
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2-large")

element = tokenize_batch_element(
    prompt="\n\nHuman: What is 2+2?\n\nAssistant:",
    chosen=" 4.",
    rejected=" 22.",
    truncation_mode="keep_end",
    tokenizer=tokenizer,
    max_length=512,
    max_prompt_length=256,
)

# element['chosen_labels'] has -100 for prompt positions
# element['chosen_input_ids'] = prompt_tokens + chosen_tokens + [EOS]

Related Pages

Implements Principle

Principle:Eric_mitchell_Direct_preference_optimization_Tokenization

Requires Environment

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment