Delete models/rope.py with huggingface_hub

models/rope.py

@@ -1,196 +0,0 @@
-"""
-RoPE (Rotary Position Embeddings) — Codificação posicional rotacional.
-
-Referência: Su et al., 2021 — "RoFormer: Enhanced Transformer with Rotary
-Position Embedding".
-
-Propriedade chave: <RoPE(q, m), RoPE(k, n)> = <RoPE(q, m-n), k>
-→ a atenção depende apenas da diferença posicional relativa.
-
-Ver docs/MATH_ANALYSIS.md seção 4.1.
-
-Autor: CNN-BiGRU Project
-"""
-from __future__ import annotations
-
-import math
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-
-
-def precompute_rope_frequencies(
-    head_dim: int,
-    max_seq_len: int,
-    base: float = 10000.0,
-    device: Optional[torch.device] = None,
-    dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """
-    Pré-computa as frequências angulares theta_i para RoPE.
-
-    Fórmula: theta_i = base^(-2i / head_dim)  para i = 0, 1, ..., head_dim/2 - 1
-
-    Args:
-        head_dim: dimensão de cada cabeça (deve ser par)
-        max_seq_len: comprimento máximo da sequência
-        base: base da exponencial (padrão 10000)
-        device: device do tensor
-        dtype: dtype do tensor
-
-    Returns:
-        cos_sin: tensor [max_seq_len, head_dim] com [cos(theta_0*m), cos(theta_1*m), ...,
-                 sin(theta_0*m), sin(theta_1*m), ...] intercalados por par
-                 (formato compatível com apply_rope).
-
-    Matematicamente:
-        Para cada posição m em [0, max_seq_len) e cada par de dimensões 2i:
-            angle = m * theta_i = m * base^(-2i/head_dim)
-        A matriz de rotação no plano (2i, 2i+1) é:
-            [[cos(angle), -sin(angle)],
-             [sin(angle),  cos(angle)]]
-    """
-    if head_dim % 2 != 0:
-        raise ValueError(f"head_dim deve ser par; recebido {head_dim}")
-
-    # theta_i = base^(-2i / head_dim), i = 0, ..., head_dim/2 - 1
-    i = torch.arange(0, head_dim, 2, device=device, dtype=dtype)  # [head_dim/2]
-    theta = 1.0 / (base ** (i / head_dim))  # [head_dim/2]
-
-    # Posições m
-    positions = torch.arange(max_seq_len, device=device, dtype=dtype)  # [max_seq_len]
-
-    # angle[m, i] = m * theta_i
-    angles = torch.outer(positions, theta)  # [max_seq_len, head_dim/2]
-
-    # Duplicar para intercalar: [cos_0, cos_1, ..., sin_0, sin_1, ...]
-    cos = torch.cos(angles)  # [max_seq_len, head_dim/2]
-    sin = torch.sin(angles)
-
-    # Formato final: [max_seq_len, head_dim] com pares (cos_i, sin_i)
-    cos_sin = torch.stack([cos, sin], dim=-1).reshape(max_seq_len, head_dim)
-    return cos_sin
-
-
-def apply_rope(
-    x: torch.Tensor,
-    cos_sin: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Aplica RoPE ao tensor x.
-
-    Args:
-        x: [batch, n_heads, seq_len, head_dim]
-        cos_sin: [seq_len, head_dim] — tensor pré-computado
-
-    Returns:
-        x_rotated: [batch, n_heads, seq_len, head_dim]
-
-    Matemática (para cada batch, head, posição m, par de dims (2i, 2i+1)):
-        x'[2i]   = x[2i]   * cos(angle) - x[2i+1] * sin(angle)
-        x'[2i+1] = x[2i]   * sin(angle) + x[2i+1] * cos(angle)
-
-    O tensor cos_sin tem formato [seq_len, head_dim] com:
-        cos_sin[m, 2i]   = cos(m * theta_i)
-        cos_sin[m, 2i+1] = sin(m * theta_i)
-    """
-    if x.dim() != 4:
-        raise ValueError(
-            f"x deve ser [batch, n_heads, seq_len, head_dim]; recebido {x.shape}"
-        )
-
-    B, H, T, D = x.shape
-    if D % 2 != 0:
-        raise ValueError(f"head_dim deve ser par; recebido {D}")
-
-    # Reshape x para [B, H, T, D/2, 2] — pares de dimensões
-    x_pairs = x.float().reshape(B, H, T, D // 2, 2)
-
-    # Reshape cos_sin para [T, D/2, 2]
-    cos_sin_t = cos_sin[:T].to(x.device).float().reshape(T, D // 2, 2)
-    cos = cos_sin_t[..., 0]  # [T, D/2]
-    sin = cos_sin_t[..., 1]
-
-    # Broadcast: [1, 1, T, D/2]
-    cos = cos.unsqueeze(0).unsqueeze(0)
-    sin = sin.unsqueeze(0).unsqueeze(0)
-
-    # Aplicar rotação
-    x_real = x_pairs[..., 0]  # [B, H, T, D/2]
-    x_imag = x_pairs[..., 1]
-    rotated_real = x_real * cos - x_imag * sin
-    rotated_imag = x_real * sin + x_imag * cos
-
-    # Reagrupar para [B, H, T, D]
-    rotated = torch.stack([rotated_real, rotated_imag], dim=-1).reshape(B, H, T, D)
-    return rotated.to(x.dtype)
-
-
-class RotaryPositionEmbedding(nn.Module):
-    """
-    Módulo PyTorch para RoPE — mantém o cache de frequências como buffer.
-
-    Uso:
-        rope = RotaryPositionEmbedding(head_dim=64, max_seq_len=512)
-        # Em forward:
-        q = rope(q)  # q: [batch, n_heads, seq_len, head_dim]
-        k = rope(k)
-    """
-
-    def __init__(self, head_dim: int, max_seq_len: int = 512, base: float = 10000.0):
-        super().__init__()
-        if head_dim % 2 != 0:
-            raise ValueError(f"head_dim deve ser par; recebido {head_dim}")
-        self.head_dim = head_dim
-        self.max_seq_len = max_seq_len
-        self.base = base
-
-        # Pré-computar e registrar como buffer (move com .to(device))
-        cos_sin = precompute_rope_frequencies(head_dim, max_seq_len, base)
-        self.register_buffer("cos_sin", cos_sin, persistent=False)
-
-    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
-        """
-        Aplica RoPE a x.
-
-        Args:
-            x: [batch, n_heads, seq_len, head_dim]
-            offset: deslocamento posicional (para geração autoregressiva com cache KV,
-                    onde a posição do novo token é len(cached) + i)
-
-        Returns:
-            x_rotated: mesmo shape que x
-        """
-        if x.dim() != 4:
-            raise ValueError(
-                f"x deve ser [batch, n_heads, seq_len, head_dim]; recebido {x.shape}"
-            )
-        T = x.size(2)
-        if offset + T > self.max_seq_len:
-            # Estender o cache dinamicamente
-            self._extend_cache(offset + T)
-        # Slice das posições relevantes
-        cos_sin = self.cos_sin[offset:offset + T]
-        return apply_rope(x, cos_sin)
-
-    def _extend_cache(self, new_max: int) -> None:
-        """Estende o cache de frequências para new_max posições."""
-        new_max_pow2 = 1
-        while new_max_pow2 < new_max:
-            new_max_pow2 *= 2
-        new_cos_sin = precompute_rope_frequencies(
-            self.head_dim, new_max_pow2, self.base,
-            device=self.cos_sin.device, dtype=self.cos_sin.dtype
-        )
-        self.cos_sin = new_cos_sin.to(self.cos_sin.device)
-        self.max_seq_len = new_max_pow2
-        # Re-registrar como buffer (substitui o anterior)
-        self.register_buffer("cos_sin", self.cos_sin, persistent=False)
-
-
-__all__ = [
-    "precompute_rope_frequencies",
-    "apply_rope",
-    "RotaryPositionEmbedding",
-]