bridge.data.datasets.preference_tokenization

Module Contents

Functions

tokenize_conversation	Tokenize one conversation into (input_ids, loss_mask), or return a drop reason.
_rendered_eos_end	End (exclusive) of a final-turn EOS followed by nothing but whitespace, or None if there is none.
_as_messages
_decode_tool_calls	HF datasets stores tool_calls as a JSON string when message schemas differ; the template needs the list.
_decode_arguments
build_pair_conversations	Build the chosen/rejected conversations for one row, or return a drop reason.

API

bridge.data.datasets.preference_tokenization.tokenize_conversation(

tokenizer,

messages: list[dict],

max_seq_length: int,

tools: Sequence[Mapping[str, Any]] | None = None,

all_assistant_turns: bool = False,

chat_template_kwargs: Mapping[str, Any] | None = None,

) → tuple[list[int], list[bool]] | str

Tokenize one conversation into (input_ids, loss_mask), or return a drop reason.

Scores the final assistant turn (header included) through its EOS; all_assistant_turns also scores every earlier assistant turn. tools are rendered into the prompt, never scored. chat_template_kwargs reach the template (e.g. truncate_history_thinking=False keeps the reasoning of assistant turns before the last user turn, which Nemotron 3 / Qwen3 drop by default).

bridge.data.datasets.preference_tokenization._rendered_eos_end(

tokenizer,

input_ids: list[int],

context_len: int,

eos_token_id: int,

) → int | None

End (exclusive) of a final-turn EOS followed by nothing but whitespace, or None if there is none.

bridge.data.datasets.preference_tokenization._as_messages(

value: str | Sequence[Mapping[str, Any]],

role: str,

) → list[dict]

bridge.data.datasets.preference_tokenization._decode_tool_calls(

message: Mapping[str, Any],

) → dict[str, Any]

HF datasets stores tool_calls as a JSON string when message schemas differ; the template needs the list.

bridge.data.datasets.preference_tokenization._decode_arguments(

call: Mapping[str, Any],

) → dict[str, Any]

bridge.data.datasets.preference_tokenization.build_pair_conversations(

row: Mapping[str, Any],

*,

chosen_key: str = 'chosen',

rejected_key: str = 'rejected',

prompt_key: str | None = None,

) → tuple[list[dict], list[dict]] | str

Build the chosen/rejected conversations for one row, or return a drop reason.

Full trajectories (no prompt_key) must first differ at an assistant turn; the rest is shared context.