src/bigru_t/model/token_compress.py

"""
token_compress — Compressores e indexadores para atenção de 2M tokens.

Pipeline (3 estágios):
1. ToMe (Token Merging) — Bolya et al., 2023
   Funde tokens similares via matching bipartido soft. Reduz N → N' com
   N' = N - r·T (r tokens por camada, T camadas). Custo O(N log N).

2. KNN Indexer (IVF-like)
   Mantém centroides via k-means. Para cada query, busca top-k tokens
   no índice (k << N'). Custo O(N' log N') para construir, O(k log N')
   por query.

3. Per-Layer Attention Pattern
   Cada camada tem seu próprio padrão de sparsidade, selecionável:
     - "full":     dense attention clássica (O(N²)) — apenas para N' pequeno
     - "window":   atenção local em janela w
     - "global":   tokens especiais (CLS, registradores) attendem a tudo
     - "block":    atenção em blocos (BLOCKWISE)
   O padrão é configurável por camada via lista `patterns`.

Complexidade total: O(N log N + T · k · N') comparado a O(T · N²) do
Transformer denso. Para N=2M, T=24, k=256: ganho ~ 50×.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# 1. ToMeCompressor
# ---------------------------------------------------------------------------
def _bipartite_soft_matching(
    metric: torch.Tensor,
    r: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Matching bipartido soft (Bolya et al., 2023 — algoritmo original).

    Passos:
      1. Ordena tokens por score (norma do embedding — proxy para PCA[0]).
      2. Separa alternadamente em dois conjuntos A (índices pares) e B
         (índices ímpares) — disjuntos, |A| = |B| = N//2.
      3. Para cada i ∈ [0, N//2), forma par candidato (a_i, b_i).
      4. Calcula similaridade s_i = ⟨a_i, b_i⟩.
      5. Seleciona os r pares com maior similaridade para fundir.
         Os demais pares (N//2 - r) mantêm ambos os tokens.

    metric: (B, N, C)
    r: nº de pares a fundir
    Retorna:
      a_idx: (B, r)            — índices de A nos pares fundidos (âncoras)
      b_idx: (B, r)            — índices de B nos pares fundidos
      ungated_idx: (B, N - 2r) — índices dos tokens não fundidos
    """
    B, N, C = metric.shape
    M = N // 2  # número de pares

    if r <= 0 or M < r:
        all_idx = torch.arange(N, device=metric.device).unsqueeze(0).expand(B, -1)
        return all_idx[:, :0], all_idx[:, :0], all_idx

    with torch.no_grad():
        scores = metric.norm(dim=-1)                                    # (B, N)
        sorted_idx = scores.argsort(dim=-1, descending=True)            # (B, N)
        # Alterna pares/ímpares: A = sorted_idx[:, 0::2], B = sorted_idx[:, 1::2].
        a_idx_all = sorted_idx[:, 0:2 * M:2]                            # (B, M)
        b_idx_all = sorted_idx[:, 1:2 * M:2]                            # (B, M)
        # Similaridade entre cada par (a_i, b_i).
        a_set = metric.gather(1, a_idx_all.unsqueeze(-1).expand(-1, -1, C))  # (B, M, C)
        b_set = metric.gather(1, b_idx_all.unsqueeze(-1).expand(-1, -1, C))  # (B, M, C)
        sim = (a_set * b_set).sum(dim=-1)                                # (B, M)
        # Top-r pares mais similares para fundir.
        top_r = sim.topk(r, dim=-1).indices                              # (B, r)
        # Máscara de pares fundidos.
        merged_mask = torch.zeros(B, M, dtype=torch.bool, device=metric.device)
        merged_mask.scatter_(1, top_r, True)
        # Índices de pares não-fundidos (constantes no batch).
        pair_idx = torch.arange(M, device=metric.device).unsqueeze(0).expand(B, -1)
        keep_pair = pair_idx[~merged_mask]                              # (B*(M-r),)
        keep_pair = keep_pair.reshape(B, M - r)
        # Token indices.
        a_merge = a_idx_all.gather(1, top_r)                            # (B, r)
        b_merge = b_idx_all.gather(1, top_r)                            # (B, r)
        a_keep = a_idx_all.gather(1, keep_pair)                         # (B, M-r)
        b_keep = b_idx_all.gather(1, keep_pair)                         # (B, M-r)
        # Ungated = a_keep ∪ b_keep (total = 2*(M-r) = N - 2r se N par).
        ungated = torch.cat([a_keep, b_keep], dim=1)                    # (B, 2(M-r))
    return a_merge, b_merge, ungated


