src/model/lfm2_small.py

"""LFM2Small: scaled-down LFM2.5-1.2B backbone for transaction sequences.

Reimplements (not subclasses) the core LFM2 architecture at ~8.3M total params.
8 layers in conv-conv-attn-conv-attn-conv-attn-conv order, preserving every
structural choice from the full 1.2B model:

  - Gated short convolution with depthwise causal Conv1d
  - Grouped query attention (4Q / 2KV, group size 2) with QK RMSNorm
  - SwiGLU MLP with auto-adjusted intermediate dimension
  - Pre-norm residual connections
  - Final RMSNorm (embedding_norm) before LM heads

Module naming matches LFM2 conventions exactly:
  layers[i].self_attn.{q_proj, k_proj, v_proj, out_proj, q_layernorm, k_layernorm}
  layers[i].conv.{in_proj, out_proj, conv}
  layers[i].feed_forward.{w1, w2, w3}
  layers[i].{operator_norm, ffn_norm}
  embedding_norm

Reference: modeling_lfm2.py in HuggingFace transformers (LiquidAI/LFM2-1.2B).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml

from src.data.schema import SchemaConfig, load_schema
from src.model.embedding import StructuredEmbedding
from src.model.heads import PerFeatureLMHeads


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass
class ModelConfig:
    """Typed config for LFM2Small. Loads from configs/model.yaml."""

    hidden_size: int = 256
    intermediate_size: int = 1024
    num_attention_heads: int = 4
    num_key_value_heads: int = 2
    num_layers: int = 8
    layer_order: list[str] = field(default_factory=lambda: [
        "conv", "conv", "attn", "conv", "attn", "conv", "attn", "conv",
    ])
    conv_kernel_size: int = 3
    block_auto_adjust_ff_dim: bool = True
    block_multiple_of: int = 256
    block_ffn_dim_multiplier: float = 1.0
    rms_norm_eps: float = 1e-6
    rope_theta: float = 1_000_000.0
    max_position_embeddings: int = 4096
    initializer_range: float = 0.02
    num_transactions: int = 64
    num_features: int = 15

    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.num_attention_heads

    @property
    def num_kv_groups(self) -> int:
        return self.num_attention_heads // self.num_key_value_heads

    @property
    def effective_intermediate_size(self) -> int:
        """MLP dim after LFM2's block_auto_adjust_ff_dim.

        With hidden=256, intermediate=1024: int(2*1024/3)=682, rounded to 768.
        """
        if not self.block_auto_adjust_ff_dim:
            return self.intermediate_size
        size = int(2 * self.intermediate_size / 3)
        size = int(self.block_ffn_dim_multiplier * size)
        return self.block_multiple_of * (
            (size + self.block_multiple_of - 1) // self.block_multiple_of
        )

    @classmethod
    def from_yaml(cls, path: str | Path) -> ModelConfig:
        with open(path) as f:
            raw = yaml.safe_load(f)
        bb = raw.get("backbone", {})
        seq = raw.get("sequence", {})
        return cls(
            hidden_size=bb.get("hidden_size", 256),
            intermediate_size=bb.get("intermediate_size", 1024),
            num_attention_heads=bb.get("num_attention_heads", 4),
            num_key_value_heads=bb.get("num_key_value_heads", 2),
            num_layers=bb.get("num_layers", 8),
            layer_order=bb.get("layer_order", [
                "conv", "conv", "attn", "conv", "attn", "conv", "attn", "conv",
            ]),
            conv_kernel_size=bb.get("conv_kernel_size", 3),
            block_auto_adjust_ff_dim=bb.get("block_auto_adjust_ff_dim", True),
            block_multiple_of=bb.get("block_multiple_of", 256),
            block_ffn_dim_multiplier=bb.get("block_ffn_dim_multiplier", 1.0),
            rms_norm_eps=bb.get("rms_norm_eps", 1e-6),
            rope_theta=bb.get("rope_theta", 1_000_000.0),
            num_transactions=seq.get("num_transactions", 64),
            num_features=seq.get("features_per_transaction", 15),
        )


# ---------------------------------------------------------------------------
# Building blocks
# ---------------------------------------------------------------------------


class RMSNorm(nn.Module):
    """Root mean square layer normalization (Lfm2RMSNorm)."""

    def __init__(self, dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (self.weight * x).to(dtype)


class RotaryEmbedding(nn.Module):
    """Rotary position embeddings. Flat token positions 0..S-1."""

    def __init__(
        self, head_dim: int, max_seq_len: int = 4096, theta: float = 1_000_000.0,
    ) -> None:
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.max_seq_len = max_seq_len

    def forward(
        self, x: torch.Tensor, position_ids: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Returns (cos, sin) each shaped (B, S, head_dim)."""
        # (1, D/2, 1) @ (B, 1, S) -> (B, D/2, S) -> (B, S, D/2)
        inv_freq = self.inv_freq[None, :, None].float().to(x.device)
        pos = position_ids[:, None, :].float()
        freqs = (inv_freq @ pos).transpose(1, 2)
        emb = torch.cat([freqs, freqs], dim=-1)  # (B, S, head_dim)
        return emb.cos().to(x.dtype), emb.sin().to(x.dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """cos/sin: (B, S, D) unsqueezed to (B, 1, S, D). q/k: (B, H, S, D)."""
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expand KV heads for GQA: (B, H_kv, S, D) -> (B, H_kv*n_rep, S, D)."""
    if n_rep == 1:
        return x
    B, H, S, D = x.shape
    return x[:, :, None, :, :].expand(B, H, n_rep, S, D).reshape(B, H * n_rep, S, D)


# ---------------------------------------------------------------------------
# Layers
# ---------------------------------------------------------------------------


class SwiGLU(nn.Module):
    """SwiGLU MLP (Lfm2MLP): w2(silu(w1(x)) * w3(x))."""

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        intermediate = config.effective_intermediate_size
        self.w1 = nn.Linear(config.hidden_size, intermediate, bias=False)
        self.w3 = nn.Linear(config.hidden_size, intermediate, bias=False)
        self.w2 = nn.Linear(intermediate, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class ShortConv(nn.Module):
    """Gated short convolution (Lfm2ShortConv).

    in_proj splits hidden -> (B_gate, C_gate, x). B_gate * x feeds a causal
    depthwise Conv1d, output gated by C_gate, then out_proj. Left-padding
    (padding=kernel-1, truncated to seqlen) ensures no future token leakage.
    """

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        h, k = config.hidden_size, config.conv_kernel_size
        self.in_proj = nn.Linear(h, 3 * h, bias=False)
        self.conv = nn.Conv1d(h, h, kernel_size=k, groups=h, bias=False, padding=k - 1)
        self.out_proj = nn.Linear(h, h, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # (B, S, D) -> project -> (B, 3D, S) -> chunk into three (B, D, S) tensors
        seqlen = hidden_states.shape[1]
        BCx = self.in_proj(hidden_states).transpose(-1, -2)
        B_gate, C_gate, x = BCx.chunk(3, dim=-2)
        # Causal depthwise conv: left-padded, truncate right to preserve causality
        conv_out = self.conv(B_gate * x)[..., :seqlen]
        y = C_gate * conv_out
        return self.out_proj(y.transpose(-1, -2).contiguous())


class Attention(nn.Module):
    """Grouped query attention with QK RMSNorm (Lfm2Attention).

    QK norms after projection and before rotary stabilize deep training.
    Present in LFM2 but absent from LLaMA-family models.
    """

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.num_kv_groups = config.num_kv_groups
        self.head_dim = config.head_dim
        self.scaling = self.head_dim ** -0.5

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)

        self.q_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        B, S, _ = hidden_states.shape

        # Project -> reshape to heads -> QK norm -> transpose to (B, H, S, D)
        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)

        q = self.q_layernorm(q).transpose(1, 2)  # (B, H, S, D)
        k = self.k_layernorm(k).transpose(1, 2)  # (B, H_kv, S, D)
        v = v.transpose(1, 2)                     # (B, H_kv, S, D)

        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        k = repeat_kv(k, self.num_kv_groups)  # (B, H, S, D)
        v = repeat_kv(v, self.num_kv_groups)

        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=self.scaling)
        return self.out_proj(attn_out.transpose(1, 2).reshape(B, S, -1).contiguous())


class DecoderLayer(nn.Module):
    """Pre-norm residual: conv or attention + SwiGLU (Lfm2DecoderLayer).

    x = x + op(operator_norm(x))       # op = conv or self_attn
    x = x + feed_forward(ffn_norm(x))
    """

    def __init__(self, config: ModelConfig, layer_idx: int) -> None:
        super().__init__()
        self.is_attention_layer = config.layer_order[layer_idx] == "attn"

        if self.is_attention_layer:
            self.self_attn = Attention(config)
        else:
            self.conv = ShortConv(config)

        self.feed_forward = SwiGLU(config)
        self.operator_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        residual = hidden_states
        if self.is_attention_layer:
            hidden_states = self.self_attn(self.operator_norm(hidden_states), position_embeddings)
        else:
            hidden_states = self.conv(self.operator_norm(hidden_states))
        hidden_states = hidden_states + residual
        return hidden_states + self.feed_forward(self.ffn_norm(hidden_states))


# ---------------------------------------------------------------------------
# Full model
# ---------------------------------------------------------------------------


class LFM2Small(nn.Module):
    """LFM2-small: structured embedding + interleaved backbone + tied LM heads.

    ~8.3M params at hidden=256 (embedding ~1.7M, backbone ~6.6M, heads tied).
    """

    def __init__(self, config: ModelConfig, schema: SchemaConfig) -> None:
        super().__init__()
        self.config = config

        assert len(config.layer_order) == config.num_layers
        assert config.hidden_size % config.num_attention_heads == 0
        assert config.num_attention_heads % config.num_key_value_heads == 0

        self.embedding = StructuredEmbedding(schema, config.hidden_size)
        self.layers = nn.ModuleList([
            DecoderLayer(config, i) for i in range(config.num_layers)
        ])
        self.rotary_emb = RotaryEmbedding(
            config.head_dim, config.max_position_embeddings, config.rope_theta,
        )
        self.embedding_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.lm_heads = PerFeatureLMHeads(self.embedding)

        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights following LFM2 conventions. Skips lm_heads (tied)."""
        for name, module in self.named_modules():
            if name.startswith("lm_heads"):
                continue
            if isinstance(module, (nn.Linear, nn.Conv1d)):
                nn.init.normal_(module.weight, std=self.config.initializer_range)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=self.config.initializer_range)

    def backbone_forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """Embedding + backbone + final norm. Returns (B, S, D).

        Use for downstream heads (fraud prediction) that skip LM logits.
        """
        hidden_states = self.embedding(token_ids)  # (B, T*F, D)
        position_ids = torch.arange(
            hidden_states.shape[1], device=hidden_states.device,
        ).unsqueeze(0)
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for layer in self.layers:
            hidden_states = layer(hidden_states, position_embeddings)

        return self.embedding_norm(hidden_states)

    def forward(self, token_ids: torch.Tensor) -> list[torch.Tensor]:
        """Token IDs -> per-feature logits for causal LM pretraining.

        Args:
            token_ids: (B, T, F) int tensor.

        Returns:
            15 tensors, each (B, T*F, vocab_size_f). Position p predicts
            position p+1; the training loop selects head[(p+1) % num_features].
        """
        return self.lm_heads(self.backbone_forward(token_ids))

    def param_count(self) -> dict[str, int]:
        """Parameter counts by component. Accounts for weight tying."""
        emb = sum(p.numel() for p in self.embedding.parameters())
        backbone = sum(p.numel() for p in self.layers.parameters())
        backbone += sum(p.numel() for p in self.embedding_norm.parameters())
        total = sum(p.numel() for p in self.parameters())
        return {"embedding": emb, "backbone": backbone, "lm_heads_tied": 0, "total_unique": total}

    @classmethod
    def from_config_files(cls, model_yaml: str | Path, schema_yaml: str | Path) -> LFM2Small:
        """Construct from YAML config files."""
        return cls(ModelConfig.from_yaml(model_yaml), load_schema(schema_yaml))