# RNN 的伪代码 — 注意这里是一个 for 循环，必须串行执行
hidden = zeros(hidden_size)
for word in sentence:      # 必须顺序处理每个词！
    hidden = rnn_cell(word, hidden)  # 当前词 + 上一步的隐藏状态
output = hidden  # 最终状态包含"全部信息"

"那只在公园里追蝴蝶的、毛色金黄的、带着红色项圈的小猫，最终还是没有抓到它。"

1. 拿着 Query（"apple"）和每个 Key 比较相似度
2. 发现 Key="apple" 完全匹配 → 得分最高
3. Key="apply" 有点像 → 得分较高
4. Key="banana" 完全不像 → 得分很低
5. 根据得分高低，加权获取对应的 Value
6. 最终结果：主要获取"苹果"的翻译，略带一点"应用"的信息

import torch

# 一个简单的例子：3个词，每个词用4维向量表示
query = torch.tensor([1.0, 0.5, 0.0, 0.3])      # "apple" 的向量
keys = torch.tensor([
    [1.0, 0.4, 0.1, 0.3],   # "apple" — 非常相似
    [0.8, 0.3, 0.2, 0.2],   # "apply" — 有点像
    [0.0, 0.1, 0.9, 0.7],   # "banana" — 完全不像
])

# 点积计算相似度
scores = torch.matmul(keys, query)
print(f"原始得分: {scores}")
# 输出类似: tensor([1.39, 1.01, 0.26])

Query（我要查的词 "apple"）:  [1.0, 0.5, 0.0, 0.3]

Key₁（"apple" 的索引）:      [1.0, 0.4, 0.1, 0.3]
Key₂（"apply" 的索引）:      [0.8, 0.3, 0.2, 0.2]
Key₃（"banana" 的索引）:     [0.0, 0.1, 0.9, 0.7]

  Q  ·  K₁  =  1.0×1.0  +  0.5×0.4  +  0.0×0.1  +  0.3×0.3
             =    1.0    +    0.2    +    0.0    +    0.09
             =    1.29

  Q  ·  K₂  =  1.0×0.8  +  0.5×0.3  +  0.0×0.2  +  0.3×0.2
             =    0.8    +    0.15   +    0.0    +    0.06
             =    1.01

  Q  ·  K₃  =  1.0×0.0  +  0.5×0.1  +  0.0×0.9  +  0.3×0.7
             =    0.0    +    0.05   +    0.0    +    0.21
             =    0.26

Score("apple") = 1.29  ← 最高！果然 "apple" 和 "apple" 最像
Score("apply") = 1.01  ← 有点像，但不如完全匹配
Score("banana") = 0.26 ← 得分很低，"banana" 和 "apple" 差别大

# 应用 Softmax
weights = torch.softmax(scores, dim=0)
print(f"注意力权重: {weights}")
# 输出类似: tensor([0.50, 0.34, 0.16])
# "apple" 获得最高权重，"banana" 最低

values = torch.tensor([
    [1.0, 0.0, 0.0],   # "苹果" 的表示
    [0.0, 1.0, 0.0],   # "应用" 的表示
    [0.0, 0.0, 1.0],   # "香蕉" 的表示
])

# 加权组合
output = torch.matmul(weights, values)
print(f"最终输出: {output}")
# 输出类似: tensor([0.50, 0.34, 0.16])
# 主要是"苹果"的信息，少量"应用"的信息

"动物没有穿过马路，因为它太累了。"

import torch
import torch.nn as nn

# 模拟输入：4个词，每个词3维嵌入
X = torch.tensor([
    [1.0, 0.0, 1.0],   # "I"
    [0.0, 1.0, 1.0],   # "love"
    [1.0, 1.0, 0.0],   # "deep"
    [0.0, 1.0, 0.0],   # "learning"
], dtype=torch.float32)

d_model = 3  # 嵌入维度

# 三个投影矩阵（实际中这些是可学习的参数）
W_Q = torch.tensor([[1, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=torch.float32)
W_K = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=torch.float32)
W_V = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], dtype=torch.float32)

