Transformer 架构与 self-attention

"""手算 self-attention：三个词、二维向量，每一步都能拿计算器核对。

故意不用随机数，Q/K/V 全是写死的小整数，这样每一个中间结果
都能自己推一遍，确认 attention 没有任何魔法。
"""
import math
import torch
import torch.nn.functional as F


def main():
    d_k = 2

    # 三个词，每个词 2 维。这里直接给出 Q/K/V，跳过 x @ Wq 那一步
    Q = torch.tensor([[1.0, 0.0],      # 词0 的查询
                      [0.0, 1.0],      # 词1 的查询
                      [1.0, 1.0]])     # 词2 的查询
    K = torch.tensor([[1.0, 0.0],      # 词0 的标签
                      [0.0, 1.0],
                      [1.0, 1.0]])
    V = torch.tensor([[10.0, 0.0],     # 词0 的内容
                      [0.0, 10.0],
                      [5.0, 5.0]])

    # ------------------------------------------------ 第一步：打分
    scores = Q @ K.T
    print('原始分数表 scores = Q @ Kᵀ')
    print(scores)
    # 手算第 0 行：词0 的 Q=[1,0] 分别点乘三个 K
    #   与 K0=[1,0] → 1*1 + 0*0 = 1
    #   与 K1=[0,1] → 1*0 + 0*1 = 0
    #   与 K2=[1,1] → 1*1 + 0*1 = 1
    # 所以第 0 行是 [1, 0, 1]，和打印结果一致

    # ------------------------------------------------ 第二步：缩放
    scaled = scores / math.sqrt(d_k)
    print('\n除以 √d_k = %.4f 之后' % math.sqrt(d_k))
    print(scaled)
    # 第 0 行变成 [0.7071, 0, 0.7071]

    # ------------------------------------------------ 第三步：softmax
    weights = F.softmax(scaled, dim=-1)
    print('\n注意力权重（每行和为 1）')
    print(weights)
    print('每行之和', weights.sum(dim=-1))
    # 手算第 0 行：exp(0.7071)=2.0281, exp(0)=1, exp(0.7071)=2.0281
    #   分母 = 2.0281 + 1 + 2.0281 = 5.0562
    #   权重 = [0.4011, 0.1978, 0.4011]

    # ------------------------------------------------ 第四步：加权求和
    out = weights @ V
    print('\n输出 = 权重 @ V')
    print(out)
    # 手算第 0 行第 0 列：0.4011*10 + 0.1978*0 + 0.4011*5 = 6.0165
    #            第 1 列：0.4011*0 + 0.1978*10 + 0.4011*5 = 3.9835
    # 两列相加正好 10 —— 因为 V 每行之和都是 10，而权重和为 1，
    # 「加权平均」的性质在这里体现得非常直白

    print('\n输出每行之和（应当都是 10）', out.sum(dim=-1))

    # ------------------------------------------------ 形状不变
    print('\n输入 Q 形状', Q.shape, '→ 输出形状', out.shape)
    # 进出同形是 attention 能无限堆层的前提：
    # 一层的输出可以直接当下一层的输入，不需要任何适配层

    # ------------------------------------------------ 对称性观察
    print('\nscores 是对称的吗:', torch.allclose(scores, scores.T))
    # 这个例子里 Q 和 K 恰好相同，所以分数表对称；
    # 真实模型里 Wq 和 Wk 是两套不同的权重，分数表一般不对称：
    # 「A 关注 B 的程度」不等于「B 关注 A 的程度」


if __name__ == '__main__':
    main()

"""为什么要除以 √d_k：不除会发生什么，用数字看清楚。

点积 q·k 是 d_k 个乘积之和。q、k 的每一维若是均值 0、方差 1 的独立随机数，
那么点积的方差就是 d_k，标准差是 √d_k —— 维度越大，分数的绝对值越大。
分数一大，softmax 就被推到饱和区：一个位置接近 1，其余接近 0，
梯度几乎为 0，这一层学不动。除以 √d_k 正好把标准差拉回 1。
"""
import math
import torch
import torch.nn.functional as F


