nemo_automodel.components.models.llama_bidirectional.model

Llama Bidirectional model for embedding and retrieval tasks.

This module provides a bidirectional variant of Llama that is auto-discovered by the ModelRegistry via the ModelClass export.

To add support for other backbones (e.g., Qwen2, Mistral), create a similar module in a new directory (e.g., qwen2_bidirectional/) with its own ModelClass export.

Module Contents

Classes

Name	Description
`LlamaBidirectionalConfig`	Configuration class for LlamaBidirectionalModel.
`LlamaBidirectionalForSequenceClassification`	Llama Bidirectional Model with a sequence classification/regression head.
`LlamaBidirectionalModel`	Llama retrieval model with a configurable attention mode.

Functions

Name	Description
`_pool`	Pool hidden states using the specified pooling method.
`_register_with_hf_auto_classes`	Register bidirectional Llama types with Hugging Face Auto classes.

Data

ModelClass

API

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

Bases: LlamaConfig

Configuration class for LlamaBidirectionalModel.

Extends LlamaConfig with additional parameters for bidirectional attention and pooling configurations.

model_type

= 'llama_bidirec'

class nemo_automodel.components.models.llama_bidirectional.model.LlamaBidirectionalForSequenceClassification(
    config
)

Bases: LlamaPreTrainedModel

Llama Bidirectional Model with a sequence classification/regression head.

This model adds a classification head on top of the bidirectional Llama model and includes configurable pooling strategies.

effective_score_temperature

float

Return the temperature applied to score logits by this model.

model

= LlamaBidirectionalModel(config)

num_labels

= config.num_labels

score

nemo_automodel.components.models.llama_bidirectional.model.LlamaBidirectionalForSequenceClassification._init_weights(
    module
)

nemo_automodel.components.models.llama_bidirectional.model.LlamaBidirectionalForSequenceClassification.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: typing.Union[transformers.cache_utils.Cache, typing.List[torch.FloatTensor]] | 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,
    kwargs = {}
) -> typing.Union[typing.Tuple, transformers.modeling_outputs.SequenceClassifierOutputWithPast]

class nemo_automodel.components.models.llama_bidirectional.model.LlamaBidirectionalModel(
    config: transformers.models.llama.configuration_llama.LlamaConfig
)

Bases: LlamaModel

Llama retrieval model with a configurable attention mode.

False preserves legacy bidirectional behavior; true uses the stock Llama causal mask.

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

Pool hidden states using the specified pooling method.

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

Enables AutoConfig, AutoModel, and AutoModelForSequenceClassification to resolve their corresponding bidirectional config and model classes.

nemo_automodel.components.models.llama_bidirectional.model.ModelClass = [LlamaBidirectionalModel, LlamaBidirectionalForSequenceClassification]