# Step 1: 投影得到 Q, K, V
Q = X @ W_Q   # [4, 3] @ [3, 3] = [4, 3]
K = X @ W_K   # [4, 3] @ [3, 3] = [4, 3]
V = X @ W_V   # [4, 3] @ [3, 3] = [4, 3]
print(f"Q:\n{Q}")
print(f"K:\n{K}")
print(f"V:\n{V}")

# Step 2: 计算注意力分数 Q · K^T
scores = Q @ K.T   # [4, 3] @ [3, 4] = [4, 4]
print(f"注意力分数矩阵:\n{scores}")
# 这是一个 4×4 的矩阵：scores[i][j] 表示第 i 个词对第 j 个词的原始注意力分数

# Step 3: 缩放（下一节会详细解释为什么要缩放）
import math
d_k = d_model  # Key 的维度
scaled_scores = scores / math.sqrt(d_k)
print(f"缩放后的分数:\n{scaled_scores}")

# Step 4: Softmax — 每一行变成概率分布
attention_weights = torch.softmax(scaled_scores, dim=-1)
print(f"注意力权重:\n{attention_weights}")
# 每一行的和都等于 1.0
print(f"每行求和验证: {attention_weights.sum(dim=-1)}")

# Step 5: 加权组合 Value
output = attention_weights @ V   # [4, 4] @ [4, 3] = [4, 3]
print(f"注意力输出:\n{output}")
# 每一行是对所有 Value 向量的加权组合

输入 X（4个词 × 3维嵌入）:        权重矩阵 W_Q（3×3）:
         d₁   d₂   d₃                  d₁'  d₂'  d₃'
  "I"   [1.0, 0.0, 1.0]           d₁ [ 1,   0,   1 ]
  "love"[0.0, 1.0, 1.0]     ×     d₂ [ 0,   1,   0 ]
  "deep"[1.0, 1.0, 0.0]           d₃ [ 1,   0,   0 ]
  "learning"[0.0, 1.0, 0.0]

Q[0][0] = X[0][0]×W_Q[0][0] + X[0][1]×W_Q[1][0] + X[0][2]×W_Q[2][0]
        = 1.0 × 1  +  0.0 × 0  +  1.0 × 1
        = 1.0 + 0.0 + 1.0
        = 2.0

Q[0][1] = 1.0 × 0  +  0.0 × 1  +  1.0 × 0  =  0.0

Q[0][2] = 1.0 × 1  +  0.0 × 0  +  1.0 × 0  =  1.0

完整的 Q 矩阵：
         d₁'  d₂'  d₃'
  "I"   [2.0, 0.0, 1.0]   ← 1×1+0×0+1×1, 1×0+0×1+1×0, 1×1+0×0+1×0
  "love"[1.0, 1.0, 0.0]   ← 0×1+1×0+1×1, 0×0+1×1+1×0, 0×1+1×0+1×0
  "deep"[1.0, 1.0, 1.0]   ← 1×1+1×0+0×1, 1×0+1×1+0×0, 1×1+1×0+0×0
  "learning"[0.0, 1.0, 0.0] ← 0×1+1×0+0×1, 0×0+1×1+0×0, 0×1+1×0+0×0

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        # 三个线性层分别生成 Q, K, V
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        Q = self.W_Q(x)   # [batch, seq_len, d_model]
        K = self.W_K(x)
        V = self.W_V(x)

        # 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_model ** 0.5)
        weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(weights, V)

        return output, weights

# 测试
attn = SelfAttention(d_model=64)
x = torch.randn(2, 10, 64)  # batch=2, seq_len=10, d_model=64
output, weights = attn(x)
print(f"输出形状: {output.shape}")    # [2, 10, 64]
print(f"权重形状: {weights.shape}")    # [2, 10, 10]

import torch

d_k_values = [1, 4, 16, 64, 256, 512]

for d_k in d_k_values:
    q = torch.randn(1, d_k)
    k = torch.randn(5, d_k)

    # 未缩放
    raw_scores = q @ k.T
    raw_probs = torch.softmax(raw_scores, dim=-1)

    # 缩放后
    scaled_scores = raw_scores / (d_k ** 0.5)
    scaled_probs = torch.softmax(scaled_scores, dim=-1)

    print(f"d_k={d_k:>3d} | 未缩放最大权重: {raw_probs.max():.4f} | "
          f"缩放后最大权重: {scaled_probs.max():.4f}")