def stats(d_k, n=20000, seed=0):
    """随机造 n 对向量，统计点积的标准差。"""
    g = torch.Generator().manual_seed(seed)
    q = torch.randn(n, d_k, generator=g)
    k = torch.randn(n, d_k, generator=g)
    dot = (q * k).sum(dim=-1)
    return dot.std().item()


def main():
    print('%-8s %-16s %-16s %s' % ('d_k', '点积标准差(实测)', '理论值 √d_k', '除以√d_k后'))
    for d_k in (4, 16, 64, 256, 1024):
        s = stats(d_k)
        print('%-8d %-16.3f %-16.3f %.3f'
              % (d_k, s, math.sqrt(d_k), s / math.sqrt(d_k)))
    # 实测值紧贴理论值 √d_k，最后一列全部回到 1 附近，
    # 说明缩放这一步不是经验参数，是被方差公式定死的

    # ---------------------------------------------------- softmax 饱和
    print('\n同一组分数，缩放前后的 softmax 对比（d_k = 64）：')
    d_k = 64
    g = torch.Generator().manual_seed(1)
    q = torch.randn(1, d_k, generator=g)
    k = torch.randn(6, d_k, generator=g)

    raw = q @ k.T                       # 不缩放
    scaled = raw / math.sqrt(d_k)       # 缩放

    print('未缩放分数 :', [round(v, 2) for v in raw[0].tolist()])
    print('未缩放权重 :', [round(v, 4) for v in F.softmax(raw, dim=-1)[0].tolist()])
    print('缩放后分数 :', [round(v, 2) for v in scaled[0].tolist()])
    print('缩放后权重 :', [round(v, 4) for v in F.softmax(scaled, dim=-1)[0].tolist()])

    # ---------------------------------------------------- 梯度对照
    # softmax 的雅可比对角元是 p(1-p)：p 越接近 0 或 1，梯度越接近 0
    for name, s in (('未缩放', raw), ('缩放后', scaled)):
        p = F.softmax(s, dim=-1)
        grad_scale = (p * (1 - p)).max().item()
        print('%s 的最大 p(1-p) = %.6f' % (name, grad_scale))
    # 未缩放那一行的数值明显更小 —— 这就是「梯度消失在 softmax 里」的样子

    print('\n一句话：√d_k 是把分数的标准差归一到 1，让 softmax 停在敏感区。')
    print('注意分母是 √d_k（每个头的维度），不是 d_model。')
    print('多头时 d_k = d_model / n_head，这两个数不一样，写错会让缩放失效。')


if __name__ == '__main__':
    main()

