nemo_automodel.components.models.d1.model

Checkpoint-compatible Omni model with a differentiable decision API.

Module Contents

Classes

Name	Description
`D1OmniModel`	Omni’s bidirectional trunk and released marker head, preserving all weights.

API

class nemo_automodel.components.models.d1.model.D1OmniModel(
    config: nemo_automodel.components.models.d1.config.D1OmniConfig,
    trust_remote_code: bool = False
)

Bases: HFCheckpointingMixin, PreTrainedModel

Omni’s bidirectional trunk and released marker head, preserving all weights.

Perception classes come from reviewed, immutable released code. The language encoder and decision API are model-owned. Training returns tensors with their autograd graph intact, unlike the released inference probability interface.

audio

= audio_class(config.audio_config, width)

base_model_prefix

= 'd1'

encoder

= D1OmniEncoder(config.text_config)

head

main_input_name

= 'input_ids'

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vision

nemo_automodel.components.models.d1.model.D1OmniModel.encode_audio(
    audio: typing.Any
) -> torch.Tensor

Encode a mono 16 kHz clip into a perception prefix.

Parameters:

audio

Any

Float Tensor or array of shape [samples], containing a mono waveform sampled at 16 kHz, as accepted by the pinned perception module.

Returns: torch.Tensor

Tensor of shape [1, prefix_sequence, hidden].

nemo_automodel.components.models.d1.model.D1OmniModel.encode_images(
    images: list[typing.Any]
) -> torch.Tensor

Encode ordered PIL images into a perception prefix.

Returns: torch.Tensor

Tensor of shape [1, prefix_sequence, hidden], concatenating all images.

nemo_automodel.components.models.d1.model.D1OmniModel.encode_text(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    media_prefix: torch.Tensor | None = None,
    media_lengths: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]

Encode text with optional precomputed perception prefixes.

Parameters:

input_ids

torch.Tensor

Long Tensor of shape [batch, sequence], with right padding.

attention_mask

torch.Tensor

Boolean Tensor of shape [batch, sequence], true for text tokens.

media_prefix

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, padded_prefix_sequence, hidden].

media_lengths

torch.Tensor | NoneDefaults to None

Optional long Tensor of shape [batch] specifying real prefix lengths.

Returns: torch.Tensor

Pair of normalized text-state Tensor of shape

nemo_automodel.components.models.d1.model.D1OmniModel.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    marker_positions: torch.Tensor,
    marker_mask: torch.Tensor,
    question_types: torch.Tensor,
    media_prefix: torch.Tensor | None = None,
    media_lengths: torch.Tensor | None = None,
    labels: torch.Tensor | None = None
) -> transformers.modeling_outputs.SequenceClassifierOutput

Return differentiable option logits in the released internal ordering.

Parameters:

input_ids

torch.Tensor

Long Tensor of shape [batch, sequence], with right padding.

attention_mask

torch.Tensor

Boolean Tensor of shape [batch, sequence], true for text tokens.

marker_positions

torch.Tensor

Long Tensor of shape [batch, options], relative to text.

marker_mask

torch.Tensor

Boolean Tensor of shape [batch, options] marking declared options.

question_types

torch.Tensor

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

media_prefix

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, padded_prefix_sequence, hidden].

media_lengths

torch.Tensor | NoneDefaults to None

Optional long Tensor of shape [batch] giving real prefix lengths.

labels

torch.Tensor | NoneDefaults to None

Optional long Tensor of shape [batch] of internal option indices.

Returns: SequenceClassifierOutput

SequenceClassifierOutput containing float32 logits Tensor of shape

nemo_automodel.components.models.d1.model.D1OmniModel.get_input_embeddings() -> torch.nn.Embedding

Return the text input embedding table.

nemo_automodel.components.models.d1.model.D1OmniModel.prepare_text_readout(
    examples: list[nemo_automodel.components.models.d1.records.D1TextExample]
) -> nemo_automodel.components.models.d1.text_readout.D1OmniTextReadout

Prepare the released marker readout for text-only adaptation.

Parameters:

examples

list[D1TextExample]

Publisher-encoded questions; marker identity and calibration remain attached to each example and are consumed when scoring.

Returns: D1OmniTextReadout

Sequence feature extractor and the released float32 marker head.

nemo_automodel.components.models.d1.model.D1OmniModel.set_input_embeddings(
    value: torch.nn.Embedding
) -> None

Replace the text embedding table.