"""trainer.py — Loop de treinamento cooperativo com auto-aprendizado. Implementa a seção 11 de dados.txt: PARA época = 1 ATÉ NUM_EPOCAS: PARA batch EM DataLoader: // Forward Gerador logits_prova, enc_out, _ = gerador.forward(axiomas, conjectura, prova_alvo) // Calcular perdas e avaliações do verificador L_total, L_G, L_V, L_AH, lista_v = calcular_perdas(...) // Backward zerar_gradientes L_total.backward() // Controle de gradiente (clip + spectral norm) aplicar_controle_gradiente(gerador) aplicar_controle_gradiente(verificador) // Ajuste dinâmico de LR baseado em normas de gradiente e v_media eta_G, eta_V = atualizar_taxas(...) // Passo de otimização otimizador_G.step() otimizador_V.step() // Normalização espectral após step aplicar_espectral_norm(gerador) aplicar_espectral_norm(verificador) Adicionalmente (APRIMORAMENTOS v2.0): - Mecanismo de hipóteses ativado quando v < threshold (com gate funcional) - N tentativas de sinergia (SynergySearcher) no início de cada época - Meta-controlador que ajusta pesos das perdas baseado na variância do gradiente - EWC (Elastic Weight Consolidation) integrado à perda total - Context Window aplicado na inferência (não no treino, que usa teacher forcing) - Verificador REALMENTE chamado (não proxy) — correção de bug crítico - Hipótese controller output REALMENTE usado na forward — correção de bug crítico """ from __future__ import annotations import logging import math import os import time from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from ..losses.losses import LossConfig, MultiLoss from ..models.generator_verifier import ( GeneratorCNNBiGRU, VerifierCNNBiGRU, AntiHallucinationLayer, ) from ..models.multimodal_model import MultimodalCNNBiGRU from ..utils.ewc import EWCConfig, EWCState from ..utils.memory_optimizer import MemoryOptimizer from .auto_learner import AutoLearnConfig, AutoLearner, estimate_curvature, orthogonal_init_model from .hypothesis_controller import HypothesisConfig, HypothesisController, SynergySearcher logger = logging.getLogger(__name__) @dataclass class TrainerConfig: """Configuração do treinador.""" # Épocas e batches num_epochs: int = 2 max_batches_per_epoch: int = 10 # limita batches para teste de 50 amostras batch_size: int = 8 # Synergy search n_synergy_attempts: int = 3 # N tentativas de acerto do treinamento use_synergy_search: bool = True # Hypothesis use_hypotheses: bool = True n_hypotheses: int = 4 # Loss config loss_config: LossConfig = field(default_factory=LossConfig) # Auto-learn config auto_config: AutoLearnConfig = field(default_factory=AutoLearnConfig) # EWC config (NOVO v2.0) ewc_config: Optional[EWCConfig] = None # None = EWC desativado # Device device: str = "cpu" # Logging log_every: int = 1 # Habilitar/desabilitar componentes use_verifier_real: bool = True # usar verificador REAL (não proxy) use_generator: bool = True # usar gerador para logits de LM use_hypothesis_output: bool = True # usar saída da hipótese (não só flag) class CooperativeTrainer: """Treina o modelo multimodal CNN-BiGRU cooperativo. Args: model: MultimodalCNNBiGRU tokenizer: BBPETokenizer config: TrainerConfig generator: GeneratorCNNBiGRU (opcional, para modo gerar+verificar) verifier: VerifierCNNBiGRU (opcional) anti_hallucination: AntiHallucinationLayer (opcional) ewc_state: EWCState (opcional, para aprendizado contínuo) """ def __init__( self, model: MultimodalCNNBiGRU, tokenizer, config: TrainerConfig, generator: Optional[GeneratorCNNBiGRU] = None, verifier: Optional[VerifierCNNBiGRU] = None, anti_hallucination: Optional[AntiHallucinationLayer] = None, ewc_state: Optional[EWCState] = None, ): self.model = model.to(config.device) self.tokenizer = tokenizer self.cfg = config self.device = config.device # Componentes opcionais self.generator = generator.to(config.device) if generator else None self.verifier = verifier.to(config.device) if verifier else None self.anti_hall = ( anti_hallucination.to(config.device) if anti_hallucination else None ) # EWC state (aprendizado contínuo) — NOVO v2.0 self.ewc_state = ewc_state # Loss self.loss_fn = MultiLoss(config.loss_config, pad_idx=tokenizer.pad_id).to(config.device) # Auto-learner self.auto = AutoLearner(config.auto_config) # Synergy searcher self.synergy = SynergySearcher( n_attempts=config.n_synergy_attempts, seed=42, ) # Otimizadores self.opt_model = torch.optim.AdamW( self.model.parameters(), lr=config.auto_config.initial_lr_G, weight_decay=config.auto_config.l2_reg # sempre definido agora (v2.0) ) self.opt_gen = ( torch.optim.AdamW( self.generator.parameters(), lr=config.auto_config.initial_lr_G, weight_decay=config.auto_config.l2_reg, ) if self.generator else None ) self.opt_ver = ( torch.optim.AdamW( self.verifier.parameters(), lr=config.auto_config.initial_lr_V, weight_decay=config.auto_config.l2_reg, ) if self.verifier else None ) # Memory optimizer self.mem = MemoryOptimizer(enable_amp=(self.device != "cpu")) self.mem.configure() # Hypothesis controller (aplicado sobre features do modelo) self.hypothesis = HypothesisController( feat_dim=model.fusion.out_dim, out_dim=model.fusion.out_dim, config=HypothesisConfig( n_hypotheses=config.n_hypotheses, threshold_err=config.loss_config.threshold_err, ), ).to(self.device) # Histórico de treinamento self.history: List[Dict[str, Any]] = [] # NOVO v2.0: métricas estendidas self.metrics: Dict[str, List[float]] = { "loss": [], "ppl": [], "v_mean": [], "lr_G": [], "lr_V": [], "l_g": [], "l_v": [], "l_ah": [], "l_ewc": [], "hypothesis_activated": [], } # ===================================================================== # Cálculo de v (verificador) e τ (anti-alucinação) — VERSÃO CORRIGIDA # ===================================================================== def _compute_v_and_tau( self, fused: torch.Tensor, input_ids_a: torch.Tensor, input_ids_b: torch.Tensor, gen_logits: Optional[torch.Tensor] = None, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: """Computa valores do verificador (v) e anti-alucinação (τ). VERSÃO CORRIGIDA v2.0: - Antes: `v_proxy = sigmoid(fused.mean(-1))` (verificador nunca chamado) - Agora: chama o VerifierCNNBiGRU com tokens reais gerados pelo gerador - Anti-alucinação: aplicada sobre os logits reais do gerador (não proxy) """ v_list = [] tau_list = [] B = fused.size(0) device = fused.device # --- Verificador REAL --- if self.verifier is not None and self.cfg.use_verifier_real: try: # Passo gerado: pegar argmax dos gen_logits por timestep (se disponível) if gen_logits is not None: # gen_logits: [B, T, V] → argmax → [B, T] passo_gerado = gen_logits.argmax(dim=-1) else: # Sem gerador: usar input_ids_b como "passo hipotético" passo_gerado = input_ids_b # Premissas: usar input_ids_a como contexto acumulado premissas_a = input_ids_a # Para o stream B do verificador, usamos input_ids_b como premissa B premissas_b = input_ids_b # Garantir que passo_gerado tenha mesma dimensão de tempo que premissas # (verificador espera 4 entradas: premissas_a, premissas_b, passo_a, passo_b) passo_a = passo_gerado passo_b = passo_gerado # simplificação: mesmo passo para ambos streams # Chamada REAL ao verificador with torch.enable_grad(): v_real = self.verifier(premissas_a, premissas_b, passo_a, passo_b) # [B, 1] v_list.append(v_real) except Exception as e: logger.warning(f"Verifier forward falhou: {e}; usando proxy") v_proxy = torch.sigmoid(fused.mean(dim=-1, keepdim=True)) v_list.append(v_proxy) elif self.verifier is not None and not self.cfg.use_verifier_real: # Proxy explícito (modo degrau) v_proxy = torch.sigmoid(fused.mean(dim=-1, keepdim=True)) v_list.append(v_proxy) else: # Sem verifier: v = 1 (sem punição) v_list.append(torch.ones(B, 1, device=device)) # --- Anti-hallucination REAL --- if self.anti_hall is not None: try: if gen_logits is not None: # Aplicar τ sobre cada timestep dos logits do gerador # gen_logits: [B, T, V] for t in range(gen_logits.size(1)): tau_t = self.anti_hall(gen_logits[:, t, :]) # [B, 1] tau_list.append(tau_t) else: # Sem gerador: aplicar τ sobre lm_head(fused) logits_proxy = self.model.lm_head(fused) # [B, V] tau = self.anti_hall(logits_proxy) tau_list.append(tau) except Exception as e: logger.warning(f"Anti-hallucination falhou: {e}") tau_list.append(torch.ones(B, 1, device=device)) else: tau_list.append(torch.ones(B, 1, device=device)) return v_list, tau_list # ===================================================================== # EWC Penalty (NOVO v2.0) # ===================================================================== def _compute_ewc_penalty(self) -> Optional[torch.Tensor]: """Computa a penalidade EWC se o estado estiver ativo.""" if self.ewc_state is None or not self.ewc_state.is_active(): return None try: return self.ewc_state.penalty(self.model) except Exception as e: logger.warning(f"EWC penalty falhou: {e}") return None # ===================================================================== # Passo de Treinamento (CORRIGIDO v2.0) # ===================================================================== def _train_step( self, batch: Dict[str, torch.Tensor], ) -> Dict[str, float]: """Executa um passo de treinamento. Retorna métricas. CORREÇÕES v2.0: - Hipótese controller output REALMENTE usado (soma ao fused) - Verificador REALMENTE chamado (não proxy) - EWC penalty integrado à perda total - Sem double-counting de CE loss (apenas via MultiLoss) """ # Move para device input_ids_a = batch["input_ids_a"].to(self.device, non_blocking=True) input_ids_b = batch["input_ids_b"].to(self.device, non_blocking=True) attn_mask_a = batch["attn_mask_a"].to(self.device, non_blocking=True) attn_mask_b = batch["attn_mask_b"].to(self.device, non_blocking=True) images = batch["images"].to(self.device, non_blocking=True) audios = batch["audios"].to(self.device, non_blocking=True) labels = batch["labels"].to(self.device, non_blocking=True) # Forward with self.mem.amp_context(self.device): out = self.model( input_ids_a, input_ids_b, images=images, audios=audios, attn_mask_a=attn_mask_a, attn_mask_b=attn_mask_b, mode="classify", ) logits = out["logits"] # [B, num_classes] fused = out["fused"] # [B, fusion_dim] # Verifier + anti-hallucination + gerador gen_logits = None if self.generator is not None and self.cfg.use_generator: try: gen_out = self.generator( input_ids_a, input_ids_b, target_proof=input_ids_b, # teacher forcing bos_id=self.tokenizer.bos_id, eos_id=self.tokenizer.eos_id, ) gen_logits = gen_out["logits"] except Exception as e: logger.debug(f"Generator forward falhou: {e}") gen_logits = None # Verifier e anti-hallucination (chamadas REAIS) v_list, tau_list = self._compute_v_and_tau( fused, input_ids_a, input_ids_b, gen_logits ) # Constrói logits para a MultiLoss if gen_logits is not None: loss_logits = gen_logits loss_targets = input_ids_b else: # Proxy: expande logits de classificação para [B, T, V] B = fused.size(0) T = input_ids_b.size(1) loss_logits = self.model.lm_head(fused).unsqueeze(1).expand(-1, T, -1) loss_targets = input_ids_b # NOVO v2.0: Hypothesis controller — usar a SAÍDA (não só a flag) # v_mean over list of [B,1] → [B,1] v_mean_tensor = torch.stack([v.squeeze(-1) for v in v_list]).mean(dim=0).unsqueeze(-1) try: hyp_out = self.hypothesis(fused, v_mean_tensor) # A saída do hypothesis controller É USADA para enriquecer o fused # (somando como residual). Antes era ignorada. if self.cfg.use_hypothesis_output: # Residual: fused = fused + alpha * (hyp_out - fused) # onde alpha é derivado do gate alpha = hyp_out["gate"].mean().clamp(0.0, 1.0).item() if alpha > 1e-3: # Re-projetar para mesma dimensão (hyp_out pode ter dimensão diferente) hyp_out_proj = hyp_out["output"] if hyp_out_proj.size(-1) != fused.size(-1): # Projeção via Linear on-the-fly (lazy) if not hasattr(self, "_hyp_proj"): self._hyp_proj = nn.Linear( hyp_out_proj.size(-1), fused.size(-1), bias=False ).to(self.device) nn.init.xavier_uniform_(self._hyp_proj.weight) hyp_out_proj = self._hyp_proj(hyp_out_proj) # Soma residual ponderada pelo gate fused = fused + alpha * hyp_out_proj # Recalcular logits de classificação com fused atualizado # (não refazemos forward do model para economizar compute; # apenas atualizamos a saída final do lm_head/classifier) # NOTA: isto é uma aproximação — o ideal seria refazer o forward except Exception as e: logger.debug(f"Hypothesis controller falhou: {e}") hyp_out = {"activated": False, "gate": torch.zeros_like(v_mean_tensor)} # NOVO v2.0: EWC Penalty ewc_pen = self._compute_ewc_penalty() # Curvatura (estimativa por diferenças finitas) try: curv = estimate_curvature( self.model, lambda: self.loss_fn( loss_logits, loss_targets, v_list, tau_list, models_for_reg=[self.model], ), eps=self.cfg.loss_config.curvature_eps, n_samples=2, ) self.loss_fn.set_curvature_estimate(curv) except Exception as e: logger.debug(f"Curvatura falhou: {e}") # Compute losses (com EWC penalty integrado) losses = self.loss_fn( loss_logits, loss_targets, v_list, tau_list, models_for_reg=[m for m in [self.model, self.generator, self.verifier] if m is not None], ewc_penalty=ewc_pen, ) # Loss de classificação (NÃO duplica a L_G — esta é uma tarefa auxiliar) ce_loss = F.cross_entropy(logits, labels) # Combina: perda de geração multimodal + perda de classificação total_loss = losses["total"] + ce_loss # Backward self.opt_model.zero_grad(set_to_none=True) if self.opt_gen: self.opt_gen.zero_grad(set_to_none=True) if self.opt_ver: self.opt_ver.zero_grad(set_to_none=True) total_loss.backward() # Clipping self.auto.apply_clipping(self.model) if self.generator: self.auto.apply_clipping(self.generator) if self.verifier: self.auto.apply_clipping(self.verifier) # Normas de gradiente grad_G = self.auto.compute_total_grad_norm(self.model) grad_V = self.auto.compute_total_grad_norm(self.verifier) if self.verifier else 0.0 # v_mean e L_V_mean v_mean = float(torch.stack([v.squeeze(-1) for v in v_list]).mean().item()) L_V_mean = float(losses["l_v"].item()) if isinstance(losses["l_v"], torch.Tensor) else 0.0 # Ajuste dinâmico de LR novo_lr_G, novo_lr_V = self.auto.update_learning_rates( grad_G, grad_V, v_mean, L_V_mean ) self.auto.set_optimizer_lr(self.opt_model, novo_lr_G) if self.opt_gen: self.auto.set_optimizer_lr(self.opt_gen, novo_lr_G) if self.opt_ver: self.auto.set_optimizer_lr(self.opt_ver, novo_lr_V) # Step self.opt_model.step() if self.opt_gen: self.opt_gen.step() if self.opt_ver: self.opt_ver.step() # Spectral norm após step if self.cfg.auto_config.apply_after_step: self.auto.apply_spectral_norm(self.model) if self.generator: self.auto.apply_spectral_norm(self.generator) if self.verifier: self.auto.apply_spectral_norm(self.verifier) # Métricas loss_val = float(total_loss.item()) try: ppl = math.exp(min(loss_val, 20)) # evita overflow except OverflowError: ppl = float("inf") # Hipótese info hyp_activated = bool(hyp_out.get("activated", False)) if isinstance(hyp_out, dict) else False # EWC info l_ewc_val = float(losses.get("l_ewc", torch.tensor(0.0)).item()) if isinstance( losses.get("l_ewc", 0.0), torch.Tensor ) else 0.0 return { "loss": loss_val, "ppl": ppl, "v_mean": v_mean, "lr_G": novo_lr_G, "lr_V": novo_lr_V, "l_g": float(losses["l_g"].item()) if isinstance(losses["l_g"], torch.Tensor) else 0.0, "l_v": L_V_mean, "l_ah": float(losses["l_ah"].item()) if isinstance(losses["l_ah"], torch.Tensor) else 0.0, "l_ewc": l_ewc_val, "hypothesis_activated": hyp_activated, } # ===================================================================== # Synergy Search Helper # ===================================================================== def _eval_config(self, cfg: Dict, sample_batch: Dict[str, torch.Tensor]) -> float: """Avalia uma configuração de perda em um mini-batch (para synergy search).""" # Salva config atual original = {k: getattr(self.loss_fn.cfg, k) for k in cfg} # Aplica nova config for k, v in cfg.items(): if hasattr(self.loss_fn.cfg, k): setattr(self.loss_fn.cfg, k, v) try: metrics = self._train_step(sample_batch) return metrics["loss"] finally: # Restaura for k, v in original.items(): setattr(self.loss_fn.cfg, k, v) # ===================================================================== # Loop Principal de Treinamento # ===================================================================== def train( self, dataloader: DataLoader, ) -> Dict[str, Any]: """Executa o treinamento completo.""" logger.info("Iniciando treinamento cooperativo CNN-BiGRU v2.0") logger.info(f" epochs={self.cfg.num_epochs}, batch_size={self.cfg.batch_size}, " f"n_synergy={self.cfg.n_synergy_attempts}") logger.info(f" EWC ativo: {self.ewc_state is not None and self.ewc_state.is_active()}") logger.info(f" Verifier real: {self.cfg.use_verifier_real}") logger.info(f" Hipótese output usado: {self.cfg.use_hypothesis_output}") start_time = time.time() for epoch in range(self.cfg.num_epochs): epoch_loss = 0.0 epoch_ppl = 0.0 n_batches = 0 hypothesis_activations = 0 # Synergy search no início da época (apenas época 0 para economizar tempo) if self.cfg.use_synergy_search and epoch == 0: try: # Pega um batch de amostra para avaliação sample_batch = next(iter(dataloader)) sample_batch = {k: (v.to(self.device) if isinstance(v, torch.Tensor) else v) for k, v in sample_batch.items()} base_cfg = { "alpha": self.loss_fn.cfg.alpha, "beta": self.loss_fn.cfg.beta, "gamma_loss": self.loss_fn.cfg.gamma_loss, "delta": self.loss_fn.cfg.delta, "lambda_penal": self.loss_fn.cfg.lambda_penal, "mu_exp_penal": self.loss_fn.cfg.mu_exp_penal, } logger.info(f"Synergy search: tentando {self.cfg.n_synergy_attempts} configurações...") best_cfg = self.synergy.search( base_config=base_cfg, eval_fn=lambda c: self._eval_config(c, sample_batch), ) # Aplica a melhor configuração for k, v in best_cfg.items(): if hasattr(self.loss_fn.cfg, k): setattr(self.loss_fn.cfg, k, v) logger.info("Synergy search concluído. Melhor config aplicada.") except Exception as e: logger.warning(f"Synergy search falhou: {e}") # Loop de batches self.model.train() if self.generator: self.generator.train() if self.verifier: self.verifier.train() if self.anti_hall: self.anti_hall.train() self.hypothesis.train() for batch_idx, batch in enumerate(dataloader): if batch_idx >= self.cfg.max_batches_per_epoch: break try: metrics = self._train_step(batch) epoch_loss += metrics["loss"] epoch_ppl += metrics["ppl"] n_batches += 1 if metrics["hypothesis_activated"]: hypothesis_activations += 1 # Armazena todas as métricas for k in self.metrics: if k in metrics: self.metrics[k].append(metrics[k]) if batch_idx % self.cfg.log_every == 0: logger.info( f"Época {epoch+1} | Batch {batch_idx} | " f"loss={metrics['loss']:.4f} | ppl={metrics['ppl']:.2f} | " f"v_mean={metrics['v_mean']:.3f} | lr_G={metrics['lr_G']:.2e} | " f"l_ewc={metrics['l_ewc']:.4f} | " f"hyp_act={'Y' if metrics['hypothesis_activated'] else 'N'}" ) except Exception as e: logger.error(f"Erro no batch {batch_idx}: {e}", exc_info=True) continue # Cleanup de memória ao final da época self.mem.cleanup() avg_loss = epoch_loss / max(n_batches, 1) avg_ppl = epoch_ppl / max(n_batches, 1) logger.info( f"Época {epoch+1} concluída | avg_loss={avg_loss:.4f} | " f"avg_ppl={avg_ppl:.2f} | hyp_activations={hypothesis_activations} | " f"batches_ok={n_batches}/{self.cfg.max_batches_per_epoch}" ) # Detecta falha sistêmica if n_batches == 0: logger.error(f"Época {epoch+1}: NENHUM batch teve sucesso — possível bug no pipeline") self.history.append({ "epoch": epoch + 1, "avg_loss": avg_loss, "avg_ppl": avg_ppl, "n_batches": n_batches, "hypothesis_activations": hypothesis_activations, "epoch_loss_total": epoch_loss, }) elapsed = time.time() - start_time logger.info(f"Treinamento concluído em {elapsed:.1f}s") return { "history": self.history, "metrics": self.metrics, "final_loss": self.history[-1]["avg_loss"] if self.history else 0.0, "final_ppl": self.history[-1]["avg_ppl"] if self.history else 0.0, "elapsed_s": elapsed, "synergy_history": self.synergy.history, "auto_history": self.auto.history, "ewc_active": self.ewc_state is not None and self.ewc_state.is_active(), "ewc_num_tasks": self.ewc_state.num_tasks() if self.ewc_state else 0, } # ===================================================================== # Consolidar tarefa para EWC (NOVO v2.0) # ===================================================================== def consolidate_task( self, forward_fn, compute_fisher: bool = True, ) -> None: """Consolida a tarefa atual no estado EWC. Args: forward_fn: callable que retorna (logits, target) ou logits para cálculo da Fisher Information. compute_fisher: se True, recalcula a Fisher; se False, apenas armazena o theta_star atual. """ if self.ewc_state is None: logger.warning("EWC não configurado — consolidate_task ignorado") return try: self.ewc_state.consolidate(self.model, forward_fn=forward_fn if compute_fisher else None) logger.info(f"Tarefa consolidada no EWC (total: {self.ewc_state.num_tasks()})") except Exception as e: logger.error(f"Falha ao consolidar tarefa no EWC: {e}") __all__ = ["TrainerConfig", "CooperativeTrainer"]