"""多头注意力的拆分与合并：形状怎么变，为什么要分头。

一句话：分头不是把模型变大，是把同样宽的向量切成几段，
让每一段各自学一种关系（谁修饰谁、谁指代谁、主谓搭配……），
最后再拼回原宽度。参数量和单头几乎一样，表达力却分了工。
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=64, n_head=8):
        super().__init__()
        assert d_model % n_head == 0, 'd_model 必须能被 n_head 整除'
        self.d_model = d_model
        self.n_head = n_head
        self.d_k = d_model // n_head          # 每个头分到的维度

        # 四个线性层：三个造 Q/K/V，一个把拼接结果再混一次
        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)

    def split_heads(self, x):
        """(batch, seq, d_model) → (batch, n_head, seq, d_k)"""
        batch, seq, _ = x.shape
        # 先把最后一维切成 n_head 段
        x = x.view(batch, seq, self.n_head, self.d_k)
        # 再把「头」提到 seq 前面，让每个头拿到一份完整的 (seq, d_k)
        return x.transpose(1, 2)

    def merge_heads(self, x):
        """(batch, n_head, seq, d_k) → (batch, seq, d_model)"""
        batch, n_head, seq, d_k = x.shape
        x = x.transpose(1, 2)                 # 头换回去
        # transpose 之后内存不连续，必须 contiguous 才能 view
        return x.contiguous().view(batch, seq, n_head * d_k)

    def forward(self, x, mask=None):
        Q = self.split_heads(self.w_q(x))
        K = self.split_heads(self.w_k(x))
        V = self.split_heads(self.w_v(x))
        print('  拆分后 Q 形状', tuple(Q.shape))

        # 注意分母是 d_k（每个头的维度），不是 d_model
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask, float('-inf'))
        weights = F.softmax(scores, dim=-1)
        print('  每个头各自的权重表', tuple(weights.shape))

        out = weights @ V                     # (batch, n_head, seq, d_k)
        out = self.merge_heads(out)
        print('  合并后形状', tuple(out.shape))
        return self.w_o(out), weights


def main():
    torch.manual_seed(0)
    batch, seq, d_model, n_head = 2, 5, 64, 8

    x = torch.randn(batch, seq, d_model)
    mha = MultiHeadAttention(d_model, n_head)
    print('输入形状', tuple(x.shape))
    out, w = mha(x)
    print('输出形状', tuple(out.shape), '→ 进出同形，可以直接堆下一层')

    # ---------------------------------------------------- 参数量对照
    single = MultiHeadAttention(d_model, n_head=1)
    p_multi = sum(p.numel() for p in mha.parameters())
    p_single = sum(p.numel() for p in single.parameters())
    print('\n8 头参数量 :', p_multi)
    print('1 头参数量 :', p_single)
    print('是否相等   :', p_multi == p_single)
    # 完全相等：4 个 (d_model × d_model) 的线性层 + 偏置，和头数无关。
    # 所以「多头」买到的是分工，不是容量。

    print('\n每个头分到的维度 d_k =', d_model // n_head)
    print('d_model 不能被 n_head 整除时 view 会直接报错，这是选参数的硬约束。')


if __name__ == '__main__':
    main()

"""位置编码：用正弦余弦给每个位置生成一个固定的「坐标指纹」。

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


def positional_encoding(max_len, d_model):
    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len).unsqueeze(1).float()     # (max_len, 1)

    # 偶数下标 0, 2, 4 ... 对应公式里的 2i
    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
    return pe


if __name__ == '__main__':
    # 课堂例子：d_model = 4，看第 2 个位置（pos=2）的编码向量
    pe = positional_encoding(max_len=6, d_model=4)
    print('pos=2 的位置编码', pe[2])
    # 手算对照：
    #   i=0: sin(2/10000^0) = sin(2)     ≈  0.909
    #        cos(2/10000^0) = cos(2)     ≈ -0.416
    #   i=1: sin(2/10000^0.5) = sin(0.02) ≈ 0.020
    #        cos(2/10000^0.5) = cos(0.02) ≈ 0.998

    # 位置编码是加到词向量上，不是拼接
    d_model, seq_len = 4, 6
    word_embed = torch.randn(seq_len, d_model)
    x = word_embed + pe[:seq_len]
    print('加完位置编码后的形状', x.shape)      # 和词向量同形，(6, 4)

"""LayerNorm 与 BatchNorm 的分界：沿哪个方向求均值方差。

BatchNorm：对同一个特征、跨整批样本求均值方差 —— 「这一列的所有人」
LayerNorm：对同一个样本、跨它自己的所有特征求 —— 「这一个人的所有科目」
Transformer 用 LayerNorm，原因全在下面第三段。
"""
import torch
import torch.nn as nn