class ToMeCompressor(nn.Module):
    """Token Merging (ToMe) — funde r tokens similares a cada forward.

    Reduz N → N - r tokens (mantém dimensão de feature).
    """

    def __init__(self, r: int = 64, use_merge_bias: bool = False):
        super().__init__()
        self.r = r
        # Opcional: pequeno gate para mistura ponderada entre par (a, b).
        # Por padrão usamos média simples (estilo ToMe original).
        self.merge_gate = nn.Parameter(torch.tensor(0.5)) if use_merge_bias else None

    def forward(self, X: torch.Tensor, return_aux: bool = False) -> torch.Tensor:
        """X: (B, N, D) -> (B, N - r, D).

        Quando return_aux=True, retorna também os índices (para debug).
        """
        B, N, D = X.shape
        M = N // 2
        # Se não há pares suficientes para fundir, retorna X inalterado.
        if self.r <= 0 or M < self.r:
            return X
        a_pair, b_idx, ungated = _bipartite_soft_matching(X, self.r)
        # Merge: média do par (a, b).
        a_set = X.gather(1, a_pair.unsqueeze(-1).expand(-1, -1, D))   # (B, r, D)
        b_set = X.gather(1, b_idx.unsqueeze(-1).expand(-1, -1, D))    # (B, r, D)
        if self.merge_gate is not None:
            w = torch.sigmoid(self.merge_gate)
            merged = w * a_set + (1.0 - w) * b_set
        else:
            merged = 0.5 * (a_set + b_set)
        # Ungated tokens (2*(M-r)).
        ungated_set = X.gather(1, ungated.unsqueeze(-1).expand(-1, -1, D))  # (B, 2(M-r), D)
        # Concatena: ungated + merged = 2(M-r) + r = 2M - r = N - r (se N par).
        out = torch.cat([ungated_set, merged], dim=1)                # (B, N - r, D)
        if return_aux:
            return out, (a_pair, b_idx, ungated)
        return out


