bridge.models.exaone.exaone45.modelling_exaone45.vision_model#

Module Contents#

Classes#

Exaone45VisionModel

EXAONE 4.5 ViT vision model.

API#

class bridge.models.exaone.exaone45.modelling_exaone45.vision_model.Exaone45VisionModel(
transformer_config: megatron.bridge.models.exaone.exaone45.modelling_exaone45.transformer_config.Exaone45TransformerConfig,
transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec,
patch_merger_spec: megatron.core.transformer.spec_utils.ModuleSpec,
pre_process: bool = True,
post_process: bool = True,
pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,
class_token_len: int = 1,
patch_dim: int = 14,
temporal_patch_size: int = 2,
spatial_merge_size: int = 2,
spatial_patch_size: int = 14,
img_h: int = 336,
img_w: int = 336,
window_size: int = 112,
)#

Bases: megatron.core.models.common.vision_module.vision_module.VisionModule

EXAONE 4.5 ViT vision model.

Parameters:

  • transformer_config (TransformerConfig) – Transformer config.

  • transformer_layer_spec (ModuleSpec) – Specifies module to use for transformer layers.

  • patch_merger_spec (ModuleSpec) – Specifies module to use for transformer layers.

Initialization

set_input_tensor(input_tensor: torch.Tensor) → None#

Sets input tensor to the model.

Parameters:

input_tensor (Tensor) – Sets the input tensor for the model.

rot_pos_emb(grid_thw)#
forward(
hidden_states: Optional[torch.Tensor],
grid_thw: torch.Tensor,
inference_params: Optional[megatron.core.InferenceParams] = None,
extra_block_kwargs: dict = None,
) → torch.Tensor#

Forward function of the EXAONE 4.5 vision model. This function passes the input tensors through the embedding layer and then the transformer.

Parameters:

  • x (torch.Tensor) – input image/video data of shape [n_tokens, n_dims]

  • grid_thw (torch.Tensor) – the size tensor indicates grid size of each image/frame

  • packed_seq_params (PackedSeqParams) – parameters to build attention mask in the backend

Returns:

output after final transformer block of shape [b, s, h].

Return type:

x (torch.Tensor)

get_packed_seq_params(
grid_thw: Optional[torch.Tensor],
cu_seqlens: Optional[torch.Tensor] = None,
)#
get_window_index(grid_thw)#