d_k=  1 | 未缩放最大权重: 0.3500 | 缩放后最大权重: 0.3500
d_k=  4 | 未缩放最大权重: 0.5200 | 缩放后最大权重: 0.3800
d_k= 16 | 未缩放最大权重: 0.7800 | 缩放后最大权重: 0.3500
d_k= 64 | 未缩放最大权重: 0.9500 | 缩放后最大权重: 0.3200
d_k=256 | 未缩放最大权重: 0.9990 | 缩放后最大权重: 0.3600
d_k=512 | 未缩放最大权重: 0.9999 | 缩放后最大权重: 0.3400

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度

        # 四个投影矩阵
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # 1. 线性投影
        Q = self.W_Q(Q)  # [batch, seq_len, d_model]
        K = self.W_K(K)
        V = self.W_V(V)

        # 2. 切分成多个头: [batch, seq, d_model] → [batch, n_heads, seq, d_k]
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention_weights = torch.softmax(scores, dim=-1)
        context = torch.matmul(attention_weights, V)

        # 4. 拼接所有头: [batch, n_heads, seq, d_k] → [batch, seq, d_model]
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 5. 输出投影
        output = self.W_O(context)

        return output, attention_weights

# 测试
mha = MultiHeadAttention(d_model=512, n_heads=8)
x = torch.randn(2, 10, 512)  # batch=2, seq_len=10
output, weights = mha(x, x, x)  # 自注意力：Q=K=V=x
print(f"输出形状: {output.shape}")     # [2, 10, 512]
print(f"权重形状: {weights.shape}")     # [2, 8, 10, 10] — 8 个头的注意力矩阵

"猫 追 狗"  →  猫在追狗
"狗 追 猫"  →  狗在追猫

# 方案 1：线性编码（有问题！）
pe = torch.arange(0, seq_len).float()
# 位置 0: [0, 0, 0, ...]
# 位置 1: [1, 1, 1, ...]
# 位置 100: [100, 100, 100, ...]

import torch
import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        # 频率递减：低维度高频，高维度低频
        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)   # 偶数维度用 sin
        pe[:, 1::2] = torch.cos(position * div_term)    # 奇数维度用 cos

        pe = pe.unsqueeze(0)   # 增加 batch 维度: [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len, :]

# 使用示例
d_model = 128
pe = PositionalEncoding(d_model)
x = torch.randn(2, 20, d_model)  # 2 个句子，每句 20 个词
output = pe(x)
print(f"输入形状: {x.shape}")     # [2, 20, 128]
print(f"输出形状: {output.shape}") # [2, 20, 128] — 加了位置编码但形状不变

输入 x
  ↓
[Multi-Head Self-Attention]  ← 全局信息交互
  ↓
[Add & Norm]                 ← 残差连接 + 层归一化
  ↓
[Feed-Forward Network]       ← 逐位置深度加工
  ↓
[Add & Norm]                 ← 残差连接 + 层归一化
  ↓
输出

import torch
import torch.nn as nn
import math

class FeedForward(nn.Module):
    """前馈网络：两层线性变换 + ReLU"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(self.relu(self.linear1(x))))


class EncoderLayer(nn.Module):
    """单个 Encoder 层"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 子层 1: Multi-Head Self-Attention + Add & Norm
        attn_output, _ = self.self_attention(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))  # 残差 + 归一化

        # 子层 2: Feed-Forward + Add & Norm
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout2(ff_output))     # 残差 + 归一化

        return x


class Encoder(nn.Module):
    """完整 Encoder：词嵌入 + 位置编码 + N 个 Encoder 层"""
    def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len=5000, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 词嵌入 + 缩放 + 位置编码
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.pos_encoding(x)
        x = self.dropout(x)

        # 通过 N 个 Encoder 层
        for layer in self.layers:
            x = layer(x, mask)

        return x

# 测试
encoder = Encoder(
    vocab_size=10000, d_model=512, n_heads=8,
    d_ff=2048, n_layers=6
)
src = torch.randint(0, 10000, (2, 20))  # batch=2, seq_len=20
output = encoder(src)
print(f"Encoder 输出: {output.shape}")  # [2, 20, 512]

完整 Encoder 数据流:

Token IDs  →  Embedding  →  × √d_model  →  + PositionalEncoding
                                                      ↓
                                              EncoderLayer × N
                                              ┌─────────────────┐
                                              │ Self-Attention   │
                                              │ Add & Norm       │
                                              │ Feed-Forward     │
                                              │ Add & Norm       │
                                              └─────────────────┘
                                                      ↓
                                              Encoder Output
                                          [batch, seq_len, d_model]

输入 x（已生成的词）
  ↓
[Masked Multi-Head Self-Attention]  ← 只看已生成的词，不能偷看未来
  ↓
[Add & Norm]
  ↓
[Multi-Head Cross-Attention]        ← Q 来自 Decoder，K/V 来自 Encoder
  ↓
[Add & Norm]
  ↓
[Feed-Forward Network]
  ↓
[Add & Norm]
  ↓
输出

def create_causal_mask(seq_len):
    """创建因果掩码：上三角为 True（将被遮住）"""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
    return mask

# 对于 seq_len=5:
# [[0, 1, 1, 1, 1],    ← 位置 0 只能看位置 0
#  [0, 0, 1, 1, 1],    ← 位置 1 能看位置 0, 1
#  [0, 0, 0, 1, 1],    ← 位置 2 能看位置 0, 1, 2
#  [0, 0, 0, 0, 1],    ← 位置 3 能看位置 0, 1, 2, 3
#  [0, 0, 0, 0, 0]]    ← 位置 4 能看位置 0, 1, 2, 3, 4

