src/easytranslate/model/positional.py

"""
位置编码模块 — Person B 负责实现

包含:
1. SinusoidalPositionalEncoding: 经典正弦余弦位置编码
2. RotaryPositionalEmbedding (RoPE): 旋转位置编码

RoPE 是当前最前沿的位置编码方案，被 LLaMA, GPT-NeoX 等大模型广泛采用。
优势: 更好地捕获相对位置信息，支持长度外推。
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


class SinusoidalPositionalEncoding(nn.Module):
    """
    经典正弦余弦位置编码 (Vaswani et al., 2017)。

    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

    TODO [Person B]: 实现以下内容:
    1. 在 __init__ 中预计算位置编码矩阵 [max_seq_len, d_model]
    2. 注册为 buffer (不参与梯度更新)
    3. forward(x): 返回 x + pe[:, :x.size(1), :]
    """

    def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, L, D]
        Returns:
            x + positional_encoding: [B, L, D]
        """
        x = x + self.pe[:, : x.size(1), :]
        return self.dropout(x)


class RotaryPositionalEmbedding(nn.Module):
    """
    旋转位置编码 (RoPE) — Su et al., 2021

    核心思想: 通过旋转变换将位置信息编码到 Q, K 向量中。
    q' = q * cos(θ) + rotate_half(q) * sin(θ)
    k' = k * cos(θ) + rotate_half(k) * sin(θ)

    TODO [Person B]: 实现以下内容:

    __init__:
    1. 计算频率: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
    2. 注册为 buffer

    _compute_rope(seq_len):
    1. 计算 position indices: t = [0, 1, ..., seq_len-1]
    2. 计算 freqs = torch.outer(t, inv_freq)
    3. 构建 cos_cached, sin_cached

    apply_rotary_pos_emb(q, k):
    1. 对 q 和 k 应用旋转变换
    2. 返回旋转后的 q', k'

    参考: https://arxiv.org/abs/2104.09864
    """

    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

        inv_freq = 1.0 / (
            base ** (torch.arange(0, dim, 2).float() / dim)
        )
        self.register_buffer("inv_freq", inv_freq)

        # 预缓存 cos/sin
        self._cached_seq_len = 0
        self._cached_cos: torch.Tensor | None = None
        self._cached_sin: torch.Tensor | None = None

    def _compute_rope(self, seq_len: int, device: torch.device):
        if seq_len > self._cached_seq_len or self._cached_cos is None:
            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
            freqs = torch.outer(t, self.inv_freq)  # [seq_len, dim//2]
            emb = torch.cat((freqs, freqs), dim=-1)  # [seq_len, dim]
            self._cached_cos = emb.cos()[None, None, :, :]  # [1, 1, seq_len, dim]
            self._cached_sin = emb.sin()[None, None, :, :]  # [1, 1, seq_len, dim]
            self._cached_seq_len = seq_len
        return self._cached_cos, self._cached_sin

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        """将 x 的后半部分取反并与前半部分交换。"""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def apply_rotary_pos_emb(
        self,
        q: torch.Tensor,  # [B, nhead, L, d_k]
        k: torch.Tensor,  # [B, nhead, L, d_k]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        对 Q, K 应用 RoPE。

        Returns:
            (q_rotated, k_rotated)
        """
        cos, sin = self._compute_rope(q.size(2), q.device)
        q_embed = (q * cos[:, :, : q.size(2), :]) + (
            self._rotate_half(q) * sin[:, :, : q.size(2), :]
        )
        k_embed = (k * cos[:, :, : k.size(2), :]) + (
            self._rotate_half(k) * sin[:, :, : k.size(2), :]
        )
        return q_embed, k_embed