File size: 7,284 Bytes
a72ca7e ce54131 a72ca7e 4b27e53 a72ca7e 4b27e53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | 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")
# 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"])
# 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)
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
|