utils.py

import torch
import spacy


def translate_sentence(model, sentence, zh_vocab, en_ivocab, device, max_len):
    """
    使用训练好的Transformer模型翻译中文句子为英文句子

    参数:
        model: 训练好的Transformer模型
        sentence: 要翻译的中文句子(字符串)
        zh_vocab: 中文词汇表(字典)
        en_ivocab: 英文反向词汇表(字典，index->token)
        device: 设备('cuda'或'cpu')
        max_len: 最大序列长度

    返回:
        翻译后的英文句子(字符串)
    """
    # 加载SpaCy中文分词器
    spacy_zh = spacy.load('zh_core_web_sm')
    # spacy_en = spacy.load('en_core_web_sm')  # 注释掉了，因为不需要英文分词

    # 使用SpaCy对输入句子进行分词
    zh_tokens = [tok.text for tok in spacy_zh(sentence)]
    # en_tokens = [tok.text for tok in spacy_en(answer)]  # 注释掉了

    # 初始化中文序列的数值表示
    zh_nums = []
    # en_nums = []  # 注释掉了

    # 将中文token转换为词汇表中的ID
    for word in zh_tokens:
        try:
            zh_nums.append(zh_vocab[word])  # 使用词汇表获取ID
        except KeyError:
            pass  # 如果单词不在词汇表中，忽略它

    # 在序列前后添加特殊标记
    zh_nums.insert(0, 1)  # 1=<SOS> (Start of Sentence)
    zh_nums.append(2)  # 2=<EOS> (End of Sentence)

    '''
    # 以下代码被注释掉了，是另一种实现方式
    # 将数值化序列填充到最大长度
    zh_res = [0] * max_len
    zh_res[:len(zh_nums)] = zh_nums

    # 构建英文序列(同样被注释掉了)
    for word in en_tokens:
        try:
            en_nums.append(en_vocab[word])
        except KeyError:
            pass

    en_nums.insert(0, 1)
    en_nums.append(2)

    en_res = [0] * max_len
    en_res[:len(en_nums)] = en_nums
    '''

    # 将数值化后的中文序列转换为PyTorch张量，并调整维度
    src = torch.tensor(zh_nums).unsqueeze(1).to(device)  # shape: (seq_len, 1)
    src = torch.transpose(src, 0, 1)  # shape: (1, seq_len)

    '''
    # 以下代码被注释掉了，是另一种实现方式
    trg = torch.tensor(en_res).unsqueeze(1).to(device)
    trg = torch.transpose(trg, 0, 1)

    # 前向传播(被注释掉了)
    output = model(src, trg)
    output = output.reshape(-1, output.shape[2])
    return tensor2sentence(output, en_ivocab, max_len)
    '''

    # 初始化输出序列(从<SOS>开始)
    outputs = [1]  # 1=<SOS>

    # 逐步生成翻译结果
    for i in range(max_len):
        # 构建当前目标序列(形状: (1, seq_len))
        trg = torch.tensor(outputs).unsqueeze(1).to(device)
        trg = torch.transpose(trg, 0, 1)  # shape: (1, seq_len)

        # 在推理模式下进行前向传播
        with torch.no_grad():
            output = model(src, trg)  # shape: (1, seq_len, trg_vocab_size)

        # 获取当前时间步的最佳预测
        best_guess = output.argmax(2)[:, -1].item()  # 取最后一个时间步的预测

        # 将预测结果添加到输出序列
        outputs.append(best_guess)

        # 如果预测到<EOS>，则停止生成
        if best_guess == 2:
            break

    # 将预测的ID序列转换为英文句子
    translated_sentence = ""
    for i in range(max_len):
        try:
            word = en_ivocab[outputs[i]]  # 通过反向词汇表获取单词
            translated_sentence += word
            if word != "<eos>":  # 如果不是结束标记，则添加空格
                translated_sentence += " "
            else:
                break
        except KeyError:
            # 如果ID不在反向词汇表中(理论上不应该发生)
            break

    # 移除末尾可能的空格
    translated_sentence = translated_sentence.strip()
    return translated_sentence


def tensor2sentence(output, en_ivocab, max_len):
    """
    将模型输出的张量转换为英文句子

    参数:
        output: 模型输出的张量(shape: (seq_len, batch_size, trg_vocab_size))
        en_ivocab: 英文反向词汇表(字典)
        max_len: 最大序列长度

    返回:
        翻译后的英文句子(字符串)
    """
    # 获取每个时间步的最佳预测(形状: (seq_len, batch_size))
    best_words = output.argmax(1)  # 假设batch_size=1

    sentence = ""
    for i in range(max_len):
        try:
            num = best_words[i].item()  # 获取预测的ID
            word = en_ivocab[num]  # 通过反向词汇表获取单词
            sentence += word
            if word != "<eos>":  # 如果不是结束标记，则添加空格
                sentence += " "
            else:
                break
        except IndexError:
            # 如果超出序列长度，则停止
            break

    return sentence.strip()  # 移除末尾可能的空格


def list2sentence(best_words, zh_ivocab, max_len):
    """
    将数值化的ID列表转换为中文句子(备用函数)

    参数:
        best_words: 数值化的ID列表(形状: (seq_len,))
        zh_ivocab: 中文反向词汇表(字典)
        max_len: 最大序列长度

    返回:
        翻译后的中文句子(字符串)
    """
    sentence = ""
    for i in range(max_len):
        try:
            num = best_words[i].item()  # 获取ID
            word = zh_ivocab[num]  # 通过反向词汇表获取单词
            sentence += word
            if word != "<eos>":  # 如果不是结束标记，则添加空格
                sentence += " "
            else:
                break
        except IndexError:
            # 如果超出序列长度，则停止
            break

    return sentence.strip()  # 移除末尾可能的空格


def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
    """
    保存模型和优化器的状态

    参数:
        state: 包含模型和优化器状态的字典
        filename: 保存的文件名(默认: "my_checkpoint.pth.tar")
    """
    print("=> Saving checkpoint")
    torch.save(state, filename)  # 保存到文件


def load_checkpoint(checkpoint, model, optimizer):
    """
    从检查点加载模型和优化器的状态

    参数:
        checkpoint: 包含模型和优化器状态的字典
        model: 要加载状态的模型
        optimizer: 要加载状态的优化器
    """
    print("=> Loading checkpoint")
    model.load_state_dict(checkpoint["state_dict"])  # 加载模型状态
    optimizer.load_state_dict(checkpoint["optimizer"])  # 加载优化器状态