v3.0: reorganiza arquivos sob cnn_bigru/ (preserva árvore de pastas)

49b8205 verified 3 months ago

22.9 kB

"""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",
]

	"""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",
	]