# coding=utf-8
"""
Bidirectional Llama encoder for the TTC-L2V sentence-embedding model.

This is the LLM2Vec modification of a decoder-only Llama: the causal attention
mask is replaced with a fully bidirectional one, turning the model into a text
encoder. It is the modern-`transformers` equivalent of the original
`LlamaEncoderModel` (which targeted transformers 4.40) and relies on the
built-in `create_bidirectional_mask` helper instead of hand-rolled mask code.

The weights are ordinary Llama weights; only the attention mask differs.
"""
from __future__ import annotations

import torch

from transformers.cache_utils import Cache
from transformers.masking_utils import create_bidirectional_mask
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.llama.modeling_llama import LlamaModel


class LlamaBiModel(LlamaModel):
    """A Llama model with bidirectional (non-causal) self-attention."""

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if position_ids is None:
            # Positions count from 0 over the (padded) sequence, matching the
            # original LLM2Vec encoder. RoPE is relative, so left padding does
            # not change the content-token hidden states.
            position_ids = torch.arange(
                inputs_embeds.shape[1], device=inputs_embeds.device
            ).unsqueeze(0)

        # The only change vs. a vanilla decoder: a bidirectional mask.
        attention_mask = create_bidirectional_mask(
            config=self.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
        )

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_embeddings=position_embeddings,
                position_ids=position_ids,
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(last_hidden_state=hidden_states)