nemo_automodel.components.models.d1.decision_head

D1 readouts with explicit option order and frozen vocabulary semantics.

Module Contents

Classes

Name	Description
`D1OmniDecisionHead`	Differentiable released Omni head; noul markers retain false/true order.
`D1OmniHeadConfig`	Released Omni marker head configuration.
`D1TokenReadout`	Train additive answer-row updates while preserving input/output tying.
`D1TokenReadoutConfig`	Select answer-token rows without untying the released vocabulary head.

API

class nemo_automodel.components.models.d1.decision_head.D1OmniDecisionHead(
    config: nemo_automodel.components.models.d1.decision_head.D1OmniHeadConfig
)

Bases: Module

Differentiable released Omni head; noul markers retain false/true order.

State-dict names match the published checkpoint’s head subtree. Text-only temperature calibration belongs outside this module so media paths can skip it.

head

scorer

type_emb

= nn.Embedding(3, width)

nemo_automodel.components.models.d1.decision_head.D1OmniDecisionHead.forward(
    hidden_states: torch.Tensor,
    valid_tokens: torch.Tensor,
    marker_positions: torch.Tensor,
    valid_markers: torch.Tensor,
    question_types: torch.Tensor
) -> torch.Tensor

Score each option marker in internal checkpoint order.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden], excluding media prefixes.

valid_tokens

torch.Tensor

Boolean Tensor of shape [batch, sequence], true for non-padding tokens.

marker_positions

torch.Tensor

Long Tensor of shape [batch, options], indexing text positions.

valid_markers

torch.Tensor

Boolean Tensor of shape [batch, options], true for declared options.

question_types

torch.Tensor

Long Tensor of shape [batch] with choice=0, score=1, noul=2.

Returns: torch.Tensor

Float32 Tensor of shape [batch, options], with padding options set to -10000.

class nemo_automodel.components.models.d1.decision_head.D1OmniHeadConfig(
    hidden_size: int,
    layers: int = 2
)

Dataclass

Released Omni marker head configuration.

Parameters:

hidden_size

int

Encoder width, a positive multiple of the 64-wide head dimension.

layers

intDefaults to 2

Number of bidirectional transformer layers in the readout.

hidden_size

int

layers

int = 2

nemo_automodel.components.models.d1.decision_head.D1OmniHeadConfig.__post_init__() -> None

nemo_automodel.components.models.d1.decision_head.D1OmniHeadConfig.build() -> nemo_automodel.components.models.d1.decision_head.D1OmniDecisionHead

Build an uncalibrated, checkpoint-compatible Omni head.

class nemo_automodel.components.models.d1.decision_head.D1TokenReadout(
    config: nemo_automodel.components.models.d1.decision_head.D1TokenReadoutConfig,
    vocabulary_weight: torch.Tensor
)

Bases: Module

Train additive answer-row updates while preserving input/output tying.

This is an adaptation of the released readout, not an independent classifier. The zero-initialized adapter preserves the max-over-alias option contract.

delta_weight

nemo_automodel.components.models.d1.decision_head.D1TokenReadout.forward(
    hidden_states: torch.Tensor,
    alias_groups: tuple[tuple[int, ...], ...]
) -> torch.Tensor

Score declared options by their maximum answer-alias logit.

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions, at each causal chain’s answer position.

alias_groups

tuple[tuple[int, ...], ...]

Adapter row indices for each option, in target order.

Returns: torch.Tensor

Float32 Tensor of shape […, options], before option softmax.

class nemo_automodel.components.models.d1.decision_head.D1TokenReadoutConfig(
    hidden_size: int,
    token_ids: tuple[int, ...]
)

Dataclass

Select answer-token rows without untying the released vocabulary head.

Parameters:

hidden_size

int

Width of the normalized final backbone states.

token_ids

tuple[int, ...]

Unique answer-alias vocabulary IDs in adapter row order.

hidden_size

int

token_ids

tuple[int, ...]

nemo_automodel.components.models.d1.decision_head.D1TokenReadoutConfig.__post_init__() -> None

nemo_automodel.components.models.d1.decision_head.D1TokenReadoutConfig.build(
    vocabulary_weight: torch.Tensor
) -> nemo_automodel.components.models.d1.decision_head.D1TokenReadout

Build an adapter from the frozen released vocabulary weight.

Parameters:

vocabulary_weight

torch.Tensor

Tensor of shape [vocabulary, hidden_size].

Returns: D1TokenReadout

Zero-initialized row adapter, with detached base rows.