nemo_automodel.components.models.llama_nemotron_vl.model

Module Contents

Classes

Name	Description
`LlamaBidirectionalConfig`	Configuration for bidirectional (non-causal) LLaMA model.
`LlamaBidirectionalModel`	Legacy Llama retrieval model supporting causal and bidirectional attention.
`LlamaNemotronVLConfig`	Base configuration for vision-language models combining vision and language components.
`LlamaNemotronVLModel`	LlamaNemotron VL model for vision-language reranking.

Functions

Name	Description
`_filter_vision_embeddings_by_image_flags`	Keep only vision embeddings marked as real images.
`_register_with_hf_auto_classes`	Register bidirectional models with HuggingFace Auto classes.
`_replace_image_token_embeddings`	Replace image placeholder token embeddings with vision embeddings.
`pool`	-
`split_model`	-

Data

ModelClass

_DYNAMIC_CACHE_ACCEPTS_CONFIG

_HAS_NATIVE_BIDIRECTIONAL_MASK

_USE_PLURAL_CACHE_PARAM

_decoder_forward_params

_dynamic_cache_init_params

logger

API

class nemo_automodel.components.models.llama_nemotron_vl.model.LlamaBidirectionalConfig(
    pooling: str = 'avg',
    temperature: float = 1.0,
    is_causal: bool = False,
    kwargs: typing.Any = {}
)

Bases: LlamaConfig

Configuration for bidirectional (non-causal) LLaMA model.

model_type

= 'llama_bidirec'

class nemo_automodel.components.models.llama_nemotron_vl.model.LlamaBidirectionalModel(
    config: nemo_automodel.components.models.llama_nemotron_vl.model.LlamaBidirectionalConfig
)

Bases: LlamaModel

Legacy Llama retrieval model supporting causal and bidirectional attention.

Supports Transformers 4.44 through 5.x with a compatibility forward path.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaBidirectionalModel._create_bidirectional_mask(
    input_embeds: torch.Tensor,
    attention_mask: torch.Tensor | None
) -> torch.Tensor | None

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaBidirectionalModel.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    cache_position: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_hidden_states: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Encode tokens using the selected attention mode.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Integer tensor of shape [batch, sequence], mutually exclusive with inputs_embeds.

attention_mask

torch.Tensor | NoneDefaults to None

Optional padding mask of shape [batch, key_sequence], including cached tokens.

position_ids

torch.LongTensor | NoneDefaults to None

Optional integer tensor of shape [batch, sequence] or [1, sequence].

past_key_values

Cache | NoneDefaults to None

Optional cache of key/value tensors of shape [batch, kv_heads, cached_sequence, head_dim].

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [batch, sequence, hidden], mutually exclusive with input_ids.

cache_position

torch.LongTensor | NoneDefaults to None

Optional integer tensor of shape [sequence].

use_cache

bool | NoneDefaults to None

Whether to retain attention keys and values.

output_hidden_states

bool | NoneDefaults to None

Whether to return each layer’s hidden states.

**kwargs

AnyDefaults to {}

Additional arguments forwarded to the stock causal decoder.

Returns: BaseModelOutputWithPast

Output with last_hidden_state of shape [batch, sequence, hidden], optional

class nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLConfig(
    vision_config = None,
    llm_config = None,
    use_backbone_lora = 0,
    use_llm_lora = 0,
    select_layer = -1,
    force_image_size = None,
    downsample_ratio = 0.5,
    template = None,
    dynamic_image_size = False,
    use_thumbnail = False,
    min_dynamic_patch = 1,
    max_dynamic_patch = 6,
    mlp_checkpoint = True,
    pre_feature_reduction = False,
    keep_aspect_ratio = False,
    vocab_size = -1,
    q_max_length: int | None = 512,
    p_max_length: int | None = 10240,
    query_prefix: str = 'query:',
    passage_prefix: str = 'passage:',
    pooling: str = 'last',
    bidirectional_attention: bool = False,
    max_input_tiles: int = 2,
    img_context_token_id: int = 128258,
    kwargs = {}
)

Bases: PretrainedConfig

Base configuration for vision-language models combining vision and language components. This serves as the foundation for LlamaNemotronVL configurations.

llm_config

= LlamaBidirectionalConfig(**llm_config)

model_type

= 'llama_nemotron_vl'

sub_configs

vision_config

= SiglipVisionConfig(**vision_config)

vocab_size

= self.llm_config.vocab_size

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLConfig.get_text_config(
    decoder: bool | None = None,
    encoder: bool | None = None
) -> transformers.configuration_utils.PretrainedConfig

Resolve the language config through the Transformers composite contract.

Parameters:

decoder

bool | NoneDefaults to None

Select the text-output configuration when true.

encoder

