src/taoTrain/models/gamma_net.py

"""GammaNet language model using GammaSpaceModel blocks as the sequence mixer."""

from __future__ import annotations

import site
import sys
from typing import Any, Optional

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

from taoTrain.config import ModelConfig
from taoTrain.core import BaseModel

from .embeddings import FactorizedEmbedding
from .registry import register_architecture


def _import_gamma_space_block():
    """Import GammaSpaceBlock, retrying with user site-packages if needed."""
    try:
        from gamma_space_model import GammaSpaceBlock
        return GammaSpaceBlock
    except ModuleNotFoundError as exc:
        user_site = site.getusersitepackages()
        if user_site and user_site not in sys.path:
            sys.path.append(user_site)
            try:
                from gamma_space_model import GammaSpaceBlock
                return GammaSpaceBlock
            except ModuleNotFoundError:
                pass
        raise ModuleNotFoundError(
            "gamma_net requires the GammaSpaceModel package. Install it with "
            "`pip install \"git+https://github.com/Taotern/GammaSpaceModel.git\"`."
        ) from exc


class GammaMixerBlock(nn.Module):
    """Thin adapter from TaoTrain config/model flow to GammaSpaceBlock."""

    def __init__(self, config: ModelConfig):
        super().__init__()
        gamma_block_cls = _import_gamma_space_block()
        self.block = gamma_block_cls(
            d_model=config.hidden_dim,
            hidden_dim=config.gamma_hidden_dim,
            dt_min=config.gamma_dt_min,
            dt_max=config.gamma_dt_max,
            dt_init=config.gamma_dt_init,
            discretization=config.gamma_discretization,
            prenorm=config.gamma_prenorm,
            residual_scale=config.gamma_residual_scale,
            dropout=config.dropout,
            activation=config.gamma_activation,
            gate=config.gamma_gate,
            use_D=config.gamma_use_D,
            kernel_mode=config.gamma_kernel_mode,
            kernel_threshold=config.gamma_kernel_threshold,
            use_output_linear=config.gamma_use_output_linear,
            gate_bias=config.gamma_gate_bias,
            input_gate=config.gamma_input_gate,
            input_gate_bias=config.gamma_input_gate_bias,
            layer_scale_init=config.gamma_layer_scale_init,
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # GammaSpaceBlock expects a padding mask of shape [batch, seq].
        mask = attention_mask
        if mask is not None and mask.dim() > 2:
            mask = None
        y, _ = self.block(x, mask=mask, return_state=False)
        return y


@register_architecture("gamma_net")
class GammaNetLM(BaseModel):
    """Causal LM with TaoNet shell and Gamma Space Model residual blocks."""

    def __init__(self, config: ModelConfig):
        super().__init__(config)

        self.vocab_size = config.vocab_size
        self.d_model = config.hidden_dim
        self.n_layers = config.num_layers
        self.dropout = config.dropout
        self.d_ff = config.hidden_dim_ff if config.hidden_dim_ff is not None else (self.d_model * 4)
        self.use_factorized_embedding = getattr(config, "use_factorized_embedding", False)
        self.d_embed_rank = getattr(config, "d_embed_rank", 96)

        if self.use_factorized_embedding:
            self.token_embedding = FactorizedEmbedding(
                self.vocab_size,
                self.d_model,
                self.d_embed_rank,
            )
        else:
            self.token_embedding = nn.Embedding(self.vocab_size, self.d_model)

        self.embedding_dropout = nn.Dropout(self.dropout)
        self.blocks = nn.ModuleList([GammaMixerBlock(config) for _ in range(self.n_layers)])
        self.final_norm = nn.LayerNorm(self.d_model)
        self.output_head = nn.Linear(self.d_model, self.vocab_size, bias=False)

        self._init_shell_weights()
        self._print_architecture()

    def _init_shell_weights(self) -> None:
        """Initialize only TaoTrain-owned LM shell weights."""
        if isinstance(self.token_embedding, nn.Embedding):
            nn.init.normal_(self.token_embedding.weight, mean=0.0, std=self.config.init_std)
        nn.init.normal_(self.output_head.weight, mean=0.0, std=self.config.init_std)

    def _print_architecture(self) -> None:
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)

        print(f"\n{'=' * 70}")
        print("MODEL ARCHITECTURE - GAMMA_NET (Gamma Space Model)")
        print(f"{'=' * 70}")
        print(f"Embedding type:      {'Factorized' if self.use_factorized_embedding else 'Standard'}")
        print(f"Gamma hidden_dim:    {self.config.gamma_hidden_dim}")
        print(f"Gamma discretization:{self.config.gamma_discretization}")
        print(f"Gamma kernel_mode:   {self.config.gamma_kernel_mode}")
        print(f"Layers:              {self.n_layers}")
        print(f"d_model:             {self.d_model}")
        print(f"Dropout:             {self.dropout}")
        print(f"Max sequence length: {self.config.max_seq_length}")
        print(f"Total Parameters:    {total_params/1e6:>8.2f}M (trainable: {trainable_params/1e6:.2f}M)")
        print(f"{'=' * 70}\n")

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
    ) -> dict[str, Any]:
        x = self.token_embedding(input_ids)
        x = self.embedding_dropout(x)

        for block in self.blocks:
            x = block(x, attention_mask=attention_mask)

        x = self.final_norm(x)
        logits = self.output_head(x)

        loss = None
        if labels is not None:
            logits_flat = logits.view(-1, logits.size(-1))
            labels_flat = labels.view(-1)
            loss = F.cross_entropy(
                logits_flat,
                labels_flat,
                reduction="mean",
                ignore_index=-100,
            )

        return {
            "logits": logits,
            "loss": loss,
        }