"""Hierarchical session encoder producing a deterministic bottleneck.""" from __future__ import annotations from typing import Any import torch import torch.nn as nn from transformers import AutoConfig, AutoModel, AutoTokenizer from .encoder_layer_pool import LayerwiseAttention from .encoder_system_fusion import SystemContextFusion from .encoder_truncate import truncate_text_head_tail DEFAULT_ENCODER = "nomic-ai/modernbert-embed-base" # ModernBERT's own instruction prefixes. Runs before 2026-08 trained with the e5 defaults # ("query: "/"passage: ") instead; their exported configs record that and override these. DEFAULT_QUERY_PREFIX = "search_query: " DEFAULT_PASSAGE_PREFIX = "search_document: " def mean_pool(hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: mask = attention_mask.unsqueeze(-1).to(hidden.dtype) return (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0) class HierarchicalSessionEncoder(nn.Module): """CLS-style session Transformer over fused user-turn vectors.""" def __init__( self, input_dim: int, session_dim: int = 1024, num_layers: int = 4, num_heads: int = 8, dropout: float = 0.1, max_turns: int = 64, ) -> None: super().__init__() self.session_dim = session_dim self.input_proj = ( nn.Identity() if input_dim == session_dim else nn.Linear(input_dim, session_dim) ) self.cls = nn.Parameter(torch.randn(1, 1, session_dim) * 0.02) self.pos = nn.Embedding(max_turns + 1, session_dim) layer = nn.TransformerEncoderLayer( d_model=session_dim, nhead=num_heads, dim_feedforward=session_dim * 4, dropout=dropout, batch_first=True, activation="gelu", ) # enable_nested_tensor=False mirrors training: the fast path takes a different kernel # for padded batches, and it is unimplemented on MPS. self.transformer = nn.TransformerEncoder( layer, num_layers=num_layers, enable_nested_tensor=False ) self.norm = nn.LayerNorm(session_dim) def forward(self, turn_embeds: torch.Tensor, turn_mask: torch.Tensor) -> torch.Tensor: """ turn_embeds: [B, T, H] turn_mask: [B, T] True = padding """ x = self.input_proj(turn_embeds) batch, turns, _ = x.shape cls = self.cls.expand(batch, -1, -1) seq = torch.cat([cls, x], dim=1) positions = torch.arange(turns + 1, device=x.device).unsqueeze(0) seq = seq + self.pos(positions) pad = torch.cat( [torch.zeros(batch, 1, dtype=torch.bool, device=x.device), turn_mask], dim=1, ) encoded = self.transformer(seq, src_key_padding_mask=pad) return self.norm(encoded[:, 0]) class CueEncoder(nn.Module): """Trained user turn encoder + frozen system encoder → session bottleneck.""" def __init__( self, *, user_model_name: str = DEFAULT_ENCODER, system_model_name: str = DEFAULT_ENCODER, bottleneck_dim: int = 1024, session_dim: int = 1024, session_layers: int = 4, session_heads: int = 8, max_turns: int = 64, max_length: int = 256, dropout: float = 0.1, query_prefix: str = DEFAULT_QUERY_PREFIX, passage_prefix: str = DEFAULT_PASSAGE_PREFIX, pretrained_backbones: bool = True, ) -> None: super().__init__() self.user_model_name = user_model_name self.system_model_name = system_model_name self.query_prefix = query_prefix self.passage_prefix = passage_prefix self.bottleneck_dim = bottleneck_dim self.max_turns = max_turns self.max_length = max_length self.user_tokenizer = AutoTokenizer.from_pretrained( user_model_name, clean_up_tokenization_spaces=False ) self.system_tokenizer = AutoTokenizer.from_pretrained( system_model_name, clean_up_tokenization_spaces=False ) if pretrained_backbones: # low_cpu_mem_usage=False: keep weights off the meta device so a later # model.to(device) doesn't fail with "Cannot copy out of meta tensor". self.user_encoder = AutoModel.from_pretrained( user_model_name, low_cpu_mem_usage=False ) self.system_encoder = AutoModel.from_pretrained( system_model_name, low_cpu_mem_usage=False ) else: # Loading a complete CUE checkpoint overwrites every backbone weight, so only # the architecture config is needed. Skips a redundant base-weight download. self.user_encoder = AutoModel.from_config( AutoConfig.from_pretrained(user_model_name) ) self.system_encoder = AutoModel.from_config( AutoConfig.from_pretrained(system_model_name) ) for param in self.system_encoder.parameters(): param.requires_grad_(False) self.system_encoder.eval() hidden = self.user_encoder.config.hidden_size num_layers = int(self.user_encoder.config.num_hidden_layers) + 1 # embeddings + layers self.layer_pool = LayerwiseAttention(num_layers) self.system_fusion = SystemContextFusion( hidden, num_heads=session_heads, dropout=dropout ) self.session_encoder = HierarchicalSessionEncoder( input_dim=hidden, session_dim=session_dim, num_layers=session_layers, num_heads=session_heads, dropout=dropout, max_turns=max_turns, ) self.bottleneck = nn.Sequential( nn.Linear(session_dim, bottleneck_dim), nn.LayerNorm(bottleneck_dim), ) # Trained with a VICReg projector; kept so joint checkpoints load without gaps. self.vicreg_proj = nn.Sequential( nn.Linear(bottleneck_dim, 256), nn.ReLU(), nn.Linear(256, 256), ) def train(self, mode: bool = True): # noqa: A003 super().train(mode) self.system_encoder.eval() return self @torch.no_grad() def _encode_system_tokens( self, texts: list[str], device: torch.device ) -> tuple[torch.Tensor, torch.Tensor]: if not texts: return ( torch.zeros(0, 1, self.system_encoder.config.hidden_size, device=device), torch.zeros(0, 1, dtype=torch.bool, device=device), ) # Head+tail before HF truncation so long assistant drafts keep both lead-in and # ending (default tokenizer truncation would keep only the start after the prefix). truncated = [ truncate_text_head_tail( t, self.system_tokenizer, max_length=self.max_length, head_tokens=128, tail_tokens=128, prefix=self.passage_prefix, ) for t in texts ] encoded = self.system_tokenizer( [f"{self.passage_prefix}{t}" for t in truncated], padding=True, truncation=True, max_length=self.max_length, return_tensors="pt", ).to(device) out = self.system_encoder(**encoded, output_hidden_states=False) tokens = out.last_hidden_state mask = encoded["attention_mask"].bool() return tokens, mask def _encode_user_queries(self, texts: list[str], device: torch.device) -> torch.Tensor: encoded = self.user_tokenizer( [f"{self.query_prefix}{t}" for t in texts], padding=True, truncation=True, max_length=self.max_length, return_tensors="pt", ).to(device) out = self.user_encoder(**encoded, output_hidden_states=True) # Use all hidden states including embeddings. layer_states = out.hidden_states pooled_layers = self.layer_pool(layer_states) return mean_pool(pooled_layers, encoded["attention_mask"]) def encode_sessions( self, sessions: list[list[dict[str, str]]], *, device: torch.device | None = None, ) -> dict[str, torch.Tensor]: """Encode a list of role/content turn lists into bottleneck embeddings.""" device = device or next(self.parameters()).device batch = len(sessions) # Collect valid user turns with optional preceding assistant text. user_texts: list[str] = [] system_texts: list[str] = [] has_system_flags: list[bool] = [] locations: list[tuple[int, int]] = [] max_turns = 0 for b, turns in enumerate(sessions): prev_system = "" turn_i = 0 for turn in turns: role = turn.get("role", "") content = turn.get("content", "") if role in {"assistant", "system"}: prev_system = content continue if role != "user" or not content.strip(): continue if turn_i >= self.max_turns: break user_texts.append(content) system_texts.append(prev_system) has_system_flags.append(bool(prev_system.strip())) locations.append((b, turn_i)) turn_i += 1 max_turns = max(max_turns, turn_i) hidden = self.user_encoder.config.hidden_size if not user_texts: zeros = torch.zeros(batch, self.bottleneck_dim, device=device) return { "bottleneck": zeros, "session": zeros, "turn_mask": torch.ones(batch, 1, dtype=torch.bool, device=device), } user_q = self._encode_user_queries(user_texts, device) # Encode systems that exist; use dummy for empty. nonempty_idx = [i for i, flag in enumerate(has_system_flags) if flag] system_tokens = torch.zeros( len(user_texts), 1, hidden, device=device, dtype=user_q.dtype ) system_mask = torch.zeros(len(user_texts), 1, dtype=torch.bool, device=device) if nonempty_idx: toks, mask = self._encode_system_tokens( [system_texts[i] for i in nonempty_idx], device ) # Pad to common length across all rows. max_s = toks.size(1) system_tokens = torch.zeros( len(user_texts), max_s, hidden, device=device, dtype=user_q.dtype ) system_mask = torch.zeros( len(user_texts), max_s, dtype=torch.bool, device=device ) system_tokens[nonempty_idx] = toks.to(user_q.dtype) system_mask[nonempty_idx] = mask has_system = torch.tensor(has_system_flags, device=device) fused = self.system_fusion(user_q, system_tokens, system_mask, has_system) turn_embeds = torch.zeros( batch, max(1, max_turns), hidden, device=device, dtype=fused.dtype ) turn_mask = torch.ones(batch, max(1, max_turns), dtype=torch.bool, device=device) for (b, t), vec in zip(locations, fused, strict=True): turn_embeds[b, t] = vec turn_mask[b, t] = False session = self.session_encoder(turn_embeds, turn_mask) bottleneck = self.bottleneck(session) return { "bottleneck": bottleneck, "session": session, "turn_mask": turn_mask, } def forward(self, sessions: list[list[dict[str, str]]]) -> dict[str, torch.Tensor]: return self.encode_sessions(sessions) def config_dict(self) -> dict[str, Any]: return { "user_model_name": self.user_model_name, "system_model_name": self.system_model_name, "bottleneck_dim": self.bottleneck_dim, "session_dim": self.session_encoder.session_dim, "session_layers": len(self.session_encoder.transformer.layers), "session_heads": self.system_fusion.cross_attn.num_heads, "max_turns": self.max_turns, "max_length": self.max_length, "query_prefix": self.query_prefix, "passage_prefix": self.passage_prefix, }