`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.
`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_pad_to_max_shape`	Return whether a processor input must pad non-packing dimensions.
`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_to_data_url`	Encode a PIL Image as a base64 `data:` URL.
`_encode_single_image_source`	Resolve and encode one image source.
`extract_input_image_sources_from_responses_messages`	Extract image sources from Responses-API messages in encounter order.
`extract_input_images_from_responses_messages`	Load images from Responses-API input messages in encounter order.
`_materialize_ragged_pixel_values`	Fold a ragged per-image `pixel_values` list into one padded tensor.
`_stack_ragged_pixel_values`	Derive imgs_sizes from unpadded shapes, then pad into one tensor.
`_restore_tensors`	Convert a processor’s list outputs back to tensors, in place.
`attach_image_model_inputs_to_message`	Attach processor-owned image tensors without replacing rollout tokens.
`encode_images_in_examples`	Replace local image paths in NeMo Gym examples with base64 data URLs.
`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`	Mark media tokens that anchor nothing, so the model keeps their embedding.

Data#

`VLLM_MULTIMODAL_DATA_KEYS`
`NATIVE_MULTIMODAL_KEYS`
`IMAGE_CONTENT_TYPES`
`VIDEO_CONTENT_TYPES`
`AUDIO_CONTENT_TYPES`
`MULTIMODAL_CONTENT_TYPES`
`NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS`
`MEDIA_TAGS`
`MEDIA_TAGS_REVERSED`
`DEFAULT_MEDIA_EXTENSIONS`
`_PLACEHOLDER_STYLE_PROCESSOR_NAMES`
`MEDIA_TAGS_TO_ALLOWED`
`MEDIA_TAG_PATTERN`
`logger`

API#

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

nemo_rl.data.multimodal_utils.NATIVE_MULTIMODAL_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.NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS#: 8

nemo_rl.data.multimodal_utils.MEDIA_TAGS#: None

nemo_rl.data.multimodal_utils.MEDIA_TAGS_REVERSED#: None

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.

class nemo_rl.data.multimodal_utils.PackedTensor( tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]], dim_to_pack: int, *, pad_to_max_shape: bool = False, _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.
pad_to_max_shape – Pad every non-packing dimension to its batch-wide maximum before concatenating. All tensors must have the same rank.

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

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: Optional[torch.device] = 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])
>>>

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_pad_to_max_shape(processor: Any, key: str) → bool#: Return whether a processor input must pad non-packing dimensions.

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_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._encode_single_image_source(source: str) → str#: Resolve and encode one image source.

nemo_rl.data.multimodal_utils.extract_input_image_sources_from_responses_messages( messages: Any, ) → list[str | PIL.Image.Image]#: Extract image sources from Responses-API messages in encounter order.

nemo_rl.data.multimodal_utils.extract_input_images_from_responses_messages( messages: Any, ) → list[PIL.Image.Image]#: Load images from Responses-API input messages in encounter order.

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

Fold a ragged per-image pixel_values list into one padded tensor.

Processors with dynamic per-image resolution return a list of CHW tensors rather than a stacked batch. imgs_sizes is derived from the unpadded shapes first, since those exact sizes are what the projector slices with; padding happens afterwards so downstream sees the single tensor its torch.Tensor contract expects.

nemo_rl.data.multimodal_utils._stack_ragged_pixel_values( processed: dict[str, Any], tiles: list[torch.Tensor], processor: Any, ) → None#: Derive imgs_sizes from unpadded shapes, then pad into one tensor.

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.attach_image_model_inputs_to_message( message: dict[str, Any], *, images: list[PIL.Image.Image], processor: Any, pad_dynamic_image_shapes: bool = False, ) → None#: Attach processor-owned image tensors without replacing rollout tokens.

nemo_rl.data.multimodal_utils.encode_images_in_examples(nemo_gym_examples: list[dict]) → list[dict]#

Replace local image paths in NeMo Gym examples with base64 data URLs.

Walks each example’s responses_create_params.input[].content[] items, collects local image references, encodes each unique source once using a bounded thread pool, and rewrites every corresponding image part with the resulting base64 data: URL. Parts whose URL already starts with http://, https://, or data: are left untouched. Malformed items (non-dict entries, missing/empty URLs, non-list input/content) are skipped without raising.

The examples are mutated in place; the same list is also returned for convenience so callers can chain the call.

Parameters:: nemo_gym_examples – List of NeMo Gym example dicts. Each example is expected to contain a responses_create_params mapping with an input list of Responses API messages.
Returns:: The same nemo_gym_examples list, with local image references rewritten to base64 data URLs in place.

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#

Mark media tokens that anchor nothing, so the model keeps their embedding.

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

The batch is duck-typed rather than annotated as BatchedDataDict: that module imports this one, so naming it here would be circular.

nemo_rl.data.multimodal_utils#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.data.multimodal_utils`#