class DecoderLayer(nn.Module):
    """单个 Decoder 层"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        # 三个子层
        self.masked_self_attention = MultiHeadAttention(d_model, n_heads)
        self.cross_attention = MultiHeadAttention(d_model, n_heads)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)
        # 三个 LayerNorm
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        # 三个 Dropout
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
        # 子层 1: Masked Self-Attention
        attn1, _ = self.masked_self_attention(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout1(attn1))

        # 子层 2: Cross-Attention（Q=decoder, K=V=encoder）
        attn2, _ = self.cross_attention(x, encoder_output, encoder_output, src_mask)
        x = self.norm2(x + self.dropout2(attn2))

        # 子层 3: Feed-Forward
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout3(ff_output))

        return x


class Decoder(nn.Module):
    """完整 Decoder"""
    def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len=5000, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len)
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.pos_encoding(x)
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(x, encoder_output, src_mask, tgt_mask)

        return x

输入: [<BOS>, 我, 爱, 自然, 语言]
目标: [我, 爱, 自然, 语言, <EOS>]

Step 1: 输入 [<BOS>]           → 预测 "我"
Step 2: 输入 [<BOS>, 我]       → 预测 "爱"
Step 3: 输入 [<BOS>, 我, 爱]   → 预测 "自然"
...
Step N: 输入 [<BOS>, 我, 爱, ... ] → 预测 <EOS> → 停止

源语言 Token IDs                                目标语言 Token IDs（右移一位）
      ↓                                                ↓
  [Embedding]                                     [Embedding]
      ↓                                                ↓
  [+ Pos Encoding]                              [+ Pos Encoding]
      ↓                                                ↓
┌─────────────────┐                         ┌──────────────────────┐
│  Encoder × N    │                         │  Decoder × N         │
│                 │    K, V                  │                      │
│  Self-Attn      │────────────────────────→│  Masked Self-Attn    │
│  Add & Norm     │                         │  Add & Norm          │
│  FFN            │                         │  Cross-Attn (Q←Dec)  │
│  Add & Norm     │                         │  Add & Norm          │
└─────────────────┘                         │  FFN                 │
                                            │  Add & Norm          │
                                            └──────────────────────┘
                                                      ↓
                                               [Linear Layer]
                                                      ↓
                                                 [Softmax]
                                                      ↓
                                             Output Probabilities

class Transformer(nn.Module):
    """完整的 Transformer 模型"""
    def __init__(
        self,
        src_vocab_size,    # 源语言词表大小
        tgt_vocab_size,    # 目标语言词表大小
        d_model=512,       # 嵌入维度
        n_heads=8,         # 注意力头数
        d_ff=2048,         # FFN 隐藏层维度
        n_encoder_layers=6,
        n_decoder_layers=6,
        max_len=5000,
        dropout=0.1,
    ):
        super().__init__()
        self.encoder = Encoder(src_vocab_size, d_model, n_heads, d_ff,
                              n_encoder_layers, max_len, dropout)
        self.decoder = Decoder(tgt_vocab_size, d_model, n_heads, d_ff,
                              n_decoder_layers, max_len, dropout)
        # 最终输出层：d_model → vocab_size
        self.output_projection = nn.Linear(d_model, tgt_vocab_size)

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        # 1. Encoder 处理源序列
        encoder_output = self.encoder(src, src_mask)

        # 2. Decoder 处理目标序列（参考 Encoder 输出）
        decoder_output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)

        # 3. 映射到词表大小的 logits
        logits = self.output_projection(decoder_output)

        return logits

    def generate(self, src, max_len=50, start_token=1, end_token=2):
        """贪心解码生成"""
        self.eval()
        with torch.no_grad():
            encoder_output = self.encoder(src)

            # 从 <BOS> 开始
            tgt = torch.full((src.size(0), 1), start_token, dtype=torch.long, device=src.device)

            for _ in range(max_len):
                # 创建因果掩码
                tgt_mask = create_causal_mask(tgt.size(1)).to(src.device)

                decoder_output = self.decoder(tgt, encoder_output, tgt_mask=tgt_mask)
                logits = self.output_projection(decoder_output[:, -1, :])  # 只取最后一个位置
                next_token = logits.argmax(dim=-1, keepdim=True)

                tgt = torch.cat([tgt, next_token], dim=1)

                # 如果所有样本都生成了 <EOS>，停止
                if (next_token == end_token).all():
                    break

            return tgt

# 创建模型
model = Transformer(
    src_vocab_size=10000,
    tgt_vocab_size=8000,
    d_model=512,
    n_heads=8,
    d_ff=2048,
    n_encoder_layers=6,
    n_decoder_layers=6,
)

# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
print(f"总参数量: {total_params:,}")  # 约 65M 参数（类似原始论文的 base 模型）

import torch.optim as optim

# 训练配置
optimizer = optim.Adam(model.parameters(), lr=1e-4, betas=(0.9, 0.98), eps=1e-9)
criterion = nn.CrossEntropyLoss(ignore_index=0)  # 忽略 <PAD> 的损失

def train_step(model, src, tgt, optimizer, criterion):
    model.train()
    optimizer.zero_grad()

    # tgt_input: 去掉最后一个 token（作为输入）
    # tgt_output: 去掉第一个 token（作为目标）
    tgt_input = tgt[:, :-1]
    tgt_output = tgt[:, 1:]

    # 创建掩码
    tgt_mask = create_causal_mask(tgt_input.size(1)).to(src.device)

    # 前向传播
    logits = model(src, tgt_input, tgt_mask=tgt_mask)

    # 计算损失
    loss = criterion(
        logits.reshape(-1, logits.size(-1)),  # [batch*seq, vocab]
        tgt_output.reshape(-1),                # [batch*seq]
    )

    # 反向传播
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

    return loss.item()

# BERT 风格的预训练
# 输入: "我 [MASK] 自然语言 [MASK]"
# 目标: 预测 [MASK] 位置的词 → "爱", "处理"

# MLM 的掩码策略（80/10/10 规则）：
# 被选中的 15% 的 token 中：
# - 80% 替换为 [MASK]
# - 10% 替换为随机 token
# - 10% 保持不变

# GPT 风格的预训练
# 输入:   "我  爱  自然  语言"
# 目标:   "爱 自然  语言  处理"
# 位置 0: 看到 [我]         → 预测 "爱"
# 位置 1: 看到 [我, 爱]     → 预测 "自然"
# 位置 2: 看到 [我, 爱, 自然] → 预测 "语言"

# T5 风格的训练
# 输入: "我 <X> 语言处理"
# 目标: "<X> 爱 自然"
# 其中 <X> 是被遮住的连续片段的标记

Decoder-Only（GPT）训练效率：
输入:   "我 爱 自然 语言 处理"
目标:   "爱 自然 语言 处理 <EOS>"
→ 每个位置都是一个训练样本！5个token = 5个训练信号

Encoder-Only（BERT）训练效率：
输入:   "我 [MASK] 自然 [MASK] 处理"
目标:   预测 "爱" 和 "语言"
→ 只有被MASK的位置有训练信号！5个token = 2个训练信号

原始 Transformer:  Multi-Head Attention (MHA)  → 每个头独立的 K/V
     ↓
LLaMA / Mistral:  Grouped-Query Attention (GQA) → 多个 Q 头共享 K/V
     ↓
DeepSeek V3:      Multi-Head Latent Attention (MLA) → 压缩 K/V 到低维再还原

DeepSeek V3:  256 个路由专家 + 1 个共享专家，每个 token 激活 Top-8
              → 671B 总参数，但每个 token 只用 37B → 10x 计算效率提升

原始 Transformer: sin/cos 绝对位置编码 → 外推能力差
     ↓
LLaMA / 大部分现代 LLM: RoPE（旋转位置编码） → 更好的相对距离建模
     ↓
+ NTK-aware / YaRN 扩展 → 支持 100K+ 上下文

import torch
import torch.nn as nn
import math


class PositionalEncoding(nn.Module):
    """位置编码：sin/cos 固定编码"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        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)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        Q = self.W_Q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = self.dropout(torch.softmax(scores, dim=-1))
        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.W_O(context)


