import torch
from d2l import torch as d2l

# #@save 注释用于标记这个函数可能会被保存或导入到其他脚本中使用
def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),
                  cmap='Reds'):
    """显示矩阵热图"""
    d2l.use_svg_display()  # 使用SVG格式显示图像，以获得更清晰的图像质量
    num_rows, num_cols = matrices.shape[0], matrices.shape[1]  # 获取矩阵数组的行数和列数
    fig, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize,
                                 sharex=True, sharey=True, squeeze=False)
    # 创建一个图形和一组子图，每个矩阵都会在子图中显示
    for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):
        for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):
            pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap)
            # 使用imshow函数绘制矩阵的热图，detach().numpy()将矩阵从PyTorch张量转换为NumPy数组
            if i == num_rows - 1:
                ax.set_xlabel(xlabel)  # 在最下方的子图设置x轴标签
            if j == 0:
                ax.set_ylabel(ylabel)  # 在最左侧的子图设置y轴标签
            if titles:
                ax.set_title(titles[j])  # 如果提供了标题，为每个子图设置标题
    fig.colorbar(pcm, ax=axes, shrink=0.6)  # 为子图集添加一个颜色条，显示颜色映射的数值范围

attention_weights = torch.eye(10).reshape((1, 1, 10, 10))
show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

n_train = 50  # 训练样本数
x_train, _ = torch.sort(torch.rand(n_train) * 5)   # 排序后的训练样本

def f(x):
    return 2 * torch.sin(x) + x**0.8

y_train = f(x_train) + torch.normal(0.0, 0.5, (n_train,))  # 训练样本的输出
x_test = torch.arange(0, 5, 0.1)  # 测试样本
y_truth = f(x_test)  # 测试样本的真实输出
n_test = len(x_test)  # 测试样本数

def plot_kernel_reg(y_hat):
    d2l.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'],
             xlim=[0, 5], ylim=[-1, 5])
    d2l.plt.plot(x_train, y_train, 'o', alpha=0.5);

#torch.repeat_interleave 是一个PyTorch函数，用于沿指定维度重复张量中的元素。
y_hat = torch.repeat_interleave(y_train.mean(), n_test)
plot_kernel_reg(y_hat)

# X_repeat的形状:(n_test,n_train),
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))
# x_train包含着键。attention_weights的形状：(n_test,n_train),
# 每一行都包含着要在给定的每个查询的值（y_train）之间分配的注意力权重
attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)
# y_hat的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights, y_train)
plot_kernel_reg(y_hat)

d2l.show_heatmaps(attention_weights.unsqueeze(0).unsqueeze(0),
                  xlabel='Sorted training inputs',
                  ylabel='Sorted testing inputs')

weights = torch.ones((2, 10)) * 0.1
values = torch.arange(20.0).reshape((2, 10))
torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))

# 每一行都包含着相同的训练输入
X_tile = x_train.repeat((n_train, 1))
# 每一行都包含着相同的训练输出
Y_tile = y_train.repeat((n_train, 1))
# 使得每行包含`n_train - 1`个元素，即除了自己之外的其他样本，形成了`keys`张量。
keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))
# 使得每行包含`n_train - 1`个元素，即除了自己之外的其他样本，形成了`values`张量。
values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))

net = NWKernelRegression()
loss = nn.MSELoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])

for epoch in range(5):
    trainer.zero_grad()
    l = loss(net(x_train, keys, values), y_train)
    l.sum().backward()
    trainer.step()
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

d2l.show_heatmaps(net.attention_weights.unsqueeze(0).unsqueeze(0),
                  xlabel='Sorted training inputs',
                  ylabel='Sorted testing inputs')

#@save
def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)  # 如果没有提供有效长度，直接在最后一个轴上应用softmax
    else:
        shape = X.shape  # 保存X的原始形状
        if valid_lens.dim() == 1:
            # 如果valid_lens是1D张量，则对每个样本的所有序列重复相同的长度
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            # 如果valid_lens是2D张量，则将其展平
            valid_lens = valid_lens.reshape(-1)
        # 对序列进行掩蔽操作，有效长度之外的位置被置为一个非常大的负数(-1e6)
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        # 将掩蔽后的X重新塑形为原始形状，并在最后一个轴上应用softmax
        return nn.functional.softmax(X.reshape(shape), dim=-1)

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))
"""
tensor([[[0.5980, 0.4020, 0.0000, 0.0000],
         [0.5548, 0.4452, 0.0000, 0.0000]],

        [[0.3716, 0.3926, 0.2358, 0.0000],
         [0.3455, 0.3337, 0.3208, 0.0000]]])
"""

#@save
class AdditiveAttention(nn.Module):
    """加性注意力"""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        # 使用线性层将键和查询转换到相同的隐藏维度
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        # 线性层用于计算加性注意力的分数
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        # Dropout层用于正则化
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        # 首先，通过W_k和W_q将键和查询映射到隐藏空间
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 扩展维度，使得queries和keys能够进行广播相加
        # queries形状变为(batch_size, 查询的个数, 1, num_hiddens)
        # keys形状变为(batch_size, 1, 键-值对的个数, num_hiddens)
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        # 应用tanh激活函数
        features = torch.tanh(features)
        # 计算加性注意力分数
        scores = self.w_v(features).squeeze(-1)
        # 应用掩蔽softmax函数，考虑到有效长度
        self.attention_weights = masked_softmax(scores, valid_lens)
        # 应用注意力权重到值上，进行加权求和
        # values形状：(batch_size, 键-值对的个数, 值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)

# 初始化输入数据
queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))
values = torch.arange(40, dtype=torch.float32).reshape(1, 10, 4).repeat(2, 1, 1)
valid_lens = torch.tensor([2, 6])

# 实例化加性注意力模型，设置为评估模式
attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8, dropout=0.1)
attention.eval()

# 通过加性注意力模型计算输出
output = attention(queries, keys, values, valid_lens)

"""
tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)
"""

#@save
class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 设置transpose_b=True为了交换keys的最后两个维度
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)

queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
attention(queries, keys, values, valid_lens)
"""
tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]])
"""

d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads  # 定义头的数量
        # 初始化点积注意力机制，其中包含dropout
        self.attention = d2l.DotProductAttention(dropout)
        # 为查询、键、值分别初始化线性变换层
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        # 最后的线性变换层，用于合并多头注意力的输出
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # 对输入的查询、键、值进行线性变换
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # 如果提供了有效长度，则对其进行处理以匹配扩展后的批次大小
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # 调用点积注意力机制，计算多头注意力的输出
        output = self.attention(queries, keys, values, valid_lens)

        # 合并多头注意力的输出
        output_concat = transpose_output(output, self.num_heads)
        # 通过最后的线性变换层
        return self.W_o(output_concat)

#@save
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads， num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()
"""
MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)
"""

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape
"""
torch.Size([2, 4, 100])
"""

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                                   num_hiddens, num_heads, 0.5)
attention.eval()
"""
MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)
"""

batch_size, num_queries, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
attention(X, X, X, valid_lens).shape
"""
torch.Size([2, 4, 100])
"""

#@save
class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        # 初始化一个Dropout层，用于在位置编码后对输出进行dropout操作以防止过拟合
        self.dropout = nn.Dropout(dropout)
        # 初始化一个位置编码矩阵P，它有max_len行和num_hiddens列
        # 这里max_len是序列的最大长度，num_hiddens是每个词元的隐藏向量维度
        self.P = torch.zeros((1, max_len, num_hiddens))
        
        # 计算位置编码的公式部分
        X = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1) / torch.pow(10000, torch.arange(0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 对偶数位置使用正弦函数进行编码
        self.P[:, :, 0::2] = torch.sin(X)
        # 对奇数位置使用余弦函数进行编码
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        # 将位置编码添加到输入X上
        # 这里X是一个三维张量，形状为(batch_size, sequence_length, num_hiddens)
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        # 对结果应用dropout并返回
        return self.dropout(X)

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
         figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])

