`nemo_rl.data.multimodal_utils`#

Module Contents#

Classes#

PackedTensor

A logical batch of rows backed by packable tensor segments.

Functions#

`uses_image_placeholder`	Return whether a processor requires explicit image placeholders.
`row_shapes_key`	Companion tag key carrying per-row shapes for `field`.
`present_multimodal_fields`	Multimodal wire fields the rollout actually wrote for this batch.
`multimodal_row_tags`	Per-sample tag rows carrying what `to_wire`’s flattening removes.
`reassemble_packed_multimodal`	In place: rebuild `PackedTensor` for packed fields, consuming companions.
`_patchify_segments`	Cut pixel segments into vision patches and pack them into one sequence.
`_shared_preprocess_spec`	Return the preprocessing setting shared by every input.
`encode_multimodal_for_wire`	The wire value for one multimodal field. Dispatched by registry membership.
`get_multimodal_keys_from_processor`	Get keys of the multimodal data that can be used as model inputs.
`get_multimodal_default_settings_from_processor`
`get_dim_to_pack_along`	Special considerations for packing certain keys from certain processors.
`get_preprocess`	Return materialization preprocessing for one processor input.
`extract_multimodal_model_inputs`	Extract packed media inputs and sequence-aligned auxiliary tensors.
`resolve_to_image`	Resolve the image path to a PIL.Image object.
`image_size_from_source`	Return `(width, height)` for an image source.
`image_to_data_url`	Encode a PIL Image as a base64 `data:` URL.
`get_responses_content_part_url`	Return a string media source from a Responses/Chat content part.
`extract_input_media_sources_from_responses_messages`	Extract tagged image and video sources in encounter order.
`media_sources_equal`	Compare tagged media by string value or object identity.
`_materialize_ragged_pixel_values`	Preserve a ragged per-image `pixel_values` list for materialization.
`_stack_ragged_pixel_values`	Derive image sizes and preserve native-shape tiles for patchification.
`_restore_tensors`	Convert a processor’s list outputs back to tensors, in place.
`_image_num_tokens_from_processed`
`attach_image_model_inputs_to_message`	Attach image tensors while preserving rollout-token/media parity.
`attach_processor_media_token_validity_mask`	Record which rollout tokens are processor-owned media placeholders.
`video_path_to_data_url`	Inline a local or `file://` video as a base64 data URL.
`get_media_from_message`	Get all media from a message log item.
`load_media_from_message`
`build_media_token_validity_mask`	Mark media-token positions in rows that carry no media of that modality.
`media_placeholder_token_id_from_chunks`	The vocabulary id these model chunks treat as a media placeholder, if any.
`chunks_accept_media_token_validity_mask`	Whether a model chunk’s forward takes an explicit media-token validity mask.
`image_counts_by_row`	How many images each row of the batch actually carries.
`attach_media_token_validity_mask`	Attach a `[B, S]` mask of media-token positions that anchor media features.

API#

nemo_rl.data.multimodal_utils.VLLM_CONTENT_KEY#: ‘vllm_content’

nemo_rl.data.multimodal_utils.VLLM_MULTI_MODAL_DATA_KEY#: ‘vllm_multi_modal_data’

nemo_rl.data.multimodal_utils.VLLM_PROMPT_KEYS#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.IMAGE_CONTENT_TYPES#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.VIDEO_CONTENT_TYPES#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.AUDIO_CONTENT_TYPES#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.MULTIMODAL_CONTENT_TYPES#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.ROLLOUT_MATCHED_MEDIA_KEY#: ‘_rollout_matched_media’

nemo_rl.data.multimodal_utils.MEDIA_TAGS#: None

nemo_rl.data.multimodal_utils.MEDIA_TAGS_REVERSED#: None

nemo_rl.data.multimodal_utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC#: b’NEMO_RL_CACHED_VIDEO_FRAMES_V1\n’

nemo_rl.data.multimodal_utils.CACHED_VIDEO_FRAME_MANIFEST_MIME#: ‘video/x-nemo-rl-cached-frames’

nemo_rl.data.multimodal_utils.DEFAULT_MEDIA_EXTENSIONS#: None

nemo_rl.data.multimodal_utils._PLACEHOLDER_STYLE_PROCESSOR_NAMES#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.MEDIA_TAGS_TO_ALLOWED#: None

