src/oxmini/model.py

"""OxMini hybrid KDA-lite/MLA-lite causal language model."""

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any

import torch
from torch import nn
from torch.nn import functional as F

from .attention_kda import KDALiteAttention
from .attention_mla import MLALiteAttention
from .config import OxMiniConfig
from .layers import RMSNorm, SwiGLU
from .mhc import MHCResidual, StreamCollapse


@dataclass
class CausalLMOutput:
    logits: torch.Tensor
    loss: torch.Tensor | None = None


class OxMiniBlock(nn.Module):
    def __init__(self, config: OxMiniConfig, attention_type: str) -> None:
        super().__init__()
        self.use_mhc = config.use_mhc
        self.norm_attn = RMSNorm(config.n_embd, config.rms_norm_eps)
        self.norm_mlp = RMSNorm(config.n_embd, config.rms_norm_eps)
        if attention_type == "kda":
            self.attention = KDALiteAttention(
                config.n_embd,
                config.n_head,
                config.dropout,
                config.bias,
                config.kda_state_norm_cap,
            )
        elif attention_type == "mla":
            self.attention = MLALiteAttention(
                config.n_embd,
                config.n_head,
                config.mla_latent_dim,
                config.dropout,
                config.bias,
            )
        else:
            raise ValueError(f"unknown attention type: {attention_type}")
        self.mlp = SwiGLU(
            config.n_embd,
            config.n_embd * config.ffn_multiplier,
            config.dropout,
            config.bias,
        )
        if self.use_mhc:
            self.attn_residual = MHCResidual(config.hc_streams, config.use_sinkhorn_mhc)
            self.mlp_residual = MHCResidual(config.hc_streams, config.use_sinkhorn_mhc)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.use_mhc:
            # Attention and MLP each have an independent routing matrix, just as
            # a standard pre-norm block has two independent residual additions.
            x = self.attn_residual(x, lambda value: self.attention(self.norm_attn(value)))
            return self.mlp_residual(x, lambda value: self.mlp(self.norm_mlp(value)))
        x = x + self.attention(self.norm_attn(x))
        return x + self.mlp(self.norm_mlp(x))

    def forward_step(
        self,
        x: torch.Tensor,
        cache: object | None,
        position: int,
        max_cache_length: int,
    ) -> tuple[torch.Tensor, object]:
        """Increment one block while preserving its KDA state or MLA K/V cache."""

        def attention_step(value: torch.Tensor) -> tuple[torch.Tensor, object]:
            normalized = self.norm_attn(value)
            if isinstance(self.attention, KDALiteAttention):
                return self.attention.step(normalized, state=cache)  # type: ignore[arg-type]
            return self.attention.step(
                normalized,
                cache=cache,  # type: ignore[arg-type]
                position=position,
                max_cache_length=max_cache_length,
            )

        if self.use_mhc:
            # Mirror MHCResidual.forward explicitly because the attention step
            # must return both its ordinary update and persistent decoder state.
            pre, post = self.attn_residual.weights()
            attention_input = torch.einsum("bsnd,n->bsd", x, pre)
            update, next_cache = attention_step(attention_input)
            x = x + update.unsqueeze(2) * post.view(
                1, 1, self.attn_residual.streams, 1
            )
            x = self.mlp_residual(x, lambda value: self.mlp(self.norm_mlp(value)))
            return x, next_cache

        update, next_cache = attention_step(x)
        x = x + update
        return x + self.mlp(self.norm_mlp(x)), next_cache


