文本翻译,机器翻译序列(Seq2seq)

# Seq2seq
# 进行机器翻译

import collections
import math
import os

import torch
from torch import nn
from d2l import torch as d2l
import matplotlib.pyplot as plt


class Seq2SeqEncoder(d2l.Encoder):
    def __init__(self, vocab_size, embed_size, num_hidden, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # 将文本数据转化成标记关系
        # 是为了将离散的词汇信息表示为连续的、连贯的向量空间中的点，以便模型能够更好地理解单词之间的语义关系。
        self.rnn = nn.GRU(embed_size, num_hidden, num_layers, dropout=dropout)
        # 创造一个rnn层负责进行编码

    def forward(self, x, *args):
        x = self.embedding(x)
        # 创建一个嵌入层，用于将输入的整数标记（词汇表中的单词索引）转换为密集的词嵌入。这是将词汇信息转化为连续的向量表示的一种常用方法。
        # 类似于one_hot
        x = x.permute(1, 0, 2)
        # 对x进行维度置换，将批量大小放在第二维，序列长度放在第一维。这是因为RNN层通常期望输入的维度顺序是(序列长度, 批量大小, 特征维度)。
        output, state = self.rnn(x)
        # 将处理后的输入x传递给GRU层进行编码。output包含了每个时间步的隐藏状态，state包含了GRU最后一个时间步的隐藏状态。
        return output, state


encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hidden=16,
                         num_layers=2)

encoder.eval()  # 让dropout=0
x = torch.zeros((4, 7), dtype=torch.long)
# batch_size为4,序列长度为7的一个矩阵
output, state = encoder(x)

print(output.shape)
# 在编码器中将x矩阵转置,序列长度与批量大小转置,方便进行rnn处理
print(state.shape)


# state中的2是层数,每一层在最后时刻的输出,4为批量大小,16为num_hidden

class Seq2SeqDecoder(d2l.Decoder):
    def __init__(self, vocab_size, embed_size, num_hidden, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqDecoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + num_hidden, num_hidden, num_layers,
                          dropout=dropout)
        # 创建了一个GRU层，用于对输入序列进行解码。GRU的输入维度是 embed_size + num_hidden，
        # 其中 embed_size 是嵌入层输出的维度，num_hidden 是上一个时间步的隐藏状态的维度。这个GRU层有 num_layers 层，用于对序列进行建模。
        self.dense = nn.Linear(num_hidden, vocab_size)
        # 创建了一个全连接层，用于将GRU的输出转换为词汇表大小的向量，以便生成模型的输出。这一层的输出维度等于词汇表的大小
        # 输出函数

    def init_state(self, enc_outputs, *args):
        return enc_outputs[1]

    def forward(self, x, state):
        x = self.embedding(x).permute(1, 0, 2)
        context = state[-1].repeat(x.shape[0], 1, 1)
        # 广播context，使其具有与X相同的num_steps
        # 将编码器的最后一个隐藏状态 state[-1] 广播到与 x 相同的时间步数，用于上下文信息。
        x_and_context = torch.cat((x, context), 2)
        output, state = self.rnn(x_and_context, state)
        # 将连接后的输入和隐藏状态传递给GRU层，计算解码器的输出和新的隐藏状态。
        output = self.dense(output).permute(1, 0, 2)
        # 通过全连接层将GRU的输出转换为模型的最终输出，然后进行维度置换以匹配预期的输出形状。

        # output的形状:(batch_size,num_steps,vocab_size)
        # state的形状:(num_layers,batch_size,num_hidden)
        return output, state


decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8, num_hidden=16,
                         num_layers=2)
decoder.eval()

state = decoder.init_state(encoder(x))
# 对编码器的最后状态进行初始化
output, state = decoder(x, state)
# 获取输出和状态
print(output.shape, state.shape)


# ""在序列中屏蔽不相关的项"""

def sequence_mask(x, valid_len, value=0):
    max_len = x.size(1)
    mask = torch.arange((max_len), dtype=torch.float32,
                        device=x.device)[None, :] < valid_len[:, None]
    # 定义一个筛选的函数,只保留从 0 到valid长度的值,其余的全部赋值成value
    x[~mask] = value
    return x


x = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(sequence_mask(x, torch.tensor([1, 2])))


# 作用是对于一维只保留前一个,二维只保留前两个,根据valid的长度来实现

# 通过扩展softmax交叉熵函数来屏蔽不相关

class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""

    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction = 'none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss


# 训练

