Upload cnn_bigru/models/generator_verifier.py with huggingface_hub

cnn_bigru/models/generator_verifier.py

@@ -0,0 +1,344 @@
+"""generator_verifier.py — Gerador + Verificador + Anti-Alucinação.
+
+Implementa o pseudocódigo da seção 5-7 de dados.txt:
+
+    - Generator: CNN-BiGRU encoder + GRU decoder com atenção Bahdanau
+    - Verifier: CNN-BiGRU + classificador sigmoid (v ∈ [0,1])
+    - Anti-hallucination layer: ativadores lógicos fuzzy de Łukasiewicz
+      NOT(x)  = 1 - sigmoid(x)
+      AND(x,y) = relu(sigmoid(x) + sigmoid(y) - 1)
+      OR(x,y)  = min(1, sigmoid(x) + sigmoid(y))
+      IMP(x,y) = min(1, 1 - sigmoid(x) + sigmoid(y))
+"""
+from __future__ import annotations
+
+import logging
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .cooperative_bigru import CooperativeCNNBiGRU
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Funções fuzzy de Łukasiewicz
+# ============================================================================
+
+def fuzzy_not(x: torch.Tensor) -> torch.Tensor:
+    """NOT(x) = 1 - sigmoid(x)."""
+    return 1.0 - torch.sigmoid(x)
+
+
+def fuzzy_and(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """AND(x, y) = relu(sigmoid(x) + sigmoid(y) - 1)."""
+    s = torch.sigmoid(x) + torch.sigmoid(y)
+    return F.relu(s - 1.0)
+
+
+def fuzzy_or(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """OR(x, y) = min(1, sigmoid(x) + sigmoid(y))."""
+    s = torch.sigmoid(x) + torch.sigmoid(y)
+    return torch.clamp(s, 0.0, 1.0)
+
+
+def fuzzy_imp(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """IMP(x, y) = min(1, 1 - sigmoid(x) + sigmoid(y))."""
+    s = 1.0 - torch.sigmoid(x) + torch.sigmoid(y)
+    return torch.clamp(s, 0.0, 1.0)
+
+
+# ============================================================================
+# Camada Anti-Alucinação
+# ============================================================================
+
+class AntiHallucinationLayer(nn.Module):
+    """Camada anti-alucinação com lógica fuzzy de Łukasiewicz.
+
+    Recebe os logits [B, V] de um passo gerado e calcula um valor de verdade
+    suave τ ∈ [0, 1] que indica o quão "lógico" é o passo.
+
+    Implementação diferenciável:
+        - Calcula sigmoid(logit) -> probabilidade soft
+        - Aplica operadores fuzzy em pares de tokens mais prováveis
+        - Combina resultados via IMP (implicação) e AND
+    """
+
+    def __init__(self, vocab_size: int, embed_dim: int = 32):
+        super().__init__()
+        self.proj = nn.Linear(vocab_size, embed_dim)
+        self.combine = nn.Linear(embed_dim * 2, 1)
+
+    def forward(self, logits: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            logits: [B, V] logits do passo gerado
+
+        Returns:
+            tau: [B, 1] valor de verdade suave em [0, 1]
+        """
+        # Probabilidades soft via sigmoid
+        probs = torch.sigmoid(logits)  # [B, V]
+
+        # Top-2 tokens mais prováveis (diferenciável via topk)
+        top2_vals, top2_idx = torch.topk(probs, k=min(2, probs.size(-1)), dim=-1)  # [B, 2]
+
+        # Logits dos top-2 tokens via gather
+        top1_logit = logits.gather(1, top2_idx[:, 0:1])  # [B, 1]
+        top2_logit = logits.gather(1, top2_idx[:, 1:2])  # [B, 1]
+
+        # Operadores fuzzy de Łukasiewicz sobre os top-2 logits
+        imp_val = fuzzy_imp(top1_logit, top2_logit)  # [B, 1]
+        and_val = fuzzy_and(top1_logit, top2_logit)  # [B, 1]
+        or_val = fuzzy_or(top1_logit, top2_logit)    # [B, 1]
+
+        # Projeção do vetor de logits completo para embedding
+        proj = self.proj(logits)  # [B, embed_dim]
+
+        # Modulariza a projeção pelos valores fuzzy (broadcasting)
+        fuzzy_factor = (imp_val + and_val + or_val) / 3.0  # [B, 1]
+        modulated = proj * fuzzy_factor  # [B, embed_dim]
+
+        # Combina: [modulated, original_proj] -> [B, 2*embed_dim]
+        combined = torch.cat((modulated, proj), dim=-1)  # [B, 2*embed_dim]
+        tau = torch.sigmoid(self.combine(combined))  # [B, 1]
+        return tau
+
+
+# ============================================================================
+# Gerador (CNN-BiGRU encoder + GRU decoder com atenção)
+# ============================================================================
+
+class GeneratorCNNBiGRU(nn.Module):
+    """Gerador: encoder CNN-BiGRU + decoder GRU com atenção Bahdanau.
+
+    Args:
+        vocab_size, embedding_dim, cnn_filters, gru_hidden: herdados do núcleo.
+        max_proof_len: comprimento máximo da sequência gerada.
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        embedding_dim: int = 64,
+        cnn_filters: int = 64,
+        gru_hidden: int = 64,
+        n_heads: int = 4,
+        dropout: float = 0.1,
+        pad_idx: int = 1,
+        max_proof_len: int = 32,
+    ):
+        super().__init__()
+        self.pad_idx = pad_idx
+        self.vocab_size = vocab_size
+        self.gru_hidden = gru_hidden
+        self.max_proof_len = max_proof_len
+
+        # Encoder cooperativo
+        self.encoder = CooperativeCNNBiGRU(
+            vocab_size=vocab_size,
+            embedding_dim=embedding_dim,
+            cnn_filters=cnn_filters,
+            gru_hidden=gru_hidden,
+            n_heads=n_heads,
+            dropout=dropout,
+            pad_idx=pad_idx,
+        )
+
+        # Decoder GRU (unidirecional)
+        self.dec_embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
+        self.dec_gru = nn.GRUCell(embedding_dim + gru_hidden * 2, gru_hidden)
+
+        # Atenção Bahdanau
+        self.attn_W = nn.Linear(gru_hidden * 2, gru_hidden, bias=False)
+        self.attn_U = nn.Linear(gru_hidden, gru_hidden, bias=False)
+        self.attn_v = nn.Linear(gru_hidden, 1, bias=False)
+
+        # Projeção final
+        self.out_proj = nn.Linear(gru_hidden + gru_hidden * 2 + embedding_dim, vocab_size)
+
+        self.bos_id = 0  # placeholder, será sobrescrito
+        self.eos_id = 2
+
+    def _attention(
+        self,
+        dec_state: torch.Tensor,  # [B, H]
+        enc_out: torch.Tensor,    # [B, T, 2H]
+        enc_mask: torch.Tensor,   # [B, T]
+    ) -> torch.Tensor:
+        """Atenção Bahdanau: retorna contexto [B, 2H]."""
+        # score = v^T tanh(W enc + U dec)
+        proj_enc = self.attn_W(enc_out)              # [B, T, H]
+        proj_dec = self.attn_U(dec_state).unsqueeze(1)  # [B, 1, H]
+        scores = self.attn_v(torch.tanh(proj_enc + proj_dec)).squeeze(-1)  # [B, T]
+        scores = scores.masked_fill(enc_mask == 0, -1e9)
+        attn = F.softmax(scores, dim=-1)             # [B, T]
+        ctx = torch.bmm(attn.unsqueeze(1), enc_out).squeeze(1)  # [B, 2H]
+        return ctx
+
+    def forward(
+        self,
+        input_ids_a: torch.Tensor,
+        input_ids_b: torch.Tensor,
+        target_proof: Optional[torch.Tensor] = None,
+        bos_id: int = 0,
+        eos_id: int = 2,
+        max_len: Optional[int] = None,
+    ) -> dict:
+        """
+        Args:
+            input_ids_a, input_ids_b: [B, T] entradas (axiomas + conjectura)
+            target_proof: [B, T_proof] para teacher forcing (treino)
+            bos_id, eos_id: IDs dos tokens especiais
+
+        Returns:
+            dict com:
+                logits: [B, T_proof, V]
+                enc_out: [B, T, 2H]
+                dec_state: [B, H]
+        """
+        self.bos_id = bos_id
+        self.eos_id = eos_id
+
+        # 1. Encoder
+        enc = self.encoder(input_ids_a, input_ids_b, return_sequences=False)
+        # Para o decoder, usamos o fused como representação final e repetimos ao longo do tempo
+        fused = enc["fused"]  # [B, 256]
+        # enc_out "virtual": usamos a mask do stream A para o tempo
+        mask_A = enc["mask_A"]
+        T_enc = mask_A.size(1)
+        enc_out = fused.unsqueeze(1).expand(-1, T_enc, -1)  # [B, T, 256]
+        # Para atenção, usamos a primeira metade como "enc_out" e a mask do stream A
+        # Enc_out tem 2H = 256, então H = 128
+        enc_mask = mask_A
+
+        B = input_ids_a.size(0)
+        device = input_ids_a.device
+
+        # 2. Inicialização do decoder
+        # Estado inicial: projeção do fused
+        h = fused[:, :self.gru_hidden]  # [B, H] (primeira metade do fused)
+
+        # 3. Token inicial = BOS
+        if target_proof is not None:
+            # Teacher forcing
+            T_dec = target_proof.size(1)
+            inputs = target_proof  # [B, T_dec]
+        else:
+            T_dec = max_len or self.max_proof_len
+            inputs = None
+
+        logits_list = []
+        cur_token = torch.full((B, 1), bos_id, dtype=torch.long, device=device)
+
+        for t in range(T_dec):
+            if target_proof is not None:
+                tok = target_proof[:, t]  # [B]
+            else:
+                tok = cur_token.squeeze(1)
+
+            emb = self.dec_embedding(tok)  # [B, D]
+            ctx = self._attention(h, enc_out, enc_mask)  # [B, 2H]
+            gru_input = torch.cat((emb, ctx), dim=-1)    # [B, D+2H]
+            h = self.dec_gru(gru_input, h)               # [B, H]
+
+            # Logits
+            logit_input = torch.cat((h, ctx, emb), dim=-1)  # [B, H+2H+D]
+            logit = self.out_proj(logit_input)              # [B, V]
+            logits_list.append(logit)
+
+            if target_proof is None:
+                cur_token = logit.argmax(dim=-1, keepdim=True)
+
+        logits = torch.stack(logits_list, dim=1)  # [B, T_dec, V]
+
+        return {
+            "logits": logits,
+            "enc_out": enc_out,
+            "dec_state": h,
+        }
+
+
+# ============================================================================
+# Verificador (CNN-BiGRU + classificador sigmoid)
+# ============================================================================
+
+class VerifierCNNBiGRU(nn.Module):
+    """Verificador: CNN-BiGRU + classificador sigmoid (v ∈ [0,1]).
+
+    Recebe premissas (concat de axiomas + conjectura) e um passo hipotético,
+    e retorna uma probabilidade de o passo estar correto.
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        embedding_dim: int = 64,
+        cnn_filters: int = 64,
+        gru_hidden: int = 64,
+        n_heads: int = 4,
+        dropout: float = 0.1,
+        pad_idx: int = 1,
+    ):
+        super().__init__()
+        self.pad_idx = pad_idx
+        self.gru_hidden = gru_hidden
+
+        self.encoder = CooperativeCNNBiGRU(
+            vocab_size=vocab_size,
+            embedding_dim=embedding_dim,
+            cnn_filters=cnn_filters,
+            gru_hidden=gru_hidden,
+            n_heads=n_heads,
+            dropout=dropout,
+            pad_idx=pad_idx,
+        )
+
+        # Classificador sigmoid
+        self.classifier = nn.Sequential(
+            nn.Linear(gru_hidden * 2 * 2, gru_hidden),  # 256*2 = 512 (fused A+B do premissa+passo)
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(gru_hidden, 1),
+        )
+
+    def forward(
+        self,
+        premissas_a: torch.Tensor,
+        premissas_b: torch.Tensor,
+        passo_a: torch.Tensor,
+        passo_b: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            premissas_a, premissas_b: [B, T] (contexto da prova até o momento)
+            passo_a, passo_b: [B, T_p] (passo gerado a ser validado)
+
+        Returns:
+            v: [B, 1] probabilidade de o passo estar correto
+        """
+        # Encode premissas
+        enc_pre = self.encoder(premissas_a, premissas_b)
+        feat_pre = enc_pre["fused"]  # [B, 256]
+
+        # Encode passo
+        enc_passo = self.encoder(passo_a, passo_b)
+        feat_passo = enc_passo["fused"]  # [B, 256]
+
+        # Concatena e classifica
+        combined = torch.cat((feat_pre, feat_passo), dim=-1)  # [B, 512]
+        logit = self.classifier(combined)  # [B, 1]
+        v = torch.sigmoid(logit)  # [B, 1]
+        return v
+
+
+__all__ = [
+    "fuzzy_not", "fuzzy_and", "fuzzy_or", "fuzzy_imp",
+    "AntiHallucinationLayer",
+    "GeneratorCNNBiGRU",
+    "VerifierCNNBiGRU",
+]