def main():
    torch.manual_seed(0)
    batch, seq_len, d_model = 2, 4, 6
    x = torch.randn(batch, seq_len, d_model)

    # ---------------------------------------------------------- LayerNorm
    ln = nn.LayerNorm(d_model)          # 参数是「最后几维的形状」
    y_ln = ln(x)
    print('LayerNorm 输出形状', tuple(y_ln.shape))

    # 手算核对：对最后一维求均值方差，逐个位置独立归一化
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, unbiased=False, keepdim=True)   # 注意是有偏方差
    manual = (x - mean) / torch.sqrt(var + ln.eps)
    print('与手算是否一致 :', torch.allclose(y_ln, manual, atol=1e-5))
    print('归一化后每个位置的均值 ≈ 0 :',
          y_ln.mean(dim=-1).abs().max().item() < 1e-5)

    # ---------------------------------------------------------- BatchNorm
    bn = nn.BatchNorm1d(d_model)
    # BatchNorm1d 要 (batch, feature, seq)，所以必须先把维度换过来
    y_bn = bn(x.transpose(1, 2)).transpose(1, 2)
    print('\nBatchNorm 输出形状', tuple(y_bn.shape))
    # 它算的是：对每个特征通道，跨 batch × seq 所有位置求一个均值方差

    # ---------------------------------------------------------- 关键差别
    print('\n=== 改变 batch 里的其他句子，会不会影响第 0 句 ===')
    x2 = x.clone()
    x2[1] = torch.randn(seq_len, d_model) * 10      # 只动第 1 句，动得很夸张

    ln_same = torch.allclose(ln(x)[0], ln(x2)[0], atol=1e-6)
    bn_same = torch.allclose(bn(x.transpose(1, 2)).transpose(1, 2)[0],
                             bn(x2.transpose(1, 2)).transpose(1, 2)[0],
                             atol=1e-6)
    print('LayerNorm 第 0 句不受影响 :', ln_same)     # True
    print('BatchNorm 第 0 句不受影响 :', bn_same)     # False

    # 这就是选 LayerNorm 的第一个理由：
    #   每个位置的归一化只用它自己的数，和同批别的句子无关，
    #   推理时来一条也好来一百条也好，结果完全一样。

    # ---------------------------------------------------------- 变长的坑
    print('\n=== 变长序列 ===')
    print('BatchNorm 的统计量会把 pad 出来的位置一起算进去，')
    print('同一句话在不同 batch 里、补的长度不同，归一化结果就不同。')
    print('LayerNorm 逐位置独立，pad 多少都不影响真实词的表示。')

    # ---------------------------------------------------------- 可学参数
    print('\nLayerNorm 的可学参数 :',
          [tuple(p.shape) for p in ln.parameters()])
    # weight 和 bias 各 d_model 个：归一化之后再放缩平移，
    # 免得把网络好不容易学到的尺度信息抹平
    print('BatchNorm 还额外维护 running_mean / running_var，')
    print('所以它区分 train() 和 eval() 两种模式；LayerNorm 两种模式行为一致。')

    # ---------------------------------------------------------- 残差
    print('\n残差那一步写作 LayerNorm(x + Sublayer(x))：')
    print('先相加再归一化。加号让梯度有一条不经过子层的直通路，')
    print('归一化把相加后被放大的数值拉回稳定区间，两者缺一不可。')


if __name__ == '__main__':
    main()

"""算力账：注意力的开销为什么随序列长度平方增长。

注意力分数表是 (seq_len × seq_len) 的方阵，每个头一张、每层一张。
序列翻一倍，这张表变四倍 —— 这是 Transformer 唯一的结构性瓶颈，
也是长上下文模型要专门做优化的原因。
本文件只用标准库，跑出来的数字就是下面公式的直接结果。
"""


def attn_matrix_elems(seq_len, n_head, n_layer):
    """所有层、所有头的注意力分数表一共多少个元素。"""
    return n_layer * n_head * seq_len * seq_len


def attn_flops(seq_len, d_model, n_layer):
    """注意力部分的乘加次数（粗估）。

    QKᵀ 是 seq×seq×d_model，权重乘 V 又是一遍，所以是 2 倍。
    多头只是把 d_model 切开分算，总量不变，所以公式里不出现 n_head。
    """
    return n_layer * 2 * seq_len * seq_len * d_model


def ffn_flops(seq_len, d_model, d_ff, n_layer):
    """前馈网络部分：和序列长度成正比，不是平方。"""
    return n_layer * 2 * seq_len * d_model * d_ff


def human(n):
    for unit in ('', 'K', 'M', 'G', 'T'):
        if abs(n) < 1000:
            return '%.2f%s' % (n, unit)
        n /= 1000.0
    return '%.2fP' % n


def main():
    d_model, n_head, n_layer, d_ff = 768, 12, 12, 3072
    bytes_per_elem = 2                 # 半精度存储，一个数占 2 字节

    print('配置：d_model=%d  n_head=%d  n_layer=%d  d_ff=%d'
          % (d_model, n_head, n_layer, d_ff))
    print()
    print('%-10s %-14s %-14s %-14s %s'
          % ('seq_len', '分数表元素数', '分数表显存', '注意力乘加', 'FFN 乘加'))
    print('-' * 74)

    for seq_len in (128, 512, 1024, 2048, 4096, 8192):
        elems = attn_matrix_elems(seq_len, n_head, n_layer)
        mem = elems * bytes_per_elem
        a = attn_flops(seq_len, d_model, n_layer)
        f = ffn_flops(seq_len, d_model, d_ff, n_layer)
        print('%-10d %-14s %-14s %-14s %s'
              % (seq_len, human(elems), human(mem) + 'B', human(a), human(f)))

    print()
    print('读法：')
    print('  seq_len 每翻一倍，分数表元素数变四倍 —— 显存先撑不住的通常是它')
    print('  注意力乘加也是平方项，FFN 是线性项')
    print('  短序列时 FFN 占大头，长序列时注意力反超，交叉点就在 seq_len = d_ff')

    # 两式相除：attn / ffn = (seq·seq·d_model) / (seq·d_model·d_ff) = seq / d_ff
    # 所以比值为 1 当且仅当 seq_len == d_ff，和 d_model、n_layer 都无关
    cross = d_ff
    print('  本配置的交叉点：seq_len ≈ %d' % cross)
    a_c, f_c = attn_flops(cross, d_model, n_layer), ffn_flops(cross, d_model, d_ff, n_layer)
    print('  验证 seq_len=%d 时，注意力 %s / FFN %s，两者相等: %s'
          % (cross, human(a_c), human(f_c), a_c == f_c))

    print()
    print('推论：')
    print('  1) 上下文窗口不是想开多大开多大，代价是平方的')
    print('  2) 短文本任务把 max_len 设小，是最省成本的一次改动')
    print('  3) 长上下文要靠稀疏注意力、滑动窗口、线性注意力等结构改造，')
    print('     而不是单纯加显存')


if __name__ == '__main__':
    main()

"""self-attention 最小实现：四行公式，逐行对着看。

Attention(Q, K, V) = softmax(Q Kᵀ / √d) V
"""
import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)

seq_len, d_model = 4, 8          # 4 个词，每个词 8 维
x = torch.randn(seq_len, d_model)

# Q/K/V 都是同一份输入乘三套不同的权重变出来的——所以叫「自」注意力
Wq = torch.randn(d_model, d_model)
Wk = torch.randn(d_model, d_model)
Wv = torch.randn(d_model, d_model)

Q = x @ Wq        # 我要找什么
K = x @ Wk        # 我这里挂着什么标签
V = x @ Wv        # 我实际能交出什么内容

# 1) 打分：每个词的 Q 去和所有词的 K 做点积，得到 (seq_len, seq_len) 的分数表
scores = Q @ K.transpose(-2, -1)

# 2) 缩放：维度越大点积越大，除以 √d 防止 softmax 被推到饱和区、梯度消失
scores = scores / math.sqrt(d_model)

# 3) 归一化：每一行变成一组和为 1 的权重
weights = F.softmax(scores, dim=-1)
print('注意力权重形状', weights.shape)          # (4, 4)
print('每行之和', weights.sum(dim=-1))          # 全是 1

# 4) 加权求和：按权重把所有词的 V 混起来，就是这个词的新表示
out = weights @ V
print('输出形状', out.shape)                    # (4, 8)，和输入同形

"""多头注意力 + Add&Norm + 前馈网络：一个完整的编码器层。

用 PyTorch 自带的模块搭，先跑通结构，再回头抠细节。
"""
import torch
import torch.nn as nn


class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, n_head=8, d_ff=2048, dropout=0.1):
        super().__init__()
        # d_model 必须能被 n_head 整除：512 / 8 = 每个头 64 维
        assert d_model % n_head == 0, 'd_model 必须能被 n_head 整除'

        self.attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_head,
                                          dropout=dropout, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        # 前馈网络：先升维再降维，给每个位置单独做一次非线性加工
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, key_padding_mask=None):
        # 子层一：多头注意力 + 残差 + 层归一化
        attn_out, attn_weights = self.attn(x, x, x,
                                           key_padding_mask=key_padding_mask)
        x = self.norm1(x + self.dropout(attn_out))    # 加号就是残差连接

        # 子层二：前馈网络 + 残差 + 层归一化
        x = self.norm2(x + self.dropout(self.ffn(x)))
        return x, attn_weights


if __name__ == '__main__':
    layer = EncoderLayer(d_model=64, n_head=8, d_ff=128)
    x = torch.randn(2, 10, 64)              # (batch, seq_len, d_model)
    out, w = layer(x)
    print('输出形状', out.shape)             # (2, 10, 64)，进出同形
    print('注意力权重形状', w.shape)          # (2, 10, 10)

    # 进出同形，所以可以直接摞 N 层
    stack = nn.ModuleList([EncoderLayer(64, 8, 128) for _ in range(6)])
    h = x
    for blk in stack:
        h, _ = blk(h)
    print('六层之后', h.shape)                # 还是 (2, 10, 64)

"""带掩码的注意力：解码器只能看已经生成的词，不能偷看未来。"""
import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)
seq_len, d_model = 5, 8
Q = K = V = torch.randn(seq_len, d_model)

scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_model)

# 上三角（不含对角线）就是「未来的位置」，把它们设成负无穷
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))

weights = F.softmax(scores, dim=-1)
print(weights.round(decimals=3))
# 第 0 行只有第 0 列有值，第 1 行前两列有值……严格的下三角
# 关键点：屏蔽要在 softmax 之前做。softmax 之后再置 0，每行就不等于 1 了
print('每行之和', weights.sum(dim=-1))

"""三种衍生架构：encoder-only / decoder-only / encoder-decoder。

同样是 Transformer 块，按「用哪半边、加不加掩码」分成三条技术路线，
今天见到的绝大多数模型都能归进其中一类。
"""
import torch
import torch.nn as nn


D_MODEL, N_HEAD, N_LAYER, D_FF, VOCAB = 128, 4, 2, 256, 1000


# ------------------------------------------------------------------ 其一
class EncoderOnly(nn.Module):
    """双向看全文，适合「读懂一句话」的任务：分类、抽取、检索。"""

    def __init__(self, n_class=2):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, D_MODEL)
        layer = nn.TransformerEncoderLayer(D_MODEL, N_HEAD, D_FF,
                                           batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, N_LAYER)
        self.head = nn.Linear(D_MODEL, n_class)

    def forward(self, tokens):
        x = self.encoder(self.embed(tokens))   # 不加 look-ahead mask
        return self.head(x.mean(dim=1))        # 整句压成一个向量再分类


# ------------------------------------------------------------------ 其二
class DecoderOnly(nn.Module):
    """只能看左边，逐词往下写。今天的对话式大模型基本都是这一类。"""

    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, D_MODEL)
        layer = nn.TransformerEncoderLayer(D_MODEL, N_HEAD, D_FF,
                                           batch_first=True)
        # 结构上复用 EncoderLayer，靠传入因果掩码把它变成解码器
        self.blocks = nn.TransformerEncoder(layer, N_LAYER)
        self.head = nn.Linear(D_MODEL, VOCAB)  # 输出维度 = 词表大小

    def forward(self, tokens):
        seq_len = tokens.size(1)
        # 上三角为 True 的布尔掩码：第 t 个位置看不到 t 之后
        causal = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool),
                            diagonal=1)
        x = self.blocks(self.embed(tokens), mask=causal)
        return self.head(x)                    # 每个位置都预测「下一个词」