class OxMiniForCausalLM(nn.Module):
    def __init__(self, config: OxMiniConfig) -> None:
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.blocks = nn.ModuleList(
            [OxMiniBlock(config, attention_type) for attention_type in config.layer_types]
        )
        self.collapse = StreamCollapse(config.hc_streams) if config.use_mhc else nn.Identity()
        self.final_norm = RMSNorm(config.n_embd, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.apply(self._init_weights)
        if config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    @property
    def num_parameters(self) -> int:
        return sum(parameter.numel() for parameter in self.parameters())

    def forward(
        self,
        input_ids: torch.Tensor,
        targets: torch.Tensor | None = None,
    ) -> CausalLMOutput:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        if input_ids.shape[1] > self.config.block_size:
            raise ValueError(
                f"sequence length {input_ids.shape[1]} exceeds block_size {self.config.block_size}"
            )
        x = self.token_embedding(input_ids)
        if self.config.use_mhc:
            # Broadcast, do not concatenate: every stream starts as the same
            # token representation and subsequently diverges through learned
            # per-sublayer post-routing coefficients.
            x = x.unsqueeze(2).expand(-1, -1, self.config.hc_streams, -1)
        for block in self.blocks:
            x = block(x)
        x = self.collapse(x)
        logits = self.lm_head(self.final_norm(x))
        loss = None
        if targets is not None:
            # Every position predicts the next character supplied in ``targets``;
            # data batching performs the one-token shift before this call.
            loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), targets.reshape(-1))
        return CausalLMOutput(logits=logits, loss=loss)

    def forward_step(
        self,
        input_ids: torch.Tensor,
        caches: list[object | None] | None = None,
        position: int = 0,
    ) -> tuple[torch.Tensor, list[object]]:
        """Decode exactly one position and return per-layer recurrent caches."""

        if input_ids.ndim != 2 or input_ids.shape[1] != 1:
            raise ValueError("forward_step expects input IDs shaped [batch, 1]")
        if caches is None:
            caches = [None] * len(self.blocks)
        if len(caches) != len(self.blocks):
            raise ValueError("cache count must match the number of transformer blocks")

        x = self.token_embedding(input_ids)
        if self.config.use_mhc:
            x = x.unsqueeze(2).expand(-1, -1, self.config.hc_streams, -1)
        next_caches: list[object] = []
        for block, cache in zip(self.blocks, caches, strict=True):
            x, next_cache = block.forward_step(
                x,
                cache=cache,
                position=position,
                max_cache_length=self.config.block_size,
            )
            next_caches.append(next_cache)
        x = self.collapse(x)
        logits = self.lm_head(self.final_norm(x))
        return logits, next_caches

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        if input_ids.ndim != 2 or input_ids.shape[1] == 0:
            raise ValueError("input_ids must be a non-empty [batch, sequence] tensor")
        was_training = self.training
        self.eval()
        generated = input_ids
        caches: list[object | None] | None = None
        logits: torch.Tensor | None = None
        # Prefill token by token so KDA's associative state and MLA's compressed
        # K/V cache are identical to the states used during incremental decode.
        for position in range(input_ids.shape[1]):
            logits, caches = self.forward_step(
                input_ids[:, position : position + 1],
                caches=caches,
                position=position,
            )
        assert logits is not None
        for generation_index in range(max_new_tokens):
            next_logits = logits[:, -1, :]
            if not torch.isfinite(next_logits).all():
                raise FloatingPointError("non-finite logits encountered during generation")
            if temperature <= 0:
                next_token = next_logits.argmax(dim=-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k is not None:
                    k = min(top_k, next_logits.shape[-1])
                    cutoff = torch.topk(next_logits, k).values[:, [-1]]
                    next_logits = next_logits.masked_fill(
                        next_logits < cutoff, float("-inf")
                    )
                probabilities = torch.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probabilities, 1, generator=generator)
            generated = torch.cat((generated, next_token), dim=1)
            if generation_index + 1 < max_new_tokens:
                logits, caches = self.forward_step(
                    next_token,
                    caches=caches,
                    position=input_ids.shape[1] + generation_index,
                )
        if was_training:
            self.train()
        return generated

    def save_pretrained(self, directory: str | Path) -> Path:
        from safetensors.torch import save_file

        directory = Path(directory)
        directory.mkdir(parents=True, exist_ok=True)
        values: dict[str, Any] = self.config.to_dict()
        values.update({"architectures": [self.__class__.__name__], "model_type": "oxmini"})
        with (directory / "config.json").open("w", encoding="utf-8") as handle:
            json.dump(values, handle, indent=2, sort_keys=True)
            handle.write("\n")
        # Clone tied tensors so safetensors sees independent storage for both
        # state-dict keys while preserving strict-load compatibility.
        state = {
            key: value.detach().cpu().clone().contiguous()
            for key, value in self.state_dict().items()
        }
        save_file(state, str(directory / "pytorch_model.safetensors"))
        return directory

    @classmethod
    def from_pretrained(
        cls,
        model_id_or_path: str | Path,
        map_location: str | torch.device = "cpu",
        revision: str | None = None,
    ) -> "OxMiniForCausalLM":
        from safetensors.torch import load_file

        path = Path(model_id_or_path)
        if not path.exists():
            from huggingface_hub import snapshot_download

            path = Path(snapshot_download(str(model_id_or_path), revision=revision))
        config = OxMiniConfig.from_file(path / "config.json")
        model = cls(config)
        state = load_file(str(path / "pytorch_model.safetensors"), device=str(map_location))
        model.load_state_dict(state)
        return model.to(map_location)