for i in range(8):
    print(f'{i}的二进制是：{i:>03b}')
"""
0的二进制是：000
1的二进制是：001
2的二进制是：010
3的二进制是：011
4的二进制是：100
5的二进制是：101
6的二进制是：110
7的二进制是：111
"""

P = P[0, :, :].unsqueeze(0).unsqueeze(0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
                  ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

#@save
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        # 第一个全连接层，将输入维度从ffn_num_input变换到ffn_num_hiddens
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        # ReLU激活函数
        self.relu = nn.ReLU()
        # 第二个全连接层，将隐藏层维度从ffn_num_hiddens变换到ffn_num_outputs
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        # X是输入的张量，形状为(batch_size, sequence_length, ffn_num_input)
        # 先通过第一个全连接层和ReLU激活函数
        # 然后通过第二个全连接层
        # 输出张量的形状为(batch_size, sequence_length, ffn_num_outputs)
        return self.dense2(self.relu(self.dense1(X)))

ffn = PositionWiseFFN(4, 4, 8)
ffn.eval()
ffn(torch.ones((2, 3, 4)))

#@save
class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        # 初始化一个Dropout层，用于在加法操作之前对Y进行随机丢弃，以减少过拟合
        self.dropout = nn.Dropout(dropout)
        # 初始化一个LayerNorm层，用于在加法操作之后对结果进行规范化
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        # 首先对Y应用dropout，然后与X进行加法操作（残差连接）
        # 最后，对加法操作的结果进行层规范化
        # X是来自前一个层的输入，Y是当前层的输出
        return self.ln(self.dropout(Y) + X)

#@save
class EncoderBlock(nn.Module):
    """Transformer编码器块"""
    
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, use_bias=False, **kwargs):
 """ 
      参数:
    - key_size: 键向量的维度大小。
    - query_size: 查询向量的维度大小。
    - value_size: 值向量的维度大小。
    - num_hiddens: 自注意力层和前馈网络输出的隐藏单元数量，也是模型中间层的维度。
    - norm_shape: 层规范化应用的维度。这通常是输入特征的维度。
    - ffn_num_input: 前馈网络的输入维度。
    - ffn_num_hiddens: 前馈网络中间层的维度。
    - num_heads: 多头注意力机制中头的数量。
    - dropout: Dropout层的丢弃比例。
    - use_bias: 是否在多头自注意力机制和前馈网络中使用偏置项。默认为False。
 """  
        super(EncoderBlock, self).__init__(**kwargs)
        # 初始化多头自注意力机制，其中包括指定的头数和大小参数
        self.attention = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout,
            use_bias)
        # 第一个残差连接和层规范化模块
        self.addnorm1 = AddNorm(norm_shape, dropout)
        # 初始化基于位置的前馈网络
        self.ffn = PositionWiseFFN(
            ffn_num_input, ffn_num_hiddens, num_hiddens)
        # 第二个残差连接和层规范化模块
        self.addnorm2 = AddNorm(norm_shape, dropout)

    def forward(self, X, valid_lens):
        # 先通过多头自注意力机制，然后应用残差连接和层规范化
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        # 接着通过基于位置的前馈网络，再次应用残差连接和层规范化
        return self.addnorm2(Y, self.ffn(Y))

X = torch.ones((2, 100, 24))
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
encoder_blk.eval()
encoder_blk(X, valid_lens).shape
# 输出依然是：torch.Size([2, 100, 24])

