nemo_automodel.components.models.inkling.multimodal

View as Markdown

Native Inkling vision, audio, and multimodal composition modules.

Module Contents

Classes

NameDescription
InklingAudioEmbeddingsEmbed and sum the independently quantized dMel codebooks.
InklingAudioModelNative Inkling audio tower.
InklingModelCompose native text, vision, and audio Inkling towers.
InklingModelOutputWithPastInkling backbone output including projected image features.
InklingVisionEncoderLayerFold time and space into channels, then project one HMLP stage.
InklingVisionModelNative hierarchical MLP vision tower.

Functions

NameDescription
_minimum_cost_scale_indicesSelect distinct ordered scales with minimum absolute-log cost.
_plan_output_scalesPlan the time, height, width, and channel sizes of every HMLP stage.
_prime_factorsReturn prime factors in ascending order.

Data

__all__

API

class nemo_automodel.components.models.inkling.multimodal.InklingAudioEmbeddings(
config: nemo_automodel.components.models.inkling.configuration.InklingAudioConfig
)

Bases: Module

Embed and sum the independently quantized dMel codebooks.

embed_audio_tokens
nemo_automodel.components.models.inkling.multimodal.InklingAudioEmbeddings.forward(
input_ids: torch.Tensor
) -> torch.Tensor

Embed discretized mel bins.

Parameters:

input_ids
torch.Tensor

Long tensor of shape [frames, mel_bins].

Returns: torch.Tensor

Tensor of shape [frames, hidden].

class nemo_automodel.components.models.inkling.multimodal.InklingAudioModel(
config: nemo_automodel.components.models.inkling.configuration.InklingAudioConfig
)

Bases: Module

Native Inkling audio tower.

embed_audio_tokens
= InklingAudioEmbeddings(config)
norm
nemo_automodel.components.models.inkling.multimodal.InklingAudioModel.forward(
audio_input_ids: torch.Tensor
) -> transformers.modeling_outputs.BaseModelOutputWithPooling

Encode audio tokens.

Parameters:

audio_input_ids
torch.Tensor

Long tensor of shape [frames, mel_bins].

Returns: BaseModelOutputWithPooling

An output whose hidden state and pooler output have shape [frames, hidden].

nemo_automodel.components.models.inkling.multimodal.InklingAudioModel.init_weights() -> None

Initialize the audio embedding and final normalization.

Bases: Module

Compose native text, vision, and audio Inkling towers.

audio_tower
= InklingAudioModel(config.audio_config)
language_model
vision_tower
= InklingVisionModel(config.vision_config)
nemo_automodel.components.models.inkling.multimodal.InklingModel._get_placeholder_mask(
input_ids: torch.Tensor | None,
inputs_embeds: torch.Tensor,
features: torch.Tensor,
token_id: int
) -> torch.Tensor

Match multimodal features to placeholder tokens.

Parameters:

input_ids
torch.Tensor | None

Optional long tensor of shape [batch, sequence].

inputs_embeds
torch.Tensor

Tensor of shape [batch, sequence, hidden].

features
torch.Tensor

Tensor of shape [placeholders, hidden].

token_id
int

Vocabulary ID used as the placeholder.

Returns: torch.Tensor

Boolean tensor of shape [batch, sequence, hidden].

nemo_automodel.components.models.inkling.multimodal.InklingModel.forward(
input_ids: torch.LongTensor | None = None,
pixel_values: torch.FloatTensor | None = None,
audio_input_ids: torch.LongTensor | None = None,
audio_input_ids_mask: torch.Tensor | None = None,
attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
inputs_embeds: torch.FloatTensor | None = None,
use_cache: bool | None = None,
kwargs: typing.Any = {}
) -> nemo_automodel.components.models.inkling.multimodal.InklingModelOutputWithPast

Merge multimodal features and run the text backbone.

Parameters:

input_ids
torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence].

pixel_values
torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [patches, time, height, width, channels].

audio_input_ids
torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [audios, frames, mel_bins].

audio_input_ids_mask
torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [audios, frames].

attention_mask
torch.Tensor | dict[str, torch.Tensor] | NoneDefaults to None

Optional padding tensor of shape [batch, total_sequence] or mapping of prepared attention masks.

position_ids
torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence].

past_key_values
InklingDynamicCache | NoneDefaults to None

Optional model-owned decoding cache.

inputs_embeds
torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [batch, sequence, hidden].

use_cache
bool | NoneDefaults to None

Whether to allocate and return a decoding cache.

**kwargs
AnyDefaults to {}

Additional text-attention arguments.

Returns: InklingModelOutputWithPast

An output whose last hidden state has shape [batch, sequence, hidden].