# ---------------------------------------------------------------------------
# 2. KNNIndexer (IVF-like, sem dependência do FAISS)
# ---------------------------------------------------------------------------
class KNNIndexer(nn.Module):
    """Indexador k-NN baseado em k-means (estilo IVF).

    Mantém uma lista invertida de tokens por centroide. Query busca
    apenas nos centroides mais próximos (n_probe) e retorna top-k tokens
    globais. Custo por query: O(n_probe · avg_list_size + n_centroids).

    Funciona em batch: para cada query, retorna índices dos top-k tokens
    e os pesos de similaridade (cosseno).
    """

    def __init__(
        self,
        dim: int,
        n_centroids: int = 256,
        n_probe: int = 8,
        k: int = 64,
        n_iter: int = 5,
    ):
        super().__init__()
        self.dim = dim
        self.n_centroids = n_centroids
        self.n_probe = n_probe
        self.k = k
        self.n_iter = n_iter
        # Centroides — buffer treinável em modo "build".
        self.register_buffer("centroids", torch.randn(n_centroids, dim) * 0.02)
        # Listas invertidas (chave = centroide idx, valor = índices de tokens).
        # Em runtime mantemos como listas Python — não persistimos em state_dict.
        self._lists: List[torch.Tensor] = []
        self._tokens: Optional[torch.Tensor] = None  # (N, D)
        self._built = False

    @torch.no_grad()
    def build(self, tokens: torch.Tensor) -> None:
        """Constrói o índice IVF sobre o banco de tokens.

        tokens: (N, D)
        """
        N, D = tokens.shape
        assert D == self.dim
        device = tokens.device
        # 1. K-means rápido (n_iter iterações de Lloyd).
        if N <= self.n_centroids:
            self.centroids = tokens.clone()
            self._lists = [torch.tensor([i], device=device) for i in range(N)]
        else:
            # Inicialização: amostragem aleatória.
            perm = torch.randperm(N, device=device)[:self.n_centroids]
            self.centroids = tokens[perm].clone()
            for _ in range(self.n_iter):
                # Atribui cada token ao centroide mais próximo.
                d2 = (tokens.unsqueeze(1) - self.centroids.unsqueeze(0)).pow(2).sum(-1)  # (N, K)
                assign = d2.argmin(dim=-1)                                              # (N,)
                # Atualiza centroides.
                for k in range(self.n_centroids):
                    mask = assign == k
                    if mask.any():
                        self.centroids[k] = tokens[mask].mean(dim=0)
            # Constói listas invertidas.
            d2 = (tokens.unsqueeze(1) - self.centroids.unsqueeze(0)).pow(2).sum(-1)
            assign = d2.argmin(dim=-1)
            self._lists = [
                (assign == k).nonzero(as_tuple=True)[0].to(device)
                for k in range(self.n_centroids)
            ]
        self._tokens = tokens
        self._built = True

    @torch.no_grad()
    def search(self, queries: torch.Tensor, k: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        """Busca top-k tokens para cada query.

        queries: (B, D)
        Retorna:
          indices: (B, k) — índices no banco _tokens
          weights: (B, k) — similaridade cosseno (em [-1, 1])
        """
        if not self._built:
            raise RuntimeError("Índice não construído. Chame build() antes.")
        k = k or self.k
        B = queries.shape[0]
        device = queries.device
        # 1. Encontra n_probe centroides mais próximos por query.
        q_norm = F.normalize(queries, dim=-1)                          # (B, D)
        c_norm = F.normalize(self.centroids, dim=-1)                   # (K, D)
        sim_c = q_norm @ c_norm.t()                                    # (B, K)
        n_probe = min(self.n_probe, self.n_centroids)
        top_c = sim_c.topk(n_probe, dim=-1).indices                    # (B, n_probe)

        # 2. Para cada query, agrega tokens das listas invertidas correspondentes.
        all_idx_list = []
        all_w_list = []
        for b in range(B):
            cand_indices = torch.cat([self._lists[c.item()] for c in top_c[b] if len(self._lists[c.item()]) > 0])
            if cand_indices.numel() == 0:
                # Fallback: busca global.
                cand_indices = torch.arange(self._tokens.shape[0], device=device)
            cand_tokens = self._tokens[cand_indices]                   # (M, D)
            q_b = q_norm[b]                                            # (D,)
            sims = (cand_tokens / cand_tokens.norm(dim=-1, keepdim=True).clamp(min=1e-8)) @ q_b  # (M,)
            kk = min(k, sims.numel())
            topk = sims.topk(kk)
            idx = cand_indices[topk.indices]                           # (kk,)
            w = topk.values                                            # (kk,)
            # Pad com zeros se necessário.
            if kk < k:
                pad_idx = torch.zeros(k - kk, dtype=torch.long, device=device)
                pad_w = torch.zeros(k - kk, device=device)
                idx = torch.cat([idx, pad_idx])
                w = torch.cat([w, pad_w])
            all_idx_list.append(idx)
            all_w_list.append(w)
        indices = torch.stack(all_idx_list, dim=0)                     # (B, k)
        weights = torch.stack(all_w_list, dim=0)                       # (B, k)
        return indices, weights


# ---------------------------------------------------------------------------
# 3. CompressionPipeline (ToMe + KNN index)
# ---------------------------------------------------------------------------
class CompressionPipeline(nn.Module):
    """Pipeline de compressão: ToMe → KNN index → retorna subconjunto + pesos."""

    def __init__(
        self,
        dim: int,
        tome_r: int = 64,
        n_centroids: int = 256,
        n_probe: int = 8,
        k: int = 64,
    ):
        super().__init__()
        self.tome = ToMeCompressor(r=tome_r)
        self.indexer = KNNIndexer(dim=dim, n_centroids=n_centroids, n_probe=n_probe, k=k)
        self._built = False

    @torch.no_grad()
    def build(self, tokens: torch.Tensor) -> None:
        """Constrói índice KNN sobre tokens pré-comprimidos por ToMe."""
        compressed = self.tome(tokens)
        self.indexer.build(compressed)
        self._built = True

    def forward(
        self,
        queries: torch.Tensor,
        keys: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Para cada query, retorna top-k keys + pesos.

        queries: (B, N_q, D)
        keys:    (B, N_k, D) — índice construído para o batch atual.

        Em modo "global" (índice precompilado via build), apenas queries são
        processadas e keys são ignoradas.

        Retorna:
          indices: (B, N_q, k) — índices das top-k keys
          weights: (B, N_q, k) — similaridades (cosseno)
          compressed_keys: (B, N_k - r, D) — keys pós-ToMe (se r > 0)
        """
        # Comprime keys via ToMe.
        compressed_keys = self.tome(keys)
        # Para cada batch element, faz busca local (não usa índice global pré-construído).
        B, Nq, D = queries.shape
        Nk_c = compressed_keys.shape[1]
        k = self.indexer.k
        # Similaridade cosseno entre queries e compressed_keys.
        q_norm = F.normalize(queries, dim=-1)                          # (B, Nq, D)
        k_norm = F.normalize(compressed_keys, dim=-1)                  # (B, Nk_c, D)
        sim = q_norm @ k_norm.transpose(-1, -2)                        # (B, Nq, Nk_c)
        topk = sim.topk(min(k, Nk_c), dim=-1)                          # (B, Nq, k)
        indices = topk.indices                                          # (B, Nq, k)
        weights = topk.values                                           # (B, Nq, k)
        return indices, weights, compressed_keys