modeling_embedding.py

from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
from transformers.modeling_outputs import BaseModelOutputWithPast

try:
    # pyrefly: ignore [missing-import]
    from .configuration_echo import EchoConfig

    # pyrefly: ignore [missing-import]
    from .modeling_echo import EchoModel, EchoPreTrainedModel
except ImportError:
    from echo_dsrn.configuration_echo import EchoConfig
    from echo_dsrn.modeling_echo import EchoModel, EchoPreTrainedModel


class EchoModelForSentenceEmbedding(EchoPreTrainedModel):
    """
    Sentence embedding adapter for Echo-DSRN.
    Extracts the recurrent state 'c' or sequences from layers and shapes them
    for sentence-transformers compatibility.
    """

    def __init__(self, config: EchoConfig):
        super().__init__(config)
        self.model = EchoModel(config)
        self.pooling_mode = getattr(config, "pooling_mode", "c_T")

        # Determine target dimension for the projection input
        if self.pooling_mode == "hybrid":
            proj_in_dim = config.hidden_size * (config.num_heads + 1)
        elif self.pooling_mode == "mean_x_out":
            proj_in_dim = config.hidden_size
        else:  # "c_T" or "mean_c_all"
            proj_in_dim = config.hidden_size * config.num_heads

        # Optional projection layer to map back to a specific target embedding dimension.
        self.project_embeddings = getattr(config, "project_embeddings", False)
        self.projection_mlp = getattr(config, "projection_mlp", False)
        if self.projection_mlp:
            target_dim = getattr(config, "embedding_dim", config.hidden_size)
            hidden_dim = getattr(config, "projection_hidden_dim", 1024)
            self.projection = nn.Sequential(
                nn.Linear(proj_in_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, target_dim, bias=False),
            )
        elif self.project_embeddings:
            target_dim = getattr(config, "embedding_dim", config.hidden_size)
            self.projection = nn.Linear(proj_in_dim, target_dim, bias=False)
        else:
            self.projection = None

        self.post_init()

    def get_input_embeddings(self):
        return self.model.embedding

    def set_input_embeddings(self, value):
        self.model.embedding = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        pooling_mode = getattr(self.config, "pooling_mode", "c_T")
        output_all_states = pooling_mode in ["mean_c_all", "hybrid"]

        # 1. Base model forward pass
        outputs = self.model(
            input_ids=input_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            output_all_states=output_all_states,
            **kwargs,
        )

        # Determine sequence length for broadcasting
        if input_ids is not None:
            seq_len = input_ids.shape[1]
        elif inputs_embeds is not None:
            seq_len = inputs_embeds.shape[1]
        else:
            seq_len = 1

        def mean_pooling(token_embeddings, mask):
            input_mask_expanded = mask.unsqueeze(-1).expand(token_embeddings.size()).float()
            sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
            sum_mask = input_mask_expanded.sum(1)
            sum_mask = torch.clamp(sum_mask, min=1e-9)
            return sum_embeddings / sum_mask

        # 2. Extract and pool representations according to pooling_mode
        if pooling_mode == "mean_c_all":
            # Extract full sequence of recurrent slow states c_all from last layer
            c_all = outputs.all_c_all[-1]  # shape: (Batch, Seq_Len, State_Dim)
            if attention_mask is not None:
                pooled = mean_pooling(c_all, attention_mask)
            else:
                pooled = c_all.mean(dim=1)
        elif pooling_mode == "mean_x_out":
            # Mean pool the final hidden state
            last_hidden_state = outputs.last_hidden_state  # shape: (Batch, Seq_Len, hidden_size)
            if attention_mask is not None:
                pooled = mean_pooling(last_hidden_state, attention_mask)
            else:
                pooled = last_hidden_state.mean(dim=1)
        elif pooling_mode == "hybrid":
            # Concatenate pooled fast states (h_all) and slow states (c_all) from last layer
            h_all = outputs.all_h_all[-1]  # shape: (Batch, Seq_Len, hidden_size)
            c_all = outputs.all_c_all[-1]  # shape: (Batch, Seq_Len, State_Dim)
            if attention_mask is not None:
                pooled_h = mean_pooling(h_all, attention_mask)
                pooled_c = mean_pooling(c_all, attention_mask)
            else:
                pooled_h = h_all.mean(dim=1)
                pooled_c = c_all.mean(dim=1)
            pooled = torch.cat(
                [pooled_h, pooled_c], dim=-1
            )  # shape: (Batch, hidden_size + State_Dim)
        else:  # "c_T" (default baseline behavior)
            past = outputs.past_key_values
            if hasattr(past, "__getitem__"):
                last_layer_state = past[-1]
            elif hasattr(past, "states"):  # EchoCache support
                last_layer_state = past.states[-1]
            else:
                raise ValueError("Could not extract recurrent state from model cache.")
            pooled = last_layer_state[1]  # shape: (Batch, State_Dim)

        # 3. Apply optional projection
        if self.projection is not None:
            embeddings = self.projection(pooled)
        else:
            embeddings = pooled

        # 4. Broadcast to shape (Batch, Seq_Len, Dim) for pooling safety
        embeddings_3d = embeddings.unsqueeze(1).expand(-1, seq_len, -1)

        if not return_dict:
            return (embeddings_3d, outputs.past_key_values)

        return BaseModelOutputWithPast(
            last_hidden_state=embeddings_3d,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )