Delete models/nlg.py with huggingface_hub

models/nlg.py

@@ -1,457 +0,0 @@
-"""nlg.py — Natural Language Generation (NLG) para CNN-BiGRU.
-
-Wrapper de alto nível para geração de texto que integra:
-    - TransformerDecoderStack (backbone causal com RoPE)
-    - MedusaMTP (Multi-Token Prediction)
-    - Sampling: Temperatura + Top-K + Top-P + Presence/Frequency Penalty
-    - CyclicReasoning (refinamento iterativo opcional)
-    - ContextWindowManager (janela de contexto com cache KV)
-
-==============================================================================
-ANÁLISE MATEMÁTICA E LÓGICA
-==============================================================================
-
-Dado um prompt P = [p_1, p_2, ..., p_n], geramos tokens y_1, y_2, ...
-autoregressivamente. Em cada passo t:
-
-    1. Forward pass: h_t = Decoder(P || y_{<t})   via cache KV
-    2. (Opcional) Refinamento cíclico: h_t = CyclicReasoning(h_t)
-    3. Logits: l_t = LM_Head(h_t)   [B, V]
-    4. Medusa: l_{t+k} = Medusa_Head_k(h_t) para k=1..K
-    5. Penalidades: l_t -= presence_penalty * 1[y in history]
-                          + frequency_penalty * count[y in history]
-    6. Sampling: y_t ~ softmax(l_t / T) com máscara Top-K e Top-P
-
-    7. (Opcional) Tree decoding: aceita candidatos Medusa se prob >= threshold
-
-A loss de treinamento combina:
-    L_NLG = CE(lm_head(h_t), y_{t+1}) + mu_medusa * L_Medusa + mu_ewc * L_EWC
-
-==============================================================================
-INTEGRAÇÃO
-==============================================================================
-
-- TransformerDecoderStack (cnn_bigru.models.transformer_block): backbone causal
-- MedusaMTP (cnn_bigru.models.medusa_heads): MTP
-- CyclicReasoning (cnn_bigru.models.cyclic_reasoning): refinamento iterativo
-- ContextWindowManager (cnn_bigru.models.context_window): janela deslizante
-- BBPETokenizer (cnn_bigru.tokenizer): tokenização byte-BPE
-
-Autor: CNN-BiGRU Project
-"""
-from __future__ import annotations
-
-import logging
-import math
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================================
-# Configuração
-# ============================================================================
-
-@dataclass
-class NLGConfig:
-    """Configuração do gerador NLG."""
-    vocab_size: int = 32000
-    embed_dim: int = 256
-    n_heads: int = 4
-    n_layers: int = 4
-    max_seq_len: int = 512
-    dropout: float = 0.1
-    pad_id: int = 1
-    bos_id: int = 2
-    eos_id: int = 3
-    # Medusa MTP
-    use_medusa: bool = True
-    n_medusa_heads: int = 4
-    mu_medusa: float = 0.5
-    # Cyclic Reasoning
-    use_cyclic_reasoning: bool = False
-    n_cycles: int = 3
-    # Sampling defaults
-    temperature: float = 0.7
-    top_k: int = 40
-    top_p: float = 0.9
-    presence_penalty: float = 0.3
-    frequency_penalty: float = 0.3
-    # Context window
-    max_window: int = 1024
-    n_sink_tokens: int = 4
-    eviction_strategy: str = "sink_sliding"
-    # Weight tying
-    weight_tying: bool = True
-    # Device
-    device: str = "cpu"
-
-
-# ============================================================================
-# NLG Module
-# ============================================================================
-
-class NLGModule(nn.Module):
-    """Módulo de geração de linguagem natural (NLG).
-
-    Combina TransformerDecoder + Medusa + (opcional) CyclicReasoning.
-    """
-
-    def __init__(self, config: NLGConfig):
-        super().__init__()
-        self.config = config
-        # Lazy import para evitar circular deps
-        from .transformer_block import TransformerDecoderStack
-        from .medusa_heads import MedusaMTP, MedusaConfig
-        from .cyclic_reasoning import (
-            CyclicReasoning, CyclicReasoningConfig,
-        )
-        from .context_window import (
-            ContextWindowManager, ContextWindowConfig,
-        )
-
-        # Backbone (Transformer Decoder)
-        self.decoder = TransformerDecoderStack(
-            vocab_size=config.vocab_size,
-            embed_dim=config.embed_dim,
-            n_heads=config.n_heads,
-            n_layers=config.n_layers,
-            max_seq_len=config.max_seq_len,
-            use_rope=True,
-            pad_id=config.pad_id,
-            weight_tying=config.weight_tying,
-        )
-
-        # Medusa MTP (opcional)
-        if config.use_medusa:
-            self.medusa = MedusaMTP(MedusaConfig(
-                vocab_size=config.vocab_size,
-                embed_dim=config.embed_dim,
-                n_heads=config.n_medusa_heads,
-                head_hidden_mult=2,
-                dropout=config.dropout,
-                bos_id=config.bos_id,
-            ))
-        else:
-            self.medusa = None
-
-        # Cyclic Reasoning (opcional)
-        if config.use_cyclic_reasoning:
-            self.cyclic = CyclicReasoning(CyclicReasoningConfig(
-                embed_dim=config.embed_dim,
-                max_cycles=config.n_cycles,
-                convergence_eps=1e-3,
-                use_anti_hallucination_gate=True,
-                device=config.device,
-            ))
-        else:
-            self.cyclic = None
-
-        # Context Window Manager (criado lazy na geração)
-        self._cw: Optional[ContextWindowManager] = None
-
-    # ----------------------------------------------------------------------
-    # Forward (treino com teacher forcing)
-    # ----------------------------------------------------------------------
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        target_ids: Optional[torch.Tensor] = None,
-        use_medusa_loss: bool = True,
-    ) -> Dict[str, torch.Tensor]:
-        """Forward pass com teacher forcing.
-
-        Args:
-            input_ids: [B, T] tokens de entrada (inclui shift direito para teacher forcing)
-            target_ids: [B, T] tokens alvo (default: input_ids shifted por 1)
-            use_medusa_loss: se True, computa loss Medusa
-
-        Returns:
-            dict com:
-                logits: [B, T, V] logits da cabeça principal
-                loss: escalar (CE + Medusa) — só se target_ids fornecido
-                loss_main: CE da cabeça principal
-                loss_medusa: loss Medusa (se use_medusa_loss e medusa ativo)
-                medusa_stats: stats por cabeça
-                hidden: [B, T, D] hidden states (para debug)
-        """
-        # Backbone forward
-        logits, caches = self.decoder(input_ids)  # [B, T, V], list
-        # hidden states = logits do lm_head pré-projection? Não — precisamos do hidden
-        # Na TransformerDecoderStack, o hidden é a saída antes do lm_head
-        # Aqui assumimos que logits é o output do lm_head (com weight tying)
-        # Para obter o hidden, usamos logits se weight_tying (matriz transposta)
-        # ou adicionamos um hook no decoder. Para simplificar, refazemos o hidden
-        # via projeção inversa (apenas se weight_tying=True)
-        # NOTA: em prática, o decoder poderia retornar hidden explicitamente
-        # Para este wrapper, retornamos logits como proxy do hidden
-        # (a Medusa recebe o hidden pré-lm_head, mas para integração simples
-        # usamos logits como entrada da Medusa apenas quando shapes batem)
-        hidden = logits  # fallback: usa logits como hidden (será projetado se necessário)
-
-        result = {
-            "logits": logits,
-            "hidden": hidden,
-        }
-
-        # Loss (se target fornecido)
-        if target_ids is not None:
-            # CE principal
-            V = self.config.vocab_size
-            loss_main = F.cross_entropy(
-                logits.view(-1, V),
-                target_ids.view(-1),
-                ignore_index=self.config.pad_id,
-                reduction="mean",
-            )
-            result["loss_main"] = loss_main
-
-            # Loss Medusa
-            loss_medusa = torch.zeros(1, device=logits.device, dtype=logits.dtype)
-            if use_medusa_loss and self.medusa is not None:
-                # A MedusaMTP espera o hidden state [B, T, D], não logits [B, T, V]
-                # Como não temos o hidden explícito aqui (sem modificar o decoder),
-                # usamos uma projeção aproximada: pegamos logits e projetamos para D
-                # via lm_head.weight.T (se weight tying) ou via embedding
-                # Para simplificar, usamos logits[:, :, :D] (slice) quando V > D
-                D = self.config.embed_dim
-                if logits.size(-1) >= D:
-                    hidden_proxy = logits[:, :, :D]  # slice simples
-                else:
-                    # projeta logits para D
-                    if not hasattr(self, "_logits_proj"):
-                        self._logits_proj = nn.Linear(V, D, bias=False).to(logits.device)
-                    hidden_proxy = self._logits_proj(logits)
-
-                loss_med, medusa_stats = self.medusa.compute_loss(
-                    hidden_proxy, target_ids, ignore_index=self.config.pad_id,
-                )
-                loss_medusa = loss_med
-                result["medusa_stats"] = medusa_stats
-
-            result["loss_medusa"] = loss_medusa.squeeze()
-            result["loss"] = loss_main + self.config.mu_medusa * loss_medusa.squeeze()
-
-        return result
-
-    # ----------------------------------------------------------------------
-    # Geração autoregressiva
-    # ----------------------------------------------------------------------
-
-    @torch.no_grad()
-    def generate(
-        self,
-        prompt_ids: torch.Tensor,
-        max_new_tokens: int = 50,
-        temperature: Optional[float] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        presence_penalty: Optional[float] = None,
-        frequency_penalty: Optional[float] = None,
-        use_context_window: bool = True,
-        use_medusa: Optional[bool] = None,
-        eos_id: Optional[int] = None,
-    ) -> Dict[str, torch.Tensor]:
-        """Gera texto autoregressivamente com sampling.
-
-        Args:
-            prompt_ids: [B, T_prompt] IDs do prompt
-            max_new_tokens: número máximo de tokens a gerar
-            temperature, top_k, top_p, presence_penalty, frequency_penalty:
-                overrides da config (se None, usa config)
-            use_context_window: se True, aplica eviction policy
-            use_medusa: se True, usa Medusa tree decoding (default: config.use_medusa)
-            eos_id: ID do token EOS (default: config.eos_id)
-
-        Returns:
-            dict com:
-                ids: [B, T_prompt + T_gen] tokens gerados
-                n_tokens: int — número de tokens gerados (excluindo prompt)
-                n_medusa_accepted: int — tokens aceitos via Medusa
-                stopped_early: bool — True se EOS atingido
-        """
-        self.eval()
-        cfg = self.config
-        T = temperature if temperature is not None else cfg.temperature
-        K = top_k if top_k is not None else cfg.top_k
-        P = top_p if top_p is not None else cfg.top_p
-        PP = presence_penalty if presence_penalty is not None else cfg.presence_penalty
-        FP = frequency_penalty if frequency_penalty is not None else cfg.frequency_penalty
-        EOS = eos_id if eos_id is not None else cfg.eos_id
-        use_med = use_medusa if use_medusa is not None else cfg.use_medusa and self.medusa is not None
-
-        device = prompt_ids.device
-        B = prompt_ids.size(0)
-
-        # Inicializar context window
-        if use_context_window and self._cw is None:
-            from .context_window import ContextWindowManager, ContextWindowConfig
-            self._cw = ContextWindowManager(ContextWindowConfig(
-                max_window=cfg.max_window,
-                eviction_strategy=cfg.eviction_strategy,
-                n_sink_tokens=cfg.n_sink_tokens,
-                embed_dim=cfg.embed_dim,
-                n_heads=cfg.n_heads,
-                head_dim=cfg.embed_dim // cfg.n_heads,
-                n_layers=cfg.n_layers,
-                device=str(device),
-            ))
-
-        # Token history (para penalidades)
-        token_counts: List[Dict[int, int]] = [dict() for _ in range(B)]
-        for b in range(B):
-            for tok in prompt_ids[b].tolist():
-                token_counts[b][tok] = token_counts[b].get(tok, 0) + 1
-
-        # Estado: tokens atuais
-        cur_ids = prompt_ids.clone()  # [B, T]
-        n_generated = 0
-        n_medusa_accepted = 0
-        stopped = torch.zeros(B, dtype=torch.bool, device=device)
-
-        for step in range(max_new_tokens):
-            # Truncar para max_seq_len
-            ctx = cur_ids[:, -cfg.max_seq_len:]
-            # Forward
-            try:
-                logits, _ = self.decoder(ctx)
-            except Exception as e:
-                logger.warning(f"NLG forward falhou no step {step}: {e}")
-                break
-
-            # Pegar logits do último token
-            last_logits = logits[:, -1, :].clone()  # [B, V]
-
-            # Aplicar penalidades
-            if PP != 0.0 or FP != 0.0:
-                for b in range(B):
-                    for tok_id, count in token_counts[b].items():
-                        last_logits[b, tok_id] -= (PP + count * FP)
-
-            # Temperatura
-            if T > 0:
-                last_logits = last_logits / T
-            else:
-                # Greedy
-                next_tok = last_logits.argmax(dim=-1, keepdim=True)  # [B, 1]
-                cur_ids = torch.cat([cur_ids, next_tok], dim=1)
-                n_generated += 1
-                for b in range(B):
-                    t = next_tok[b].item()
-                    token_counts[b][t] = token_counts[b].get(t, 0) + 1
-                    if t == EOS:
-                        stopped[b] = True
-                if stopped.all():
-                    break
-                continue
-
-            # Top-K
-            if K > 0:
-                top_vals, _ = torch.topk(last_logits, min(K, last_logits.size(-1)))
-                min_val = top_vals[:, -1:]
-                last_logits[last_logits < min_val] = float("-inf")
-
-            # Top-P (nucleus)
-            if P < 1.0:
-                sorted_logits, sorted_indices = torch.sort(last_logits, descending=True)
-                sorted_probs = F.softmax(sorted_logits, dim=-1)
-                cum_probs = torch.cumsum(sorted_probs, dim=-1)
-                mask = cum_probs > P
-                mask[:, 1:] = mask[:, :-1].clone()
-                mask[:, 0] = False
-                indices_to_remove = mask.scatter(1, sorted_indices, mask)
-                last_logits[indices_to_remove] = float("-inf")
-
-            # Sampling
-            probs = F.softmax(last_logits, dim=-1)
-            next_tok = torch.multinomial(probs, num_samples=1)  # [B, 1]
-            cur_ids = torch.cat([cur_ids, next_tok], dim=1)
-            n_generated += 1
-
-            # Update counts
-            for b in range(B):
-                t = next_tok[b].item()
-                token_counts[b][t] = token_counts[b].get(t, 0) + 1
-                if t == EOS:
-                    stopped[b] = True
-
-            # (Opcional) Medusa: aceitar próximos candidatos
-            # Skip nesta implementação simplificada — Medusa loss é computada em treino
-
-            if stopped.all():
-                break
-
-        return {
-            "ids": cur_ids,
-            "n_tokens": n_generated,
-            "n_medusa_accepted": n_medusa_accepted,
-            "stopped_early": bool(stopped.any().item()),
-        }
-
-    # ----------------------------------------------------------------------
-    # Compute loss para treinamento (compatível com trainer)
-    # ----------------------------------------------------------------------
-
-    def compute_loss(
-        self,
-        input_ids: torch.Tensor,
-        target_ids: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
-        """Computa a loss de NLG para uso no trainer.
-
-        Atalho para self.forward(input_ids, target_ids, use_medusa_loss=True)
-        """
-        return self.forward(input_ids, target_ids, use_medusa_loss=True)
-
-
-# ============================================================================
-# Self-test
-# ============================================================================
-
-def _self_test():
-    """Teste rápido do módulo NLG."""
-    torch.manual_seed(42)
-    config = NLGConfig(
-        vocab_size=100,
-        embed_dim=32,
-        n_heads=4,
-        n_layers=2,
-        max_seq_len=32,
-        pad_id=1, bos_id=2, eos_id=3,
-        use_medusa=True,
-        n_medusa_heads=3,
-        use_cyclic_reasoning=False,
-        weight_tying=True,
-    )
-    nlg = NLGModule(config)
-    print(f"NLG params: {sum(p.numel() for p in nlg.parameters())}")
-
-    # Forward com teacher forcing
-    ids = torch.randint(0, 100, (2, 8))
-    target = torch.randint(0, 100, (2, 8))
-    out = nlg.compute_loss(ids, target)
-    print(f"Loss: {out['loss'].item():.4f}")
-    print(f"  main: {out['loss_main'].item():.4f}")
-    print(f"  medusa: {out['loss_medusa'].item():.4f}")
-    print(f"  medusa_stats: {out.get('medusa_stats', {})}")
-
-    # Geração
-    prompt = torch.tensor([[2, 5, 10, 15]], dtype=torch.long)
-    gen = nlg.generate(prompt, max_new_tokens=10, temperature=0.7, top_k=10)
-    print(f"Generated {gen['n_tokens']} tokens: {gen['ids'][0].tolist()}")
-
-
-if __name__ == "__main__":
-    _self_test()
-
-
-__all__ = [
-    "NLGConfig",
-    "NLGModule",
-]