cue / encoder_model.py

Upload folder using huggingface_hub

6c05b39 verified 3 months ago

12.5 kB

"""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,
        }

	"""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,
	}