model.py

"""
GPT pequeno com duas variantes de residual:
  - "baseline"  : soma residual fixa  (padrão Transformer)
  - "attnres"   : Block Attention Residuals (atenção sobre estados anteriores no bloco)

Arquitetura:
  d_model    = 256
  n_layers   = 8
  n_heads    = 8
  d_ff       = 1024
  max_seq    = 512
  vocab_size = 16384
  ~25M parâmetros

Pequeno o suficiente para treinar rápido na RTX 3080 e mostrar diferenças
de convergência entre as variantes.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field


@dataclass
class GPTConfig:
    vocab_size : int   = 16_384
    max_seq    : int   = 512
    d_model    : int   = 512      # 256→512: ~4× mais parâmetros por camada
    n_layers   : int   = 12       # 8→12: mais profundidade, mais estados para AttnRes explorar
    n_heads    : int   = 8
    d_ff       : int   = 2_048    # 1024→2048
    dropout    : float = 0.1
    variant    : str   = "baseline"
    block_size : int   = 4

# Tamanhos de referência:
#   small  : d=256, L=8,  ff=1024  → ~10M params  (experimento rápido)
#   medium : d=512, L=12, ff=2048  → ~50M params  (treino sério, ~6h na 3080)
#   large  : d=768, L=16, ff=3072  → ~120M params (cabível na 3080 com seq=256)


# ── Blocos base ────────────────────────────────────────────────────────────────

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        assert cfg.d_model % cfg.n_heads == 0
        self.n_heads = cfg.n_heads
        self.d_head  = cfg.d_model // cfg.n_heads
        self.qkv     = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.out     = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.drop    = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        q, k, v  = self.qkv(x).split(D, dim=-1)

        def reshape(t):
            return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        q, k, v = reshape(q), reshape(k), reshape(v)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.drop.p if self.training else 0.0)
        y = y.transpose(1, 2).contiguous().view(B, T, D)
        return self.out(y)


class FeedForward(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_ff, bias=False),
            nn.GELU(),
            nn.Linear(cfg.d_ff, cfg.d_model, bias=False),
            nn.Dropout(cfg.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ── Residual: variante clássica ────────────────────────────────────────────────

class FixedResidual(nn.Module):
    """x + sublayer(x)  —  residual padrão."""
    def forward(self, x: torch.Tensor, states: list, out: torch.Tensor) -> torch.Tensor:
        return x + out


# ── Residual: variante AttnRes ─────────────────────────────────────────────────

class AttnResidual(nn.Module):
    """
    Substitui a soma residual fixa por uma agregação aprendida sobre
    hidden states anteriores no bloco.

    Design "zero-init gate" (inspirado em ReZero/Zero-Init):
      - alpha inicializado em 0 → começa idêntico ao residual clássico
      - modelo abre o canal AttnRes gradualmente conforme aprende
      - gradientes fluem normalmente desde o primeiro step

    Fórmula:
      h = x + alpha * (attn_aggregate(states) - x) + out
        = x + out          quando alpha=0  (residual clássico)
        = attn_agg + out   quando alpha=1  (AttnRes completo)
    """
    def __init__(self, d_model: int):
        super().__init__()
        d_proj = max(d_model // 8, 16)
        self.q_proj = nn.Linear(d_model, d_proj, bias=False)
        self.k_proj = nn.Linear(d_model, d_proj, bias=False)
        self.scale  = d_proj ** -0.5
        # Gate: sigmoid(-4) ≈ 0.018 → começa quase como residual clássico
        # O modelo abre o canal AttnRes gradualmente conforme aprende
        self.alpha  = nn.Parameter(torch.full((1,), -4.0))

        nn.init.normal_(self.q_proj.weight, std=0.02)
        nn.init.normal_(self.k_proj.weight, std=0.02)

    def forward(self, x: torch.Tensor, states: list[torch.Tensor], out: torch.Tensor) -> torch.Tensor:
        if not states:
            return x + out

        # Empilha estados anteriores: (B, T, N, D)
        stacked = torch.stack(states, dim=2)
        B, T, N, D = stacked.shape

        q = self.q_proj(x)                                       # (B, T, d_proj)
        k = self.k_proj(stacked.reshape(B * T, N, D))            # (B*T, N, d_proj)
        k = k.reshape(B, T, N, -1)

        scores  = torch.einsum("btd,btnd->btn", q, k) * self.scale  # (B, T, N)
        weights = torch.softmax(scores, dim=-1)                       # (B, T, N)

        attn_res = torch.einsum("btn,btnd->btd", weights, stacked)    # (B, T, D)

        # Interpola entre residual clássico (x) e residual por atenção (attn_res)
        # alpha=0 → x + out  (clássico);  alpha → 1 → attn_res + out  (AttnRes)
        alpha    = torch.sigmoid(self.alpha)
        residual = x + alpha * (attn_res - x)

        return residual + out


# ── Bloco Transformer ─────────────────────────────────────────────────────────

class TransformerBlock(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.ln1  = nn.LayerNorm(cfg.d_model)
        self.attn = CausalSelfAttention(cfg)
        self.ln2  = nn.LayerNorm(cfg.d_model)
        self.ff   = FeedForward(cfg)

        if cfg.variant == "attnres":
            self.res_a = AttnResidual(cfg.d_model)
            self.res_m = AttnResidual(cfg.d_model)
        else:
            self.res_a = FixedResidual()
            self.res_m = FixedResidual()

        # Preenchido pelo GPTModel durante o forward
        self.block_states: list[torch.Tensor] = []

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.res_a(x, self.block_states, self.attn(self.ln1(x)))
        x = self.res_m(x, self.block_states, self.ff(self.ln2(x)))
        return x


# ── Modelo GPT ────────────────────────────────────────────────────────────────

class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg

        self.embed   = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.max_seq,    cfg.d_model)
        self.drop    = nn.Dropout(cfg.dropout)
        self.blocks  = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
        self.ln_f    = nn.LayerNorm(cfg.d_model)
        self.head    = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)

        # Weight tying: embedding e head compartilham pesos
        self.head.weight = self.embed.weight

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.02)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, std=0.02)

    def _reset_block_states(self):
        """Limpa buffers de estado antes de cada forward."""
        for i, block in enumerate(self.blocks):
            block.block_states = []

    def _update_block_states(self, layer_idx: int, state: torch.Tensor):
        """Adiciona o estado atual ao buffer do bloco correto.

        Sem detach: mantém o grafo de computação para que q_proj e k_proj
        recebam gradientes e possam aprender a selecionar estados relevantes.
        """
        if self.cfg.variant != "attnres":
            return
        block_idx = layer_idx // self.cfg.block_size
        next_idx  = layer_idx + 1
        if (next_idx < self.cfg.n_layers and
                next_idx // self.cfg.block_size == block_idx):
            self.blocks[next_idx].block_states = (
                self.blocks[layer_idx].block_states + [state]
            )

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
        B, T = idx.shape
        assert T <= self.cfg.max_seq, f"Sequência {T} > max_seq {self.cfg.max_seq}"

        pos = torch.arange(T, device=idx.device)
        x   = self.drop(self.embed(idx) + self.pos_emb(pos))

        self._reset_block_states()

        for i, block in enumerate(self.blocks):
            x = block(x)
            self._update_block_states(i, x)

        x    = self.ln_f(x)
        logits = self.head(x)  # (B, T, V)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, self.cfg.vocab_size),
                targets.reshape(-1),
                ignore_index=-1,
            )

        return logits, loss

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new: int, temperature: float = 0.8,
                 top_k: int = 40) -> torch.Tensor:
        for _ in range(max_new):
            idx_cond = idx[:, -self.cfg.max_seq:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs  = F.softmax(logits, dim=-1)
            next_t = torch.multinomial(probs, num_samples=1)
            idx    = torch.cat([idx, next_t], dim=1)
        return idx

    def n_params(self) -> int:
        return sum(p.numel() for p in self.parameters())


# ── Utilitários ────────────────────────────────────────────────────────────────

def make_model(variant: str = "baseline", **overrides) -> GPT:
    cfg = GPTConfig(variant=variant, **overrides)
    return GPT(cfg)


if __name__ == "__main__":
    for v in ("baseline", "attnres"):
        m = make_model(v)
        x = torch.randint(0, 16384, (2, 64))
        logits, loss = m(x[:, :-1], x[:, 1:])
        print(f"{v:10s}  params: {m.n_params():,}  logits: {logits.shape}  loss: {loss.item():.4f}")