nemo_rl.data.multimodal_utils.MEDIA_TAG_PATTERN#: ‘compile(…)’

nemo_rl.data.multimodal_utils.logger#: ‘getLogger(…)’

nemo_rl.data.multimodal_utils.uses_image_placeholder(processor: Any) → bool#

Return whether a processor requires explicit image placeholders.

Parameters:: processor – Multimodal processor to classify.
Returns:: Whether the processor expands image placeholders through __call__ rather than tokenized apply_chat_template.

nemo_rl.data.multimodal_utils.PER_TOKEN_MULTIMODAL_FIELDS#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.PACKED_MULTIMODAL_FIELDS#: ‘frozenset(…)’

nemo_rl.data.multimodal_utils.ROW_SHAPES_SUFFIX#: ‘__row_shapes’

nemo_rl.data.multimodal_utils.row_shapes_key(field: str) → str#: Companion tag key carrying per-row shapes for field.

nemo_rl.data.multimodal_utils.ROW_GEOMETRY_SHAPES#: ‘shapes’

nemo_rl.data.multimodal_utils.ROW_GEOMETRY_PREPROCESS_MODE#: ‘preprocess_mode’

nemo_rl.data.multimodal_utils.ROW_GEOMETRY_PREPROCESS_KWARGS#: ‘preprocess_kwargs’

nemo_rl.data.multimodal_utils.WIRE_MULTIMODAL_FIELDS#: None

nemo_rl.data.multimodal_utils.present_multimodal_fields( meta: nemo_rl.data_plane.interfaces.KVBatchMeta, ) → list[str]#

Multimodal wire fields the rollout actually wrote for this batch.

Intersecting with meta.fields is required, not defensive: the noop adapter and the TQ contract both raise on a fetch for a field that was never written, and text-only runs write none of these. Sorted for a deterministic field list across ranks.

nemo_rl.data.multimodal_utils.multimodal_row_tags( multimodal: dict, sample_count: int, ) → list[dict[str, Any]] | None#

Per-sample tag rows carrying what to_wire’s flattening removes.

KVBatchMeta.tags is the transport’s channel for per-sample primitives: it is aligned 1:1 with sample_ids and projected automatically by subset/slice/concat, so a worker holding a shard gets its own rows without anyone re-keying them. The data plane never interprets the contents.

Carries shapes (per-row, and unrecoverable once to_wire flattens) and the field’s preprocessing settings. Deliberately not a pad target: the width padding lands at is scratch that the model discards. Mcore consumes pre-patchified pixels without spatial dimensions, while AutoModel pads at materialization, so nothing batch-wide has to be agreed across shards.

nemo_rl.data.multimodal_utils.reassemble_packed_multimodal( fields: dict, tags: list[dict[str, Any]] | None = None, ) → None#

In place: rebuild PackedTensor for packed fields, consuming companions.

Called by the read path once its columns are materialized. Leaves anything that is not a wire-form packed field untouched, so it is safe to call on any column dict.

Raises:: ValueError – A packed field arrived without its shapes companion. Every producer mints the two together (multimodal_row_tags beside encode_multimodal_for_wire) and every KVBatchMeta transform projects tags alongside sample_ids, so a missing companion means the transport dropped it. Reconstructing anyway would hand the model 1-D pixels and train image-blind without an error.

nemo_rl.data.multimodal_utils._patchify_segments( segments: list[torch.Tensor], *, patch_dim: int, ) → torch.Tensor#

Cut pixel segments into vision patches and pack them into one sequence.

Each [N, channels, H, W] segment is processed at its native resolution into a [C_i, P²] block, where C_i is N * rows * columns and P² is the flattened patch width (channels * patch_dim**2). Blocks are packed along dimension zero, then a batch dimension is added to produce [1, total_C, P²]. Already-patchified segments in that final layout are accepted so repeated materialization is safe.

nemo_rl.data.multimodal_utils._shared_preprocess_spec( from_packed_tensors: list[PackedTensor], ) → dict[str, Any]#: Return the preprocessing setting shared by every input.