nemo_automodel.components.models.inkling.multimodal.InklingModel.get_audio_features(
audio_input_ids: torch.Tensor,
audio_input_ids_mask: torch.Tensor | None = None
) -> transformers.modeling_outputs.BaseModelOutputWithPooling

Project discretized audio into text hidden space.

Parameters:

audio_input_ids
torch.Tensor

Long tensor of shape [audios, frames, mel_bins].

audio_input_ids_mask
torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [audios, frames].

Returns: BaseModelOutputWithPooling

An output whose hidden state has shape [valid_frames, text_hidden].

nemo_automodel.components.models.inkling.multimodal.InklingModel.get_image_features(
pixel_values: torch.Tensor,
kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPooling

Project image/video patches into text hidden space.

Parameters:

pixel_values
torch.Tensor

Tensor of shape [patches, time, height, width, channels].

**kwargs
AnyDefaults to {}

Reserved for the common multimodal calling convention.

Returns: BaseModelOutputWithPooling

An output whose pooler output has shape [patches, text_hidden].

nemo_automodel.components.models.inkling.multimodal.InklingModel.get_input_embeddings() -> torch.nn.Module

Return the text token-embedding module.

nemo_automodel.components.models.inkling.multimodal.InklingModel.init_weights(
buffer_device: torch.device
) -> None

Initialize every tower for checkpoint-free construction.

class nemo_automodel.components.models.inkling.multimodal.InklingModelOutputWithPast(
image_hidden_states: torch.FloatTensor | None = None
)
Dataclass

Bases: BaseModelOutputWithPast

Inkling backbone output including projected image features.

image_hidden_states
FloatTensor | None = None
class nemo_automodel.components.models.inkling.multimodal.InklingVisionEncoderLayer(
input_dim: int,
output_dim: int,
temporal_fold: int,
spatial_fold: int,
add_norm: bool,
dtype: torch.dtype
)

Bases: Module

Fold time and space into channels, then project one HMLP stage.

layer_norm
projection
nemo_automodel.components.models.inkling.multimodal.InklingVisionEncoderLayer._fold_timespace_to_depth(
hidden_states: torch.Tensor
) -> torch.Tensor

Fold time and two spatial axes into the channel dimension.

Parameters:

hidden_states
torch.Tensor

Tensor of shape [patches, time, height, width, channels].

Returns: torch.Tensor

Tensor of shape “[patches, time / temporal_fold, height / spatial_fold,

nemo_automodel.components.models.inkling.multimodal.InklingVisionEncoderLayer.forward(
hidden_states: torch.Tensor
) -> torch.Tensor

Apply one hierarchical vision projection.

Parameters:

hidden_states
torch.Tensor

Tensor of shape [patches, time, height, width, channels].

Returns: torch.Tensor

Tensor of shape [patches, folded_time, folded_height, folded_width, output_channels].

nemo_automodel.components.models.inkling.multimodal.InklingVisionEncoderLayer.init_weights(
init_std: float
) -> None

Initialize the projection and optional normalization.

class nemo_automodel.components.models.inkling.multimodal.InklingVisionModel(
config: nemo_automodel.components.models.inkling.configuration.InklingVisionConfig
)

Bases: Module

Native hierarchical MLP vision tower.

encoder_layers
= nn.ModuleList()
final_norm
nemo_automodel.components.models.inkling.multimodal.InklingVisionModel.forward(
pixel_values: torch.Tensor,
kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPooling

Encode preprocessed image/video patches.

Parameters:

pixel_values
torch.Tensor

Tensor of shape [patches, time, height, width, channels].

**kwargs
AnyDefaults to {}

Reserved for the common multimodal calling convention.

Returns: BaseModelOutputWithPooling

An output whose pooler output has shape [patches, text_hidden].

nemo_automodel.components.models.inkling.multimodal.InklingVisionModel.init_weights() -> None

Initialize every vision projection and normalization.

nemo_automodel.components.models.inkling.multimodal._minimum_cost_scale_indices(
cost_matrix: torch.Tensor
) -> torch.LongTensor

Select distinct ordered scales with minimum absolute-log cost.

Parameters:

cost_matrix
torch.Tensor

Tensor of shape [layers_plus_one, available_scales].

Returns: torch.LongTensor

Long tensor of shape [layers_plus_one] indexing distinct scales.

nemo_automodel.components.models.inkling.multimodal._plan_output_scales(
config: nemo_automodel.components.models.inkling.configuration.InklingVisionConfig
) -> torch.LongTensor

Plan the time, height, width, and channel sizes of every HMLP stage.

nemo_automodel.components.models.inkling.multimodal._prime_factors(
number: int
) -> list[int]

Return prime factors in ascending order.

nemo_automodel.components.models.inkling.multimodal.__all__ = ['InklingModel', 'InklingModelOutputWithPast']