dat_lm.py

"""Decoder-only LM built from dual-attention blocks."""

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence

from .dat_config import DatLMConfig
from .dat_core import (
    DisentangledRelationalCrossAttention,
    RelationalAttention,
    RelationalCrossAttention,
)
from .transformer_core import SelfAttention
from .dat_symbols import (
    PositionalSymbolRetriever,
    RelationalSymbolicAttentionRetriever,
    RelativePositionalSymbolRetriever,
    SymbolicAttentionRetriever,
)
from .masks import build_decoder_attention_mask
from .mlm import MaskClassifier
from .transformer_components import FeedForward, PositionalEncoding, PositionalInfo, RMSNorm


class DatDecoderBlock(nn.Module):
    def __init__(self, config: DatLMConfig):
        super().__init__()
        self.norm_first = config.norm_first
        # input_dim is the block state width; hidden_dim is the total DAT width
        # used with total_n_heads to derive the shared SA/RA head_dim.
        self.sensory_attention = SelfAttention(
            input_dim=config.hidden_dim,
            n_heads=config.n_heads_sa,
            hidden_dim=config.hidden_dim,
            total_n_heads=config.total_n_heads,
            dropout=config.dropout,
            supports_relative=config.pe_type == "relative",
            use_bias_qkv=config.use_bias_qkv,
            use_bias_out=config.use_bias_out,
        )
        self.relational_attention = _build_relational_attention(config)
        if config.share_attn_params:
            self.sensory_attention.q_proj = self.relational_attention.q_proj
            self.sensory_attention.k_proj = self.relational_attention.k_proj
        self.feed_forward = FeedForward(
            input_dim=config.hidden_dim,
            hidden_dim=config.resolved_ffn_hidden_dim,
            dropout=config.dropout,
            activation=config.ffn_activation,
            use_bias=config.use_bias_ffn,
        )
        self.norm1 = _create_norm(config.norm_type, config.hidden_dim)
        self.norm2 = _create_norm(config.norm_type, config.hidden_dim)
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self,
        x: torch.Tensor,
        symbol_retriever: nn.Module,
        mask: torch.Tensor,
        pos_info: Optional[PositionalInfo],
    ) -> torch.Tensor:
        if self.norm_first:
            normed_x = self.norm1(x)
            # Retrieve symbols at the block boundary so pre-norm attention and
            # input-dependent symbol retrieval use the same representation.
            symbols = symbol_retriever(normed_x)
            sensory_output, _ = self.sensory_attention(normed_x, mask=mask, pos_info=pos_info)
            relational_output, _ = self.relational_attention(
                normed_x,
                symbols,
                mask=mask,
                pos_info=pos_info,
            )
            x = x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1))
            x = x + self.dropout(self.feed_forward(self.norm2(x)))
            return x

        symbols = symbol_retriever(x)
        sensory_output, _ = self.sensory_attention(x, mask=mask, pos_info=pos_info)
        relational_output, _ = self.relational_attention(
            x,
            symbols,
            mask=mask,
            pos_info=pos_info,
        )
        x = self.norm1(x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1)))
        x = self.norm2(x + self.dropout(self.feed_forward(x)))
        return x


