typedecide.training.dataset

Turning decisions into single-token-supervised examples. Background: Answer-token-only loss. Importing this module does not import torch.

typedecide.training.dataset

Turning Decisions into the one supervised token a typed readout ever reads.

The invariant this whole module exists to hold:

input_ids = encode(full_prompt(state, criterion)) + [answer_token]
labels    = [-100] * len(encode(full_prompt(state, criterion))) + [answer_token]
answer_token = prompt.answer_slots(tokenizer, prompt, n_options)[0][answer_index]

Every token except the last is masked. The browser never samples a sentence: it reads the logits at the position right after "Answer:" and softmaxes over that criterion's answer letters, so that position is the only one worth a gradient.

build_example is a pure function over the prompt.Tokenizer protocol -- .encode and nothing else -- so the invariant is testable against a stub with no torch and no model download. DecisionDataset is the thin indexable wrapper the trainer hands to transformers.Trainer; it deliberately does not subclass torch.utils.data.Dataset, because that base class contributes no behaviour and subclassing it would drag torch into import time for the sake of a name.

IGNORE_INDEX module-attribute

IGNORE_INDEX = -100

PREAMBLE module-attribute

PREAMBLE = f'{INSTRUCTION}\n\nEvidence:\n'

DecisionDataset

DecisionDataset(
    decisions: Sequence[Decision],
    tokenizer: Tokenizer,
    max_length: int,
    *,
    randomise_option_order: bool = False,
    seed: int = 0
)

An indexable view of decisions as masked, single-token-supervised examples.

Duck-types torch.utils.data.Dataset (__len__ and __getitem__ are all transformers.Trainer asks for) without importing torch, which keeps the masking invariant testable with a stub tokenizer.

Option-order randomisation is applied per epoch. set_epoch re-permutes every decision's options from seed and the epoch number, and the trainer calls it from an on_epoch_begin callback. Per epoch is affordable here because a permutation is a tuple rebuild, not a re-tokenisation, and it is strictly better than once up front: a model that sees each decision in a different order each epoch cannot learn "the answer is A". The manifest records which scheme ran.

Source code in src/typedecide/training/dataset.py

def __init__(
    self,
    decisions: Sequence[Decision],
    tokenizer: Tokenizer,
    max_length: int,
    *,
    randomise_option_order: bool = False,
    seed: int = 0,
) -> None:
    if not decisions:
        raise TrainingError("No decisions to train on; the dataset is empty.")
    self._source: tuple[Decision, ...] = tuple(decisions)
    self._tokenizer = tokenizer
    self._max_length = max_length
    self._randomise = randomise_option_order
    self._seed = seed
    self._epoch = 0
    self._rows: tuple[Decision, ...] = self._source
    if self._randomise:
        self.set_epoch(0)

epoch property

epoch: int

decisions property

decisions: tuple[Decision, ...]

The decisions as currently ordered -- after this epoch's permutation.

set_epoch

set_epoch(epoch: int) -> None

Re-permute the options for epoch. A no-op when randomisation is off.

Source code in src/typedecide/training/dataset.py

def set_epoch(self, epoch: int) -> None:
    """Re-permute the options for `epoch`. A no-op when randomisation is off."""
    self._epoch = epoch
    if not self._randomise:
        return
    # Mixing the epoch into the seed keeps every epoch different and the whole run
    # reproducible from the single `seed` in the config.
    self._rows = tuple(reorder_options(self._source, seed=self._seed * 1_000_003 + epoch))

preflight

preflight() -> dict[str, dict[str, int]]

Resolve the answer slots once per distinct criterion, before training starts.

Returns {criterion key: {letter: token id}} for the manifest, and raises TrainingError naming the offending decision if any criterion has no stable answer slot. Catching that now rather than at step 4,000 of a twelve-hour run is the entire point.

Source code in src/typedecide/training/dataset.py