# ------------------------------------------------------------------ 其三
class EncoderDecoder(nn.Module):
    """一边读原文一边写译文，适合序列到序列：翻译、摘要。"""

    def __init__(self):
        super().__init__()
        self.src_embed = nn.Embedding(VOCAB, D_MODEL)
        self.tgt_embed = nn.Embedding(VOCAB, D_MODEL)
        self.transformer = nn.Transformer(
            d_model=D_MODEL, nhead=N_HEAD,
            num_encoder_layers=N_LAYER, num_decoder_layers=N_LAYER,
            dim_feedforward=D_FF, batch_first=True)
        self.head = nn.Linear(D_MODEL, VOCAB)

    def forward(self, src, tgt):
        tgt_len = tgt.size(1)
        tgt_mask = torch.triu(torch.ones(tgt_len, tgt_len, dtype=torch.bool),
                              diagonal=1)
        out = self.transformer(self.src_embed(src), self.tgt_embed(tgt),
                               tgt_mask=tgt_mask)
        return self.head(out)


def main():
    torch.manual_seed(0)
    batch, seq = 2, 6
    src = torch.randint(1, VOCAB, (batch, seq))
    tgt = torch.randint(1, VOCAB, (batch, seq - 1))

    m1 = EncoderOnly(n_class=3)
    print('encoder-only    输出', tuple(m1(src).shape), '→ (batch, 类别数)')

    m2 = DecoderOnly()
    print('decoder-only    输出', tuple(m2(src).shape), '→ (batch, seq, 词表)')

    m3 = EncoderDecoder()
    print('encoder-decoder 输出', tuple(m3(src, tgt).shape), '→ (batch, 目标长, 词表)')

    print()
    for name, model in (('encoder-only', m1), ('decoder-only', m2),
                        ('encoder-decoder', m3)):
        print('%-16s 参数量 %d' % (name, sum(p.numel() for p in model.parameters())))
    # encoder-decoder 最大：两套块 + 一层额外的交叉注意力

    print()
    print('怎么选：')
    print('  要「理解」→ encoder-only，能双向看全文，但不会自己写')
    print('  要「生成」→ decoder-only，天然适配逐词续写与对话')
    print('  要「转换」→ encoder-decoder，源和目标是两种序列时最直接')


if __name__ == '__main__':
    main()

"""Transformer 骨架模板：位置编码 + N 层编码器 + 任务头。"""
import math
import torch
import torch.nn as nn


class PositionalEncoding(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)
        pe[:, 1::2] = torch.cos(position * div_term)
        # register_buffer：跟着模型存取和搬设备，但不参与梯度更新
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        # x: (batch, seq_len, d_model)，按实际长度截取后相加
        return x + self.pe[:, :x.size(1)]


class MyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=256, n_head=8, n_layer=4,
                 d_ff=1024, n_class=2, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_head, dim_feedforward=d_ff,
            dropout=dropout, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layer)

        # TODO: 整句分类取 [CLS] 或平均池化；逐词任务把 n_class 换成 vocab_size
        self.head = nn.Linear(d_model, n_class)

    def forward(self, tokens, key_padding_mask=None):
        # 论文里嵌入要乘 √d_model，让它和位置编码量级相当
        x = self.embedding(tokens) * math.sqrt(self.d_model)
        x = self.pos(x)
        x = self.encoder(x, src_key_padding_mask=key_padding_mask)

        # TODO: 换成 x[:, 0] 取首位，或 x.max(dim=1) 取最大池化
        pooled = x.mean(dim=1)
        return self.head(pooled)


