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

49b8205 verified 3 months ago

6.57 kB

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

	"""
	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_0m), cos(theta_1m), ...,
	sin(theta_0m), sin(theta_1m), ...] 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",
	]