def preflight(self) -> dict[str, dict[str, int]]:
    """Resolve the answer slots once per distinct criterion, before training starts.

    Returns `{criterion key: {letter: token id}}` for the manifest, and raises
    `TrainingError` naming the offending decision if any criterion has no stable
    answer slot. Catching that now rather than at step 4,000 of a twelve-hour run
    is the entire point.
    """
    from ..schema import LETTERS

    table: dict[str, dict[str, int]] = {}
    seen: set[tuple[str, tuple[str, ...]]] = set()
    for decision in self._source:
        criterion = decision.criterion
        signature = (criterion.question, criterion.option_ids)
        if signature in seen:
            continue
        seen.add(signature)
        slots, separator = answer_slots_for(self._tokenizer, decision)
        resolved = dict(zip(LETTERS, slots, strict=False))
        previous = table.get(criterion.key)
        if previous is not None and previous != resolved:
            log.warning(
                "Criterion %s resolves to different answer-slot ids for different "
                "option sets; the manifest records the last one seen.",
                criterion.key,
            )
        table[criterion.key] = resolved
        log.debug(
            "Criterion %s: answer slots %s with separator %r.",
            criterion.key, resolved, separator,
        )
    return table

AnswerTokenCollator

AnswerTokenCollator(pad_token_id: int)

Collates build_example outputs into the tensors Trainer expects.

Written rather than borrowed from transformers so the padding rule above is stated in this repository, where the invariant it protects also lives.

Source code in src/typedecide/training/dataset.py

def __init__(self, pad_token_id: int) -> None:
    self.pad_token_id = pad_token_id

pad_token_id instance-attribute

pad_token_id = pad_token_id

build_example

build_example(
    tokenizer: Tokenizer,
    decision: Decision,
    max_length: int,
) -> dict[str, list[int]]

One training example: the whole prompt masked out, plus one supervised letter.

Returns input_ids, attention_mask and labels, all the same length and all at most max_length long.

Truncation, when the prompt does not fit, removes tokens from the left of the state. It never touches the tail -- the criterion, the lettered options and "Answer:" -- because the answer slot is defined by the characters immediately before it, and trimming the tail would move the supervised position onto a token that is not an answer letter at all. The instruction preamble is kept at the front where it can be, so what is dropped is the oldest evidence, the way a human skimming a long record would drop it.

Raises TrainingError, naming the decision, if the row is unlabelled, if the tokenizer has no stable answer slot for this prompt, or if the criterion alone is already longer than max_length.

Source code in src/typedecide/training/dataset.py

def build_example(
    tokenizer: Tokenizer, decision: Decision, max_length: int
) -> dict[str, list[int]]:
    """One training example: the whole prompt masked out, plus one supervised letter.

    Returns `input_ids`, `attention_mask` and `labels`, all the same length and all at
    most `max_length` long.

    Truncation, when the prompt does not fit, removes tokens **from the left of the
    state**. It never touches the tail -- the criterion, the lettered options and
    "Answer:" -- because the answer slot is defined by the characters immediately
    before it, and trimming the tail would move the supervised position onto a token
    that is not an answer letter at all. The instruction preamble is kept at the front
    where it can be, so what is dropped is the oldest evidence, the way a human
    skimming a long record would drop it.

    Raises `TrainingError`, naming the decision, if the row is unlabelled, if the
    tokenizer has no stable answer slot for this prompt, or if the criterion alone is
    already longer than `max_length`.
    """
    if not decision.labelled:
        raise TrainingError(
            f"Decision {decision.id!r} has no answer_id, so there is no token to train "
            "on. Filter to labelled decisions before training, or label this row."
        )
    if max_length < 2:
        raise TrainingError(
            f"max_length must leave room for at least the answer token, got {max_length}."
        )

    criterion = decision.criterion
    prompt = full_prompt(decision.state, criterion)

    try:
        slots, _separator = answer_slots_for(tokenizer, decision)
        answer_token = slots[decision.answer_index]
    except SchemaError as error:  # answer_id not in this criterion; schema caught it
        raise TrainingError(f"Decision {decision.id!r}: {error}") from error

    prompt_ids = list(tokenizer.encode(prompt, add_special_tokens=False))
    if not prompt_ids:
        raise TrainingError(f"Decision {decision.id!r} encoded to zero prompt tokens.")

    budget = max_length - 1  # the answer token is not negotiable
    if len(prompt_ids) > budget:
        prompt_ids = _truncate_state_from_left(tokenizer, decision, prompt_ids, budget)

    input_ids = [*prompt_ids, answer_token]
    labels = [IGNORE_INDEX] * len(prompt_ids) + [answer_token]
    return {
        "input_ids": input_ids,
        "attention_mask": [1] * len(input_ids),
        "labels": labels,
    }

answer_slots_for

answer_slots_for(
    tokenizer: Tokenizer, decision: Decision
) -> tuple[list[int], str]

The answer-letter token ids at this decision's prompt boundary.

A thin wrapper over prompt.answer_slots whose only job is to attach the decision id to the failure. A row whose answer slot cannot be resolved is a row whose loss would land on the wrong token, so it is reported, never quietly dropped.

