Delete models/nlp.py with huggingface_hub

models/nlp.py

@@ -1,654 +0,0 @@
-"""nlp.py — Natural Language Processing (NLP) para CNN-BiGRU.
-
-Módulo de processamento de linguagem natural que fornece tarefas clássicas
-de NLP em cima do backbone cooperativo CNN-BiGRU:
-
-    - Sequence Classification (sentiment, topic, NLI)
-    - Token Classification (NER, POS tagging)
-    - Span Detection (QA extractive)
-    - Sentence Pair Classification (paraphrase, entailment)
-    - Embeddings para retrieval semântico
-
-==============================================================================
-ANÁLISE MATEMÁTICA E LÓGICA
-==============================================================================
-
-Dado o output do CooperativeCNNBiGRU (representações fused [B, 256] ou
-sequências temporais [B, T, 128] por stream), cada tarefa NLP aplica uma
-cabeça específica:
-
-1. Sequence Classification:
-        h_cls = SelfAttentionSummary(seq)  [B, D]
-        logits = Linear(D, num_classes)    [B, C]
-        loss = CrossEntropy(logits, y)
-
-2. Token Classification:
-        logits = Linear(D, num_labels)     [B, T, L]
-        loss = CrossEntropy(logits, y)  (ignore padding)
-
-3. Span Detection (QA):
-        start_logits = Linear(D, 1)        [B, T]
-        end_logits = Linear(D, 1)          [B, T]
-        loss = CE(start_logits, start_pos) + CE(end_logits, end_pos)
-
-4. Sentence Pair Classification:
-        Usa o dual-stream A/B do CooperativeCNNBiGRU diretamente
-        (já que o modelo é naturalmente dual-stream).
-        logits = Linear(D_fused, num_classes)
-
-5. Embeddings para retrieval:
-        h = mean_pool(seq)                 [B, D]
-        h = normalize(h, p=2)              [B, D]
-        similarity = cosine(h_query, h_doc)
-        loss = InfoNCE (contrastive)
-
-==============================================================================
-INTEGRAÇÃO
-==============================================================================
-
-- CooperativeCNNBiGRU (cnn_bigru.models.cooperative_bigru): backbone
-- SelfAttentionSummary: já incluído no cooperative_bigru
-- BBPETokenizer: tokenização
-
-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 NLPConfig:
-    """Configuração do módulo NLP."""
-    # Dimensões (precisam bater com o backbone)
-    embed_dim: int = 64           # embedding de tokens
-    cnn_filters: int = 64
-    gru_hidden: int = 64          # 64 por direção -> 128 por stream
-    feat_per_stream: int = 128    # = 2 * gru_hidden
-    feat_fused: int = 256         # = 2 * feat_per_stream
-    n_heads: int = 4
-    dropout: float = 0.1
-    pad_idx: int = 1
-    # Tarefas
-    num_classes_seq: int = 3      # sequence classification
-    num_labels_tok: int = 7       # token classification (NER-style)
-    # Pooling
-    pooling: str = "self_attn"    # "self_attn" | "mean" | "max" | "cls"
-    # Loss
-    label_smoothing: float = 0.0
-    # Device
-    device: str = "cpu"
-
-
-# ============================================================================
-# Sequence Classification Head
-# ============================================================================
-
-class SequenceClassificationHead(nn.Module):
-    """Cabeça para classificação de sequência inteira.
-
-    Aplica pooling (self-attention, mean, max ou cls) sobre a sequência
-    temporal e projeta para num_classes logits.
-
-    Args:
-        embed_dim: dimensão de entrada (e.g., 256 para fused, 128 por stream)
-        num_classes: número de classes de saída
-        dropout: prob de dropout
-        pooling: estratégia de pooling
-        n_heads: cabeças para self-attention pooling
-    """
-
-    def __init__(
-        self,
-        embed_dim: int,
-        num_classes: int,
-        dropout: float = 0.1,
-        pooling: str = "self_attn",
-        n_heads: int = 4,
-    ):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.pooling = pooling
-        self.num_classes = num_classes
-
-        # Pooling
-        if pooling == "self_attn":
-            # Self-attention pooling simplificado (sem CLS token)
-            self.attn_query = nn.Linear(embed_dim, 1)
-        elif pooling == "cls":
-            # CLS token learnable
-            self.cls = nn.Parameter(torch.zeros(1, 1, embed_dim))
-            nn.init.normal_(self.cls, std=0.02)
-        # mean / max: não precisa de params
-
-        self.dropout = nn.Dropout(dropout)
-        self.classifier = nn.Linear(embed_dim, num_classes)
-        nn.init.xavier_uniform_(self.classifier.weight)
-        nn.init.zeros_(self.classifier.bias)
-
-    def forward(
-        self,
-        seq: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            seq: [B, T, D] sequência de hidden states
-            mask: [B, T] máscara de padding (1=real, 0=pad)
-
-        Returns:
-            logits: [B, C]
-            pooled: [B, D] representação pooled
-        """
-        if seq.dim() == 2:
-            # [B, D] — já é pooled (e.g., fused)
-            pooled = seq
-        else:
-            B, T, D = seq.shape
-
-            if self.pooling == "self_attn":
-                # Attention pooling: softmax over T
-                scores = self.attn_query(seq).squeeze(-1)  # [B, T]
-                if mask is not None:
-                    scores = scores.masked_fill(mask == 0, -1e9)
-                # Lidar com linhas totalmente mascaradas
-                all_neg = (scores <= -1e9).all(dim=-1, keepdim=True)
-                scores = torch.where(all_neg, torch.zeros_like(scores), scores)
-                weights = F.softmax(scores, dim=-1).unsqueeze(-1)  # [B, T, 1]
-                pooled = (seq * weights).sum(dim=1)  # [B, D]
-
-            elif self.pooling == "mean":
-                if mask is not None:
-                    m = mask.unsqueeze(-1).float()
-                    pooled = (seq * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
-                else:
-                    pooled = seq.mean(dim=1)
-
-            elif self.pooling == "max":
-                if mask is not None:
-                    seq_masked = seq.masked_fill(mask.unsqueeze(-1) == 0, -1e9)
-                    pooled, _ = seq_masked.max(dim=1)
-                    # Lidar com sequências totalmente mascaradas (retorna 0)
-                    all_neg = (pooled <= -1e9).all(dim=-1, keepdim=True)
-                    pooled = torch.where(all_neg, torch.zeros_like(pooled), pooled)
-                else:
-                    pooled, _ = seq.max(dim=1)
-
-            elif self.pooling == "cls":
-                cls = self.cls.expand(B, -1, -1)  # [B, 1, D]
-                seq_ext = torch.cat([cls, seq], dim=1)  # [B, T+1, D]
-                pooled = seq_ext[:, 0, :]  # pega CLS
-
-            else:
-                raise ValueError(f"pooling desconhecido: {self.pooling}")
-
-        pooled = self.dropout(pooled)
-        logits = self.classifier(pooled)
-        return logits, pooled
-
-
-# ============================================================================
-# Token Classification Head
-# ============================================================================
-
-class TokenClassificationHead(nn.Module):
-    """Cabeça para classificação de tokens (NER, POS tagging).
-
-    Aplica Linear(D, num_labels) sobre cada posição da sequência.
-
-    Args:
-        embed_dim: dimensão de entrada
-        num_labels: número de rótulos (e.g., 7 para BIO tagging)
-        dropout: prob de dropout
-    """
-
-    def __init__(self, embed_dim: int, num_labels: int, dropout: float = 0.1):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.num_labels = num_labels
-        self.dropout = nn.Dropout(dropout)
-        self.classifier = nn.Linear(embed_dim, num_labels)
-        nn.init.xavier_uniform_(self.classifier.weight)
-        nn.init.zeros_(self.classifier.bias)
-
-    def forward(self, seq: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            seq: [B, T, D] sequência de hidden states
-        Returns:
-            logits: [B, T, L]
-        """
-        return self.classifier(self.dropout(seq))
-
-
-# ============================================================================
-# Span Detection Head (QA Extractive)
-# ============================================================================
-
-class SpanDetectionHead(nn.Module):
-    """Cabeça para detecção de span (QA extractive).
-
-    Prediz posições de início e fim do span de resposta.
-
-    Args:
-        embed_dim: dimensão de entrada
-        dropout: prob de dropout
-    """
-
-    def __init__(self, embed_dim: int, dropout: float = 0.1):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.dropout = nn.Dropout(dropout)
-        self.start_classifier = nn.Linear(embed_dim, 1)
-        self.end_classifier = nn.Linear(embed_dim, 1)
-        for layer in (self.start_classifier, self.end_classifier):
-            nn.init.xavier_uniform_(layer.weight)
-            nn.init.zeros_(layer.bias)
-
-    def forward(self, seq: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            seq: [B, T, D]
-        Returns:
-            start_logits: [B, T]
-            end_logits: [B, T]
-        """
-        seq = self.dropout(seq)
-        start_logits = self.start_classifier(seq).squeeze(-1)  # [B, T]
-        end_logits = self.end_classifier(seq).squeeze(-1)
-        return start_logits, end_logits
-
-
-# ============================================================================
-# Embedding Head (para retrieval semântico)
-# ============================================================================
-
-class EmbeddingHead(nn.Module):
-    """Cabeça para gerar embeddings normalizados (retrieval semântico).
-
-    Aplica mean-pool com máscara + LayerNorm + normalização L2.
-
-    Args:
-        embed_dim: dimensão de entrada e saída
-        use_layer_norm: se True, aplica LayerNorm antes da normalização L2
-    """
-
-    def __init__(self, embed_dim: int, use_layer_norm: bool = True):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.use_layer_norm = use_layer_norm
-        if use_layer_norm:
-            self.ln = nn.LayerNorm(embed_dim)
-
-    def forward(
-        self,
-        seq: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-        normalize: bool = True,
-    ) -> torch.Tensor:
-        """
-        Args:
-            seq: [B, T, D]
-            mask: [B, T] (1=real, 0=pad)
-            normalize: se True, normaliza com L2
-        Returns:
-            embedding: [B, D]
-        """
-        if mask is not None:
-            m = mask.unsqueeze(-1).float()
-            pooled = (seq * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
-        else:
-            pooled = seq.mean(dim=1)
-
-        if self.use_layer_norm:
-            pooled = self.ln(pooled)
-
-        if normalize:
-            pooled = F.normalize(pooled, p=2, dim=-1)
-
-        return pooled
-
-
-# ============================================================================
-# NLP Module (wrapper integrando todas as cabeças)
-# ============================================================================
-
-class NLPModule(nn.Module):
-    """Módulo NLP integrando múltiplas tarefas no backbone CNN-BiGRU.
-
-    Args:
-        config: configuração NLP
-        backbone: modelo CooperativeCNNBiGRU (opcional — se None, cria um novo)
-    """
-
-    def __init__(
-        self,
-        config: NLPConfig,
-        backbone: Optional[nn.Module] = None,
-    ):
-        super().__init__()
-        self.config = config
-
-        # Backbone (cria se não fornecido)
-        if backbone is None:
-            from .cooperative_bigru import CooperativeCNNBiGRU
-            backbone = CooperativeCNNBiGRU(
-                vocab_size=32000,  # placeholder, sobrescrito via set_vocab_size
-                embedding_dim=config.embed_dim,
-                cnn_filters=config.cnn_filters,
-                gru_hidden=config.gru_hidden,
-                n_heads=config.n_heads,
-                dropout=config.dropout,
-                pad_idx=config.pad_idx,
-            )
-        self.backbone = backbone
-
-        # Cabeças (uma por tarefa)
-        self.seq_head = SequenceClassificationHead(
-            embed_dim=config.feat_fused,
-            num_classes=config.num_classes_seq,
-            dropout=config.dropout,
-            pooling=config.pooling,
-            n_heads=config.n_heads,
-        )
-        self.tok_head = TokenClassificationHead(
-            embed_dim=config.feat_per_stream,  # por stream (usa seq_A ou seq_B)
-            num_labels=config.num_labels_tok,
-            dropout=config.dropout,
-        )
-        self.span_head = SpanDetectionHead(
-            embed_dim=config.feat_per_stream,
-            dropout=config.dropout,
-        )
-        self.embed_head = EmbeddingHead(
-            embed_dim=config.feat_per_stream,
-            use_layer_norm=True,
-        )
-
-    # ----------------------------------------------------------------------
-    # Tarefas
-    # ----------------------------------------------------------------------
-
-    def sequence_classification(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-    ) -> Dict[str, torch.Tensor]:
-        """Classificação de sequência (usando representação fused).
-
-        Args:
-            input_ids_a, input_ids_b: [B, T]
-
-        Returns:
-            dict com logits [B, C] e pooled [B, D]
-        """
-        out = self.backbone(input_ids_a, input_ids_b, return_sequences=False)
-        fused = out["fused"]  # [B, 256]
-        logits, pooled = self.seq_head(fused, mask=None)
-        return {"logits": logits, "pooled": pooled, "fused": fused}
-
-    def token_classification(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        stream: str = "A",
-    ) -> torch.Tensor:
-        """Classificação de tokens (NER/POS).
-
-        Args:
-            input_ids_a, input_ids_b: [B, T]
-            stream: "A" ou "B" — qual stream usar para token classification
-
-        Returns:
-            logits: [B, T, L]
-        """
-        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
-        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
-        if seq is None:
-            # Fallback: usar fused (broadcast para T)
-            fused = out["fused"]  # [B, 256]
-            T = input_ids_a.size(1)
-            seq = fused.unsqueeze(1).expand(-1, T, -1)  # [B, T, 256]
-            # Projeta para feat_per_stream se necessário
-            if fused.size(-1) != self.config.feat_per_stream:
-                if not hasattr(self, "_proj_tok"):
-                    self._proj_tok = nn.Linear(
-                        fused.size(-1), self.config.feat_per_stream, bias=False
-                    ).to(fused.device)
-                seq = self._proj_tok(seq)
-        return self.tok_head(seq)
-
-    def span_detection(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        stream: str = "A",
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Detecção de span (QA extractive).
-
-        Returns:
-            (start_logits, end_logits) — ambos [B, T]
-        """
-        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
-        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
-        if seq is None:
-            fused = out["fused"]
-            T = input_ids_a.size(1)
-            seq = fused.unsqueeze(1).expand(-1, T, -1)
-            if fused.size(-1) != self.config.feat_per_stream:
-                if not hasattr(self, "_proj_span"):
-                    self._proj_span = nn.Linear(
-                        fused.size(-1), self.config.feat_per_stream, bias=False
-                    ).to(fused.device)
-                seq = self._proj_span(seq)
-        return self.span_head(seq)
-
-    def embed(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        stream: str = "A",
-        normalize: bool = True,
-    ) -> torch.Tensor:
-        """Gera embedding normalizado para retrieval semântico.
-
-        Returns:
-            embedding: [B, D]
-        """
-        out = self.backbone(input_ids_a, input_ids_b, return_sequences=True)
-        seq = out.get("seq_A") if stream == "A" else out.get("seq_B")
-        mask = out.get("mask_A") if stream == "A" else out.get("mask_B")
-        if seq is None:
-            # Fallback: fused
-            return self.embed_head(
-                out["fused"].unsqueeze(1), mask=None, normalize=normalize
-            )
-        if seq.size(-1) != self.config.feat_per_stream:
-            if not hasattr(self, "_proj_emb"):
-                self._proj_emb = nn.Linear(
-                    seq.size(-1), self.config.feat_per_stream, bias=False
-                ).to(seq.device)
-            seq = self._proj_emb(seq)
-        return self.embed_head(seq, mask=mask, normalize=normalize)
-
-    # ----------------------------------------------------------------------
-    # Loss functions para cada tarefa
-    # ----------------------------------------------------------------------
-
-    def sequence_classification_loss(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        labels: torch.Tensor,
-    ) -> torch.Tensor:
-        """Computa CE loss para sequence classification."""
-        out = self.sequence_classification(input_ids_a, input_ids_b)
-        return F.cross_entropy(
-            out["logits"], labels,
-            label_smoothing=self.config.label_smoothing,
-        )
-
-    def token_classification_loss(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        labels: torch.Tensor,
-        stream: str = "A",
-    ) -> torch.Tensor:
-        """Computa CE loss para token classification.
-
-        Args:
-            labels: [B, T] rótulos por token (use -100 ou pad_idx para ignorar)
-        """
-        logits = self.token_classification(input_ids_a, input_ids_b, stream=stream)
-        # Ignora posições com label = -100 ou mask = 0
-        mask = (input_ids_a != self.config.pad_idx).float()
-        ignore_mask = (labels == -100) | (labels == self.config.pad_idx)
-        labels_clamped = labels.clone()
-        labels_clamped[ignore_mask] = -100
-        return F.cross_entropy(
-            logits.view(-1, self.config.num_labels_tok),
-            labels_clamped.view(-1),
-            ignore_index=-100,
-        )
-
-    def span_detection_loss(
-        self,
-        input_ids_a: torch.Tensor,
-        input_ids_b: torch.Tensor,
-        start_positions: torch.Tensor,
-        end_positions: torch.Tensor,
-        stream: str = "A",
-    ) -> torch.Tensor:
-        """Computa loss para span detection.
-
-        Args:
-            start_positions, end_positions: [B] índices do span (long)
-        """
-        start_logits, end_logits = self.span_detection(
-            input_ids_a, input_ids_b, stream=stream
-        )
-        # Ignora posições fora do span (mask)
-        loss_start = F.cross_entropy(start_logits, start_positions, ignore_index=-100)
-        loss_end = F.cross_entropy(end_logits, end_positions, ignore_index=-100)
-        return (loss_start + loss_end) / 2.0
-
-    def contrastive_loss(
-        self,
-        query_ids_a: torch.Tensor,
-        query_ids_b: torch.Tensor,
-        pos_ids_a: torch.Tensor,
-        pos_ids_b: torch.Tensor,
-        neg_ids_a: torch.Tensor,
-        neg_ids_b: torch.Tensor,
-        temperature: float = 0.07,
-    ) -> torch.Tensor:
-        """InfoNCE loss para retrieval semântico.
-
-        Args:
-            query, pos, neg: tensores [B, T] para query, positivo e negativo
-        """
-        q = self.embed(query_ids_a, query_ids_b, normalize=True)
-        p = self.embed(pos_ids_a, pos_ids_b, normalize=True)
-        n = self.embed(neg_ids_a, neg_ids_b, normalize=True)
-
-        # Similaridade coseno (já normalizado -> produto escalar)
-        sim_pos = (q * p).sum(dim=-1, keepdim=True) / temperature  # [B, 1]
-        sim_neg = (q * n).sum(dim=-1, keepdim=True) / temperature  # [B, 1]
-
-        # InfoNCE: maximizar sim_pos - sim_neg
-        logits = torch.cat([sim_pos, sim_neg], dim=-1)  # [B, 2]
-        labels = torch.zeros(q.size(0), dtype=torch.long, device=q.device)  # 0 = pos
-        return F.cross_entropy(logits, labels)
-
-
-# ============================================================================
-# Self-test
-# ============================================================================
-
-def _self_test():
-    """Teste rápido do módulo NLP."""
-    torch.manual_seed(42)
-    from .cooperative_bigru import CooperativeCNNBiGRU
-    config = NLPConfig(
-        embed_dim=32,
-        cnn_filters=32,
-        gru_hidden=32,
-        feat_per_stream=64,
-        feat_fused=128,
-        n_heads=4,
-        dropout=0.1,
-        pad_idx=1,
-        num_classes_seq=3,
-        num_labels_tok=5,
-        pooling="self_attn",
-    )
-    backbone = CooperativeCNNBiGRU(
-        vocab_size=100,
-        embedding_dim=config.embed_dim,
-        cnn_filters=config.cnn_filters,
-        gru_hidden=config.gru_hidden,
-        n_heads=config.n_heads,
-        dropout=config.dropout,
-        pad_idx=config.pad_idx,
-    )
-    nlp = NLPModule(config, backbone=backbone)
-
-    B, T = 2, 8
-    ids_a = torch.randint(2, 100, (B, T))
-    ids_b = torch.randint(2, 100, (B, T))
-
-    # Sequence classification
-    out = nlp.sequence_classification(ids_a, ids_b)
-    print(f"Seq cls: logits {out['logits'].shape}, pooled {out['pooled'].shape}")
-    labels = torch.randint(0, 3, (B,))
-    loss_seq = nlp.sequence_classification_loss(ids_a, ids_b, labels)
-    print(f"  loss={loss_seq.item():.4f}")
-
-    # Token classification
-    tok_logits = nlp.token_classification(ids_a, ids_b)
-    print(f"Tok cls: {tok_logits.shape}")
-    tok_labels = torch.randint(0, 5, (B, T))
-    loss_tok = nlp.token_classification_loss(ids_a, ids_b, tok_labels)
-    print(f"  loss={loss_tok.item():.4f}")
-
-    # Span detection
-    start_logits, end_logits = nlp.span_detection(ids_a, ids_b)
-    print(f"Span: start {start_logits.shape}, end {end_logits.shape}")
-    sp = torch.randint(0, T, (B,))
-    ep = torch.randint(0, T, (B,))
-    loss_span = nlp.span_detection_loss(ids_a, ids_b, sp, ep)
-    print(f"  loss={loss_span.item():.4f}")
-
-    # Embedding
-    emb = nlp.embed(ids_a, ids_b)
-    print(f"Embed: {emb.shape}, norm: {emb.norm(dim=-1).mean().item():.4f}")
-
-    print("NLP self-test OK")
-
-
-if __name__ == "__main__":
-    _self_test()
-
-
-__all__ = [
-    "NLPConfig",
-    "SequenceClassificationHead",
-    "TokenClassificationHead",
-    "SpanDetectionHead",
-    "EmbeddingHead",
-    "NLPModule",
-]