Delete models/transformer_block.py with huggingface_hub

models/transformer_block.py

@@ -1,401 +0,0 @@
-"""
-Transformer Decoder Block — CausalSelfAttention + TransformerBlock.
-
-Implementa o bloco decodificador causal conforme dados.txt seções 6-7:
-  - Multi-head self-attention com máscara triangular inferior
-  - Scaled dot-product: (q @ k.T) / sqrt(d_head)
-  - QKV projection unificada
-  - Pre-LN: x = x + attn(ln1(x)); x = x + ffn(ln2(x))
-  - FFN: Linear(d, 4d) → GELU → Linear(4d, d) → Dropout
-  - Weight tying opcional entre embedding e LM head
-
-Inclui:
-  - RoPE (Rotary Position Embeddings) opcional
-  - Cache KV opcional para geração autoregressiva eficiente
-  - Padding mask (ignora tokens [PAD])
-  - Causal mask (não olha para o futuro)
-
-Autor: CNN-BiGRU Project
-"""
-from __future__ import annotations
-
-import logging
-import math
-from dataclasses import dataclass, field
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .rope import RotaryPositionEmbedding
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================================
-# Configuração
-# ============================================================================
-
-@dataclass
-class TransformerBlockConfig:
-    """Configuração de um bloco Transformer Decoder."""
-    embed_dim: int = 256
-    n_heads: int = 4
-    ff_dim: Optional[int] = None  # default: 4 * embed_dim
-    dropout: float = 0.1
-    use_rope: bool = True
-    max_seq_len: int = 512
-    rope_base: float = 10000.0
-    # Limiar para attention dropout
-    attn_dropout: float = 0.1
-    # Bias nas projeções lineares
-    qkv_bias: bool = False
-    out_bias: bool = False
-    ffn_bias: bool = False
-
-
-# ============================================================================
-# Causal Self-Attention
-# ============================================================================
-
-class CausalSelfAttention(nn.Module):
-    """
-    Mecanismo de Atenção Multi-Head com Máscara Causal (Triangular Inferior).
-
-    Conforme dados.txt linhas 553-592.
-
-    Fórmula matemática:
-        Attention(Q, K, V) = softmax( (Q @ K^T) / sqrt(d_k) + mask ) @ V
-
-    Com RoPE opcional, Q e K são rotacionados antes do produto escalar.
-    """
-
-    def __init__(self, config: TransformerBlockConfig):
-        super().__init__()
-        self.config = config
-        self.embed_dim = config.embed_dim
-        self.n_heads = config.n_heads
-        if self.embed_dim % self.n_heads != 0:
-            raise ValueError(
-                f"embed_dim ({self.embed_dim}) deve ser divisível por "
-                f"n_heads ({self.n_heads})"
-            )
-        self.head_dim = self.embed_dim // self.n_heads
-
-        # Projeção QKV unificada (3 * embed_dim) — ganho de performance
-        self.qkv_projection = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
-        self.out_projection = nn.Linear(self.embed_dim, self.embed_dim, bias=config.out_bias)
-
-        self.attn_dropout = nn.Dropout(config.attn_dropout)
-        self.resid_dropout = nn.Dropout(config.dropout)
-
-        # RoPE (opcional)
-        self.use_rope = config.use_rope
-        if self.use_rope:
-            self.rope = RotaryPositionEmbedding(
-                head_dim=self.head_dim,
-                max_seq_len=config.max_seq_len,
-                base=config.rope_base,
-            )
-        else:
-            self.rope = None
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        kv_cache_k: Optional[torch.Tensor] = None,
-        kv_cache_v: Optional[torch.Tensor] = None,
-        position_offset: int = 0,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            x: [batch, seq_len, embed_dim]
-            padding_mask: [batch, seq_len] — 1 para token real, 0 para PAD
-                          (se None, não aplica)
-            kv_cache_k, kv_cache_v: cache KV de passos anteriores
-                [batch, n_heads, cached_len, head_dim] ou None
-            position_offset: offset posicional (para uso com cache KV)
-
-        Returns:
-            out: [batch, seq_len, embed_dim]
-            new_k: [batch, n_heads, total_len, head_dim] — K atualizado para cache
-            new_v: [batch, n_heads, total_len, head_dim] — V atualizado
-        """
-        B, T, C = x.size()
-
-        # Projeção QKV unificada
-        qkv = self.qkv_projection(x)  # [B, T, 3*C]
-        q, k, v = qkv.split(self.embed_dim, dim=2)  # cada um: [B, T, C]
-
-        # Reshape para multi-head: [B, n_heads, T, head_dim]
-        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-
-        # Aplicar RoPE a Q e K (não a V)
-        if self.use_rope and self.rope is not None:
-            q = self.rope(q, offset=position_offset)
-            k = self.rope(k, offset=position_offset)
-
-        # Concatenar com cache KV (se fornecido)
-        if kv_cache_k is not None and kv_cache_v is not None:
-            k = torch.cat([kv_cache_k, k], dim=2)  # [B, n_heads, cached+T, head_dim]
-            v = torch.cat([kv_cache_v, v], dim=2)
-        new_k, new_v = k, v
-
-        total_len = k.size(2)
-
-        # Atenção por produto escalar escalonado
-        # scores: [B, n_heads, T, total_len]
-        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
-
-        # Máscara causal: o token na posição i só pode olhar para posições <= i
-        # Em modo cache: pos i corresponde à posição absoluta (offset + i)
-        # O token atual (offset+i) não pode olhar para futuros (offset+i+1, ...)
-        # Como K contém [cached..., current_T], a máscara deve permitir que
-        # cada q[j] (j=0..T-1) veja k[0..cached+j]
-        if T > 0:
-            # Índices absolutos de Q: [offset, offset+T-1]
-            # Índices absolutos de K: [0, total_len-1]
-            q_pos = torch.arange(position_offset, position_offset + T, device=x.device)  # [T]
-            k_pos = torch.arange(0, total_len, device=x.device)  # [total_len]
-            # permitido se k_pos <= q_pos
-            causal_mask = k_pos.unsqueeze(0) <= q_pos.unsqueeze(1)  # [T, total_len]
-            # Aplicar: onde causal_mask==False → -inf
-            attn_scores = attn_scores.masked_fill(
-                ~causal_mask.unsqueeze(0).unsqueeze(0),  # [1, 1, T, total_len]
-                float('-inf')
-            )
-
-        # Máscara de padding (se fornecida)
-        if padding_mask is not None:
-            # padding_mask: [B, T_query] — mas a máscara deve afetar a chave (K)
-            # Se um token K é PAD, sua atenção deve ser zero
-            # Como o cache K pode conter tokens do passado que NÃO são PAD,
-            # a máscara de padding deve cobrir todo o K (cached + atual)
-            # Para simplicidade, assumimos que tokens cached não são PAD
-            # e usamos padding_mask apenas para o K atual
-            if padding_mask.dim() == 2:
-                # [B, T] → [B, 1, 1, T]
-                pad_mask_k = (padding_mask > 0).unsqueeze(1).unsqueeze(1)
-            else:
-                pad_mask_k = (padding_mask > 0).unsqueeze(1).unsqueeze(1)
-            # Construir máscara completa sobre K
-            if kv_cache_k is not None:
-                cached_len = kv_cache_k.size(2)
-                # Tokens cached: considerados válidos (não PAD)
-                cached_part = torch.ones(
-                    B, 1, 1, cached_len, device=x.device, dtype=pad_mask_k.dtype
-                )
-                pad_mask_full = torch.cat([cached_part, pad_mask_k], dim=3)
-            else:
-                pad_mask_full = pad_mask_k
-            attn_scores = attn_scores.masked_fill(~pad_mask_full, float('-inf'))
-
-        # Softmax + dropout
-        # Lidar com linhas totalmente -inf (todas as chaves são PAD)
-        # Substituir -inf por 0 para evitar NaN
-        all_neg = (attn_scores <= -1e9).all(dim=-1, keepdim=True)
-        attn_scores = torch.where(all_neg, torch.zeros_like(attn_scores), attn_scores)
-
-        attn_weights = F.softmax(attn_scores, dim=-1)
-        attn_weights = self.attn_dropout(attn_weights)
-
-        # Fusão dos contextos
-        out = torch.matmul(attn_weights, v)  # [B, n_heads, T, head_dim]
-        out = out.transpose(1, 2).contiguous().view(B, T, C)
-
-        out = self.resid_dropout(self.out_projection(out))
-        return out, new_k, new_v
-
-
-# ============================================================================
-# Transformer Block
-# ============================================================================
-
-class TransformerBlock(nn.Module):
-    """
-    Bloco decodificador padrão com Pre-LN (conforme dados.txt linhas 594-612).
-
-    Estrutura:
-        x = x + attn(ln1(x))
-        x = x + ffn(ln2(x))
-    FFN: Linear(d, 4d) → GELU → Linear(4d, d) → Dropout
-    """
-
-    def __init__(self, config: TransformerBlockConfig):
-        super().__init__()
-        self.config = config
-        self.ln1 = nn.LayerNorm(config.embed_dim)
-        self.attn = CausalSelfAttention(config)
-        self.ln2 = nn.LayerNorm(config.embed_dim)
-        ff_dim = config.ff_dim or (4 * config.embed_dim)
-        self.ffn = nn.Sequential(
-            nn.Linear(config.embed_dim, ff_dim, bias=config.ffn_bias),
-            nn.GELU(),
-            nn.Linear(ff_dim, config.embed_dim, bias=config.ffn_bias),
-            nn.Dropout(config.dropout),
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        kv_cache_k: Optional[torch.Tensor] = None,
-        kv_cache_v: Optional[torch.Tensor] = None,
-        position_offset: int = 0,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            x: [batch, seq_len, embed_dim]
-            padding_mask: [batch, seq_len] opcional
-            kv_cache_k, kv_cache_v: cache KV opcional
-            position_offset: offset posicional para RoPE
-
-        Returns:
-            out: [batch, seq_len, embed_dim]
-            new_k: cache K atualizado
-            new_v: cache V atualizado
-        """
-        # Pre-LN: normaliza antes de passar para a sub-camada
-        normed = self.ln1(x)
-        attn_out, new_k, new_v = self.attn(
-            normed, padding_mask=padding_mask,
-            kv_cache_k=kv_cache_k, kv_cache_v=kv_cache_v,
-            position_offset=position_offset,
-        )
-        x = x + attn_out
-
-        # FFN
-        x = x + self.ffn(self.ln2(x))
-        return x, new_k, new_v
-
-
-# ============================================================================
-# Stack de Transformer Blocks (Decoder completo)
-# ============================================================================
-
-class TransformerDecoderStack(nn.Module):
-    """
-    Empilhamento de N TransformerBlocks (decoder-only, estilo GPT).
-
-    Conforme dados.txt linhas 614-648.
-
-    Suporta weight tying (compartilhamento de pesos entre token_embedding e lm_head).
-    """
-
-    def __init__(
-        self,
-        vocab_size: int,
-        embed_dim: int = 256,
-        n_heads: int = 4,
-        n_layers: int = 2,
-        ff_dim: Optional[int] = None,
-        dropout: float = 0.1,
-        max_seq_len: int = 512,
-        use_rope: bool = True,
-        pad_id: int = 0,
-        weight_tying: bool = True,
-    ):
-        super().__init__()
-        self.vocab_size = vocab_size
-        self.embed_dim = embed_dim
-        self.max_seq_len = max_seq_len
-        self.pad_id = pad_id
-        self.weight_tying = weight_tying
-
-        block_config = TransformerBlockConfig(
-            embed_dim=embed_dim,
-            n_heads=n_heads,
-            ff_dim=ff_dim,
-            dropout=dropout,
-            use_rope=use_rope,
-            max_seq_len=max_seq_len,
-        )
-
-        self.token_embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_id)
-        # Sem positional embedding separado quando RoPE está ativo
-        if not use_rope:
-            self.position_embedding = nn.Embedding(max_seq_len, embed_dim)
-        else:
-            self.position_embedding = None
-
-        self.blocks = nn.ModuleList([
-            TransformerBlock(block_config) for _ in range(n_layers)
-        ])
-        self.ln_f = nn.LayerNorm(embed_dim)
-        self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False)
-
-        # Weight tying
-        if weight_tying:
-            self.lm_head.weight = self.token_embedding.weight
-
-    def forward(
-        self,
-        idx: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        kv_caches: Optional[list] = None,
-        position_offset: int = 0,
-        return_hidden: bool = False,
-    ) -> torch.Tensor:
-        """
-        Args:
-            idx: [batch, seq_len] — token IDs
-            padding_mask: [batch, seq_len] — 1 para real, 0 para PAD
-            kv_caches: lista de (k, v) por camada, ou None
-            position_offset: offset posicional (para uso com cache)
-            return_hidden: se True, retorna também o hidden state antes do lm_head
-
-        Returns:
-            logits: [batch, seq_len, vocab_size]
-            hidden (opcional): [batch, seq_len, embed_dim]
-        """
-        B, T = idx.size()
-        if T > self.max_seq_len:
-            raise ValueError(
-                f"Sequência de tamanho {T} excede o limite máximo configurado de "
-                f"{self.max_seq_len}."
-            )
-
-        # Embedding
-        x = self.token_embedding(idx)
-        if not self.weight_tying and self.position_embedding is not None:
-            positions = torch.arange(
-                position_offset, position_offset + T,
-                dtype=torch.long, device=idx.device
-            ).unsqueeze(0)
-            x = x + self.position_embedding(positions)
-
-        # Passar pelos blocos
-        new_caches = []
-        for i, block in enumerate(self.blocks):
-            cache_k = None
-            cache_v = None
-            if kv_caches is not None and i < len(kv_caches):
-                cache_k, cache_v = kv_caches[i]
-            x, new_k, new_v = block(
-                x, padding_mask=padding_mask,
-                kv_cache_k=cache_k, kv_cache_v=cache_v,
-                position_offset=position_offset,
-            )
-            new_caches.append((new_k, new_v))
-
-        # LayerNorm final
-        x = self.ln_f(x)
-
-        # LM head
-        logits = self.lm_head(x)
-
-        if return_hidden:
-            return logits, x, new_caches
-        return logits, new_caches
-
-
-__all__ = [
-    "TransformerBlockConfig",
-    "CausalSelfAttention",
-    "TransformerBlock",
-    "TransformerDecoderStack",
-]