Source code in src/typedecide/training/dataset.py

def answer_slots_for(tokenizer: Tokenizer, decision: Decision) -> tuple[list[int], str]:
    """The answer-letter token ids at this decision's prompt boundary.

    A thin wrapper over `prompt.answer_slots` whose only job is to attach the decision
    id to the failure. A row whose answer slot cannot be resolved is a row whose loss
    would land on the wrong token, so it is reported, never quietly dropped.
    """
    prompt = full_prompt(decision.state, decision.criterion)
    try:
        return answer_slots(tokenizer, prompt, len(decision.criterion.options))
    except PromptError as error:
        raise TrainingError(
            f"Decision {decision.id!r} (criterion {decision.criterion.key!r}): {error}"
        ) from error

reorder_options

reorder_options(
    decisions: Sequence[Decision], *, seed: int
) -> list[Decision]

Permute each decision's options, preserving answer_id.

Defers to typedecide.data.augment.randomise_option_order, so the package has exactly one implementation of the permutation and one seeding scheme. (An inline fallback used to live here for the days before data/ existed; it seeded by row position rather than by decision id, so had it ever run, the same seed would have produced a different order. It was unreachable and is gone.)

Source code in src/typedecide/training/dataset.py

def reorder_options(decisions: Sequence[Decision], *, seed: int) -> list[Decision]:
    """Permute each decision's options, preserving `answer_id`.

    Defers to `typedecide.data.augment.randomise_option_order`, so the package has
    exactly one implementation of the permutation and one seeding scheme. (An inline
    fallback used to live here for the days before `data/` existed; it seeded by row
    position rather than by decision id, so had it ever run, the same seed would have
    produced a different order. It was unreachable and is gone.)
    """
    from ..data.augment import randomise_option_order

    return list(randomise_option_order(decisions, seed=seed))

pad_batch

pad_batch(
    features: Sequence[dict[str, list[int]]],
    pad_token_id: int,
) -> dict[str, list[list[int]]]

Right-pad a batch of examples to a common length.

Right-padding is correct for training: the real tokens keep positions 0..n-1, which is what the default position_ids give, and the padding is masked out of attention and carries -100 labels so it contributes no loss. (Batched scoring is the opposite case and must left-pad with explicit position ids; that belongs to evaluation/, not here.)

Source code in src/typedecide/training/dataset.py

def pad_batch(
    features: Sequence[dict[str, list[int]]], pad_token_id: int
) -> dict[str, list[list[int]]]:
    """Right-pad a batch of examples to a common length.

    Right-padding is correct for *training*: the real tokens keep positions
    0..n-1, which is what the default `position_ids` give, and the padding is masked
    out of attention and carries `-100` labels so it contributes no loss. (Batched
    *scoring* is the opposite case and must left-pad with explicit position ids;
    that belongs to `evaluation/`, not here.)
    """
    if not features:
        raise TrainingError("Cannot collate an empty batch.")
    width = max(len(f["input_ids"]) for f in features)
    batch: dict[str, list[list[int]]] = {"input_ids": [], "attention_mask": [], "labels": []}
    for feature in features:
        pad = width - len(feature["input_ids"])
        batch["input_ids"].append(feature["input_ids"] + [pad_token_id] * pad)
        batch["attention_mask"].append(feature["attention_mask"] + [0] * pad)
        batch["labels"].append(feature["labels"] + [IGNORE_INDEX] * pad)
    return batch

labelled_only

labelled_only(
    decisions: Iterable[Decision],
) -> tuple[list[Decision], list[str]]

Split decisions into the trainable ones and the ids of the unlabelled rest.

The caller decides what to do about the unlabelled ones; this function never discards a row without handing back its id, because "n went in, fewer came out" with no names attached is how a bad dataset survives a run.

Source code in src/typedecide/training/dataset.py

def labelled_only(decisions: Iterable[Decision]) -> tuple[list[Decision], list[str]]:
    """Split decisions into the trainable ones and the ids of the unlabelled rest.

    The caller decides what to do about the unlabelled ones; this function never
    discards a row without handing back its id, because "n went in, fewer came out"
    with no names attached is how a bad dataset survives a run.
    """
    keep: list[Decision] = []
    dropped: list[str] = []
    for decision in decisions:
        if decision.labelled:
            keep.append(decision)
        else:
            dropped.append(decision.id)
    return keep, dropped