#@save
class TransformerEncoder(d2l.Encoder):
    """Transformer编码器"""
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, use_bias=False, **kwargs):
        """
        初始化Transformer编码器。

        参数:
        - vocab_size: 词汇表大小，用于嵌入层。
        - key_size: 键向量的维度。
        - query_size: 查询向量的维度。
        - value_size: 值向量的维度。
        - num_hiddens: 自注意力层和前馈网络的输出维度，即模型的隐藏单元数。
        - norm_shape: 层规范化应用的维度，通常是模型的隐藏层维度。
        - ffn_num_input: 前馈网络的输入维度。
        - ffn_num_hiddens: 前馈网络的隐藏层维度。
        - num_heads: 多头自注意力机制中的头数。
        - num_layers: 编码器中的层数。
        - dropout: Dropout比例。
        - use_bias: 是否在自注意力和前馈网络中使用偏置项，默认为False。

        其他参数:
        - **kwargs: 捕获未明确列出的关键字参数。
        """
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens  # 存储隐藏单元数
        # 初始化嵌入层，将词汇表索引转换为固定大小的密集向量
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        # 初始化位置编码，为每个位置的词元添加位置信息
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        # 创建编码器层的堆叠
        self.blks = nn.Sequential()
        for i in range(num_layers):
            # 逐层添加编码器块
            self.blks.add_module("block"+str(i),
                EncoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, use_bias))

    def forward(self, X, valid_lens, *args):
        """
        前向传播函数。

        参数:
        - X: 输入序列的张量。
        - valid_lens: 输入序列中每个元素的有效长度。

        返回:
        - 经过Transformer编码器处理的序列张量。
        """
        # 对嵌入向量进行缩放并添加位置编码
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            # 逐个编码器块处理
            X = blk(X, valid_lens)
            # 存储每个编码器块的注意力权重
            self.attention_weights[i] = blk.attention.attention.attention_weights
        return X

