nemo_automodel.components.models.biencoder.biencoder_model#

Module Contents#

Classes#

NeMoAutoModelBiencoder

NeMo AutoModel Biencoder with custom kernel support.

Data#

logger

API#

nemo_automodel.components.models.biencoder.biencoder_model.logger#

‘getLogger(…)’

class nemo_automodel.components.models.biencoder.biencoder_model.NeMoAutoModelBiencoder#

Bases: nemo_automodel._transformers.auto_model._BaseNeMoAutoModelClass

NeMo AutoModel Biencoder with custom kernel support.

This class extends _BaseNeMoAutoModelClass to provide biencoder functionality with NeMo AutoModel architecture, including support for Liger kernels and SDPA patching optimizations. The parent class handles all kernel patching (_patch_liger_kernel, _patch_attention), while this class only customizes the model initialization to use BiencoderModel.

classmethod from_pretrained(
pretrained_model_name_or_path,
share_encoder=True,
add_linear_pooler=False,
out_dimension=None,
do_gradient_checkpointing=False,
train_n_passages: int = 1,
eval_negative_size: int = 0,
pooling: str = 'avg',
l2_normalize: bool = True,
t: float = 1.0,
use_liger_kernel: bool = True,
use_sdpa_patching: bool = True,
sdpa_method: Optional[List[torch.nn.attention.SDPBackend]] = None,
**kwargs,
)#

Load a biencoder model from pretrained weights.

This method uses BiencoderModel.build to initialize the model, then applies kernel patching from the parent class (_patch_liger_kernel, _patch_attention). If patching fails, the method retries with adjusted parameters.

Parameters:

  • pretrained_model_name_or_path – Path to pretrained model or model identifier

  • share_encoder – Whether to share encoder weights between query and passage

  • add_linear_pooler – Whether to add a linear pooler layer

  • out_dimension – Output dimension for linear pooler

  • do_gradient_checkpointing – Whether to enable gradient checkpointing

  • train_n_passages – Number of passages per query during training

  • eval_negative_size – Number of negative samples during evaluation

  • pooling – Pooling strategy (‘avg’, ‘cls’, ‘last’, etc.)

  • l2_normalize – Whether to L2 normalize embeddings

  • t – Temperature for scaling similarity scores

  • use_liger_kernel – Whether to apply Liger kernel optimizations

  • use_sdpa_patching – Whether to apply SDPA patching

  • sdpa_method – SDPA backend methods to use

  • **kwargs – Additional arguments passed to BiencoderModel.build

Returns:

BiencoderModel instance with loaded and patched weights

.. rubric:: Notes

If kernel patching fails, the partially constructed model is deleted and the method recurses once with use_liger_kernel=False or use_sdpa_patching=False