class nemo_rl.data.multimodal_utils.PackedTensor( tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]], dim_to_pack: int, *, preprocess_mode: Optional[str] = None, preprocess_kwargs: Optional[dict[str, Any]] = None, _row_offsets: Optional[list[int]] = None, _segment_indices: Optional[list[int]] = None, _segment_provenance: Optional[list[bytes]] = None, )#

A logical batch of rows backed by packable tensor segments.

The default representation is intentionally the legacy one: every entry in tensors is one logical row and no deduplication metadata is allocated. enable_deduplication adds stable provenance to the physical segments. Operations that combine or slice dedup-enabled values then use a CSR-like logical-row mapping:

_row_offsets partitions the flattened logical segment references.
_segment_indices maps each logical segment reference to tensors.
_segment_provenance is stable across deepcopy/pickle and is the only evidence used to re-intern physical segments.

Prompt identity is deliberately absent: belonging to the same prompt group makes media a candidate for sharing, but never proves media equality.

Worked example. Prompt A has one image and prompt B has two, each already merged by :meth:merge_segments into a single logical row::

A: tensors = [tA]        _row_offsets = [0, 1]  _segment_indices = [0]
B: tensors = [tB0, tB1]  _row_offsets = [0, 2]  _segment_indices = [0, 1]

GRPO then expands each prompt into G generations. The rows are deep-copied while _prepare_multimodal_sharing aliases the media leaves, so

Meth:

concat re-interns them by provenance. For G=3::

tensors = [tA, tB0, tB1] # 3 physical segments _row_offsets = [0, 1, 2, 3, 5, 7, 9] # 6 logical rows _segment_indices = [0, 0, 0, 1, 2, 1, 2, 1, 2] len() = 6

Six logical rows over nine segment references and three physical tensors, so physical memory is flat in G. Row 4 reads as _segment_indices[_row_offsets[4]:_row_offsets[5]] == [1, 2], i.e. both of prompt B’s images, without copying either.

Initialization

Wrap per-item tensors for concatenation along dim_to_pack.

Parameters:

tensors – A tensor or list of per-item tensors. List entries may be None for items without this modality.
dim_to_pack – Dimension along which as_tensor concatenates.
preprocess_mode – Optional preprocessing applied by as_tensor. Supported values are pad_to_max_shape and patchify. Patchify requires dim_to_pack=0 and changes 4-D inputs into a 3-D [1, total_patches, P²] tensor.
preprocess_kwargs – Extra arguments for preprocess_mode. Patchify requires patch_dim.

__setstate__(state: dict[str, Any]) → None#: Restore both current and pre-deduplication pickled instances.

broadcast_parts() → tuple[nemo_rl.data.multimodal_utils.PackedTensor, list[Optional[tuple[int, ...]]], str, str, list[torch.Tensor]]#

Split semantic state from physical data for replica broadcast.

The tensor-free header retains every other instance attribute, so new semantic state does not require a parallel broadcast schema update.

rebuild_from_broadcast_parts( shapes: list[Optional[tuple[int, ...]]], payload: torch.Tensor, ) → nemo_rl.data.multimodal_utils.PackedTensor#: Attach received physical data to a header from broadcast_parts.

property _preprocess_spec: dict[str, Any]#: Return keyword arguments that preserve preprocessing in a copy.

property deduplication_enabled: bool#: Whether this value carries stable physical-segment provenance.

logical_segment_counts_by_row() → list[int]#: Return the number of non-empty media segments in each logical row.

iter_logical_segments()#: Yield physical tensor segments in logical row/segment order.

enable_deduplication() → nemo_rl.data.multimodal_utils.PackedTensor#: Assign stable provenance lazily without changing logical contents.

_row_segment_indices(row: int) → list[int]#

__deepcopy__( memo: dict[int, Any], ) → nemo_rl.data.multimodal_utils.PackedTensor#: Share immutable media segments only for an explicitly enabled value.

as_tensor( device: torch.device | None = None, mode: str | None = None, ) → Optional[torch.Tensor]#

__len__() → int#

to( device: str | torch.device, ) → nemo_rl.data.multimodal_utils.PackedTensor#

Move physical segments in place, retaining provenance.

A device move is value-preserving, so provenance stays valid and two copies of the same segment still re-intern. This is the opposite of

