nemo_automodel.components.models.ministral_bidirectional

Bidirectional Ministral3 model for embedding and retrieval tasks.

Submodules

Package Contents

Classes

Name	Description
`Ministral3BidirectionalConfig`	Configuration for Ministral3BidirectionalModel with pooling and temperature settings.
`Ministral3BidirectionalModel`	Ministral3 retrieval model with a configurable attention mode.
`Mistral3BidirectionalConfig`	Configuration for Mistral3BidirectionalModel with pooling and temperature settings.
`Mistral3BidirectionalModel`	Mistral3 retrieval model with a configurable Ministral3 text tower.
`Mistral3VLBidirectionalForSequenceClassification`	Bidirectional Mistral3 VLM with a sequence-scoring head.

API

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

Bases: Ministral3Config

Configuration for Ministral3BidirectionalModel with pooling and temperature settings.

model_type

= 'ministral3_bidirec'

class nemo_automodel.components.models.ministral_bidirectional.model.Ministral3BidirectionalModel(
    config: nemo_automodel.components.models.ministral_bidirectional.model.Ministral3BidirectionalConfig
)

Bases: Ministral3Model

Ministral3 retrieval model with a configurable attention mode.

The historical class identity is retained for legacy checkpoint compatibility.

nemo_automodel.components.models.ministral_bidirectional.model.Ministral3BidirectionalModel.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,
    use_cache: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Forward pass with the attention mode stored in the model config.

Causal mode delegates to the Hugging Face parent implementation. Non-causal mode uses the retrieval-specific bidirectional mask.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Integer token IDs [batch, sequence], exclusive with inputs_embeds.

attention_mask

torch.Tensor | NoneDefaults to None

Optional padding mask [batch, sequence], with zero for padding.

position_ids

torch.LongTensor | NoneDefaults to None

Integer positions [batch, sequence] or [1, sequence].

past_key_values

Cache | NoneDefaults to None

Optional Transformers cache in its native layout.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Floating-point embeddings [batch, sequence, hidden].

use_cache

bool | NoneDefaults to None

Whether to populate the returned cache.

cache_position

torch.LongTensor | NoneDefaults to None

Integer cache positions [sequence].

**kwargs

AnyDefaults to {}

Additional Hugging Face forward options.

Returns: BaseModelOutputWithPast

Output with floating-point last_hidden_state [batch, sequence, hidden],

class nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalConfig(
    pooling: str = 'avg',
    temperature: float | None = None,
    score_temperature: float | None = None,
    kwargs: typing.Any = {}
)

Bases: Mistral3Config

Configuration for Mistral3BidirectionalModel with pooling and temperature settings.

model_type

= 'mistral3_bidirec'

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalConfig.__post_init__(
    kwargs: typing.Any = {}
) -> None

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalConfig.to_dict() -> dict[str, typing.Any]

Serialize pooling under its public configuration key.

class nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel(
    config: nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalConfig
)

Bases: Mistral3Model

Mistral3 retrieval model with a configurable Ministral3 text tower.

The historical class identity is retained while the text tower honors either policy.

_sentence_transformer_input_mode

= 'structured_multimodal'

language_model

= Ministral3BidirectionalModel(config.text_config)

retrieval_processor_target

str

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel._dummy_vision_sum() -> torch.Tensor

Run the smallest valid Pixtral image through the sharded vision tower.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel._nemo_apply_liger_kernel(
    liger_kernel_transformers
) -> None

Apply compatible Liger kernels to the Ministral text tower only.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel.forward(
    input_ids: torch.LongTensor | None = None,
    pixel_values: torch.FloatTensor | None = None,
    image_sizes: 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,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    run_dummy_vision: bool | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, ...] | transformers.models.mistral3.modeling_mistral3.Mistral3ModelOutputWithPast

Run text or multimodal inputs through one FSDP root forward.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Tensor of shape [batch, sequence] containing text and image placeholder token IDs.

pixel_values

torch.FloatTensor | NoneDefaults to None

Tensor of shape [images, channels, height, width].

image_sizes

torch.LongTensor | NoneDefaults to None

Tensor of shape [images, 2] containing image height and width.

attention_mask

torch.Tensor | NoneDefaults to None

Tensor of shape [batch, sequence].

position_ids

torch.LongTensor | NoneDefaults to None

Tensor of shape [batch, sequence].

past_key_values

Cache | NoneDefaults to None

Optional cache with per-layer key and value tensors of shape [batch, kv_heads, cached_sequence, head_dim], where kv_heads is the number of key/value heads and cached_sequence is the stored token count. The cache is updated in place when used.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Tensor of shape [batch, sequence, hidden].

use_cache

bool | NoneDefaults to None

Whether to return updated key/value cache state.

output_attentions

bool | NoneDefaults to None

Whether to return attention tensors.

output_hidden_states

bool | NoneDefaults to None

Whether to return intermediate hidden states.

return_dict

bool | NoneDefaults to None

Whether to return a structured model output.

run_dummy_vision

bool | NoneDefaults to None

