nemo_automodel.components.models.bagel.attention_masks

Packed multimodal attention mask helpers.

Module Contents

Functions

Name	Description
`_repeat_document_labels`	Return a per-token document id for packed samples.
`_repeat_split_labels`	Return per-token split labels used by the packed block-mask predicate.
`create_sparse_mask`	Create the block-mask predicate for packed multimodal attention.

API

nemo_automodel.components.models.bagel.attention_masks._repeat_document_labels(
    document_lens: list[int],
    device: torch.device
) -> torch.Tensor

Return a per-token document id for packed samples.

nemo_automodel.components.models.bagel.attention_masks._repeat_split_labels(
    split_lens: list[int],
    attn_modes: list[str],
    device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]

Return per-token split labels used by the packed block-mask predicate.

nemo_automodel.components.models.bagel.attention_masks.create_sparse_mask(
    document_lens: list[int],
    split_lens: list[int],
    attn_modes: list[str],
    device: torch.device
) -> typing.Any

Create the block-mask predicate for packed multimodal attention.