transformer

transformer.ipynb

@@ -0,0 +1,1077 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "89fbc397",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Parameter containing:\n",
+      "tensor([[ 1.9787,  1.9409,  1.7920, -1.0682,  0.3885,  0.2064,  0.8532, -0.4364],\n",
+      "        [-0.5978,  0.4850, -0.7542,  1.8672, -0.5355,  0.0457,  1.4320,  2.0734],\n",
+      "        [-2.5576,  1.0819,  0.5857, -1.9211, -0.0412,  0.8866, -0.5337, -0.6320],\n",
+      "        [-0.6629, -0.2139,  0.1620, -1.6161, -0.0565, -0.1995,  1.3500,  0.2705],\n",
+      "        [ 1.0395, -1.0146,  0.0549,  0.6911, -1.1399,  0.0717, -0.2826,  1.1461],\n",
+      "        [ 1.0596,  0.7725, -0.0256,  1.5319,  0.3158, -0.1126,  0.7849, -0.9682],\n",
+      "        [ 0.7389,  1.2391, -0.9184,  0.3582, -0.9273,  0.4013, -0.7880, -0.4279],\n",
+      "        [ 1.3871, -0.0653,  1.4355,  0.3321,  0.2264,  0.8935, -0.8899,  0.8854],\n",
+      "        [-0.2120,  1.5278, -0.6495,  1.0855, -0.3273,  0.5673,  0.8668, -1.8364],\n",
+      "        [ 1.8725,  0.7562,  1.3969, -0.0723, -0.0850,  1.1399, -1.2521, -0.3105]],\n",
+      "       requires_grad=True)\n",
+      "tensor([[-0.5978,  0.4850, -0.7542,  1.8672, -0.5355,  0.0457,  1.4320,  2.0734]],\n",
+      "       grad_fn=<EmbeddingBackward0>)\n",
+      "tensor([[[-0.5978,  0.4850, -0.7542,  1.8672, -0.5355,  0.0457,  1.4320,\n",
+      "           2.0734],\n",
+      "         [-2.5576,  1.0819,  0.5857, -1.9211, -0.0412,  0.8866, -0.5337,\n",
+      "          -0.6320],\n",
+      "         [-0.6629, -0.2139,  0.1620, -1.6161, -0.0565, -0.1995,  1.3500,\n",
+      "           0.2705]],\n",
+      "\n",
+      "        [[ 1.0395, -1.0146,  0.0549,  0.6911, -1.1399,  0.0717, -0.2826,\n",
+      "           1.1461],\n",
+      "         [ 1.0596,  0.7725, -0.0256,  1.5319,  0.3158, -0.1126,  0.7849,\n",
+      "          -0.9682],\n",
+      "         [ 1.9787,  1.9409,  1.7920, -1.0682,  0.3885,  0.2064,  0.8532,\n",
+      "          -0.4364]]], grad_fn=<EmbeddingBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import math\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "# 演示场景：假设词表只有 10 个 token，每个 token 用 8 维向量表示\n",
+    "emb = nn.Embedding(num_embeddings=10, embedding_dim=8)\n",
+    "\n",
+    "# nn.Embedding 内部维护的就是那张 (10, 8) 的查找表，随机初始化\n",
+    "print(emb.weight)   # torch.Size([10, 8])\n",
+    "\n",
+    "# 查单个 token：索引 1 对应的 8 维向量\n",
+    "print(emb(torch.tensor([1])))   # torch.Size([1, 8])\n",
+    "\n",
+    "# 查一个 batch：两句话，每句 3 个 token\n",
+    "# 输出形状 (2, 3, 8) —— (batch, seq_len, d_model)\n",
+    "two_sentences = torch.tensor([\n",
+    "    [1, 2, 3],\n",
+    "    [4, 5, 0],   # 0 预留给 padding（填充位），让两句话对齐长度\n",
+    "])\n",
+    "print(emb(two_sentences))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "34fddfbb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[ 3.5261, -0.3545,  1.0208, -3.6344,  2.4014, -5.0555,  2.1495,\n",
+      "          -2.4776],\n",
+      "         [-1.6115, -3.6160,  2.7583, -2.6699,  2.2102,  0.6809, -2.4801,\n",
+      "           4.3560],\n",
+      "         [-4.0357, -2.4071,  5.4715,  0.0622,  3.9720,  1.1684,  4.5511,\n",
+      "          -5.1458]],\n",
+      "\n",
+      "        [[ 0.6346, -1.3376, -0.1375,  4.1304,  0.4468,  1.8783, -2.5956,\n",
+      "           2.6552],\n",
+      "         [ 2.4246, -3.5624,  0.2341, -2.2205, -0.3354, -4.0794, -4.2326,\n",
+      "          -2.3740],\n",
+      "         [ 1.4035, -0.7918, -5.0030, -0.8911,  8.1170,  2.1314, -4.8738,\n",
+      "          -2.7081]]], grad_fn=<MulBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "class Embeddings(nn.Module):\n",
+    "    \"\"\"词嵌入层：查表 + 按 sqrt(d_model) 缩放\"\"\"\n",
+    "\n",
+    "    def __init__(self, vocab_size: int, d_model: int):\n",
+    "        super().__init__()\n",
+    "        # lut = look-up table，本质就是那张 (vocab_size, d_model) 的参数矩阵\n",
+    "        self.lut = nn.Embedding(vocab_size, d_model)\n",
+    "        self.d_model = d_model\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x 形状: (batch, seq_len)，装的是 token 索引\n",
+    "        # 输出形状: (batch, seq_len, d_model)\n",
+    "        # 注意这里乘了 sqrt(d_model)，原因见 2.2 节\n",
+    "        return self.lut(x) * math.sqrt(self.d_model)\n",
+    "\n",
+    "two_sentences = torch.tensor([\n",
+    "    [1, 2, 3],\n",
+    "    [4, 5, 0],   # 0 预留给 padding（填充位），让两句话对齐长度\n",
+    "])\n",
+    "emb_layer = Embeddings(vocab_size=10, d_model=8)\n",
+    "emb_out = emb_layer(two_sentences)\n",
+    "print(emb_out)   # torch.Size([2, 3, 8])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "80a8b462",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[5.0949, 6.8103, 1.0000, 8.6680, 7.6680],\n",
+      "        [5.0000, 7.0000, 1.0000, 9.0000, 8.0000],\n",
+      "        [5.5379, 5.9242, 1.0000, 7.1174, 6.1174]])\n",
+      "tensor([[5.0000, 7.0000, 1.0000, 9.0000, 8.0000],\n",
+      "        [5.0949, 6.8103, 1.0000, 8.6680, 7.6680],\n",
+      "        [5.5379, 5.9242, 1.0000, 7.1174, 6.1174]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "# 第一组：三行向量按原始顺序排列\n",
+    "q1 = torch.tensor([[0, 1, 2, 1, 5],\n",
+    "                   [3, 9, 3, 8, 2],\n",
+    "                   [2, 3, 1, 5, 6]],dtype=torch.float)\n",
+    "k1 = torch.tensor([[3, 2, 1, 5, 5],\n",
+    "                   [1, 5, 5, 0, 1],\n",
+    "                   [0, 8, 9, 6, 1]],dtype=torch.float)\n",
+    "v1 = torch.tensor([[7, 3, 1, 2, 1],\n",
+    "                   [4, 7, 1, 1, 9],\n",
+    "                   [5, 7, 1, 9, 8]],dtype=torch.float)\n",
+    "\n",
+    "# 第二组：把前两行对调（相当于把句子里的词换了顺序）\n",
+    "q2 = torch.tensor([[3, 9, 3, 8, 2],\n",
+    "                  [0, 1, 2, 1, 5],\n",
+    "                  [2, 3, 1, 5, 6]],dtype=torch.float)\n",
+    "k2 = torch.tensor([[1, 5, 5, 0, 1],\n",
+    "                  [3, 2, 1, 5, 5],\n",
+    "                  [0, 8, 9, 6, 1]],dtype=torch.float)\n",
+    "v2 = torch.tensor([[4, 7, 1, 1, 9],\n",
+    "                  [7, 3, 1, 2, 1],\n",
+    "                  [5, 7, 1, 9, 8]],dtype=torch.float)\n",
+    "\n",
+    "softmax = torch.nn.Softmax(dim=1)\n",
+    "\n",
+    "def simple_attention(q, k, v):\n",
+    "    \"\"\"最朴素的注意力：QK^T -> softmax -> 加权 V\"\"\"\n",
+    "    scores = torch.matmul(q, k.transpose(-1, -2))\n",
+    "    weights = softmax(scores)\n",
+    "    return torch.matmul(weights, v)\n",
+    "\n",
+    "out1 = simple_attention(q1, k1, v1)\n",
+    "out2 = simple_attention(q2, k2, v2)\n",
+    "\n",
+    "# 对比会发现：out2 就是 out1 的行跟着对调了一下\n",
+    "# 注意力分数本身毫无变化 —— 模型根本不知道谁在前谁在后\n",
+    "print(out1)\n",
+    "print(out2) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "42d3f04b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00,\n",
+      "           1.0000e+00,  0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  9.9833e-02,  9.9500e-01,  9.9998e-03,\n",
+      "           9.9995e-01,  1.0000e-03,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  1.9867e-01,  9.8007e-01,  1.9999e-02,\n",
+      "           9.9980e-01,  2.0000e-03,  1.0000e+00]]])\n",
+      "torch.Size([2, 3, 8])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# 位置编码层\n",
+    "## 位置编码的实现其实很简单，直接对照着公式去敲代码就可以，\n",
+    "# 下面这个代码只是其中一种实现方式；\n",
+    "## 从理解来讲，需要注意的就是偶数和奇数在公式上有一个共同部分；\n",
+    "## pos代表的是单词在句子中的索引，这点需要注意；\n",
+    "# 比如max_len是128个，那么索引就是从0，1，2，...,127\n",
+    "##假设我的d_model是512，2i那个符号中i从0取到了255，\n",
+    "# 那么2i对应取值就是0,2,4...510,那么2i+1对应取值就是1,3,5...511\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "class PositionalEncoding(nn.Module):\n",
+    "    def __init__(self, d_model,max_len=5000,dropout=0.1) -> None:\n",
+    "        super().__init__()\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "        \n",
+    "        pe = torch.zeros(max_len,d_model)\n",
+    "        \n",
+    "        pos = torch.arange(0,max_len).unsqueeze(1)\n",
+    "        div_term = torch.pow(10000, torch.arange(0, d_model, 2).float() / d_model)\n",
+    "        \n",
+    "        ## 这里需要注意的是pe[:, 0::2]这个用法，就是从0开始到最后面，\n",
+    "        # 步长为2，其实代表的就是偶数位置\n",
+    "        pe[:, 0::2] = torch.sin(pos / div_term)\n",
+    "        ##这里需要注意的是pe[:, 1::2]这个用法，就是从1开始到最后面，\n",
+    "        # 步长为2，其实代表的就是奇数位置\n",
+    "        pe[:, 1::2] = torch.cos(pos / div_term)\n",
+    "\n",
+    "        #升维\n",
+    "        pe = pe.unsqueeze(0)\n",
+    "        #将位置编码矩阵设置为缓冲区\n",
+    "        # (这些缓冲区不会被视为模型的可训练参数（即不会被优化器更新），\n",
+    "        # 但会作为模型状态的一部分保存和加载)\n",
+    "        #相当于把位置信息保存在了一个常量数组中\n",
+    "        self.register_buffer(\"pe\",pe)\n",
+    "    \n",
+    "    def forward(self,x):\n",
+    "        print(self.pe[:,:x.size(1)])\n",
+    "        x = x+ self.pe[:,:x.size(1)]\n",
+    "        return self.dropout(x)\n",
+    "\n",
+    "pos = PositionalEncoding(8)\n",
+    "pos_out = pos(emb_out)\n",
+    "print(pos_out.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a8c53ac0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "OrderedDict([('weight', tensor([[-0.1045,  0.4786],\n",
+      "        [-0.7842,  0.9056]])), ('running_mean', tensor([0., 0.]))])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "class MyModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(MyModel, self).__init__()\n",
+    "        self.weight = nn.Parameter(torch.randn(2, 2))  # 可训练参数\n",
+    "         # 注册缓冲区 这个张量会随模型的状态字典 (state_dict) 一起保存和加载\n",
+    "         #保存不需要训练但需持久化的张量。\n",
+    "        self.register_buffer('running_mean', torch.zeros(2)) \n",
+    "\n",
+    "model = MyModel()\n",
+    "print(model.state_dict())  # 输出包含 weight 和 running_mean"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "80177eae",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[ 0.1789,  0.9565,  0.8954],\n",
+      "         [-0.7689,  0.2724, -0.8498],\n",
+      "         [ 0.3944,  0.5488,  1.0585]]])\n",
+      "tensor([[[False, False,  True]]])\n",
+      "torch.Size([1, 1, 3])\n",
+      "tensor([[[ 1.7893e-01,  9.5654e-01, -1.0000e+09],\n",
+      "         [-7.6893e-01,  2.7236e-01, -1.0000e+09],\n",
+      "         [ 3.9439e-01,  5.4877e-01, -1.0000e+09]]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#mask_fill方法 将张量中特定位置的元素替换为指定的值\n",
+    "import torch\n",
+    "scores = torch.randn(1, 3, 3)\n",
+    "print(scores)\n",
+    "\n",
+    "inputs = torch.tensor([\n",
+    "    [4, 5, 0]\n",
+    "])\n",
+    "\n",
+    "mask = (inputs == 0).unsqueeze(1)\n",
+    "print(mask)\n",
+    "print(mask.shape)\n",
+    "\n",
+    "scores = scores.masked_fill(mask, -1e9)\n",
+    "print(scores)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "41cfc85e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 1, 3])\n",
+      "屏蔽前:\n",
+      " tensor([[[-1.0046,  0.1224, -1.2707],\n",
+      "         [ 0.5761,  0.2447,  0.5002],\n",
+      "         [-1.1239,  0.3064, -2.3147]]])\n",
+      "屏蔽后:\n",
+      " tensor([[[-1.0046e+00,  1.2237e-01, -1.0000e+09],\n",
+      "         [ 5.7612e-01,  2.4473e-01, -1.0000e+09],\n",
+      "         [-1.1239e+00,  3.0637e-01, -1.0000e+09]]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "#注意力机制增加掩码\n",
+    "def get_padding_mask(x, padding_idx=0):\n",
+    "    \"\"\"生成 padding 掩码：pad 位置为 True（待屏蔽）\n",
+    "\n",
+    "    输入 x: (batch, seq_len)\n",
+    "    输出:   (batch, 1, seq_len) —— 第 1 维留给注意力分数的行做广播\n",
+    "    \"\"\"\n",
+    "    return (x == padding_idx).unsqueeze(1)\n",
+    "\n",
+    "mask = get_padding_mask(inputs)\n",
+    "print(mask.shape)   # torch.Size([2, 1, 3])\n",
+    "\n",
+    "# 看看 masked_fill 的实际效果\n",
+    "scores = torch.randn(1, 3, 3)   # 假装是两个样本的注意力分数\n",
+    "print(\"屏蔽前:\\n\", scores)\n",
+    "scores = scores.masked_fill(mask, -1e9)\n",
+    "print(\"屏蔽后:\\n\", scores)\n",
+    "# 第二句的最后一个位置全被压成 -1e9，\n",
+    "# softmax 之后它对任何位置的注意力权重都 ≈ 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3780dcb3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([2, 3, 8])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import math\n",
+    "import torch\n",
+    "\n",
+    "def attention(q, k, v, mask=None, dropout=None):\n",
+    "    \"\"\"缩放点积注意力（Scaled Dot-Product Attention）\n",
+    "\n",
+    "    参数:\n",
+    "        q, k, v: 形状 (batch, n_head, seq_len, d_k) 的查询/键/值张量\n",
+    "        mask:    布尔掩码，True 的位置会被屏蔽（不参与注意力分配）\n",
+    "        dropout: 对注意力权重做随机失活，训练时防过拟合\n",
+    "    返回:\n",
+    "        上下文向量，形状与 v 一致\n",
+    "    \"\"\"\n",
+    "    # d_k 是每个头的维度（不是总维度 d_model！）\n",
+    "    d_k = k.size(-1)\n",
+    "\n",
+    "    # 第一步：Q·K^T 算原始相关性分数，形状 (batch, n_head, seq_len, seq_len)\n",
+    "    # 第 (i, j) 个元素 = 第 i 个位置\"问\"第 j 个位置的匹配程度\n",
+    "    # 第二步：除以 sqrt(d_k) 缩放，防止大维度下 softmax 饱和\n",
+    "    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)\n",
+    "\n",
+    "    # 第三步（可选）：掩码。把不该看的位置填成 -1e9（约等于负无穷）\n",
+    "    # softmax(-1e9) ≈ 0，这些位置分到的注意力权重趋近于零，等于被\"蒙住眼睛\"\n",
+    "    if mask is not None:\n",
+    "        scores = scores.masked_fill(mask, -1e9)\n",
+    "\n",
+    "    # 第四步：softmax 在最后一维归一化，每行权重之和为 1\n",
+    "    # 至此，分数变成了\"每个位置该分多少注意力\"的概率分布\n",
+    "    p_attn = torch.softmax(scores, dim=-1)\n",
+    "\n",
+    "    if dropout is not None:\n",
+    "        p_attn = dropout(p_attn)\n",
+    "\n",
+    "    # 第五步：用权重对 V 加权求和，得到融合了全局上下文的新表示\n",
+    "    # 输出形状 (batch, n_head, seq_len, d_k)\n",
+    "    return torch.matmul(p_attn, v)\n",
+    "\n",
+    "\n",
+    "# 用第二章的位置编码输出做个冒烟测试：q=k=v（自注意力就是\"自己问自己\"）\n",
+    "result = attention(pos_out, pos_out, pos_out)\n",
+    "print(result.shape)   # 与输入形状一致"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "865d6691",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[-1.3416, -0.4472,  0.4472,  1.3416],\n",
+      "        [ 0.0000,  0.0000,  0.0000,  0.0000]],\n",
+      "       grad_fn=<NativeLayerNormBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "# 定义张量x\n",
+    "# x是一个2×4大小的张量，代表了2个样本，每个样本有4个特征\n",
+    "# 当使用层归一化对x进行计算时，会对x的每一行的4个特征进行归一化\n",
+    "x = torch.tensor([[1.0, 2.0, 3.0, 4.0],\n",
+    "                  [1.0, 1.0, 1.0, 1.0]])\n",
+    "\n",
+    "# 定义层归一化实例\n",
+    "# 传入的参数4，表示层归一化的输入维度是4\n",
+    "# 也就是对应某一层的4个神经元的输出\n",
+    "layer_norm = nn.LayerNorm(4)\n",
+    "\n",
+    "# 计算层归一化的结果\n",
+    "output = layer_norm(x)\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "4fcad02a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "γ和β的参数值：\n",
+      "gamma: Parameter containing:\n",
+      "tensor([1., 1., 1., 1.], requires_grad=True)\n",
+      "beta: Parameter containing:\n",
+      "tensor([0., 0., 0., 0.], requires_grad=True)\n",
+      "γ和β的尺寸：\n",
+      "gamma.shape: torch.Size([4])\n",
+      "beta.shape: torch.Size([4])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "# 定义张量x\n",
+    "# x是一个2×4大小的张量，代表了2个样本，每个样本有4个特征\n",
+    "# 当使用层归一化对x进行计算时，会对x的每一行的4个特征进行归一化\n",
+    "x = torch.tensor([[1.0, 2.0, 3.0, 4.0],\n",
+    "                  [1.0, 1.0, 1.0, 1.0]])\n",
+    "\n",
+    "# 定义层归一化实例\n",
+    "# 传入的参数4，表示层归一化的输入维度是4\n",
+    "# 也就是对应某一层的4个神经元的输出\n",
+    "layer_norm = nn.LayerNorm(4)\n",
+    "\n",
+    "# 它们分别对应了γ和β两个参数\n",
+    "print(\"γ和β的参数值：\")\n",
+    "print(\"gamma:\", layer_norm.weight)\n",
+    "print(\"beta:\", layer_norm.bias)\n",
+    "\n",
+    "# 打印γ和β的尺寸\n",
+    "print(\"γ和β的尺寸：\")\n",
+    "print(\"gamma.shape:\", layer_norm.weight.shape)\n",
+    "print(\"beta.shape:\", layer_norm.bias.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "20785542",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([2, 3, 8])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class MultiHeadedAttention(nn.Module):\n",
+    "    \"\"\"多头注意力 + 残差连接 + LayerNorm（Encoder/Decoder 通用积木）\"\"\"\n",
+    "\n",
+    "    def __init__(self, d_model: int, n_head: int, dropout: float = 0.1):\n",
+    "        super().__init__()\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "\n",
+    "        # d_model 必须能被头数整除，每头分到 d_token 维\n",
+    "        # 例: d_model=512, n_head=8 -> 每头 64 维\n",
+    "        assert d_model % n_head == 0\n",
+    "        self.d_token = d_model // n_head\n",
+    "        self.n_head = n_head\n",
+    "\n",
+    "        # 注意这里 Q/K/V 的投影矩阵都是 (d_model, d_model)：\n",
+    "        # 实现上不为每个头单独建小矩阵，而是用一个大矩阵投影后再\"切分\"，\n",
+    "        # 数学上等价于 h 个小矩阵，但矩阵运算效率高得多\n",
+    "        self.W_Q = nn.Linear(d_model, d_model, bias=False)\n",
+    "        self.W_K = nn.Linear(d_model, d_model, bias=False)\n",
+    "        self.W_V = nn.Linear(d_model, d_model, bias=False)\n",
+    "\n",
+    "        # 所有头的输出拼接后，再过一次线性层做跨头信息融合\n",
+    "        self.linear = nn.Linear(d_model, d_model, bias=False)\n",
+    "        self.norm = nn.LayerNorm(d_model)\n",
+    "\n",
+    "    def forward(self, q, k, v, mask=None):\n",
+    "        # 残差连接要用的\"原始输入\"，先存一份（残差见第四章）\n",
+    "        res = q\n",
+    "\n",
+    "        # ---- 分头 ----\n",
+    "        # 投影后形状: (batch, seq_len, d_model)\n",
+    "        # view 拆成 (batch, seq_len, n_head, d_token)\n",
+    "        # transpose 换成 (batch, n_head, seq_len, d_token)\n",
+    "        # 把头数提到第 1 维，后面就可以对\"每个头\"并行做矩阵乘\n",
+    "        q = self.W_Q(q).view(q.shape[0], q.shape[1], self.n_head, self.d_token).transpose(1, 2)\n",
+    "        k = self.W_K(k).view(k.shape[0], k.shape[1], self.n_head, self.d_token).transpose(1, 2)\n",
+    "        v = self.W_V(v).view(v.shape[0], v.shape[1], self.n_head, self.d_token).transpose(1, 2)\n",
+    "\n",
+    "        # 掩码升一维以对齐广播: (batch, 1, seq_len) -> (batch, 1, 1, seq_len)\n",
+    "        # 这样同一个掩码会作用到所有头上\n",
+    "        if mask is not None:\n",
+    "            mask = mask.unsqueeze(1)\n",
+    "\n",
+    "        # ---- 每个头各自算注意力 ----\n",
+    "        context = attention(q, k, v, mask, self.dropout)\n",
+    "\n",
+    "        # ---- 拼头 + 融合 ----\n",
+    "        # transpose 回来: (batch, seq_len, n_head, d_token)\n",
+    "        # reshape 拼回: (batch, seq_len, d_model)\n",
+    "        concat = context.transpose(1, 2).reshape(q.shape[0], -1, self.n_head * self.d_token)\n",
+    "        output = self.linear(concat)\n",
+    "\n",
+    "        # ---- 残差 + 层归一化（第四章详解）----\n",
+    "        return self.norm(res + output)\n",
+    "\n",
+    "\n",
+    "# 测试：8 维模型、2 个头\n",
+    "mha = MultiHeadedAttention(d_model=8, n_head=2)\n",
+    "mha_out = mha(pos_out, pos_out, pos_out)\n",
+    "print(mha_out.shape)   # 输入什么形状，输出还是什么形状"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "6bb81078",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([2, 3, 8])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class FeedForwardNet(nn.Module):\n",
+    "    \"\"\"位置前馈网络：升维 -> 激活 -> 降维，外加残差 + LayerNorm\n",
+    "\n",
+    "    注意\"position-wise\"的含义：每个位置的向量独立过同一个 FFN，\n",
+    "    位置之间在这里没有任何信息交流（交流已经在注意力里做完了）。\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):\n",
+    "        super().__init__()\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "        # w1: 升维，d_model -> d_ff（原论文 d_ff = 4 * d_model）\n",
+    "        self.w1 = nn.Linear(d_model, d_ff)\n",
+    "        # w2: 降维，d_ff -> d_model，必须压回去才能接残差\n",
+    "        self.w2 = nn.Linear(d_ff, d_model)\n",
+    "        self.relu = nn.ReLU()   # 现代实践可换成 nn.GELU()\n",
+    "        self.norm = nn.LayerNorm(d_model)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # 存一份输入，给残差用\n",
+    "        res = x\n",
+    "        # 升维 + 非线性：在更宽的空间里筛选、加工特征\n",
+    "        x = self.relu(self.w1(x))\n",
+    "        # 降维 + dropout：浓缩回 d_model，顺带正则化\n",
+    "        x = self.dropout(self.w2(x))\n",
+    "        # 残差相加 + 层归一化，与注意力子层的收尾方式完全一致\n",
+    "        return self.norm(x + res)\n",
+    "\n",
+    "\n",
+    "# 接在多头注意力后面测试\n",
+    "ffn = FeedForwardNet(d_model=8, d_ff=32)\n",
+    "ffn_out = ffn(mha_out)\n",
+    "print(ffn_out.shape)   # 形状依旧不变，Transformer 的维度守恒铁律"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "6d765e4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class EncoderLayer(nn.Module):\n",
+    "    \"\"\"编码器子层：多头自注意力 + 前馈网络（各自带残差和归一化）\"\"\"\n",
+    "\n",
+    "    def __init__(self, d_model: int, n_head: int, d_ff: int, dropout: float = 0.1):\n",
+    "        super().__init__()\n",
+    "        self.mha = MultiHeadedAttention(d_model, n_head, dropout)\n",
+    "        self.ffn = FeedForwardNet(d_model, d_ff, dropout)\n",
+    "\n",
+    "    def forward(self, x, mask=None):\n",
+    "        # 自注意力的典型特征：q、k、v 全是同一个 x（自己问自己）\n",
+    "        output = self.mha(x, x, x, mask)\n",
+    "        return self.ffn(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "f3ddb9b2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00,\n",
+      "           1.0000e+00,  0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  9.9833e-02,  9.9500e-01,  9.9998e-03,\n",
+      "           9.9995e-01,  1.0000e-03,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  1.9867e-01,  9.8007e-01,  1.9999e-02,\n",
+      "           9.9980e-01,  2.0000e-03,  1.0000e+00]]])\n",
+      "torch.Size([2, 3, 8])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class Encoder(nn.Module):\n",
+    "    \"\"\"完整编码器：Embedding -> 位置编码 -> N 个 EncoderLayer\"\"\"\n",
+    "\n",
+    "    def __init__(self, vocab_size, d_model, n_head, d_ff, N=6, dropout=0.1):\n",
+    "        super().__init__()\n",
+    "        self.emb = Embeddings(vocab_size, d_model)\n",
+    "        self.pos = PositionalEncoding(d_model=d_model, dropout=dropout)\n",
+    "        # nn.ModuleList 专门用来装\"一组结构相同的层\"，\n",
+    "        # 用普通 list 的话 PyTorch 注册不上参数，模型会漏学\n",
+    "        self.layers = nn.ModuleList([\n",
+    "            EncoderLayer(d_model, n_head, d_ff, dropout)\n",
+    "            for _ in range(N)\n",
+    "        ])\n",
+    "\n",
+    "    def forward(self, x, mask=None):\n",
+    "        x = self.emb(x)    # 索引 -> 向量\n",
+    "        x = self.pos(x)    # 注入位置信息\n",
+    "        for layer in self.layers:\n",
+    "            x = layer(x, mask)   # 逐层提炼语义\n",
+    "        return x\n",
+    "\n",
+    "\n",
+    "# 冒测试\n",
+    "src_inputs = torch.tensor([[1, 2, 3], [4, 5, 0]])\n",
+    "src_mask = get_padding_mask(src_inputs, 0)\n",
+    "encoder = Encoder(vocab_size=10, d_model=8, n_head=2, d_ff=32)\n",
+    "memory = encoder(src_inputs, src_mask)\n",
+    "print(memory.shape)   # (2, 3, 8)：源序列的语义表示，等下喂给 Decoder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "f0b080dc",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[1., 0., 0.],\n",
+      "        [1., 1., 0.],\n",
+      "        [1., 1., 1.]])\n",
+      "tensor([[0., 1., 1.],\n",
+      "        [0., 0., 1.],\n",
+      "        [0., 0., 0.]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "#保留矩阵的下三角部分，包括主对角线，而将上三角部分填充为零\n",
+    "tril = torch.tril(torch.ones((3, 3)))\n",
+    "print(tril)\n",
+    "print(1-tril)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "f9d74a70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "def get_subsequent_mask(size: int):\n",
+    "    \"\"\"生成三角掩码：需要屏蔽的\"未来位置\"为 True\n",
+    "\n",
+    "    size 为序列长度，输出形状 (1, size, size)，1 是广播用的 batch 维\n",
+    "    用 1 减去下三角矩阵，就把\"上三角（未来）\"标成了 1（待屏蔽）\n",
+    "    \"\"\"\n",
+    "    mask_shape = (1, size, size)\n",
+    "    return 1 - torch.tril(torch.ones(mask_shape)).byte()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "62364706",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[False,  True,  True],\n",
+      "         [False, False,  True],\n",
+      "         [False, False, False]],\n",
+      "\n",
+      "        [[False,  True,  True],\n",
+      "         [False, False,  True],\n",
+      "         [False, False,  True]]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Decoder 的最终掩码 = padding mask 和三角掩码的叠加（按位或）\n",
+    "# 既要遮住填充位，也要遮住未来位置\n",
+    "tgt_inputs = torch.tensor([[1, 2, 3], [4, 5, 0]])\n",
+    "pad_mask = get_padding_mask(tgt_inputs, 0)   # (2, 1, 3)\n",
+    "sub_mask = get_subsequent_mask(3)            # (1, 3, 3)\n",
+    "tgt_mask = (sub_mask | pad_mask) != 0        # 广播后 (2, 3, 3)，True = 屏蔽\n",
+    "print(tgt_mask)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "5ce2a55d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DecoderLayer(nn.Module):\n",
+    "    \"\"\"解码器子层：掩码自注意力 + 交叉注意力 + 前馈网络\"\"\"\n",
+    "\n",
+    "    def __init__(self, d_model, n_head, d_ff, dropout=0.1):\n",
+    "        super().__init__()\n",
+    "        # 自注意力：理解目标序列自身（带三角掩码，防偷看未来）\n",
+    "        self.self_mha = MultiHeadedAttention(d_model, n_head, dropout)\n",
+    "        # 交叉注意力：q 来自 Decoder，k/v 来自 Encoder 的输出 memory\n",
+    "        self.cross_mha = MultiHeadedAttention(d_model, n_head, dropout)\n",
+    "        self.ffn = FeedForwardNet(d_model, d_ff, dropout)\n",
+    "\n",
+    "    def forward(self, x, mask, memory, src_mask):\n",
+    "        # mask:     目标序列的掩码（padding + 三角叠加）\n",
+    "        # src_mask: 源序列的 padding mask（交叉注意力用，不用三角——原文随便看）\n",
+    "        x = self.self_mha(x, x, x, mask)\n",
+    "        x = self.cross_mha(x, memory, memory, src_mask)\n",
+    "        return self.ffn(x)\n",
+    "\n",
+    "\n",
+    "class Decoder(nn.Module):\n",
+    "    \"\"\"完整解码器：Embedding -> 位置编码 -> N 个 DecoderLayer\"\"\"\n",
+    "\n",
+    "    def __init__(self, vocab_size, d_model, n_head, d_ff, N=6, dropout=0.1):\n",
+    "        super().__init__()\n",
+    "        self.emb = Embeddings(vocab_size, d_model)\n",
+    "        self.pos = PositionalEncoding(d_model=d_model, dropout=dropout)\n",
+    "        self.layers = nn.ModuleList([\n",
+    "            DecoderLayer(d_model, n_head, d_ff, dropout)\n",
+    "            for _ in range(N)\n",
+    "        ])\n",
+    "\n",
+    "    def forward(self, x, mask, memory, src_mask):\n",
+    "        x = self.emb(x)\n",
+    "        x = self.pos(x)\n",
+    "        for layer in self.layers:\n",
+    "            x = layer(x, mask, memory, src_mask)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "b8dd8e0c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Generator(nn.Module):\n",
+    "    \"\"\"输出层：高维向量 -> 词表概率分布（本质是个分类网络）\"\"\"\n",
+    "\n",
+    "    def __init__(self, d_model: int, vocab_size: int):\n",
+    "        super().__init__()\n",
+    "        self.linear = nn.Linear(d_model, vocab_size)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return torch.softmax(self.linear(x), dim=-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "916621a5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "总参数量: 44,112,904\n",
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
+      "           0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
+      "           1.0366e-04,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
+      "           2.0733e-04,  1.0000e+00]]])\n",
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
+      "           0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
+      "           1.0366e-04,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
+      "           2.0733e-04,  1.0000e+00],\n",
+      "         [ 1.4112e-01, -9.8999e-01,  2.4509e-01,  ...,  1.0000e+00,\n",
+      "           3.1099e-04,  1.0000e+00]]])\n",
+      "torch.Size([2, 4, 8])\n",
+      "tensor([[0, 0, 6, 0],\n",
+      "        [6, 6, 6, 6]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class Transformer(nn.Module):\n",
+    "    \"\"\"完整的 Encoder-Decoder Transformer\"\"\"\n",
+    "\n",
+    "    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N=6, dropout=0.1):\n",
+    "        super().__init__()\n",
+    "        # 源语言和目标语言各有自己的词表，所以 Embedding 也是两套\n",
+    "        self.encoder = Encoder(src_vocab_size, d_model, n_head, d_ff, N, dropout)\n",
+    "        self.decoder = Decoder(tgt_vocab_size, d_model, n_head, d_ff, N, dropout)\n",
+    "        self.generator = Generator(d_model, tgt_vocab_size)\n",
+    "\n",
+    "    def forward(self, src_x, src_mask, tgt_x, tgt_mask):\n",
+    "        # 第一步：Encoder 把源序列读成 memory\n",
+    "        memory = self.encoder(src_x, src_mask)\n",
+    "        # 第二步：Decoder 结合 memory 和已生成的前文，产出每个位置的表示\n",
+    "        output = self.decoder(tgt_x, tgt_mask, memory, src_mask)\n",
+    "        # 第三步：投影到目标词表，得到概率分布\n",
+    "        return self.generator(output)\n",
+    "\n",
+    "\n",
+    "def make_model(src_vocab_size, tgt_vocab_size, d_model=512, n_head=8, d_ff=2048, N=6, dropout=0.1):\n",
+    "    \"\"\"建模型 + Xavier 初始化（初始化详见第六章）\"\"\"\n",
+    "    model = Transformer(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N, dropout)\n",
+    "    for p in model.parameters():\n",
+    "        if p.dim() > 1:   # 只初始化二维以上的权重矩阵，一维的 bias/LayerNorm 参数跳过\n",
+    "            nn.init.xavier_uniform_(p)\n",
+    "    return model\n",
+    "\n",
+    "\n",
+    "# 完整跑一遍：玩具级词表 + 论文标准超参\n",
+    "model = make_model(src_vocab_size=6, tgt_vocab_size=8)\n",
+    "total_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"总参数量: {total_params:,}\")   # 四千多万，小模型也有大身板\n",
+    "\n",
+    "# 构造输入\n",
+    "src_inputs = torch.tensor([[1, 2, 3], [4, 5, 0]])\n",
+    "src_mask = get_padding_mask(src_inputs, 0)\n",
+    "tgt_inputs = torch.tensor([[1, 2, 3, 4], [4, 5, 0, 0]])\n",
+    "tgt_pad_mask = get_padding_mask(tgt_inputs, 0)\n",
+    "tgt_mask = (get_subsequent_mask(4) | tgt_pad_mask) != 0\n",
+    "\n",
+    "predict = model(src_inputs, src_mask, tgt_inputs, tgt_mask)\n",
+    "print(predict.shape)   # (2, 4, 8)：2 句话 × 4 个位置 × 8 词表概率\n",
+    "print(torch.argmax(predict, dim=-1))   # 每个位置概率最大的 token 索引"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "353981b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn as nn\n",
+    "\n",
+    "def make_model(src_vocab_size, tgt_vocab_size, d_model=512, n_head=8, d_ff=2048, N=6, dropout=0.1):\n",
+    "    \"\"\"建模 + Xavier 初始化，一步到位\"\"\"\n",
+    "    model = Transformer(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N, dropout)\n",
+    "    for p in model.parameters():\n",
+    "        # 只初始化矩阵型参数（权重），跳过 bias、LayerNorm 的 γ/β 这类一维参数\n",
+    "        # 一维参数按 PyTorch 默认值走就挺好\n",
+    "        if p.dim() > 1:\n",
+    "            nn.init.xavier_uniform_(p)\n",
+    "    return model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "14ecca39",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Total number of parameters: 44112904\n",
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
+      "           0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
+      "           1.0366e-04,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
+      "           2.0733e-04,  1.0000e+00]]])\n",
+      "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
+      "           0.0000e+00,  1.0000e+00],\n",
+      "         [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
+      "           1.0366e-04,  1.0000e+00],\n",
+      "         [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
+      "           2.0733e-04,  1.0000e+00],\n",
+      "         [ 1.4112e-01, -9.8999e-01,  2.4509e-01,  ...,  1.0000e+00,\n",
+      "           3.1099e-04,  1.0000e+00]]])\n",
+      "torch.Size([2, 4, 8])\n",
+      "tensor([[[0.0869, 0.0226, 0.0461, 0.0466, 0.1743, 0.1472, 0.3472, 0.1291],\n",
+      "         [0.0679, 0.0243, 0.0554, 0.0495, 0.1762, 0.1299, 0.1985, 0.2983],\n",
+      "         [0.0679, 0.0344, 0.0777, 0.0284, 0.1801, 0.1206, 0.2998, 0.1910],\n",
+      "         [0.0982, 0.0298, 0.1479, 0.0563, 0.1244, 0.1876, 0.1037, 0.2520]],\n",
+      "\n",
+      "        [[0.1642, 0.0531, 0.0051, 0.0924, 0.0232, 0.4650, 0.0696, 0.1274],\n",
+      "         [0.2598, 0.0652, 0.0097, 0.0568, 0.0264, 0.4170, 0.0596, 0.1056],\n",
+      "         [0.2501, 0.0407, 0.0053, 0.1077, 0.0189, 0.3562, 0.0986, 0.1225],\n",
+      "         [0.0827, 0.0811, 0.0028, 0.0639, 0.0204, 0.6040, 0.0658, 0.0794]]],\n",
+      "       grad_fn=<SoftmaxBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "src_vocab_size = 6\n",
+    "tgt_vocab_size = 8\n",
+    "d_model = 512\n",
+    "n_head = 8\n",
+    "d_ff = 2048\n",
+    "N = 6\n",
+    "dropout = 0.1\n",
+    "\n",
+    "model = make_model(src_vocab_size, tgt_vocab_size, d_model, n_head, d_ff, N, dropout)\n",
+    "# print(model)\n",
+    "# 计算模型的总参数数量\n",
+    "total_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"Total number of parameters: {total_params}\")\n",
+    "\n",
+    "# 输入数据\n",
+    "src_inputs = torch.tensor([\n",
+    "    [1, 2, 3],\n",
+    "    [4, 5, 0],\n",
+    "])\n",
+    "src_mask = get_padding_mask(src_inputs, 0)\n",
+    "\n",
+    "tgt_inputs = torch.tensor([\n",
+    "    [1, 2, 3, 4],\n",
+    "    [4, 5, 0, 0],\n",
+    "])\n",
+    "# 处理mask\n",
+    "tgt_pad_mask = get_padding_mask(tgt_inputs, 0)\n",
+    "subsequent_mask = get_subsequent_mask(4)\n",
+    "tgt_mask = tgt_pad_mask | subsequent_mask\n",
+    "tgt_mask = tgt_mask != 0\n",
+    "\n",
+    "predict = model(src_inputs, src_mask, tgt_inputs, tgt_mask)\n",
+    "print(predict.shape)\n",
+    "print(predict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c8b1100",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "af29d879",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "BLEU = 15.82\n"
+     ]
+    }
+   ],
+   "source": [
+    "# pip install sacrebleu\n",
+    "import sacrebleu\n",
+    "\n",
+    "# 场景：给\"菜谱描述生成\"模型的输出打分\n",
+    "# 参考译文（每个样本可以给多个参考，更稳）\n",
+    "refs = [\n",
+    "    ['这道菜外酥里嫩，一口下去满嘴香。', '这道菜外皮酥脆，内里鲜嫩多汁。'],\n",
+    "    ['汤底用老母鸡熬了八个小时。',     '高汤以老母鸡慢炖八小时而成。'],\n",
+    "    ['这家店的服务态度特别好。',         '该店服务热情周到。'],\n",
+    "]\n",
+    "# 模型的实际输出\n",
+    "hyps = [\n",
+    "    '这道菜外皮酥脆，里面鲜嫩多汁。',\n",
+    "    '汤底是用老母鸡炖了八个小时。',\n",
+    "    '这家店服务特别棒。',\n",
+    "]\n",
+    "\n",
+    "# tokenize='zh' 启用中文分词模式；BLEU 满分 100\n",
+    "bleu = sacrebleu.corpus_bleu(hyps, refs, tokenize='zh')\n",
+    "print(f\"BLEU = {bleu.score:.2f}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gitcode",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}