class DatDecoderLM(nn.Module):
    def __init__(self, config: DatLMConfig):
        super().__init__()
        self.config = config
        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.embedding_dropout = nn.Dropout(config.dropout)
        position_dim = (
            config.hidden_dim if config.pe_type in {"sinusoidal", "learned", "none"}
            else config.head_dim
        )
        self.position_encoder = PositionalEncoding(
            embedding_dim=position_dim,
            pe_type=config.pe_type,
            max_len=config.max_seq_len,
            theta=config.rope_theta,
            max_rel_pos=config.max_rel_pos,
            init_range=config.init_range,
        )
        self.symbol_retrievers = _build_symbol_retrievers(config)
        self.layers = nn.ModuleList(DatDecoderBlock(config) for _ in range(config.n_layers))
        self.final_norm = _create_norm(config.norm_type, config.hidden_dim)
        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        if config.mlm_head_enabled:
            self.mlm_head = MaskClassifier(
                hidden_dim=config.hidden_dim,
                vocab_size=config.vocab_size,
                norm_type=config.norm_type,
                ffn_activation=config.ffn_activation,
                use_bias_ffn=config.use_bias_ffn,
                dropout=config.dropout,
                word_embedding=self.lm_head.weight if not config.tie_lm_head else self.token_embeddings.weight,
            )
        self._init_weights()

    @property
    def device(self) -> torch.device:
        return self.token_embeddings.weight.device

    def _init_weights(self) -> None:
        if self.config.init_scheme == "xavier_uniform":
            nn.init.xavier_uniform_(
                self.token_embeddings.weight,
                gain=nn.init.calculate_gain("linear"),
            )
            if self.config.tie_lm_head:
                self.lm_head.weight = self.token_embeddings.weight
            else:
                nn.init.xavier_uniform_(
                    self.lm_head.weight,
                    gain=nn.init.calculate_gain("linear"),
                )
            return

        if self.config.init_scheme == "normal_0_02_scaled_projection":
            self._init_normal_0_02_scaled_projection()
            if self.config.tie_lm_head:
                self.lm_head.weight = self.token_embeddings.weight
            return

        raise ValueError(f"Unsupported init_scheme: {self.config.init_scheme}")

    def _init_normal_0_02_scaled_projection(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

        for module in self.modules():
            if isinstance(module, RelationalAttention):
                nn.init.normal_(module.wr_proj, mean=0.0, std=0.02)
            elif isinstance(module, SymbolicAttentionRetriever):
                nn.init.normal_(module.template_features, mean=0.0, std=1.0)
                nn.init.normal_(module.symbol_library, mean=0.0, std=1.0)
            elif isinstance(module, RelativePositionalSymbolRetriever) and not module.rope:
                nn.init.xavier_uniform_(module.position_encoder.rel_pos_embeddings_table.weight)

        scaled_std = 0.02 / math.sqrt(2 * self.config.n_layers)
        ffn_scaled_suffix = (
            "feed_forward.w_up.weight"
            if self.config.ffn_activation == "swiglu"
            else "feed_forward.linear2.weight"
        )
        for name, parameter in self.named_parameters():
            # Exclude the MLM head from depth-scaled init: its feed_forward
            # down-projection shares the same suffix as decoder layers but
            # should use the regular 0.02 initialization.
            if name.startswith("mlm_head"):
                continue
            if name.endswith("o_proj.weight") or name.endswith(ffn_scaled_suffix):
                nn.init.normal_(parameter, mean=0.0, std=scaled_std)

    def _build_attention_mask(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        bidirectional: bool = False,
    ) -> torch.Tensor:
        return build_decoder_attention_mask(
            input_ids=input_ids,
            pad_token_id=self.config.pad_token_id,
            eos_token_id=self.config.eos_token_id,
            sequence_boundary_policy=self.config.sequence_boundary_policy,
            attention_mask=attention_mask,
            segment_boundary_token_id=self.config.segment_boundary_token_id,
            bidirectional=bidirectional,
        )

    def encode_for_objective(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        bidirectional: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if input_ids.dim() != 2:
            raise ValueError(f"input_ids must be rank-2 [batch, seq], got shape {tuple(input_ids.shape)}")
        if input_ids.shape[1] > self.config.max_seq_len:
            raise ValueError(
                f"Sequence length {input_ids.shape[1]} exceeds max_seq_len {self.config.max_seq_len}"
            )


        mask = self._build_attention_mask(input_ids, attention_mask, bidirectional=bidirectional)
        token_embeddings = self.token_embeddings(input_ids)
        hidden_states = token_embeddings
        pos_info = self.position_encoder.get_positional_info(input_ids.shape[1], input_ids.device)
        if pos_info.apply_to_embeddings:
            if pos_info.embeddings is None:
                raise ValueError(f"Embedding-level pe_type {pos_info.pe_type} did not provide embeddings.")
            hidden_states = hidden_states + pos_info.embeddings.unsqueeze(0)
        hidden_states = self.embedding_dropout(hidden_states)

        for symbol_retriever, layer in zip(self.symbol_retrievers, self.layers):
            hidden_states = layer(
                hidden_states,
                symbol_retriever=symbol_retriever,
                mask=mask,
                pos_info=pos_info,
            )

        hidden_states = self.final_norm(hidden_states)
        return token_embeddings, hidden_states

    def forward(
        self,
        input_ids: torch.Tensor,
        targets: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        _, hidden_states = self.encode_for_objective(input_ids, attention_mask=attention_mask)
        logits = self.lm_head(hidden_states)

        loss = None
        if targets is not None:
            if targets.shape != input_ids.shape:
                raise ValueError(
                    f"targets shape must match input_ids shape, got {tuple(targets.shape)} "
                    f"vs {tuple(input_ids.shape)}"
                )
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=-1,
            )

        return logits, loss

    def forward_mlm(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Encode bidirectionally and apply the MLM head.

        Args:
            input_ids: Masked input token ids [batch, seq].
            attention_mask: Validity mask [batch, seq].

        Returns:
            MLM logits at every position [batch, seq, vocab].
        """
        if not self.config.mlm_head_enabled:
            raise ValueError(
                "forward_mlm requires mlm_head_enabled=True; "
                "the MLM head is not instantiated."
            )
        _, hidden_states = self.encode_for_objective(
            input_ids, attention_mask=attention_mask, bidirectional=True
        )
        return self.mlm_head(hidden_states)

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        do_sample: bool = True,
        eos_token_id: Optional[int] = None,
    ) -> torch.Tensor:
        if temperature <= 0.0:
            raise ValueError(f"temperature must be positive, got {temperature}")
        if max_new_tokens < 0:
            raise ValueError(f"max_new_tokens must be non-negative, got {max_new_tokens}")
        if top_k is not None and top_k <= 0:
            raise ValueError(f"top_k must be positive when provided, got {top_k}")
        if eos_token_id is None:
            eos_token_id = self.config.eos_token_id

        if attention_mask is None:
            generated_sequences = [row.clone() for row in input_ids]
        else:
            current_attention_mask = attention_mask.bool()
            generated_sequences = []
            for row, row_mask in zip(input_ids, current_attention_mask):
                generated_sequence = row[row_mask]
                if generated_sequence.numel() == 0:
                    raise ValueError("Each input row must contain at least one unmasked token for generation.")
                generated_sequences.append(generated_sequence)

        finished = torch.tensor(
            [sequence[-1].item() == eos_token_id for sequence in generated_sequences],
            dtype=torch.bool,
            device=input_ids.device,
        )

        for _ in range(max_new_tokens):
            context_sequences = [sequence[-self.config.max_seq_len :] for sequence in generated_sequences]
            context_ids = pad_sequence(
                context_sequences,
                batch_first=True,
                padding_value=self.config.pad_token_id,
            )
            context_mask = pad_sequence(
                [
                    torch.ones(sequence.shape[0], dtype=torch.bool, device=input_ids.device)
                    for sequence in context_sequences
                ],
                batch_first=True,
                padding_value=False,
            )
            logits, _ = self(context_ids, attention_mask=context_mask)
            last_positions = context_mask.long().sum(dim=1) - 1
            batch_indices = torch.arange(logits.shape[0], device=logits.device)
            next_token_logits = logits[batch_indices, last_positions, :] / temperature

            if top_k is not None:
                k = min(top_k, next_token_logits.size(-1))
                top_values, _ = torch.topk(next_token_logits, k=k)
                cutoff = top_values[:, -1].unsqueeze(-1)
                next_token_logits = next_token_logits.masked_fill(next_token_logits < cutoff, float("-inf"))

            if do_sample:
                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            for row_index in range(next_token.shape[0]):
                if finished[row_index]:
                    continue
                generated_sequences[row_index] = torch.cat(
                    [generated_sequences[row_index], next_token[row_index]],
                )
                if next_token[row_index, 0].item() == eos_token_id:
                    finished[row_index] = True

            if torch.all(finished):
                break

        return pad_sequence(
            generated_sequences,
            batch_first=True,
            padding_value=self.config.pad_token_id,
        )


def _build_symbol_retrievers(config: DatLMConfig) -> nn.ModuleList:
    retriever = _build_symbol_retriever(config)
    if config.shared_symbol_retriever:
        return nn.ModuleList([retriever] * config.n_layers)
    return nn.ModuleList(_build_symbol_retriever(config) for _ in range(config.n_layers))


def _build_symbol_retriever(config: DatLMConfig) -> nn.Module:
    if config.symbol_retrieval == "symbolic":
        return SymbolicAttentionRetriever(
            hidden_dim=config.hidden_dim,
            symbol_dim=config.resolved_symbol_dim,
            n_symbols=config.resolved_n_symbols,
            n_heads=config.resolved_symbolic_attn_n_heads,
            dropout=config.dropout,
            use_bias=config.symbolic_use_bias,
        )
    if config.symbol_retrieval == "positional":
        return PositionalSymbolRetriever(
            symbol_dim=config.resolved_symbol_dim,
            max_len=config.max_seq_len,
            sinusoidal=config.positional_symbols_sinusoidal,
        )
    if config.symbol_retrieval == "relative":
        max_rel_distance = config.max_rel_pos if config.max_rel_pos is not None else max(1, config.max_seq_len // 2)
        return RelativePositionalSymbolRetriever(
            symbol_dim=config.resolved_symbol_dim,
            max_rel_distance=max_rel_distance,
            rope=config.relative_symbols_rope,
            theta=config.rope_theta,
        )
    if config.symbol_retrieval == "relsymbolic":
        return RelationalSymbolicAttentionRetriever(
            hidden_dim=config.hidden_dim,
            symbol_dim=config.resolved_symbol_dim,
            rel_n_heads=config.relsymbolic_rel_n_heads,
            symbolic_attn_n_heads=config.relsymbolic_symbolic_attn_n_heads,
            n_symbols=config.resolved_n_symbols,
            neighborhood_size=config.relsymbolic_neighborhood_size,
            include_self=config.relsymbolic_include_self,
            normalize_rels=config.relsymbolic_normalize_rels,
            dropout=config.relsymbolic_dropout,
            trainable_symbols=config.relsymbolic_trainable_symbols,
            rel_scale=config.relsymbolic_rel_scale,
            symbolic_attn_scale=config.relsymbolic_symbolic_attn_scale,
            use_bias=config.relsymbolic_use_bias,
        )
    raise ValueError(f"Unsupported symbol_retrieval: {config.symbol_retrieval}")


def _build_relational_attention(config: DatLMConfig) -> nn.Module:
    use_relative_symbols = config.symbol_retrieval == "relative"
    if config.ra_type == "ra":
        return RelationalAttention(
            hidden_dim=config.hidden_dim,
            symbol_dim=config.resolved_symbol_dim,
            n_heads=config.n_heads_ra,
            total_n_heads=config.total_n_heads,
            n_relations=config.resolved_ra_n_relations,
            dropout=config.dropout,
            rel_activation=config.ra_rel_activation,
            symmetric_rels=config.ra_symmetric_rels,
            use_relative_positional_symbols=use_relative_symbols,
            use_bias_qkv=config.use_bias_qkv,
            use_bias_out=config.use_bias_out,
        )
    if config.ra_type == "rca":
        return RelationalCrossAttention(
            hidden_dim=config.hidden_dim,
            symbol_dim=config.resolved_symbol_dim,
            n_heads=config.n_heads_ra,
            total_n_heads=config.total_n_heads,
            dropout=config.dropout,
            activation=config.ra_rel_activation,
            use_relative_positional_symbols=use_relative_symbols,
            use_bias_qkv=config.use_bias_qkv,
            use_bias_out=config.use_bias_out,
        )
    if config.ra_type == "disrca":
        return DisentangledRelationalCrossAttention(
            hidden_dim=config.hidden_dim,
            symbol_dim=config.resolved_symbol_dim,
            n_heads=config.n_heads_ra,
            total_n_heads=config.total_n_heads,
            dropout=config.dropout,
            rel_activation=config.ra_rel_activation,
            use_relative_positional_symbols=use_relative_symbols,
            use_bias_qkv=config.use_bias_qkv,
            use_bias_out=config.use_bias_out,
        )
    raise ValueError(f"Unsupported ra_type: {config.ra_type}")


def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
    if norm_type == "layernorm":
        return nn.LayerNorm(hidden_dim)
    if norm_type == "rmsnorm":
        return RMSNorm(hidden_dim)
    raise ValueError(f"Unsupported norm_type: {norm_type}")