Meth:: to_dtype, which changes values and therefore mints fresh provenance. The asymmetry is deliberate: :meth:concat re-interns on provenance alone, so a value-changing operation must not keep it.

The corollary is that two values sharing provenance on different devices would merge to whichever appears first in concat. That is currently unreachable – BatchedDataDict.to and get_multimodal_dict only touch top-level values, while message-level segments are nested inside message_log and worker-side device moves happen post-serialization on independent copies – and it would surface as a loud mixed-device torch.cat error rather than silent corruption. Keep it that way: do not move a subset of segments.

to_dtype( dtype: torch.dtype, ) → nemo_rl.data.multimodal_utils.PackedTensor#

Return an independent wrapper without expanding logical segments.

Dtype conversion creates new physical tensor values, so deduplicated inputs receive new provenance. When the dtype already matches, immutable tensor segments and their provenance remain shared, but mutable wrapper state is copied. The logical row-to-segment mapping is preserved exactly.

Non-floating-point segments are returned unchanged. Integer media metadata (grid sizes, frame counts) is index data, so casting it to a float dtype would silently corrupt it.

slice( indices: Union[list[int], torch.Tensor], ) → nemo_rl.data.multimodal_utils.PackedTensor#

classmethod empty_like( other: nemo_rl.data.multimodal_utils.PackedTensor, ) → nemo_rl.data.multimodal_utils.PackedTensor#: Return empty logical rows matching other.

classmethod empty_rows_like( other: nemo_rl.data.multimodal_utils.PackedTensor, num_rows: int, ) → nemo_rl.data.multimodal_utils.PackedTensor#: Return num_rows logical rows containing no media segments.

classmethod concat( from_packed_tensors: list[nemo_rl.data.multimodal_utils.PackedTensor], ) → nemo_rl.data.multimodal_utils.PackedTensor#

Concatenate a list of PackedTensor objects into a single PackedTensor.

The underlying tensors from the PackedTensors are combined into a single list of tensors and used to create a new PackedTensor.

Each batch must have the same dim_to_pack.

Example:

>>> import torch
>>> from nemo_rl.data.multimodal_utils import PackedTensor
>>> p1 = PackedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6])], dim_to_pack=0)
>>> p2 = PackedTensor([torch.tensor([7, 8, 9])], dim_to_pack=0)
>>> p3 = PackedTensor.concat([p1, p2])
>>> p3.tensors
[tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9])]
>>> p3.as_tensor()
tensor([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>>

classmethod merge_segments( from_packed_tensors: list[nemo_rl.data.multimodal_utils.PackedTensor], ) → nemo_rl.data.multimodal_utils.PackedTensor#

Merge message-turn values, collapsing compact inputs to one row.

The legacy path retains one logical row per physical segment; its caller materializes those segments into one conversation row later.

classmethod flattened_concat( from_packed_tensors: list[nemo_rl.data.multimodal_utils.PackedTensor], ) → nemo_rl.data.multimodal_utils.PackedTensor#

Given a list of PackedTensor objects, flattens each PackedTensor and then concatenates them into a single PackedTensor.

Each PackedTensor is first flattened by packing along the PackedTensor’s dim_to_pack dimension. Then, the resulting flattened tensors are used to create a new PackedTensor.

This is different from PackedTensor.concat which simply extends the underlying list of tensors. This is important because the slice and __len__ methods operate on the underlying list of tensors. Note, however, that calling as_tensor on the resulting PackedTensor will result in the same tensor as concat.

Each batch must have the same dim_to_pack.

Example:

>>> import torch
>>> from nemo_rl.data.multimodal_utils import PackedTensor
>>> p1 = PackedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6])], dim_to_pack=0)
>>> p2 = PackedTensor([torch.tensor([7, 8, 9])], dim_to_pack=0)
>>> p3 = PackedTensor.flattened_concat([p1, p2])
>>> p3.tensors
[tensor([1, 2, 3, 4, 5, 6]), tensor([7, 8, 9])]
>>> p3.as_tensor()
tensor([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>>

_row_segments() → list[list[torch.Tensor]]#

Per logical row, its non-empty physical segments.

One entry per logical row. Under deduplication a logical row maps to several shared physical segments (_row_offsets / _segment_indices), so iterating self.tensors directly would yield the physical segment count instead of the batch size. _row_segment_indices degrades to [row] for the legacy one-tensor-per-row layout.

static _shapes_of( row_segments: list[list[torch.Tensor]], ) → list[list[list[int]]]#

shapes[row][segment] for an already-walked row/segment list.

Shared by :meth:row_shapes and :meth:to_wire so the shape encoding cannot drift between the tag minter and the payload encoder.

row_shapes() → list[list[list[int]]]#

The shapes half of :meth:to_wire, without encoding the payload.

Callers that only need the geometry – :func:multimodal_row_tags – use this so they do not pay to_wire’s torch.cat a second time.

to_wire() → tuple[Optional[torch.Tensor], list[list[list[int]]]]#

Encode as a flattened torch.jagged tensor plus its row shapes.

Returns (None, []) when every logical row is empty so the caller can skip the field entirely. Otherwise:

nested — one row per logical row, each row the 1-D concat of that row’s segments. Flattening is what makes torch.jagged total: rows vary only in dim 0, so ragged trailing dims and mixed rank both encode, and nothing is padded to make a container accept them. Empty rows become zero-length placeholders.
shapes — shapes[row][segment] is that segment’s true shape. Required because TQ derives per_sample_shapes from the value it is handed, so flat rows make it record flat lengths. The caller ships this on KVBatchMeta.tags; see

func:

nemo_rl.data.multimodal_utils.multimodal_row_tags.

Only dim_to_pack=0 is supported today; other values would need ragged_idx on the nested tensor and a matching transpose in :meth:from_wire.

classmethod from_wire( nested: torch.Tensor, shapes: list[list[list[int]]], *, preprocess_mode: Optional[str] = None, preprocess_kwargs: Optional[dict[str, Any]] = None, ) → Optional[nemo_rl.data.multimodal_utils.PackedTensor]#

Reconstruct from the value produced by :meth:to_wire.

Returns None for an empty batch (no rows); the caller should skip the field entirely rather than instantiate an empty PackedTensor.

shapes is the companion returned by :meth:to_wire, carried on KVBatchMeta.tags (see :func:multimodal_row_tags). Each flat row is split by segment numel and reshaped back to its true shape. It is required, not optional: reconstructing without it yields 1-D segments, which train image-blind without erroring.

A zero-length row means the sample had no media – it becomes None (not a (0, ...) tensor) so an image-free shard reconstructs as legacy does: as_tensor returns None and logical_segment_counts_by_row reports 0 rather than 1.

The preprocessing settings are restored onto the rebuilt value, not materialized. Segments come back at their true shapes and stay separate via the CSR row map; :meth:as_tensor preprocesses at use time.

Mirrors :meth:to_wire; both assume dim_to_pack=0.

nemo_rl.data.multimodal_utils.encode_multimodal_for_wire( k: str, v: Union[nemo_rl.data.multimodal_utils.PackedTensor, torch.Tensor], ) → Optional[torch.Tensor]#

The wire value for one multimodal field. Dispatched by registry membership.

Returns None when the field has nothing to ship (an all-empty packed batch), in which case the caller omits the column entirely. The wire key is always k: one wire field per logical field.

Payload only. Per-token fields ride rectangular; packed fields ride as one flattened torch.jagged value. The geometry :meth:PackedTensor.from_wire needs to undo that flattening – per-row segment shapes plus the preprocessing settings – is minted separately by

Func:: multimodal_row_tags and shipped on KVBatchMeta.tags. TQ cannot derive it, because it reads per_sample_shapes off the flattened rows it is handed.

nemo_rl.data.multimodal_utils.UNDECLARED_MULTIMODAL_MODEL_INPUTS#: (‘imgs_sizes’, ‘num_frames’, ‘pixel_values_flat’, ‘image_num_patches’)

nemo_rl.data.multimodal_utils.get_multimodal_keys_from_processor(processor) → list[str]#

Get keys of the multimodal data that can be used as model inputs.

This will be used in the data_processor function to determine which keys to use as model inputs.

nemo_rl.data.multimodal_utils.get_multimodal_default_settings_from_processor( processor, ) → dict[str, dict[str, Any]]#

nemo_rl.data.multimodal_utils.get_dim_to_pack_along(processor, key: str) → int#

Special considerations for packing certain keys from certain processors.

In most cases, the packed items are along dim 0

nemo_rl.data.multimodal_utils.get_preprocess( processor: Any, key: str, ) → dict[str, Any]#: Return materialization preprocessing for one processor input.

nemo_rl.data.multimodal_utils.extract_multimodal_model_inputs( processor: Any, processed: dict[str, Any], ) → dict[str, nemo_rl.data.multimodal_utils.PackedTensor | torch.Tensor]#: Extract packed media inputs and sequence-aligned auxiliary tensors.

nemo_rl.data.multimodal_utils.resolve_to_image( image_path_or_image: str | PIL.Image.Image, ) → PIL.Image.Image#

Resolve the image path to a PIL.Image object.

image_path can be either:

path to local file
url to image
base64 encoded image

nemo_rl.data.multimodal_utils.image_size_from_source( source: str | PIL.Image.Image, ) → tuple[int, int]#: Return (width, height) for an image source.

nemo_rl.data.multimodal_utils.image_to_data_url(image: PIL.Image.Image, fmt: str = 'PNG') → str#

Encode a PIL Image as a base64 data: URL.

Parameters:

image – PIL image to encode.
fmt – PIL image format used for serialization (e.g. "PNG", "JPEG"). The value is also lowercased and embedded in the MIME type of the returned URL.

Returns:

A data:image/<fmt>;base64,<payload> URL suitable for embedding in an OpenAI Responses input_image content part.

nemo_rl.data.multimodal_utils.get_responses_content_part_url( part: dict[str, Any], *keys: str, ) → str#: Return a string media source from a Responses/Chat content part.

nemo_rl.data.multimodal_utils.extract_input_media_sources_from_responses_messages( messages: Any, ) → list[tuple[str, Any]]#: Extract tagged image and video sources in encounter order.

nemo_rl.data.multimodal_utils.media_sources_equal( left: tuple[str, Any], right: tuple[str, Any], ) → bool#: Compare tagged media by string value or object identity.

nemo_rl.data.multimodal_utils._materialize_ragged_pixel_values( processed: dict[str, Any], processor: Any, ) → dict[str, Any]#

Preserve a ragged per-image pixel_values list for materialization.

Tiles remain separate until PackedTensor.as_tensor materializes them.

nemo_rl.data.multimodal_utils._stack_ragged_pixel_values( processed: dict[str, Any], tiles: list[torch.Tensor], processor: Any, ) → None#: Derive image sizes and preserve native-shape tiles for patchification.

nemo_rl.data.multimodal_utils._restore_tensors(processed: dict[str, Any]) → None#

Convert a processor’s list outputs back to tensors, in place.

return_tensors=None makes the processor hand back every output as plain Python lists, not only the ragged pixel_values that mode was requested for. Downstream expects tensors – input_ids in particular is rank-checked – so restore the rest of the batch to what return_tensors="pt" would have produced. Values that resist conversion (genuinely ragged per-image metadata) are left for their own handling.

nemo_rl.data.multimodal_utils._image_num_tokens_from_processed( processed: dict[str, Any], ) → list[int]#

nemo_rl.data.multimodal_utils.attach_image_model_inputs_to_message( message: dict[str, Any], *, images: list[PIL.Image.Image], processor: Any, pad_dynamic_image_shapes: bool = False, expected_num_tokens_per_image: collections.abc.Sequence[int] | None = None, ) → None#: Attach image tensors while preserving rollout-token/media parity.

nemo_rl.data.multimodal_utils.attach_processor_media_token_validity_mask( message: dict[str, Any], processor: Any, ) → None#: Record which rollout tokens are processor-owned media placeholders.

nemo_rl.data.multimodal_utils._VIDEO_EXT_TO_MIME#: None

nemo_rl.data.multimodal_utils.video_path_to_data_url(video_path: str) → str#: Inline a local or file:// video as a base64 data URL.

nemo_rl.data.multimodal_utils.get_media_from_message( message: dict[str, Any], ) → dict[str, list[Any]]#: Get all media from a message log item.

nemo_rl.data.multimodal_utils.load_media_from_message( message: dict[str, Any], processor=None, multimodal_load_kwargs: Optional[dict[str, dict[str, Any]]] = None, ) → dict[str, list[Any]]#

nemo_rl.data.multimodal_utils.build_media_token_validity_mask( input_ids: torch.Tensor, media_token_id: int, media_counts_by_row: collections.abc.Sequence[int], base_mask: Optional[torch.Tensor] = None, ) → Optional[torch.Tensor]#

Mark media-token positions in rows that carry no media of that modality.

A media token is an ordinary vocabulary entry with its own embedding row. It only means “a projected feature belongs here” when media is attached; in a text-only row the same id is whatever the author wrote, and counting it as a placeholder makes the model demand a feature that does not exist.

Rows that do carry media keep every position valid, so a real placeholder/feature disagreement there is still reported rather than silently masked away.

Parameters:

input_ids – [B, S] token ids, one row per sample.
media_token_id – Vocabulary id the model treats as a media placeholder.
media_counts_by_row –
Attached media items per row, e.g. from

meth:

PackedTensor.logical_segment_counts_by_row.
base_mask – Optional [B, S] validity mask to refine, so masks for several modalities can be combined.

Returns:

A [B, S] bool mask, or None when nothing needs masking and the caller should leave the model’s own derivation untouched.

nemo_rl.data.multimodal_utils.media_placeholder_token_id_from_chunks( chunks: collections.abc.Sequence[Any], ) → Optional[int]#: The vocabulary id these model chunks treat as a media placeholder, if any.

nemo_rl.data.multimodal_utils.chunks_accept_media_token_validity_mask( chunks: collections.abc.Sequence[Any], ) → bool#

Whether a model chunk’s forward takes an explicit media-token validity mask.

Checked against the signature rather than a class flag so a model that does not know about the mask never receives it: such a forward would absorb it into **kwargs and silently ignore it, which looks identical to the mask having been applied. Failing to send it is visible; sending it into a void is not.

nemo_rl.data.multimodal_utils.image_counts_by_row( batch: Any, num_rows: int, ) → Optional[list[int]]#

How many images each row of the batch actually carries.

Returns None when the batch describes its images in a way this cannot read, so the caller leaves the model’s own derivation alone rather than guessing a count and masking against it.

nemo_rl.data.multimodal_utils.attach_media_token_validity_mask( batch: Any, media_token_id: Optional[int], ) → None#

Attach a [B, S] mask of media-token positions that anchor media features.

Builds the mask while rows still are samples. Sequence packing later concatenates those rows into one THD sequence, after which no per-row information can be recovered, so the packing step carries this through the same transform as input_ids rather than deriving it downstream.

Parameters:

batch – Batch with [B, S] input_ids.
media_token_id – Media placeholder token ID, or None if the model does not accept the mask, in which case any existing mask is removed. The Megatron policy worker passes None when _model_accepts_media_token_validity_mask is False.

`VLLM_CONTENT_KEY`
`VLLM_MULTI_MODAL_DATA_KEY`
`VLLM_PROMPT_KEYS`
`IMAGE_CONTENT_TYPES`
`VIDEO_CONTENT_TYPES`
`AUDIO_CONTENT_TYPES`
`MULTIMODAL_CONTENT_TYPES`
`ROLLOUT_MATCHED_MEDIA_KEY`
`MEDIA_TAGS`
`MEDIA_TAGS_REVERSED`
`CACHED_VIDEO_FRAME_MANIFEST_MAGIC`
`CACHED_VIDEO_FRAME_MANIFEST_MIME`
`DEFAULT_MEDIA_EXTENSIONS`
`_PLACEHOLDER_STYLE_PROCESSOR_NAMES`
`MEDIA_TAGS_TO_ALLOWED`
`MEDIA_TAG_PATTERN`
`logger`
`PER_TOKEN_MULTIMODAL_FIELDS`
`PACKED_MULTIMODAL_FIELDS`
`ROW_SHAPES_SUFFIX`
`ROW_GEOMETRY_SHAPES`
`ROW_GEOMETRY_PREPROCESS_MODE`
`ROW_GEOMETRY_PREPROCESS_KWARGS`
`WIRE_MULTIMODAL_FIELDS`
`UNDECLARED_MULTIMODAL_MODEL_INPUTS`
`_VIDEO_EXT_TO_MIME`

nemo_rl.data.multimodal_utils#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.data.multimodal_utils`#