bridge.data.datasets.preference_pair

Module Contents

Classes

PreferencePairDataset

Preference-pair dataset that tokenizes rows on fetch.

Data

logger

API

bridge.data.datasets.preference_pair.logger

‘getLogger(…)’

class bridge.data.datasets.preference_pair.PreferencePairDataset(

source: collections.abc.Sequence[collections.abc.Mapping[str, Any]],

tokenizer,

max_seq_length: int,

pad_token_id: int | None = None,

pad_seq_length_to_mult: int = 1,

chosen_key: str = 'chosen',

rejected_key: str = 'rejected',

prompt_key: str | None = None,

tools_key: str | None = None,

chat_template_kwargs: collections.abc.Mapping[str, Any] | None = None,

ref_logprobs: collections.abc.Mapping[int, collections.abc.Mapping[str, Any]] | None = None,

)

Bases: torch.utils.data.Dataset

Preference-pair dataset that tokenizes rows on fetch.

Fields may be message lists or plain strings; nothing is tokenized at init, and each __getitem__ runs the chat template on that pair. __len__ must stay fixed for the sampler, so invalid pairs are not dropped - they are emitted as two-token stubs with loss_multiplier=0.0 and contribute nothing to the loss. pair_id is the source row index, which is what ref_logprobs (the scorer’s output) is keyed by.

Initialization

__len__() → int

_stub_record(idx: int, reason: str) → dict[str, Any]

__getitem__(idx: int) → dict[str, Any]