nemo_automodel.components.models.llama_bidirectional

Llama Bidirectional model for embedding and retrieval tasks.

Submodules

Package Contents

Classes

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

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.

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

Legacy Llama retrieval model with configurable attention.

The model defaults to bidirectional attention for embedding and retrieval workloads. Setting config.is_causal enables standard causal attention. Both modes use the Hugging Face forward path and its tensor/output contract.

The model is auto-discovered by ModelRegistry via the ModelClass export, enabling it to be loaded via NeMoAutoModelBiEncoder.from_pretrained().