Whether to execute a zero-weighted vision path for text-only inputs. Defaults to enabled in training. Sharded vision layers always participate during evaluation, even when false, to preserve collective ordering across ranks with different modalities.

**kwargs

AnyDefaults to {}

Additional arguments forwarded to the stock Mistral3 model.

Returns: tuple[torch.Tensor, ...] | Mistral3ModelOutputWithPast

A Mistral3ModelOutputWithPast with these fields:

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel

classmethod

Load stock Mistral3 weights into the bidirectional VL layout.

Caller-provided key mappings take precedence over this model’s default stock-checkpoint conversion.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel.get_hf_export_config() -> transformers.models.mistral3.configuration_mistral3.Mistral3Config

Return the stock Transformers config used for portable inference export.

Subclasses must explicitly re-assert _export_as_stock_model = True. This prevents inherited export from silently discarding custom forward behavior that a stock Mistral3Model cannot reproduce.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalModel.supports_config(
    config: transformers.PretrainedConfig
) -> bool

classmethod

Select this implementation only for Ministral3 text towers.

Parameters:

config

PretrainedConfig

Checkpoint configuration before any retrieval conversion.

Returns: bool

Whether the checkpoint has a supported text architecture.

class nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification(
    config: nemo_automodel.components.models.ministral_bidirectional.model.Mistral3BidirectionalConfig
)

Bases: Mistral3PreTrainedModel

Bidirectional Mistral3 VLM with a sequence-scoring head.

_sentence_transformer_input_mode

= 'structured_multimodal'

base_model_prefix

= 'model'

effective_score_temperature

float

Return the temperature applied to score logits by this model.

model

= Mistral3BidirectionalModel(config)

num_labels

= config.num_labels

score

state_dict_adapter

= Mistral3RerankerStateDictAdapter()

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification._get_consolidated_hf_metadata_exporter(
    tokenizer: object,
    original_model_path: str | None
) -> nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerMetadataExporter

Select the portable reranker export for direct and distributed checkpoint saves.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification._init_weights(
    module: torch.nn.Module
) -> None

Initialize the scoring head in its configured dtype.

Parameters:

module

nn.Module

Module whose parameters are initialized in place.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification._nemo_apply_liger_kernel(
    liger_kernel_transformers
) -> None

Apply compatible Liger kernels to the nested Mistral text tower.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification.forward(
    input_ids: torch.LongTensor | None = None,
    pixel_values: torch.FloatTensor | None = None,
    image_sizes: 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,
    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,
    run_dummy_vision: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.SequenceClassifierOutputWithPast | tuple[torch.Tensor, ...]

Score each multimodal question-document sequence.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Tensor of shape [batch, sequence] containing text and image placeholder token IDs.

pixel_values

torch.FloatTensor | NoneDefaults to None

Tensor of shape [images, channels, height, width].

image_sizes

torch.LongTensor | NoneDefaults to None

Tensor of shape [images, 2] containing image height and width before padding.

attention_mask

torch.Tensor | NoneDefaults to None

Tensor of shape [batch, sequence].

position_ids

torch.LongTensor | NoneDefaults to None

Tensor of shape [batch, sequence].

past_key_values

Cache | NoneDefaults to None

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Tensor of shape [batch, sequence, hidden].

labels

torch.LongTensor | NoneDefaults to None

Tensor of shape [query_batch] containing the relevant candidate index for each query. Accepted for compatibility but intentionally unused because the recipe computes the grouped cross-entropy loss.

use_cache

bool | NoneDefaults to None

Whether to return updated key/value cache state.

output_attentions

bool | NoneDefaults to None

Whether to return attention tensors.

output_hidden_states

bool | NoneDefaults to None

Whether to return intermediate hidden states.

return_dict

bool | NoneDefaults to None

Whether to return a structured model output.

run_dummy_vision

bool | NoneDefaults to None

Whether text-only inputs should run dummy vision. Sharded vision always participates during evaluation to keep collectives aligned across ranks with different modalities.

**kwargs

AnyDefaults to {}

Additional arguments forwarded to Mistral3BidirectionalModel.

Returns: SequenceClassifierOutputWithPast | tuple[torch.Tensor, ...]

A SequenceClassifierOutputWithPast with these fields:

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification.from_pretrained(
    pretrained_model_name_or_path: str | os.PathLike | None,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification

classmethod

Load stock Mistral3 VLM weights into the reranker layout.

Stock conditional-generation checkpoints nest the text backbone under model.language_model.model. This reranker owns that backbone directly under model.language_model, so checkpoint loading must remove the extra model segment.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification.save_pretrained(
    save_directory: str | os.PathLike,
    kwargs: typing.Any = {}
) -> None

Save trained weights using the shared Transformers/vLLM reranker head name.

nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification.supports_config(
    config: transformers.PretrainedConfig
) -> bool

classmethod

Return whether the checkpoint uses the supported Ministral3 text tower.

Parameters:

config

PretrainedConfig

Checkpoint configuration before retrieval conversion.

Returns: bool

Whether the checkpoint can use this scoring backbone.