#位置编码 (Positional Encoding)：给没有顺序的矩阵注入"位置感"

import torch
import torch.nn as nn
import math

class SinCosPositionalEncoding(nn.Module):
    """原始 Transformer 论文中的正弦/余弦绝对位置编码"""
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        """
        Args:
            d_model: 每个词的向量维度（论文中为 512）
            max_len: 预设的最大序列长度
            dropout: 防止过拟合
        """
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 初始化一个全 0 的位置编码矩阵，形状 (max_len, d_model)
        position_enc = torch.zeros(max_len, d_model)

        # 位置索引序列 [0, 1, 2, ..., max_len-1] → 形状 (max_len, 1)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        # 频率衰减因子：不同维度使用不同的变化频率
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        # 偶数维用 sin，奇数维用 cos，实现不同位置向量的唯一性
        position_enc[:, 0::2] = torch.sin(position * div_term)
        position_enc[:, 1::2] = torch.cos(position * div_term)

        # 扩展为 (1, max_len, d_model)，方便与 batch 输入相加
        position_enc = position_enc.unsqueeze(0)

        # 注册为 buffer：不参与训练，但保存到模型文件中
        self.register_buffer("position_enc", position_enc)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 输入词向量，形状 (batch_size, seq_len, d_model)
        Returns:
            加上位置编码后的向量
        """
        # 只取当前序列长度对应的位置编码
        x = x + self.position_enc[:, :x.size(1), :]
        return self.dropout(x)

import torch
import torch.nn as nn

class LearnedPositionalEncoding(nn.Module):
    """端到端可学习的绝对位置编码"""
    def __init__(self, d_model: int, max_len: int = 5000):
        """
        Args:
            d_model: 向量维度
            max_len: 最大序列长度（推理时不能超过此值！）
        """
        super().__init__()
        # 使用 Embedding 存储可学习的位置向量
        self.position_emb = nn.Embedding(max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 输入词向量，形状 (batch_size, seq_len, d_model)
        Returns:
            加上位置编码后的向量
        """
        batch_size, seq_len, _ = x.shape

        # 生成位置索引，并在 batch 维度上复制
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1)

        # 取出对应的位置向量并相加
        position_vectors = self.position_emb(positions)
        return x + position_vectors

import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
    """现代 LLM 标配的 RoPE 旋转位置编码（核心简化版）"""
    def __init__(self, head_dim: int, base: int = 10000):
        """
        Args:
            head_dim: 每个注意力头的向量维度
            base: 频率衰减基数（沿用原始 Transformer 的 10000）
        """
        super().__init__()
        # 预计算频率衰减因子
        self.inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))

    def forward(self, x: torch.Tensor, seq_len: int = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: Q/K 向量，形状 (batch_size, num_heads, seq_len, head_dim)
            seq_len: 当前序列长度
        Returns:
            cos_emb, sin_emb: 用于旋转的余弦和正弦向量
        """
        if seq_len is None:
            seq_len = x.size(2)

        # 生成位置序列
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)

        # 计算每个位置、每个维度对的旋转角度
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)

        # 扩展成与 Q/K 一样的形状
        cos_emb = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
        sin_emb = torch.cat([freqs.sin(), freqs.sin()], dim=-1)

        return cos_emb, sin_emb