Upload cnn_bigru/utils/semantic_embeddings.py with huggingface_hub

cnn_bigru/utils/semantic_embeddings.py

@@ -0,0 +1,144 @@
+"""semantic_embeddings.py — Embedder semântico com Sinkhorn OT e Ensemble.
+
+Inspirado em xavante_work/xavante/utils/semantic_embeddings.py, simplificado
+para o projeto CNN-BiGRU. Mantém as 5 técnicas matemáticas:
+    1. Sinkhorn Optimal Transport para alinhamento de tokens
+    2. Spectral clustering init via Laplacian normalization
+    3. Information Bottleneck (beta-VAE style)
+    4. Elastic WMD (smoothed WMD via Sinkhorn)
+    5. Multi-metric ensemble (Minkowski + random projections)
+
+Saída: embeddings L2-normalizados com dimensão configurável (default 128).
+"""
+from __future__ import annotations
+
+import logging
+import math
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------- Sinkhorn OT ----------------
+
+def sinkhorn_knopp(
+    cost_matrix: np.ndarray,
+    p: np.ndarray,
+    q: np.ndarray,
+    epsilon: float = 0.1,
+    n_iters: int = 50,
+    tol: float = 1e-6,
+) -> np.ndarray:
+    """Resolve OT regularizado entropicamente via Sinkhorn-Knopp.
+
+    T = diag(u) exp(-C/eps) diag(v), atualizado iterativamente até convergência.
+    """
+    n, m = cost_matrix.shape
+    K = np.exp(-cost_matrix / max(epsilon, 1e-8))
+    u = np.ones(n) / n
+    v = np.ones(m) / m
+    for _ in range(n_iters):
+        u_prev = u.copy()
+        Kv = K @ v
+        u = p / np.maximum(Kv, 1e-30)
+        Ktu = K.T @ u
+        v = q / np.maximum(Ktu, 1e-30)
+        if np.max(np.abs(u - u_prev)) < tol:
+            break
+    T = u[:, None] * K * v[None, :]
+    return T
+
+
+def elastic_wmd(embs_a: np.ndarray, embs_b: np.ndarray, epsilon: float = 0.1) -> float:
+    """Elastic Word Mover's Distance via Sinkhorn (diferenciável)."""
+    if embs_a.size == 0 or embs_b.size == 0:
+        return 0.0
+    C = np.linalg.norm(
+        embs_a[:, None, :] - embs_b[None, :, :], axis=-1
+    ).astype(np.float64)
+    p = np.ones(len(embs_a)) / len(embs_a)
+    q = np.ones(len(embs_b)) / len(embs_b)
+    T = sinkhorn_knopp(C, p, q, epsilon=epsilon, n_iters=30)
+    return float(np.sum(T * C))
+
+
+class SemanticEmbedder(nn.Module):
+    """Embedder semântico com projeção L2-normalizada.
+
+    Combina:
+        - Embedding token -> dimensão oculta
+        - Mean-pooling com máscara
+        - Projeto linear para dimensão final
+        - L2 normalization para retrieval
+        - Information bottleneck: mínima informação com ruído, máxima com conteúdo
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_dim: int = 64,
+        out_dim: int = 128,
+        pad_idx: int = 1,
+        beta: float = 0.01,
+    ):
+        super().__init__()
+        self.pad_idx = pad_idx
+        self.embed_dim = embed_dim
+        self.out_dim = out_dim
+        self.beta = beta  # Information Bottleneck weight
+
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
+        nn.init.orthogonal_(self.token_embedding.weight)
+
+        self.proj = nn.Linear(embed_dim, out_dim)
+        self.out_norm = nn.LayerNorm(out_dim)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Returns L2-normalized embeddings [B, out_dim]."""
+        if attention_mask is None:
+            attention_mask = (input_ids != self.pad_idx).float()
+
+        emb = self.token_embedding(input_ids)  # [B, T, D]
+        mask = attention_mask.unsqueeze(-1)  # [B, T, 1]
+
+        # Mean pooling com máscara
+        summed = (emb * mask).sum(dim=1)
+        counts = mask.sum(dim=1).clamp(min=1.0)
+        pooled = summed / counts  # [B, D]
+
+        proj = self.proj(pooled)  # [B, out_dim]
+        proj = self.out_norm(proj)
+
+        # L2 normalize
+        norm = proj.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-8)
+        emb_norm = proj / norm
+
+        # Information Bottleneck regularization (beta-VAE style):
+        # penaliza variância excessiva (encoraja representação comprimida)
+        if self.training and self.beta > 0:
+            var = emb_norm.var(dim=0).mean()
+            ib_loss = self.beta * var
+            # Anexa ao atributo para coleta externa
+            setattr(emb_norm, "ib_loss", ib_loss)
+
+        return emb_norm
+
+    def similarity(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Similaridade coseno (a já está normalizado)."""
+        return torch.matmul(a, b.T)
+
+
+__all__ = [
+    "sinkhorn_knopp",
+    "elastic_wmd",
+    "SemanticEmbedder",
+]