`nemo_automodel.components.models.bagel.attention_masks`#

Packed multimodal attention mask helpers.

Module Contents#

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

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._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.create_sparse_mask( document_lens: list[int], split_lens: list[int], attn_modes: list[str], device: torch.device, ) → Any#: Create the block-mask predicate for packed multimodal attention.