nemo_automodel.components.loss.decision_batch

Validate identified question scores against supervision before padding once.

Model adapters attach identities in their actual output order. This boundary can reject mislabeled ordering and widths; it cannot infer identities from anonymous logit values or detect a caller attaching incorrect metadata to those values.

Module Contents

Classes

Name	Description
`DecisionBatch`	Validated, record-major inputs to the decision objective.
`DecisionTargetProtocol`	Read-only supervision required by this component.
`DecisionTargets`	Padded supervision and the original record/question/option identities.
`QuestionLogits`	Scores carrying their model-output identity through batch assembly.

Functions

Name	Description
`_validated_rows`	Check identified ragged rows against the prepared supervision order.
`assemble_decision_batch`	Validate identified scores and pad them using the supervision’s mask.

API

class nemo_automodel.components.loss.decision_batch.DecisionBatch(
    logits: torch.Tensor,
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets
)

Dataclass

Validated, record-major inputs to the decision objective.

logits

Tensor

Tensor of shape [questions, max_options] with -inf padding and autograd connected to every input question’s scores.

targets

DecisionTargets

Supervision tensors with DecisionTargets layouts; its option_mask is the single authoritative mask for these logits.

class nemo_automodel.components.loss.decision_batch.DecisionTargetProtocol()

Protocol

Read-only supervision required by this component.

label

int

Gold option index, or -100 when unlabelled.

option_ids

tuple[str, ...]

Option identities in model-output order.

question_id

str

Question identity within its record.

question_type

int

Loss question type: 0=noul, 1=choice, 2=score.

rewards

tuple[float, ...] | None

Optional utilities aligned with the option identities.

class nemo_automodel.components.loss.decision_batch.DecisionTargets(
    labels: torch.Tensor,
    option_mask: torch.Tensor,
    question_types: torch.Tensor,
    rewards: torch.Tensor | None = None,
    reward_mask: torch.Tensor | None = None,
    ref_logprobs: torch.Tensor | None = None,
    weights: torch.Tensor | None = None,
    record_ids: tuple[str, ...] = (),
    questions: tuple[tuple[nemo_automodel.components.loss.decision_batch.DecisionTargetProtocol, ...], ...] = ()
)

Dataclass

Padded supervision and the original record/question/option identities.

labels

Tensor

Long tensor of shape [questions]; gold option index or -100.

option_mask

Tensor

Bool tensor of shape [questions, max_options]; True for a real option. The assembly boundary uses this same mask for logits.

question_types

Tensor

Long tensor of shape [questions]; 0=noul, 1=choice, 2=score.

questions

tuple[tuple[DecisionTargetProtocol, ...], ...] = ()

Original targets grouped by record, retaining option order. Direct tensor-only construction is valid for isolated loss math; assembly requires identities prepared by from_questions.

record_ids

tuple[str, ...] = ()

Input record identifiers in record-major order.

ref_logprobs

Tensor | None = None

Optional float tensor of shape [questions, max_options].

reward_mask

Tensor | None = None

Optional bool tensor of shape [questions]; explicit rewards are available on True rows, with labelled one-hot rewards otherwise.

rewards

Tensor | None = None

Optional float tensor of shape [questions, max_options].

weights

Tensor | None = None

Optional float tensor of shape [questions].

nemo_automodel.components.loss.decision_batch.DecisionTargets.from_questions(
    record_ids: collections.abc.Sequence[str],
    questions: collections.abc.Sequence[collections.abc.Sequence[nemo_automodel.components.loss.decision_batch.DecisionTargetProtocol]],
    question_weights: collections.abc.Mapping[str, float] | None = None,
    reference_logprobs: collections.abc.Sequence[collections.abc.Sequence[nemo_automodel.components.loss.decision_batch.QuestionLogits]] | None = None,
    device: torch.device | str | None = None
) -> nemo_automodel.components.loss.decision_batch.DecisionTargets

classmethod

Prepare supervision before model forwards, preserving identities.

This preparation allows a distributed optimizer window to reduce its supervision denominator without retaining every microbatch’s graph.

Parameters:

record_ids

Sequence[str]

Input record IDs, aligned with the outer question axis.

questions

Sequence[Sequence[DecisionTargetProtocol]]

Targets in each record’s expected question/option order.

question_weights

Mapping[str, float] | NoneDefaults to None

Optional per-question weights; missing IDs use 1.

reference_logprobs

Sequence[Sequence[QuestionLogits]] | NoneDefaults to None

Optional identified reference tensors of shape [options] per question, in the same nested record/question order. They must be finite; stored copies are detached from autograd.

device

torch.device | str | NoneDefaults to None

Device for the supervision tensors. Defaults to CPU.

Returns: DecisionTargets

Supervision with layouts documented by DecisionTargets. Padding

nemo_automodel.components.loss.decision_batch.DecisionTargets.to(
    device: torch.device | str
) -> nemo_automodel.components.loss.decision_batch.DecisionTargets

Return supervision on device, retaining immutable identity metadata.

Parameters:

device

torch.device | str

Destination for every tensor-bearing field described above.

Returns: DecisionTargets

Targets with the same tensor shapes as this instance. Unchanged-device

class nemo_automodel.components.loss.decision_batch.QuestionLogits(
    record_id: str,
    question_id: str,
    option_ids: tuple[str, ...],
    values: torch.Tensor
)

Dataclass

Scores carrying their model-output identity through batch assembly.

option_ids

tuple[str, ...]

Allowed option identifiers in the score tensor’s order.

question_id

str

Question identifier within the input record.

record_id

str

Input record identifier.

values

Tensor

Floating tensor of shape [options], containing raw logits or, only when used as a reference policy, log-probabilities.

nemo_automodel.components.loss.decision_batch._validated_rows(
    scores: collections.abc.Sequence[collections.abc.Sequence[nemo_automodel.components.loss.decision_batch.QuestionLogits]],
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets
) -> list[torch.Tensor]

Check identified ragged rows against the prepared supervision order.

Parameters:

scores

Sequence[Sequence[QuestionLogits]]

Nested record/question score tensors of shape [options] each.

targets

DecisionTargets

Prepared supervision with the DecisionTargets tensor layouts.

Returns: list[torch.Tensor]

Original score tensors of shape [options], flattened in record-major

nemo_automodel.components.loss.decision_batch.assemble_decision_batch(
    logits: collections.abc.Sequence[collections.abc.Sequence[nemo_automodel.components.loss.decision_batch.QuestionLogits]],
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets
) -> nemo_automodel.components.loss.decision_batch.DecisionBatch

Validate identified scores and pad them using the supervision’s mask.

Parameters:

logits

Sequence[Sequence[QuestionLogits]]

Per-record/question identified tensors of shape [options]. All rows must have the same dtype/device and follow the target identities.

targets

DecisionTargets

Supervision with DecisionTargets layouts, prepared through from_questions and on the same device as logits.

Returns: DecisionBatch

DecisionBatch with record-major logits [questions, max_options] and