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`	Llama retrieval model with a configurable attention mode.

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.