encoder = TransformerEncoder(
    200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
encoder.eval()
encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape
#输出:torch.Size([2, 100, 24])

class DecoderBlock(nn.Module):
    """解码器中第i个块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, i, **kwargs):
        """
        初始化解码器块。

        参数:
        - key_size, query_size, value_size: 键、查询、值向量的维度。
        - num_hiddens: 多头注意力机制和前馈网络的隐藏层维度。
        - norm_shape: 层规范化的维度。
        - ffn_num_input, ffn_num_hiddens: 前馈网络的输入和隐藏层维度。
        - num_heads: 多头自注意力的头数。
        - dropout: Dropout层的丢弃比例。
        - i: 解码器块的索引，用于在解码过程中跟踪状态。
        """
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        # 第一个多头自注意力层，用于处理目标序列自身的依赖关系。
        self.attention1 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        # 第一个残差连接和层规范化。
        self.addnorm1 = AddNorm(norm_shape, dropout)
        # 第二个多头自注意力层，用于处理编码器的输出和当前目标序列的依赖关系。
        self.attention2 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        # 第二个残差连接和层规范化。
        self.addnorm2 = AddNorm(norm_shape, dropout)
        # 基于位置的前馈网络。
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        # 第三个残差连接和层规范化。
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def forward(self, X, state):
        """
        前向传播函数。

        参数:
        - X: 输入序列的张量。
        - state: 包含编码器输出和其他信息的状态，用于解码器的计算。

        返回:
        - 解码器块的输出和更新后的状态。
        """
        enc_outputs, enc_valid_lens = state[0], state[1]  # 编码器的输出和有效长度。
        # 处理解码器的自注意力状态更新，区分训练和预测模式。
        if state[2][self.i] is None:
            key_values = X
        else:
            key_values = torch.cat((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values

        # 自注意力机制
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # 编码器-解码器注意力机制
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        # 前馈网络
        return self.addnorm3(Z, self.ffn(Z)), state

decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, 0)
decoder_blk.eval()
X = torch.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
decoder_blk(X, state)[0].shape

class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, **kwargs):
        """
        初始化Transformer解码器。

        参数:
        - vocab_size: 词汇表的大小，用于嵌入层和最终的输出层。
        - key_size, query_size, value_size: 在多头自注意力机制中，键、查询和值的维度。
        - num_hiddens: 内部嵌入的维度，也是多头自注意力和前馈网络的输出维度。
        - norm_shape: 层规范化应用的维度，通常是模型的隐藏层维度。
        - ffn_num_input, ffn_num_hiddens: 前馈网络的输入和隐藏层的维度。
        - num_heads: 多头自注意力中的头数。
        - num_layers: 解码器中的层数。
        - dropout: Dropout层的丢弃率。
        """
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        # 词嵌入层，将输入的词汇索引转换为密集的向量表示。
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        # 位置编码层，为每个元素的嵌入添加位置信息。
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        # 解码器块的堆叠，每个块包括两个多头自注意力层和一个前馈网络。
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),
                DecoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, i))
        # 最终的线性层，将解码器的输出转换为词汇表大小的向量，用于预测下一个词元。
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        """
        初始化解码器的状态，包括编码器的输出和有效长度。
        """
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]

    def forward(self, X, state):
        """
        前向传播函数。

        参数:
        - X: 输入序列的张量。
        - state: 包含编码器输出和其他信息的状态。

        返回:
        - 解码器的输出和更新后的状态。
        """
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            # 存储当前块的解码器自注意力权重和编码器-解码器注意力权重。
            self._attention_weights[0][i] = blk.attention1.attention.attention_weights
            self._attention_weights[1][i] = blk.attention2.attention.attention_weights
        # 将解码器的输出通过一个线性层转换为预测下一个词元的概率分布。
        return self.dense(X), state

    @property
    def attention_weights(self):
        """
        返回解码器内部的注意力权重。
        """
        return self._attention_weights

# 设置模型参数
num_hiddens, num_layers, dropout = 32, 2, 0.1  # 模型隐藏单元数，层数，以及dropout比率
batch_size, num_steps = 64, 10  # 批量大小和序列长度（步数）
lr, num_epochs = 0.005, 200  # 学习率和训练周期
device = d2l.try_gpu()  # 训练使用的设备，尽量使用GPU

# 设置前馈网络和多头注意力的参数
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4  # 前馈网络输入和隐藏层维度，注意力头数
key_size, query_size, value_size = 32, 32, 32  # 设置键、查询、值的维度（对于Transformer通常相等）
norm_shape = [32]  # 层规范化应用的维度

# 加载机器翻译数据集
train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

# 初始化编码器
encoder = TransformerEncoder(
    len(src_vocab),  # 源语言词汇表大小
    key_size, query_size, value_size,
    num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
    num_heads, num_layers, dropout)

# 初始化解码器
decoder = TransformerDecoder(
    len(tgt_vocab),  # 目标语言词汇表大小
    key_size, query_size, value_size,
    num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
    num_heads, num_layers, dropout)

# 组合编码器和解码器为一个完整的模型
net = d2l.EncoderDecoder(encoder, decoder)

# 训练模型
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
#loss 0.030, 5202.9 tokens/sec on cuda:0

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')
    
"""
go . => va !,  bleu 1.000
i lost . => j'ai perdu .,  bleu 1.000
he's calm . => il est calme .,  bleu 1.000
i'm home . => je suis chez moi .,  bleu 1.000
"""

enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
    -1, num_steps))
enc_attention_weights.shape

# torch.Size([2, 4, 10, 10])

d2l.show_heatmaps(
    enc_attention_weights.cpu(), xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))

dec_attention_weights_2d = [head[0].tolist()
                            for step in dec_attention_weight_seq
                            for attn in step for blk in attn for head in blk]
dec_attention_weights_filled = torch.tensor(
    pd.DataFrame(dec_attention_weights_2d).fillna(0.0).values)
dec_attention_weights = dec_attention_weights_filled.reshape((-1, 2, num_layers, num_heads, num_steps))
dec_self_attention_weights, dec_inter_attention_weights = \
    dec_attention_weights.permute(1, 2, 3, 0, 4)
dec_self_attention_weights.shape, dec_inter_attention_weights.shape

# Plusonetoincludethebeginning-of-sequencetoken
d2l.show_heatmaps(
    dec_self_attention_weights[:, :, :, :len(translation.split()) + 1],
    xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(7, 3.5))

d2l.show_heatmaps(
    dec_inter_attention_weights, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))