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,
        _flattened_segment_mask,
        _pool_hidden_states,
    )
except ImportError:
    from echo_dsrn.configuration_echo import EchoConfig
    from echo_dsrn.modeling_echo import (
        EchoModel,
        EchoPreTrainedModel,
        _flattened_segment_mask,
        _pool_hidden_states,
    )


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.
    """

    _supports_attention_backend = True

    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")
        # Support explicit override via kwargs, else use pooling_mode
        explicit = kwargs.pop("output_all_states", None)
        output_all_states = (
            explicit if explicit is not None else (pooling_mode in ["mean_c_all", "hybrid"])
        )

        # vLLM's Transformers backend runs every sequence of a step as one
        # flattened [1, N] forward (no attention_mask, position_ids restarting
        # per sequence).  The DSRN recurrence cannot reset mid-forward, so
        # each segment runs as its own forward with fresh state.
        new_seq = _flattened_segment_mask(position_ids, input_ids, attention_mask)
        if new_seq is not None:
            return self._forward_flattened_segments(
                input_ids,
                position_ids,
                new_seq,
                return_dict,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                **kwargs,
            )

        # 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,
            attention_mask=attention_mask,
            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

        # 2. Pool representations according to pooling_mode (shared helper)
        pooled = _pool_hidden_states(outputs, pooling_mode, attention_mask)

        # 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)

        result = BaseModelOutputWithPast(
            last_hidden_state=embeddings_3d,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
        # Propagate all_c_all / all_h_all from the raw model output
        if hasattr(outputs, "all_c_all"):
            result.all_c_all = outputs.all_c_all
        if hasattr(outputs, "all_h_all"):
            result.all_h_all = outputs.all_h_all
        return result

    def _forward_flattened_segments(
        self,
        input_ids,
        position_ids,
        new_seq,
        return_dict,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        """Run each flattened-batch segment as its own forward.

        vLLM concatenates all sequences scheduled in a step into one
        ``[1, N]`` forward with ``position_ids`` restarting at each sequence
        start.  The DSRN recurrence carries state across the whole forward
        (its boundary handling freezes, it does not reset), so the segments
        cannot share one scan.  Running each segment independently reproduces
        the single-request semantics exactly; the pooled vectors are then
        stitched back into the ``[1, N]`` layout vLLM expects.
        """
        seg_ids = torch.cumsum(new_seq.long(), dim=1) - 1  # (1, N)
        num_segs = int(seg_ids.max().item()) + 1
        seg_outputs = []
        for i in range(num_segs):
            sel = seg_ids[0] == i
            seg_outputs.append(
                self.forward(
                    input_ids=input_ids[:, sel],
                    position_ids=position_ids[:, sel],
                    **kwargs,
                )
            )

        if not return_dict:
            return (
                torch.cat([o[0] for o in seg_outputs], dim=1),
                seg_outputs[-1][1],
            )

        result = BaseModelOutputWithPast(
            last_hidden_state=torch.cat([o.last_hidden_state for o in seg_outputs], dim=1),
            past_key_values=seg_outputs[-1].past_key_values,
            hidden_states=(
                [
                    torch.cat([o.hidden_states[layer_idx] for o in seg_outputs], dim=1)
                    for layer_idx in range(len(seg_outputs[0].hidden_states))
                ]
                if seg_outputs[0].hidden_states is not None
                else None
            ),
            attentions=(
                [
                    torch.cat([o.attentions[layer_idx] for o in seg_outputs], dim=1)
                    for layer_idx in range(len(seg_outputs[0].attentions))
                ]
                if seg_outputs[0].attentions is not None
                else None
            ),
        )
        if hasattr(seg_outputs[0], "all_c_all") and seg_outputs[0].all_c_all is not None:
            result.all_c_all = [
                torch.cat([o.all_c_all[layer_idx] for o in seg_outputs], dim=1)
                for layer_idx in range(len(seg_outputs[0].all_c_all))
            ]
        if hasattr(seg_outputs[0], "all_h_all") and seg_outputs[0].all_h_all is not None:
            result.all_h_all = [
                torch.cat([o.all_h_all[layer_idx] for o in seg_outputs], dim=1)
                for layer_idx in range(len(seg_outputs[0].all_h_all))
            ]
        return result