File size: 5,169 Bytes
49b8205 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """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,
return_ib_loss: bool = False,
) -> torch.Tensor:
"""Returns L2-normalized embeddings [B, out_dim].
Args:
input_ids: [B, T] token IDs
attention_mask: [B, T] opcional (1=real, 0=PAD)
return_ib_loss: se True, retorna também a loss IB (Information Bottleneck)
para que o trainer possa adicioná-la à perda total.
Returns:
emb_norm: [B, out_dim] L2-normalizado
(opcional) ib_loss: escalar (tensor 0-dim) com a loss IB
"""
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.beta > 0:
var = emb_norm.var(dim=0).mean()
ib_loss = self.beta * var
else:
ib_loss = torch.zeros((), device=emb_norm.device)
if return_ib_loss:
return emb_norm, ib_loss
# Para compatibilidade, anexa ao tensor (mas trainer deve usar return_ib_loss=True)
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",
]
|