Delete src/bigru_t/model/unified_model.py

src/bigru_t/model/unified_model.py

@@ -1,523 +0,0 @@
-"""unified_model.py — UnifiedModel: orquestra todos os componentes (Lemas 1-4).
-
-Arquitetura BiGRU_T_version:
-
-    x (batch, T, input_dim)
-        │
-        ├──→ u8cell_T_1  ─→ h_1 ─┐
-        ├──→ u8cell_T_2  ─→ h_2 ─┤       Lema 1: alpha = softmax(logits / T)
-        ├──→ ...                  ├─→ H = stack([alpha_k * h_k])  → OrqCell
-        └──→ u8cell_T_K  ─→ h_K ─┘                │
-                                                   ↓
-                                                  o (batch, d_cache)
-                                                   │
-                                    ┌──────────────┴──────────────┐
-                                    ↓                              ↓
-                                TrainT(o)                     HypT(o, stop_grad=True)
-                                    │                              │
-                                    ↓                              ↓
-                                y_hat                           delta
-                                    │                              │
-                                    └──────→ y_final = y_hat + delta * use_hyp  ←── Lema 3
-
-Lema 2: gradient_surgery(g_main, g_hyp) → projeção ortogonal
-Lema 4: MetaConfigurator ajusta T (temperatura) e τ (limiar de ativação da hipótese)
-
-Diferenças vs. v13.9.2 GRURingV139:
-- 1 modelo agregador (OrqCell) em vez de 4 camadas cooperativas empilhadas
-- Seleção de módulos por softmax (Lema 1) em vez de Confidence Gate
-- Hipótese global (HypT) em vez de 16 hipóteses paralelas com LoRA
-- Quantização W8A8 fake-quant em todas as Linear (Lema 3) em vez de
-  W8A8 QOperator só em inferência
-- Tau aprendível (Lema 4) em vez de trust_lr fixo
-
-Adaptação para Language Modeling:
-- input_dim = d_model (embedding dim)
-- output_dim = vocab_size
-- Adiciona camada de embedding (token IDs → embeddings)
-- Adiciona positional encoding (sinusoidal)
-"""
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass, field
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .u8cell_t import u8cell_T
-from .orq_cell import OrqCell
-from .train_t import TrainT
-from .hyp_t import HypT
-from .module_selector import ModuleSelector
-# V5: aplicar_w8a8_to_model substitui apply_w8a8 (bug fix crítico)
-from ..quantization.quantized_linear import QuantizedLinear, apply_w8a8_to_model
-from ..reasoning.circular_reasoning_wasserstein import CircularReasoningWasserstein
-
-
-@dataclass
-class UnifiedModelConfig:
-    """Configuração do UnifiedModel.
-
-    Campos agrupados por componente:
-        - Tokenizer / embedding: vocab_size, d_model, max_seq_len, pad_token_id
-        - u8cell_T: input_dim (default = d_model), bigru_hidden, d_transformer, nhead_tu, d_ff_tu, output_dim_u8cell
-        - Module selector: max_modules, lambda_ent
-        - OrqCell: cache_len, d_cache, nhead_orq, d_ff_orq
-        - TrainT: trainT_dim, nhead_train, d_ff_train, num_layers_train
-        - HypT: hypT_dim, nhead_hyp, d_ff_hyp, num_layers_hyp
-        - Quantization: num_bits (W8A8)
-        - Misc: dropout
-
-    Nota: `input_dim` é automaticamente sincronizado com `d_model` em
-    __post_init__ se não for explicitamente especificado. Isso garante
-    que a saída do embedding (dim=d_model) matches a entrada do BiGRU4
-    (input_dim).
-    """
-    # Embedding
-    vocab_size: int = 16384
-    d_model: int = 128
-    max_seq_len: int = 64
-    pad_token_id: int = 1
-
-    # u8cell_T
-    input_dim: int = 0  # 0 = usar d_model automaticamente (ver __post_init__)
-    bigru_hidden: int = 32
-    d_transformer: int = 64
-    nhead_tu: int = 4
-    d_ff_tu: int = 128
-    output_dim_u8cell: int = 64
-
-    # Module selector (Lema 1)
-    max_modules: int = 64
-    lambda_ent: float = 0.01
-
-    # OrqCell
-    cache_len: int = 16
-    d_cache: int = 128
-    nhead_orq: int = 8
-    d_ff_orq: int = 256
-
-    # TrainT
-    trainT_dim: int = 128
-    nhead_train: int = 4
-    d_ff_train: int = 256
-    num_layers_train: int = 2
-
-    # HypT
-    # V4: num_layers_hyp aumentado de 2 → 8 (mais capacidade para correção delta)
-    # V5: num_layers_hyp=8 é o PADRÃO PERMANENTE (não mudar)
-    hypT_dim: int = 128
-    nhead_hyp: int = 4
-    d_ff_hyp: int = 256
-    num_layers_hyp: int = 8
-
-    # Quantization (Lema 3)
-    num_bits: int = 8
-
-    # Circular Reasoning (Teorema 4 — contração W₂)
-    use_circular_reasoning: bool = True
-    circular_n_cycles: int = 3
-    circular_contraction_target: float = 0.9
-    circular_init_gate: float = 0.05
-
-    # Misc
-    dropout: float = 0.1
-
-    def __post_init__(self):
-        """Sincroniza input_dim com d_model se não for explicitamente especificado."""
-        if self.input_dim <= 0:
-            self.input_dim = self.d_model
-        # Validações de consistência
-        if self.d_transformer % self.nhead_tu != 0:
-            raise ValueError(
-                f"d_transformer ({self.d_transformer}) must be divisible by nhead_tu ({self.nhead_tu})"
-            )
-        if self.d_cache % self.nhead_orq != 0:
-            raise ValueError(
-                f"d_cache ({self.d_cache}) must be divisible by nhead_orq ({self.nhead_orq})"
-            )
-        if self.trainT_dim % self.nhead_train != 0:
-            raise ValueError(
-                f"trainT_dim ({self.trainT_dim}) must be divisible by nhead_train ({self.nhead_train})"
-            )
-        if self.hypT_dim % self.nhead_hyp != 0:
-            raise ValueError(
-                f"hypT_dim ({self.hypT_dim}) must be divisible by nhead_hyp ({self.nhead_hyp})"
-            )
-
-
-class PositionalEncoding(nn.Module):
-    """Positional encoding sinusoidal (Vaswani et al., 2017).
-
-    Reaproveita a filosofia do v13.9.2 (que usa RoPE), mas como o BiGRU_T
-    processa a sequência via BiGRU (não via attention causal), o PE sinusoidal
-    é suficiente e mais simples.
-    """
-
-    def __init__(self, d_model: int, max_len: int = 512):
-        super().__init__()
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """x: (batch, T, d_model) → x + pe[:, :T]"""
-        return x + self.pe[:, : x.size(1)]
-
-
-class UnifiedModel(nn.Module):
-    """Modelo BiGRU_T_version — orquestra Lemas 1-4.
-
-    Args:
-        config: UnifiedModelConfig
-
-    Forward:
-        x: (batch, T) token IDs OU (batch, T, input_dim) embeddings
-           (detetado por ndim: 2 = token IDs, 3 = embeddings)
-        temperature: float (default 1.0, sobrescrito pelo MetaConfigurator)
-        use_hypothesis: bool (default False, sobrescrito pelo trainer se loss > tau)
-        stop_grad_hyp: bool (default True, Lema 3)
-
-    Returns:
-        (y_hat, delta) — ambos (batch, output_dim). y_final = y_hat + delta * use_hypothesis.
-    """
-
-    def __init__(self, config: Optional[UnifiedModelConfig] = None):
-        super().__init__()
-        self.config = config or UnifiedModelConfig()
-        c = self.config
-
-        # Embedding de tokens (para LM)
-        self.token_embedding = nn.Embedding(c.vocab_size, c.d_model, padding_idx=c.pad_token_id)
-        self.positional_encoding = PositionalEncoding(c.d_model, c.max_seq_len * 2)
-        self.emb_dropout = nn.Dropout(c.dropout)
-
-        # Module selector (Lema 1)
-        self.module_selector = ModuleSelector(c.max_modules, lambda_ent=c.lambda_ent)
-
-        # Lista de módulos u8cell_T
-        self.u8cells = nn.ModuleList([
-            u8cell_T(
-                input_dim=c.input_dim,
-                hidden_dim_bigru=c.bigru_hidden,
-                d_transformer=c.d_transformer,
-                nhead=c.nhead_tu,
-                d_ff=c.d_ff_tu,
-                output_dim=c.output_dim_u8cell,
-                dropout=c.dropout,
-            )
-            for _ in range(c.max_modules)
-        ])
-
-        # Orquestrador
-        self.orq_cell = OrqCell(
-            cache_len=c.cache_len,
-            d_cache=c.d_cache,
-            d_model_in=c.output_dim_u8cell,
-            nhead=c.nhead_orq,
-            d_ff=c.d_ff_orq,
-            dropout=c.dropout,
-        )
-
-        # Cabeça principal
-        self.train_T = TrainT(
-            d_input=c.d_cache,
-            d_model=c.trainT_dim,
-            nhead=c.nhead_train,
-            d_ff=c.d_ff_train,
-            output_dim=c.vocab_size,  # LM: prevê próximo token
-            num_layers=c.num_layers_train,
-            dropout=c.dropout,
-        )
-
-        # Cabeça de hipótese (Lema 3)
-        self.hyp_T = HypT(
-            d_input=c.d_cache,
-            d_model=c.hypT_dim,
-            nhead=c.nhead_hyp,
-            d_ff=c.d_ff_hyp,
-            output_dim=c.vocab_size,
-            num_layers=c.num_layers_hyp,
-            dropout=c.dropout,
-        )
-
-        # Limiar de punição (τ) — atualizado pelo MetaConfigurator (Lema 4)
-        # Init = 1.0 (sem hipótese até a loss cair abaixo de 1.0 — raro no início)
-        self.register_buffer("tau", torch.tensor(float("inf")))
-
-        # Circular Reasoning (Teorema 4 — contração W₂ sobre a representação agregada)
-        # Aplicado APÓS OrqCell e ANTES de TrainT/HypT. Refina 'o' em n_cycles
-        # iterações de Φ(s) = LayerNorm(s + gate * RefineHead(s)).
-        self.circular_reasoning = (
-            CircularReasoningWasserstein(
-                d_model=c.d_cache,
-                n_cycles=c.circular_n_cycles,
-                contraction_target=c.circular_contraction_target,
-                init_gate=c.circular_init_gate,
-            )
-            if c.use_circular_reasoning
-            else None
-        )
-
-        # V5: Parâmetros APRENDÍVEIS do modelo (auto-ajuste próprio)
-        # λ_ent: peso da regularização de entropia do module_selector (Lema 1)
-        # Parametrização via exp (sempre positivo)
-        # init: λ_ent = c.lambda_ent = 0.01
-        # O residual_gate do CircularReasoning JÁ é Parameter (n_cycles params)
-        # e é gerenciado diretamente pelo MetaConfigurator.
-        self.log_lambda_ent = nn.Parameter(
-            torch.tensor(math.log(max(1e-6, c.lambda_ent)), dtype=torch.float32)
-        )
-
-        # V5: Aplicar W8A8 em todas as camadas lineares (exceto Embedding, que é nn.Embedding)
-        # BUG FIX CRÍTICO V5: model.apply(fn) NÃO substitui módulos (descarta retorno).
-        # Usar apply_w8a8_to_model (substituição in-place recursiva).
-        n_quantized = apply_w8a8_to_model(self)
-        # Nota: nn.GRU não é afetada (não é Linear); Embedding não é afetada (nn.Embedding)
-        # Log em runtime (não criticar se n_quantized=0 — modelo pode não ter Linear)
-        self._n_w8a8_layers = n_quantized
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        temperature: float = 1.0,
-        use_hypothesis: bool = False,
-        stop_grad_hyp: bool = True,
-        return_aux: bool = False,
-        use_circular: Optional[bool] = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Forward do UnifiedModel.
-
-        Args:
-            x: (batch, T) token IDs OU (batch, T, input_dim) embeddings
-            temperature: temperatura do softmax do Lema 1
-            use_hypothesis: se True, calcula delta (cabeça HypT ativa)
-            stop_grad_hyp: se True, detach a entrada da HypT (Lema 3)
-            return_aux: se True, retorna também entropy_reg, alpha, e info circular
-            use_circular: override p/ ativar/desativar CircularReasoning no forward
-                (default: usa self.config.use_circular_reasoning)
-
-        Returns:
-            (y_hat, delta) por padrão.
-            (y_hat, delta, aux_dict) se return_aux=True.
-
-        Fluxo completo (V3):
-            x → embedding → [L1: module_selector α] → K×u8cell_T → OrqCell → o
-              → [CRW: Φ^k(o) por n_cycles] → o* (contração W₂)
-              → TrainT(o*) = y_hat
-              → HypT(o*, stop_grad) = delta (se use_hypothesis)
-        """
-        # Embedding se necessário
-        if x.dim() == 2:
-            # x: (batch, T) token IDs → (batch, T, d_model)
-            x = self.token_embedding(x)
-            x = self.positional_encoding(x)
-            x = self.emb_dropout(x)
-        # Agora x: (batch, T, input_dim)
-
-        # V5: do_circular computed early (was previously duplicated below)
-        do_circular = (
-            use_circular if use_circular is not None
-            else self.config.use_circular_reasoning
-        )
-
-        # Lema 1: seleciona módulos via softmax com temperatura
-        alpha, entropy_reg = self.module_selector(temperature=temperature)
-
-        # Calcula saídas de todos os módulos e pondera por alpha
-        # V4: coleta métricas de diversificação BiGRU (paralelismo real)
-        H_list = []
-        H_unweighted = []  # V4: saídas h_k SEM ponderar, para diversidade INTER-módulos
-        bigru_metrics_per_module = []  # lista de dicts (1 por u8cell_T)
-        for k, cell in enumerate(self.u8cells):
-            if return_aux:
-                # V4: forward com return_aux para capturar diversidade BiGRU
-                h_k, bigru_aux = cell(x, return_aux=True)
-                bigru_metrics_per_module.append(bigru_aux)
-            else:
-                h_k = cell(x)  # (batch, output_dim_u8cell)
-            H_unweighted.append(h_k)
-            h_k_weighted = h_k * alpha[k]
-            H_list.append(h_k_weighted)
-        H = torch.stack(H_list, dim=1)  # (batch, max_modules, output_dim_u8cell)
-        # V4: stack das saídas não-ponderadas para diversidade INTER-módulos
-        H_unweighted_stack = torch.stack(H_unweighted, dim=1)  # (batch, K, output_dim_u8cell)
-
-        # Orquestrador: produz representação agregada
-        o = self.orq_cell(H)  # (batch, d_cache)
-
-        # Circular Reasoning (Teorema 4): refine 'o' via contração W₂
-        circular_info = None
-        # do_circular já foi computado no início do forward (V5)
-        if do_circular and self.circular_reasoning is not None:
-            # CRW aceita (B, D) ou (B, T, D); passamos (B, D) — LayerNorm normaliza última dim
-            crw_out = self.circular_reasoning(o)
-            o = crw_out["output"]  # representação refinada s*
-            circular_info = {
-                "converged": crw_out["converged"],
-                "w2_final": crw_out["w2_final"],
-                "n_cycles_used": crw_out["n_cycles_used"],
-                "converged_step": crw_out["converged_step"],
-                "contraction_valid": crw_out["contraction_valid"],
-                "q_empirical_avg": (
-                    sum(crw_out["q_empirical"]) / max(1, len(crw_out["q_empirical"]))
-                    if crw_out["q_empirical"] else 1.0
-                ),
-            }
-
-        # Predição principal
-        y_hat_main = self.train_T(o)  # (batch, vocab_size)
-
-        # Predição de hipótese (se desejado)
-        if use_hypothesis:
-            delta = self.hyp_T(o, stop_grad=stop_grad_hyp)
-        else:
-            delta = torch.zeros_like(y_hat_main)
-
-        if return_aux:
-            # V4: agrega métricas BiGRU de todos os K u8cell_T módulos
-            bigru_diversity_summary = self._aggregate_bigru_diversity(
-                bigru_metrics_per_module, alpha, H_unweighted_stack
-            )
-            aux = {
-                "entropy_reg": entropy_reg,
-                "alpha": alpha.detach(),
-                "o": o.detach(),
-                "circular_info": circular_info,
-                "bigru_diversity": bigru_diversity_summary,  # V4
-            }
-            return y_hat_main, delta, aux
-        return y_hat_main, delta
-
-    def _aggregate_bigru_diversity(
-        self,
-        metrics_per_module: list[dict],
-        alpha: torch.Tensor,
-        H_unweighted: torch.Tensor,
-    ) -> dict:
-        """Agrega métricas BiGRU de todos os K u8cell_T módulos.
-
-        Args:
-            metrics_per_module: lista de K dicts (cada um = saída de u8cell_T._compute_bigru_diversity_metrics)
-            alpha: (K,) — pesos do ModuleSelector
-            H_unweighted: (batch, K, output_dim_u8cell) — saídas h_k SEM ponderar
-
-        Returns:
-            dict com médias agregadas sobre os K módulos (ponderadas por alpha):
-              - mean_cosine_diversity: diversidade média inter-BiGRU (ponderada por alpha)
-              - mean_gate_norm_ratio: variação média das normas
-              - mean_activation_variance: variância média
-              - mean_parallelism_score: score médio de paralelismo
-              - module_level_cosine_diversity: 1 - mean(cos_sim entre h_k dos K módulos)
-                  (diversidade INTER-módulos, não intra-módulo)
-              - per_module_summary: lista de K dicts com métricas essenciais
-        """
-        if not metrics_per_module:
-            return {
-                "mean_cosine_diversity": 0.0,
-                "mean_gate_norm_ratio": 0.0,
-                "mean_activation_variance": 0.0,
-                "mean_parallelism_score": 0.0,
-                "module_level_cosine_diversity": 0.0,
-                "per_module_summary": [],
-            }
-
-        # Ponderar por alpha_k (módulos mais ativos pesam mais)
-        alpha_np = alpha.detach().cpu().numpy()
-        alpha_norm = alpha_np / (alpha_np.sum() + 1e-8)
-
-        cd_list = [m["cosine_diversity"] for m in metrics_per_module]
-        gnr_list = [m["gate_norm_ratio"] for m in metrics_per_module]
-        av_list = [m["activation_variance"] for m in metrics_per_module]
-        ps_list = [m["parallelism_score"] for m in metrics_per_module]
-
-        mean_cd = sum(w * v for w, v in zip(alpha_norm, cd_list))
-        mean_gnr = sum(w * v for w, v in zip(alpha_norm, gnr_list))
-        mean_av = sum(w * v for w, v in zip(alpha_norm, av_list))
-        mean_ps = sum(w * v for w, v in zip(alpha_norm, ps_list))
-
-        # V4: Diversidade INTER-módulos usando H_unweighted (batch, K, output_dim_u8cell)
-        # Média sobre batch → (K, output_dim_u8cell), depois cosine sim entre pares
-        if H_unweighted is not None and H_unweighted.size(1) >= 2:
-            module_reps = H_unweighted.detach().mean(dim=0)  # (K, output_dim_u8cell)
-            K = module_reps.size(0)
-            inter_cos_sims = []
-            for i in range(K):
-                for j in range(i + 1, K):
-                    cs = F.cosine_similarity(
-                        module_reps[i].unsqueeze(0),
-                        module_reps[j].unsqueeze(0),
-                        dim=-1,
-                    )
-                    inter_cos_sims.append(float(cs.item()))
-            mean_inter_cos = sum(inter_cos_sims) / len(inter_cos_sims)
-            module_level_cd = 1.0 - mean_inter_cos
-        else:
-            module_level_cd = 0.0
-
-        return {
-            "mean_cosine_diversity": float(mean_cd),
-            "mean_gate_norm_ratio": float(mean_gnr),
-            "mean_activation_variance": float(mean_av),
-            "mean_parallelism_score": float(mean_ps),
-            "module_level_cosine_diversity": float(module_level_cd),
-            "per_module_summary": [
-                {
-                    "cosine_diversity": cd_list[i],
-                    "gate_norm_ratio": gnr_list[i],
-                    "parallelism_score": ps_list[i],
-                    "per_bigru_norms": metrics_per_module[i]["per_bigru_norms"],
-                }
-                for i in range(len(metrics_per_module))
-            ],
-        }
-
-    def get_active_modules(self, threshold: float = 0.01) -> int:
-        """Retorna nº efetivo de módulos ativos (para auto-configuração)."""
-        return self.module_selector.get_active_modules(threshold)
-
-    # =====================================================================
-    # V5: Propriedades para parâmetros aprendíveis (auto-ajuste do modelo)
-    # =====================================================================
-    @property
-    def lambda_ent_value(self) -> float:
-        """λ_ent atual = exp(log_lambda_ent) — peso da reg. de entropia."""
-        return float(torch.exp(self.log_lambda_ent.detach()).item())
-
-    @property
-    def init_gate_value(self) -> float:
-        """init_gate atual = mean(circular_reasoning.residual_gate).
-
-        V5: O gate do CircularReasoning é Parameter interno (n_cycles params).
-        O MetaConfigurator gerencia esse Parameter diretamente.
-        """
-        if self.circular_reasoning is not None:
-            return float(self.circular_reasoning.residual_gate.detach().mean().item())
-        return 0.0
-
-    def count_parameters(self) -> dict:
-        """Conta parâmetros totais e treináveis."""
-        total = sum(p.numel() for p in self.parameters())
-        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        return {
-            "total": total,
-            "total_M": total / 1e6,
-            "trainable": trainable,
-            "trainable_M": trainable / 1e6,
-            "n_w8a8_layers": getattr(self, "_n_w8a8_layers", 0),
-        }
-
-
-def create_unified_model(config: Optional[UnifiedModelConfig] = None) -> tuple[UnifiedModel, UnifiedModelConfig]:
-    """Factory: cria UnifiedModel e retorna (model, config)."""
-    cfg = config or UnifiedModelConfig()
-    model = UnifiedModel(cfg)
-    return model, cfg