if __name__ == '__main__':
    model = MyTransformer(vocab_size=10000, n_class=4)
    tokens = torch.randint(0, 10000, (2, 16))
    print(model(tokens).shape)              # torch.Size([2, 4])
    print('参数量', sum(p.numel() for p in model.parameters()))

"""两种掩码：padding mask 和 look-ahead mask，分别挡什么。

padding mask   —— 一批句子长短不一，短的补 0 凑齐；补出来的位置没有语义，
                  不能让别的词去关注它。按「哪些位置是 pad」生成。
look-ahead mask —— 解码器逐词生成时不能偷看后面的词。按「位置先后」生成，
                  和内容无关，永远是一个上三角。
两者可以同时存在，取逻辑或。
"""
import math
import torch
import torch.nn.functional as F


PAD = 0          # 约定 0 号索引是填充符


def build_padding_mask(tokens):
    """(batch, seq) 的索引 → (batch, 1, 1, seq) 的布尔掩码，True 表示要挡。"""
    mask = (tokens == PAD)                 # 逐元素比较，pad 的位置是 True
    return mask[:, None, None, :]          # 补出 head 和 query 两个维度，靠广播铺开


def build_look_ahead_mask(seq_len):
    """上三角（不含对角线）为 True：第 t 步不能看 t 之后的位置。"""
    return torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1)


def main():
    # 三句话，实际长度分别是 4、2、5，右侧用 0 补齐到 5
    tokens = torch.tensor([[7, 3, 9, 4, PAD],
                           [5, 8, PAD, PAD, PAD],
                           [1, 2, 3, 4, 6]])
    batch, seq_len = tokens.shape
    print('token 矩阵\n', tokens)

    # ------------------------------------------------------ padding mask
    pad_mask = build_padding_mask(tokens)
    print('\npadding mask 形状', tuple(pad_mask.shape))
    print('第 1 句要挡的位置', pad_mask[1, 0, 0].tolist())
    # [False, False, True, True, True] —— 后三个补位被挡

    # ------------------------------------------------------ look-ahead
    la_mask = build_look_ahead_mask(seq_len)
    print('\nlook-ahead mask（True = 挡住）')
    print(la_mask.int())
    # 第 0 行只留第 0 列，第 1 行留前两列……严格下三角

    # ------------------------------------------------------ 合并
    combined = pad_mask | la_mask           # 广播成 (batch, 1, seq, seq)
    print('\n合并后形状', tuple(combined.shape))

    # ------------------------------------------------------ 作用于分数
    torch.manual_seed(0)
    d_k = 8
    Q = torch.randn(batch, 1, seq_len, d_k)
    K = torch.randn(batch, 1, seq_len, d_k)
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)

    masked = scores.masked_fill(combined, float('-inf'))
    weights = F.softmax(masked, dim=-1)

    print('\n第 1 句、第 3 个位置的注意力权重：')
    print([round(v, 4) for v in weights[1, 0, 3].tolist()])
    # 只有前两个位置有值，其余全是 0：既没看未来，也没看补位

    print('每行之和（应当都是 1）', weights[1, 0].sum(dim=-1).tolist())

    # ------------------------------------------------------ 顺序的坑
    print('\n屏蔽必须在 softmax 之前做：')
    after = F.softmax(scores, dim=-1).masked_fill(combined, 0.0)
    print('softmax 之后再置 0，每行之和 =',
          [round(v, 4) for v in after[1, 0].sum(dim=-1).tolist()])
    # 明显不等于 1 —— 权重没有重新归一化，剩下位置的比重被整体压小了

    # ------------------------------------------------------ 整行被挡
    print('\n注意：第 1 句第 4 个位置本身就是 pad，它作为 query 时整行都被挡，')
    print('softmax(全 -inf) 会得到 nan。工程上要么在 loss 里忽略这些位置，')
    print('要么让对角线始终可见。PyTorch 的 nn.MultiheadAttention 用')
    print('key_padding_mask 参数接收这类掩码，内部已处理这种边界情况。')


if __name__ == '__main__':
    main()