bool | NoneDefaults to None

Select the text-input configuration when true.

Returns: PretrainedConfig

The language decoder config for decoder or unrestricted requests;

class nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel(
    config: nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLConfig,
    vision_model: transformers.PreTrainedModel | None = None,
    language_model: transformers.PreTrainedModel | None = None
)

Bases: PreTrainedModel

LlamaNemotron VL model for vision-language reranking. Combines a vision encoder (SigLIP) with a bidirectional language model (LLaMA) for cross-modal reranking tasks.

_no_split_modules

= ['LlamaDecoderLayer']

downsample_ratio

= config.downsample_ratio

main_input_name

= 'pixel_values'

mlp1

num_image_token

= int((grid_size * config.downsample_ratio) ** 2)

patch_size

= 14

processor

select_layer

= config.select_layer

template

= config.template

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel._embed_batch(
    inputs: typing.Dict[str, typing.Any],
    pool_type: str | None = None
)

Encodes the inputs into a tensor of embeddings. Args: inputs: A dictionary of inputs to the model. You can prepare the inputs using the processor.process_queries and processor.process_documents methods. pool_type: The type of pooling to use. If None, the pooling type is set to the pooling type configured in the model. Returns: A tensor of embeddings.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.build_collator(
    processor = None,
    kwargs = {}
)

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.encode_documents(
    images: typing.List[typing.Any] | None = None,
    texts: typing.List[str] | None = None,
    kwargs = {}
)

Encodes the input document images and texts into a tensor of embeddings. Args: images: A list of PIL.Image of document pages images. texts: A list of document page texts. Returns: A tensor of embeddings.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.encode_queries(
    queries: typing.List[str],
    kwargs = {}
)

Encodes the input queries into a tensor of embeddings. Args: queries: A list of queries. Returns: A tensor of embeddings.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.extract_feature(
    pixel_values
)

Extract and project vision features to language model space.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.forward(
    pixel_values: torch.FloatTensor = None,
    input_ids: torch.LongTensor = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    image_flags: torch.LongTensor | None = None,
    past_key_values: typing.List[torch.FloatTensor] | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    num_patches_list: typing.List[torch.Tensor] | None = None,
    run_dummy_vision: bool | None = None
) -> typing.Union[typing.Tuple, transformers.modeling_outputs.CausalLMOutputWithPast]

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.get_decoder() -> transformers.PreTrainedModel

Return the language tower used as the retrieval text decoder.

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.get_input_embeddings()

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.get_output_embeddings()

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.pixel_shuffle(
    x,
    scale_factor = 0.5
)

nemo_automodel.components.models.llama_nemotron_vl.model.LlamaNemotronVLModel.post_loss(
    loss,
    inputs
)

nemo_automodel.components.models.llama_nemotron_vl.model._filter_vision_embeddings_by_image_flags(
    vit_embeds: torch.Tensor,
    image_flags: torch.Tensor | None
) -> torch.Tensor

Keep only vision embeddings marked as real images.

nemo_automodel.components.models.llama_nemotron_vl.model._register_with_hf_auto_classes()

This is needed so that AutoModel.from_config(LlamaBidirectionalConfig) works inside LlamaForSequenceClassification.init.

nemo_automodel.components.models.llama_nemotron_vl.model._replace_image_token_embeddings(
    input_embeds: torch.Tensor,
    input_ids: torch.Tensor,
    vit_embeds: torch.Tensor,
    img_context_token_id: int
) -> torch.Tensor

Replace image placeholder token embeddings with vision embeddings.

nemo_automodel.components.models.llama_nemotron_vl.model.pool(
    last_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    pool_type: str
) -> torch.Tensor

nemo_automodel.components.models.llama_nemotron_vl.model.split_model(
    model_path,
    device
)

nemo_automodel.components.models.llama_nemotron_vl.model.ModelClass = [LlamaNemotronVLModel]

nemo_automodel.components.models.llama_nemotron_vl.model._DYNAMIC_CACHE_ACCEPTS_CONFIG = 'config' in _dynamic_cache_init_params

nemo_automodel.components.models.llama_nemotron_vl.model._HAS_NATIVE_BIDIRECTIONAL_MASK = True

nemo_automodel.components.models.llama_nemotron_vl.model._USE_PLURAL_CACHE_PARAM = 'past_key_values' in _decoder_forward_params

nemo_automodel.components.models.llama_nemotron_vl.model._decoder_forward_params = inspect.signature(LlamaDecoderLayer.forward).parameters

nemo_automodel.components.models.llama_nemotron_vl.model._dynamic_cache_init_params = inspect.signature(DynamicCache.__init__).parameters

nemo_automodel.components.models.llama_nemotron_vl.model.logger = logging.get_logger(__name__)