class FeedForward(nn.Module):
    """前馈网络"""
    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        return self.net(x)


class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = self.norm1(x + self.dropout1(self.self_attn(x, x, x, mask)))
        x = self.norm2(x + self.dropout2(self.ffn(x)))
        return x


class DecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        x = self.norm1(x + self.dropout1(self.self_attn(x, x, x, tgt_mask)))
        x = self.norm2(x + self.dropout2(self.cross_attn(x, enc_output, enc_output, src_mask)))
        x = self.norm3(x + self.dropout3(self.ffn(x)))
        return x


class Transformer(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, d_model=512, n_heads=8,
                 d_ff=2048, n_enc_layers=6, n_dec_layers=6,
                 max_len=5000, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Embeddings + Positional Encoding
        self.src_embed = nn.Embedding(src_vocab, d_model)
        self.tgt_embed = nn.Embedding(tgt_vocab, d_model)
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)

        # Encoder & Decoder stacks
        self.encoder_layers = nn.ModuleList(
            [EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_enc_layers)]
        )
        self.decoder_layers = nn.ModuleList(
            [DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_dec_layers)]
        )

        # Output projection
        self.output_proj = nn.Linear(d_model, tgt_vocab)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def encode(self, src, src_mask=None):
        x = self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))
        for layer in self.encoder_layers:
            x = layer(x, src_mask)
        return x

    def decode(self, tgt, enc_output, src_mask=None, tgt_mask=None):
        x = self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))
        for layer in self.decoder_layers:
            x = layer(x, enc_output, src_mask, tgt_mask)
        return x

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        enc_output = self.encode(src, src_mask)
        dec_output = self.decode(tgt, enc_output, src_mask, tgt_mask)
        return self.output_proj(dec_output)


# ===== 工具函数 =====
def create_causal_mask(size):
    """创建因果掩码（下三角为 1，上三角为 0）"""
    return torch.tril(torch.ones(size, size)).unsqueeze(0).unsqueeze(0)

def create_padding_mask(seq, pad_idx=0):
    """创建 padding 掩码"""
    return (seq != pad_idx).unsqueeze(1).unsqueeze(2)

# 创建模型
model = Transformer(src_vocab=5000, tgt_vocab=5000, d_model=256, n_heads=8, d_ff=1024, n_enc_layers=3, n_dec_layers=3)

# 模拟输入
src = torch.randint(1, 5000, (4, 20))  # batch=4, src_len=20
tgt = torch.randint(1, 5000, (4, 15))  # batch=4, tgt_len=15

# 创建掩码
src_mask = create_padding_mask(src)
tgt_mask = create_causal_mask(tgt.size(1)) & create_padding_mask(tgt)

# 前向传播
logits = model(src, tgt, src_mask, tgt_mask)
print(f"输出 logits 形状: {logits.shape}")  # [4, 15, 5000]

# 参数统计
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数量: {total:,}")
print(f"可训练参数量: {trainable:,}")