nemo_automodel.components.models.ministral_bidirectional.reranker_export

Portable checkpoint export for the Mistral3 pooled reranker.

Module Contents

Classes

Name	Description
`Mistral3RerankerMetadataExporter`	Snapshot the trained head on every rank, then write portable metadata on the writer rank.
`Mistral3RerankerStateDictAdapter`	Preserve backbone weights and expose the head at the native vLLM classifier path.

API

class nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerMetadataExporter(
    model: nemo_automodel.components.models.ministral_bidirectional.model.Mistral3VLBidirectionalForSequenceClassification
)

Snapshot the trained head on every rank, then write portable metadata on the writer rank.

Serialized head values are preserved, but consumers may cast them to the backbone dtype at load time. Stock Sentence Transformers Dense scoring therefore need not be bitwise identical to FP32 projection.

dense_weight

Tensor | None = None

processor

PixtralProcessor | None = None

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerMetadataExporter.save(
    hf_metadata_dir: str,
    tokenizer: object,
    original_model_path: str | None,
    v4_compatible: bool
) -> None

Save a Transformer/Pooling/Dense chain using the snapshot from this checkpoint step.

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerMetadataExporter.save_model_assets(
    directory: str | pathlib.Path
) -> None

Write stock configs and the small standalone Transformers scoring adapter.

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

Validate the export and collectively snapshot the small trained head before rank-zero I/O.

class nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerStateDictAdapter()

Bases: StateDictAdapter

Preserve backbone weights and expose the head at the native vLLM classifier path.

checkpoint_key_aliases

dict[str, str]

Map the current head name to the name used by older training checkpoints.

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    kwargs: typing.Any = {}
) -> list[tuple[str, torch.Tensor]]

Rename one parameter without gathering or copying it.

Parameters:

fqn

str

Native parameter name.

tensor

torch.Tensor

Parameter in its native layout; the scoring head has shape [labels, hidden].

**kwargs

AnyDefaults to {}

Unused adapter options.

Returns: list[tuple[str, torch.Tensor]]

One HF name and the original tensor with unchanged shape, dtype, device and placement.

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerStateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, torch.Tensor]
) -> dict[str, str]

Keep the trained head’s dtype so its exported Dense copy describes the same weights.

Parameters:

state_dict

dict[str, torch.Tensor]

HF tensors, including language_model.score.weight of shape [labels, hidden].

Returns: dict[str, str]

The scoring head’s current dtype, overriding a potentially lower-precision source checkpoint.

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerStateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs: typing.Any = {}
) -> dict[str, torch.Tensor]

Restore the native head name, accepting legacy score.weight checkpoints too.

Parameters:

hf_state_dict

dict[str, torch.Tensor]

HF tensors in native layouts; the head has shape [labels, hidden].

device_mesh

DeviceMesh | NoneDefaults to None

Unused mesh; this adapter preserves placements and does not gather.

**kwargs

AnyDefaults to {}

Unused adapter options.

Returns: dict[str, torch.Tensor]

Native names mapped to the same tensors, preserving device, dtype and distributed placements.

nemo_automodel.components.models.ministral_bidirectional.reranker_export.Mistral3RerankerStateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    kwargs: typing.Any = {}
) -> dict[str, torch.Tensor]

Rename the native scoring head without changing tensor values.

Parameters:

state_dict

dict[str, torch.Tensor]

Native parameters in their original layouts; score.weight has shape [labels, hidden].

**kwargs

AnyDefaults to {}

Unused adapter options.

Returns: dict[str, torch.Tensor]

HF names mapped to the same tensors, with the head at language_model.score.weight.