def train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):
    def xavier_init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        if type(m) == nn.GRU:
            for param in m._flat_weights_names:
                if "weight" in param:
                    nn.init.xavier_uniform_(m._parameters[param])

    net.apply(xavier_init_weights)
    net.to(device)

    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    loss = MaskedSoftmaxCELoss()
    net.train()

    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs])
    # 构建图像

    for epoch in range(num_epochs):
        timer = d2l.Timer()
        metric = d2l.Accumulator(2)  # 训练损失总和，词元数量
        for batch in data_iter:
            optimizer.zero_grad()
            x, x_valid_len, y, y_valid_len = [x.to(device) for x in batch]
            bos = torch.tensor([tgt_vocab['<bos>']] * y.shape[0],
                               device=device).reshape(-1, 1)
            dec_input = torch.cat([bos, y[:, :-1]], 1)
            y_hat, _ = net(x, dec_input, x_valid_len)
            l = loss(y_hat, y, y_valid_len)
            l.sum().backward()  # 损失函数的标量进行“反向传播”
            d2l.grad_clipping(net, 1)
            num_tokens = y_valid_len.sum()
            optimizer.step()

            with torch.no_grad():
                metric.add(l.sum(), num_tokens)

        if (epoch + 1) % 10 == 0:
            animator.add(epoch + 1, (metric[0] / metric[1],))
        print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} '
              f'tokens/sec on {str(device)}')

    plt.show()


embed_size, num_hidden, num_layers, dropout = 32, 32, 2, 0.1
batch_size, num_steps = 64, 10
lr, num_epochs, device = 0.005, 300, d2l.try_gpu()

# def read_data_nmt():
#     """载入⼊“英语－法语”数据集。"""
#     data_dir = d2l.download_extract('fra-eng')
#     with open(os.path.join(data_dir, 'fra.txt'), 'r',encoding='UTF-8') as f:
#         return f.read()
# raw_text = read_data_nmt()


train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hidden, num_layers,
                         dropout)
decoder = Seq2SeqDecoder(len(tgt_vocab), embed_size, num_hidden, num_layers,
                         dropout)
net = d2l.EncoderDecoder(encoder, decoder)
train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)


# 预测

def predict_seq2seq(net, src_sentence, src_vocab, tgt_vocab, num_steps,
                    device, save_attention_weights=False):
    """序列到序列模型的预测"""
    # 在预测时将net设置为评估模式

    net.eval()

    src_tokens = src_vocab[src_sentence.lower().split(' ')] + [src_vocab['<eos>']]
    enc_valid_len = torch.tensor([len(src_tokens)], device=device)
    src_tokens = d2l.truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])

    # 添加批量轴

    enc_x = torch.unsqueeze(
        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)

    # enc_X是把句子中的 之前添加的pad和bos删除

    enc_outputs = net.encoder(enc_x, enc_valid_len)
    dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)

    # 拿出编码器的输出和最后状态

    # 添加批量轴
    dec_x = torch.unsqueeze(torch.tensor(
        [tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)
    # 没有bos的序列句子
    output_seq, attention_weight_seq = [], []

    for _ in range(num_steps):

        # 取出 输出和状态,作为下一步的输入 和 更新状态
        y, dec_state = net.decoder(dec_x, dec_state)

        dec_x = y.argmax(dim=2)  # 出去输出,即最大预测概率
        pred = dec_x.squeeze(dim=0).type(torch.int32).item()

        # 保存注意力权重（稍后讨论）
        if save_attention_weights:
            attention_weight_seq.append(net.decoder.attention_weights)

        # 一旦序列结束词元被预测，输出序列的生成就完成了
        if pred == tgt_vocab['<eos>']:
            break
        output_seq.append(pred)

    return ' '.join(tgt_vocab.to_tokens(output_seq)), attention_weight_seq


def bleu(pred_seq, label_seq, k):  # @save
    """计算BLEU"""
    pred_tokens, label_tokens = pred_seq.split(' '), label_seq.split(' ')
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    score = math.exp(min(0, 1 - len_label / len_pred))
    for n in range(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        for i in range(len_label - n + 1):
            label_subs[' '.join(label_tokens[i: i + n])] += 1
        for i in range(len_pred - n + 1):
            if label_subs[' '.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[' '.join(pred_tokens[i: i + n])] -= 1
        score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    return score


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, attention_weight_seq = predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device)
    print(f'{eng} => {translation}, bleu {bleu(translation, fra, k=2):.3f}')