datawhale - 基于术语词典干预的机器翻译挑战赛（一）

文章目录

[torchtext 库是干什么用的？](#torchtext 库是干什么用的？)
[TranslationDataset 类](#TranslationDataset 类)
[定义 Seq2Seq模型](#定义 Seq2Seq模型)
- - Encoder
  - Decoder
  - [Seq2Seq 类](#Seq2Seq 类)
[load_terminology_dictionary 函数](#load_terminology_dictionary 函数)
- - - 示例用法
[train 函数](#train 函数)
主程序代码
模型评价
- [load_sentences 函数](#load_sentences 函数)
- [translate_sentence 函数](#translate_sentence 函数)
- [evaluate_bleu 函数](#evaluate_bleu 函数)
- 主程序
测试集上进行推理
- [inference 函数](#inference 函数)
- 主程序

torchtext 库是干什么用的？

torchtext 是一个用于处理文本数据的库，它是 PyTorch 生态系统的一部分。这个库主要用于简化文本数据的预处理和加载过程，使得在深度学习模型中使用文本数据变得更加容易。以下是 torchtext 库的一些主要功能：

数据加载 ：torchtext 提供了方便的 API 来加载和处理各种文本数据集，包括常见的数据集格式如 CSV、TSV 等。
文本预处理：它包含了一系列的文本预处理工具，如分词、词干提取、词性标注等，这些工具可以帮助你将原始文本转换为模型可以理解的格式。
词汇表管理 ：torchtext 可以自动构建词汇表，并将文本转换为数值表示（如词嵌入），这对于深度学习模型来说是必不可少的。
数据迭代器：它提供了数据迭代器，可以方便地将数据分批加载到模型中进行训练，支持多线程和多进程加载，提高了数据加载的效率。
与 PyTorch 集成 ：torchtext 与 PyTorch 深度集成，可以直接将处理好的数据输入到 PyTorch 模型中，简化了整个深度学习流程。

TranslationDataset 类

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchtext.data.utils import get_tokenizer
from collections import Counter
import random
from torch.utils.data import Subset, DataLoader
import time

# 定义数据集类
# 修改TranslationDataset类以处理术语
class TranslationDataset(Dataset):
    def __init__(self, filename, terminology):
        self.data = []
        with open(filename, 'r', encoding='utf-8') as f:
            for line in f:
                en, zh = line.strip().split('\t')
                self.data.append((en, zh))

class TranslationDataset(Dataset):：定义一个继承自Dataset类的子类TranslationDataset。
def __init__(self, filename, terminology):：初始化方法，接受两个参数：filename（数据文件名）和terminology（术语词典）。
self.data = []：初始化一个空列表self.data，用于存储数据。
with open(filename, 'r', encoding='utf-8') as f:：打开文件，使用UTF-8编码读取。
for line in f:：逐行读取文件内容。
en, zh = line.strip().split('\t')：去掉行末的换行符，并按制表符分割成英文和中文。
self.data.append((en, zh))：将英文和中文对添加到self.data列表中。

python 复制代码

        self.terminology = terminology
        
        # 创建词汇表，注意这里需要确保术语词典中的词也被包含在词汇表中
        self.en_tokenizer = get_tokenizer('basic_english')
        self.zh_tokenizer = list  # 使用字符级分词

self.terminology = terminology：将传入的术语词典赋值给self.terminology。
self.en_tokenizer = get_tokenizer('basic_english')：使用basic_english分词器对英文进行分词。
self.zh_tokenizer = list：使用字符级分词，即将中文句子拆分成单个字符。

python 复制代码

        en_vocab = Counter(self.terminology.keys())  # 确保术语在词汇表中
        zh_vocab = Counter()
        
        print("en_vocab的值为：", en_vocab)
        print("zn_vocab的值为：", zn_vocab)

en_vocab = Counter(self.terminology.keys())：初始化英文词汇表，确保术语词典中的词在词汇表中。
zh_vocab = Counter()：初始化中文词汇表。
print("en_vocab的值为：", en_vocab)：打印英文词汇表的值。
print("zn_vocab的值为：", zn_vocab)：打印中文词汇表的值。

python 复制代码

        for en, zh in self.data:
            en_vocab.update(self.en_tokenizer(en))
            zh_vocab.update(self.zh_tokenizer(zh))

for en, zh in self.data:：遍历数据集中的每一对英文和中文。
en_vocab.update(self.en_tokenizer(en))：更新英文词汇表，统计每个词的出现频率。
zh_vocab.update(self.zh_tokenizer(zh))：更新中文词汇表，统计每个字符的出现频率。

python 复制代码

        # 添加术语到词汇表
        self.en_vocab = ['<pad>', '<sos>', '<eos>'] + list(self.terminology.keys()) + [word for word, _ in en_vocab.most_common(10000)]
        self.zh_vocab = ['<pad>', '<sos>', '<eos>'] + [word for word, _ in zh_vocab.most_common(10000)]

self.en_vocab = ['<pad>', '<sos>', '<eos>'] + list(self.terminology.keys()) + [word for word, _ in en_vocab.most_common(10000)]：构建英文词汇表，包含特殊标记（<pad>, <sos>, <eos>）、术语词典中的词以及出现频率最高的10000个词。
self.zh_vocab = ['<pad>', '<sos>', '<eos>'] + [word for word, _ in zh_vocab.most_common(10000)]：构建中文词汇表，包含特殊标记（<pad>, <sos>, <eos>）以及出现频率最高的10000个字符。

python 复制代码

        self.en_word2idx = {word: idx for idx, word in enumerate(self.en_vocab)}
        self.zh_word2idx = {word: idx for idx, word in enumerate(self.zh_vocab)}

self.en_word2idx = {word: idx for idx, word in enumerate(self.en_vocab)}：构建英文词到索引的映射字典。
self.zh_word2idx = {word: idx for idx, word in enumerate(self.zh_vocab)}：构建中文词到索引的映射字典。

python 复制代码

    def __len__(self):
        return len(self.data)

def __len__(self):：定义数据集的长度方法。
return len(self.data)：返回数据集的长度。

python 复制代码

    def __getitem__(self, idx):
        en, zh = self.data[idx]
        en_tensor = torch.tensor([self.en_word2idx.get(word, self.en_word2idx['<sos>']) for word in self.en_tokenizer(en)] + [self.en_word2idx['<eos>']])
        zh_tensor = torch.tensor([self.zh_word2idx.get(word, self.zh_word2idx['<sos>']) for word in self.zh_tokenizer(zh)] + [self.zh_word2idx['<eos>']])
        return en_tensor, zh_tensor

def __getitem__(self, idx):：定义获取数据项的方法。
en, zh = self.data[idx]：获取指定索引的数据项。
en_tensor = torch.tensor([self.en_word2idx.get(word, self.en_word2idx['<sos>']) for word in self.en_tokenizer(en)] + [self.en_word2idx['<eos>']])：将英文句子转换为索引张量，并在末尾添加<eos>标记。
zh_tensor = torch.tensor([self.zh_word2idx.get(word, self.zh_word2idx['<sos>']) for word in self.zh_tokenizer(zh)] + [self.zh_word2idx['<eos>']])：将中文句子转换为索引张量，并在末尾添加<eos>标记。
return en_tensor, zh_tensor：返回转换后的英文和中文张量。

python 复制代码

def collate_fn(batch):
    en_batch, zh_batch = [], []
    for en_item, zh_item in batch:
        en_batch.append(en_item)
        zh_batch.append(zh_item)
    
    # 对英文和中文序列分别进行填充
    en_batch = nn.utils.rnn.pad_sequence(en_batch, padding_value=0, batch_first=True)
    zh_batch = nn.utils.rnn.pad_sequence(zh_batch, padding_value=0, batch_first=True)
    
    return en_batch, zh_batch

def collate_fn(batch):：定义一个用于处理批量数据的函数。
en_batch, zh_batch = [], []：初始化两个空列表，用于存储批量数据。
for en_item, zh_item in batch:：遍历批量数据中的每一对英文和中文张量。
en_batch.append(en_item)：将英文张量添加到en_batch列表中。
zh_batch.append(zh_item)：将中文张量添加到zh_batch列表中。
en_batch = nn.utils.rnn.pad_sequence(en_batch, padding_value=0, batch_first=True)：对英文序列进行填充，使其长度一致。
zh_batch = nn.utils.rnn.pad_sequence(zh_batch, padding_value=0, batch_first=True)：对中文序列进行填充，使其长度一致。
return en_batch, zh_batch：返回填充后的英文和中文批量数据。

定义 Seq2Seq模型

Encoder

python 复制代码

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        # src shape: [batch_size, src_len]
        embedded = self.dropout(self.embedding(src))
        # embedded shape: [batch_size, src_len, emb_dim]
        outputs, hidden = self.rnn(embedded)
        # outputs shape: [batch_size, src_len, hid_dim]
        # hidden shape: [n_layers, batch_size, hid_dim]
        return outputs, hidden

初始化 ：
- input_dim：输入词汇表的大小。
- emb_dim：嵌入层的维度。
- hid_dim：隐藏层的维度。
- n_layers：RNN的层数。
- dropout：Dropout的概率。
前向传播 ：
- 输入src的形状为[batch_size, src_len]。
- 通过嵌入层和Dropout层，得到embedded，形状为[batch_size, src_len, emb_dim]。
- 通过GRU层，得到outputs和hidden，形状分别为[batch_size, src_len, hid_dim]和[n_layers, batch_size, hid_dim]。

Decoder

python 复制代码

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
        self.fc_out = nn.Linear(hid_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input, hidden):
        # input shape: [batch_size, 1]
        # hidden shape: [n_layers, batch_size, hid_dim]
        
        embedded = self.dropout(self.embedding(input))
        # embedded shape: [batch_size, 1, emb_dim]
        
        output, hidden = self.rnn(embedded, hidden)
        # output shape: [batch_size, 1, hid_dim]
        # hidden shape: [n_layers, batch_size, hid_dim]
        
        prediction = self.fc_out(output.squeeze(1))
        # prediction shape: [batch_size, output_dim]
        
        return prediction, hidden

初始化 ：
- output_dim：输出词汇表的大小。
- emb_dim：嵌入层的维度。
- hid_dim：隐藏层的维度。
- n_layers：RNN的层数。
- dropout：Dropout的概率。
前向传播 ：
- 输入input的形状为[batch_size, 1]，hidden的形状为[n_layers, batch_size, hid_dim]。
- 通过嵌入层和Dropout层，得到embedded，形状为[batch_size, 1, emb_dim]。
- 通过GRU层，得到output和hidden，形状分别为[batch_size, 1, hid_dim]和[n_layers, batch_size, hid_dim]。
- 通过全连接层，得到prediction，形状为[batch_size, output_dim]。

Seq2Seq 类

python 复制代码

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        # src shape: [batch_size, src_len]
        # trg shape: [batch_size, trg_len]
        
        batch_size = src.shape[0]
        trg_len = trg.shape[1]
        trg_vocab_size = self.decoder.output_dim

        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)
        
        _, hidden = self.encoder(src)
        
        input = trg[:, 0].unsqueeze(1)  # Start token
        
        for t in range(1, trg_len):
            output, hidden = self.decoder(input, hidden)
            outputs[:, t, :] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[:, t].unsqueeze(1) if teacher_force else top1.unsqueeze(1)

        return outputs

初始化 ：
- encoder：编码器实例。
- decoder：解码器实例。
- device：设备（CPU或GPU）。
前向传播 ：
- 输入src的形状为[batch_size, src_len]，trg的形状为[batch_size, trg_len]。
- 初始化outputs，形状为[batch_size, trg_len, trg_vocab_size]。
- 通过编码器得到hidden。
- 初始化解码器的输入为trg[:, 0].unsqueeze(1)，即目标序列的起始标记。
- 循环解码目标序列的每个时间步：
  - 通过解码器得到output和hidden。
  - 将output存储到outputs中。
  - 根据teacher_forcing_ratio决定是否使用教师强制（即使用目标序列的下一个词作为输入，还是使用解码器的预测结果）。

load_terminology_dictionary 函数

函数定义:
python 复制代码
```
def load_terminology_dictionary(dict_file):
```
- dict_file: 这是一个字符串参数，表示包含术语词典的文件路径。
初始化术语词典:
python 复制代码
```
terminology = {}
```
- 创建一个空的字典 terminology，用于存储从文件中读取的术语。
打开文件并读取内容:
python 复制代码
```
with open(dict_file, 'r', encoding='utf-8') as f:
    for line in f:
        en_term, ch_term = line.strip().split('\t')
        terminology[en_term] = ch_term
```
- 使用 with open 语句打开文件，确保文件在使用后正确关闭。
- 文件以只读模式 ('r') 打开，并指定编码为 utf-8。
- 逐行读取文件内容，每行包含一个英文术语和对应的中文术语，用制表符 (\t) 分隔。
- line.strip() 用于去除行末的换行符和其他空白字符。
- split('\t') 将行按制表符分隔成两个部分，分别赋值给 en_term 和 ch_term。
- 将英文术语作为键，中文术语作为值，添加到 terminology 字典中。
返回术语词典:
python 复制代码
```
return terminology
```
- 函数返回加载后的术语词典。

示例用法

假设有一个文件 terminology.txt，内容如下：

apple  苹果
banana  香蕉
orange  橙子

调用 load_terminology_dictionary 函数：

python 复制代码

terminology = load_terminology_dictionary('terminology.txt')
print(terminology)

输出结果将是：

python 复制代码

{'apple': '苹果', 'banana': '香蕉', 'orange': '橙子'}

train 函数

python 复制代码

def train(model, iterator, optimizer, criterion, clip):

定义一个名为 train 的函数，该函数接受五个参数：model（模型）、iterator（数据迭代器）、optimizer（优化器）、criterion（损失函数）和 clip（梯度裁剪的阈值）。

python 复制代码

    model.train()

将模型设置为训练模式。这会启用诸如 dropout 和 batch normalization 等训练特有的操作。

python 复制代码

    epoch_loss = 0

初始化 epoch_loss 变量为 0，用于累积整个 epoch 的损失。

python 复制代码

    for i, (src, trg) in enumerate(iterator):

使用 enumerate 遍历数据迭代器 iterator，每次迭代获取一个批次的数据 (src, trg)，其中 src 是源序列，trg 是目标序列。

python 复制代码

        src, trg = src.to(device), trg.to(device)

将源序列和目标序列移动到指定的设备（如 GPU）上。

python 复制代码

        optimizer.zero_grad()

将优化器的梯度清零，以防止梯度累积。

python 复制代码

        output = model(src, trg)

将源序列和目标序列输入模型，得到模型的输出。

python 复制代码

        output_dim = output.shape[-1]

获取输出张量的最后一个维度的大小，即输出序列的词汇表大小。

python 复制代码

        output = output[:, 1:].contiguous().view(-1, output_dim)

去掉输出序列的第一个时间步（通常是起始标记），并将剩余部分展平成二维张量，形状为 (batch_size * (sequence_length - 1), output_dim)。

python 复制代码

        trg = trg[:, 1:].contiguous().view(-1)

去掉目标序列的第一个时间步，并将剩余部分展平成一维张量，形状为 (batch_size * (sequence_length - 1))。

python 复制代码

        loss = criterion(output, trg)

计算损失，使用定义的损失函数 criterion。

python 复制代码

        loss.backward()

反向传播，计算梯度。

python 复制代码

        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)

对模型的梯度进行裁剪，防止梯度爆炸。裁剪的阈值由参数 clip 指定。

python 复制代码

        optimizer.step()

更新模型的参数。

python 复制代码

        epoch_loss += loss.item()

将当前批次的损失值累加到 epoch_loss 中。

python 复制代码

    return epoch_loss / len(iterator)

返回整个 epoch 的平均损失。

主程序代码

python 复制代码

if __name__ == '__main__':

这是一个常见的Python惯用法，用于判断当前模块是否是主程序入口。如果是，则执行后续代码。后同。

python 复制代码

    start_time = time.time()  # 开始计时

记录程序开始执行的时间，用于后续计算总运行时间。

python 复制代码

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

检查是否有可用的CUDA设备（即GPU），如果有则使用GPU，否则使用CPU。

python 复制代码

    terminology = load_terminology_dictionary('../dataset/en-zh.dic')

加载术语字典，该字典可能包含特定领域的术语翻译。

python 复制代码

    dataset = TranslationDataset('../dataset/train.txt', terminology=terminology)

创建一个 TranslationDataset 对象，该对象用于加载和处理训练数据。

python 复制代码

    N = 1000  # int(len(dataset) * 1)  # 或者你可以设置为数据集大小的一定比例，如 int(len(dataset) * 0.1)
    subset_indices = list(range(N))
    subset_dataset = Subset(dataset, subset_indices)

选择数据集的前1000个样本进行训练。subset_indices 是一个包含前1000个索引的列表，subset_dataset 是原始数据集的一个子集。

python 复制代码

    train_loader = DataLoader(subset_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)

创建一个 DataLoader 对象，用于批量加载数据。batch_size 设置为32，shuffle 设置为 True 以打乱数据顺序，collate_fn 是一个自定义函数，用于处理批次数据的拼接。

python 复制代码

    INPUT_DIM = len(dataset.en_vocab)
    OUTPUT_DIM = len(dataset.zh_vocab)
    ENC_EMB_DIM = 256
    DEC_EMB_DIM = 256
    HID_DIM = 512
    N_LAYERS = 2
    ENC_DROPOUT = 0.5
    DEC_DROPOUT = 0.5

定义模型参数，包括输入和输出维度、嵌入维度、隐藏层维度、层数和 dropout 率。

python 复制代码

    enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)
    dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)
    model = Seq2Seq(enc, dec, device).to(device)

初始化编码器和解码器，并将其组合成一个 Seq2Seq 模型。然后将模型移动到指定的设备（GPU或CPU）。

python 复制代码

    optimizer = optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss(ignore_index=dataset.zh_word2idx['<pad>'])

定义优化器和损失函数。优化器使用 Adam，损失函数使用交叉熵损失，并忽略目标序列中的填充标记 <pad>。

python 复制代码

    N_EPOCHS = 10
    CLIP = 1

设置训练的 epoch 数和梯度裁剪的阈值。

python 复制代码

    for epoch in range(N_EPOCHS):
        train_loss = train(model, train_loader, optimizer, criterion, CLIP)
        print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')

进行10个 epoch 的训练。每个 epoch 结束后，打印当前 epoch 的训练损失。

python 复制代码

    torch.save(model.state_dict(), './translation_model_GRU.pth')

在所有 epoch 结束后，保存模型的参数到文件 translation_model_GRU.pth。

python 复制代码

    end_time = time.time()  # 结束计时
    elapsed_time_minute = (end_time - start_time)/60
    print(f"Total running time: {elapsed_time_minute:.2f} minutes")

记录程序结束时间，计算并打印总运行时间（以分钟为单位）。

模型评价

python 复制代码

import torch
from sacrebleu.metrics import BLEU
from typing import List
# 假设已经定义了TranslationDataset, Encoder, Decoder, Seq2Seq类

导入必要的库和模块。torch 用于深度学习，sacrebleu.metrics.BLEU 用于评估翻译质量，typing.List 用于类型注解。

load_sentences 函数

python 复制代码

def load_sentences(file_path: str) -> List[str]:
    with open(file_path, 'r', encoding='utf-8') as f:
        return [line.strip() for line in f]

定义函数 load_sentences，用于从文件中加载句子。file_path 是文件路径，函数返回一个包含所有句子的列表。

translate_sentence 函数

python 复制代码

def translate_sentence(sentence: str, model: Seq2Seq, dataset: TranslationDataset, terminology, device: torch.device, max_length: int = 50):

定义一个函数 translate_sentence，用于翻译单个句子。参数包括：
- sentence: 要翻译的句子。
- model: 预训练的 Seq2Seq 模型。
- dataset: TranslationDataset 对象，包含词汇表等信息。
- terminology: 术语词典。
- device: 计算设备（CPU 或 GPU）。
- max_length: 翻译句子的最大长度，默认为 50。

python 复制代码

    model.eval()

将模型设置为评估模式，关闭 dropout 和 batch normalization 等训练特有的操作。

python 复制代码

    tokens = dataset.en_tokenizer(sentence)

使用 dataset 中的 en_tokenizer 对输入句子进行分词，得到一个 token 列表。

python 复制代码

    tensor = torch.LongTensor([dataset.en_word2idx.get(token, dataset.en_word2idx['<sos>']) for token in tokens]).unsqueeze(0).to(device)  # [1, seq_len]

将 token 列表转换为索引列表，并转换为 PyTorch 张量。如果 token 不在词汇表中，则使用 <sos> 的索引。然后将张量增加一个维度并移动到指定的设备。

python 复制代码

    with torch.no_grad():
        _, hidden = model.encoder(tensor)

关闭梯度计算，以节省内存和提高速度。然后使用编码器对输入张量进行编码，得到隐藏状态 hidden。

python 复制代码

    translated_tokens = []

初始化一个空列表 translated_tokens，用于存储翻译的 token。

python 复制代码

    input_token = torch.LongTensor([[dataset.zh_word2idx['<sos>']]]).to(device)  # [1, 1]

初始化解码器的输入 token，即 <sos> 的索引，并将其转换为张量并移动到指定的设备。

python 复制代码

    for _ in range(max_length):

开始一个循环，最多迭代 max_length 次。

python 复制代码

        output, hidden = model.decoder(input_token, hidden)

使用解码器对当前输入 token 和隐藏状态进行解码，得到输出 output 和新的隐藏状态 hidden。

python 复制代码

        top_token = output.argmax(1)

从输出中选择概率最高的 token 索引。

python 复制代码

        translated_token = dataset.zh_vocab[top_token.item()]

将索引转换为对应的 token。

python 复制代码

        if translated_token == '<eos>':
            break

如果翻译的 token 是 <eos>，则结束循环。

python 复制代码

        if translated_token in terminology.values():
            for en_term, ch_term in terminology.items():
                if translated_token == ch_term:
                    translated_token = en_term
                    break

如果翻译的 token 在术语词典中，则使用术语词典中的对应词替换。

python 复制代码

        translated_tokens.append(translated_token)

将翻译的 token 添加到 translated_tokens 列表中。

python 复制代码

        input_token = top_token.unsqueeze(1)  # [1, 1]

更新输入 token 为当前选择的 token 索引。

python 复制代码

    return ''.join(translated_tokens)

将翻译的 token 列表连接成一个字符串并返回。

evaluate_bleu 函数

python 复制代码

def evaluate_bleu(model: Seq2Seq, dataset: TranslationDataset, src_file: str, ref_file: str, terminology, device: torch.device):

定义一个函数 evaluate_bleu，用于评估模型的 BLEU 分数。参数包括：
- model: 预训练的 Seq2Seq 模型。
- dataset: TranslationDataset 对象，包含词汇表等信息。
- src_file: 包含源语言句子的文件路径。
- ref_file: 包含参考翻译句子的文件路径。
- terminology: 术语词典。
- device: 计算设备（CPU 或 GPU）。

python 复制代码

    model.eval()

将模型设置为评估模式，关闭 dropout 和 batch normalization 等训练特有的操作。

python 复制代码

    src_sentences = load_sentences(src_file)

调用 load_sentences 函数加载源语言句子，返回一个包含所有源句子的列表。

python 复制代码

    ref_sentences = load_sentences(ref_file)

调用 load_sentences 函数加载参考翻译句子，返回一个包含所有参考翻译句子的列表。

python 复制代码

    translated_sentences = []

初始化一个空列表 translated_sentences，用于存储翻译的句子。

python 复制代码

    for src in src_sentences:

遍历源语言句子列表。

python 复制代码

        translated = translate_sentence(src, model, dataset, terminology, device)

调用 translate_sentence 函数翻译当前的源句子，返回翻译后的句子。

python 复制代码

        translated_sentences.append(translated)

将翻译后的句子添加到 translated_sentences 列表中。

python 复制代码

    bleu = BLEU()

创建一个 BLEU 对象，用于计算 BLEU 分数。

python 复制代码

    score = bleu.corpus_score(translated_sentences, [ref_sentences])

调用 corpus_score 方法计算翻译句子的 BLEU 分数。translated_sentences 是模型生成的翻译句子列表，[ref_sentences] 是参考翻译句子列表。

python 复制代码

    return score

返回计算得到的 BLEU 分数。

主程序

python 复制代码

if __name__ == '__main__':

python 复制代码

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

检查是否有可用的CUDA设备（即GPU），如果有则使用GPU，否则使用CPU。

python 复制代码

    terminology = load_terminology_dictionary('../dataset/en-zh.dic')

加载术语词典，该字典可能包含特定领域的术语翻译。

python 复制代码

    dataset = TranslationDataset('../dataset/train.txt', terminology)

创建一个 TranslationDataset 对象，该对象用于加载和处理训练数据，并传递术语词典。

python 复制代码

    INPUT_DIM = len(dataset.en_vocab)
    OUTPUT_DIM = len(dataset.zh_vocab)
    ENC_EMB_DIM = 256
    DEC_EMB_DIM = 256
    HID_DIM = 512
    N_LAYERS = 2
    ENC_DROPOUT = 0.5
    DEC_DROPOUT = 0.5

定义模型参数，包括输入和输出维度、嵌入维度、隐藏层维度、层数和 dropout 率。

python 复制代码

    enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)
    dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)
    model = Seq2Seq(enc, dec, device).to(device)

初始化编码器和解码器，并将其组合成一个 Seq2Seq 模型。然后将模型移动到指定的设备（GPU或CPU）。

python 复制代码

    model.load_state_dict(torch.load('./translation_model_GRU.pth'))

加载预训练的模型参数。

python 复制代码

    bleu_score = evaluate_bleu(model, dataset, '../dataset/dev_en.txt', '../dataset/dev_zh.txt', terminology=terminology, device=device)

调用 evaluate_bleu 函数评估模型的 BLEU 分数。传递的参数包括模型、数据集、源语言文件路径、参考翻译文件路径、术语词典和设备。

python 复制代码

    print(f'BLEU-4 score: {bleu_score.score:.2f}')

打印评估得到的 BLEU-4 分数，保留两位小数。

测试集上进行推理

inference 函数

python 复制代码

def inference(model: Seq2Seq, dataset: TranslationDataset, src_file: str, save_dir: str, terminology, device: torch.device):

定义一个函数 inference，用于进行模型推理并将结果保存到文件中。参数包括：
- model: 预训练的 Seq2Seq 模型。
- dataset: TranslationDataset 对象，包含词汇表等信息。
- src_file: 包含源语言句子的文件路径。
- save_dir: 保存翻译结果的文件路径。
- terminology: 术语词典。
- device: 计算设备（CPU 或 GPU）。

python 复制代码

    model.eval()

将模型设置为评估模式，关闭 dropout 和 batch normalization 等训练特有的操作。

python 复制代码

    src_sentences = load_sentences(src_file)

调用 load_sentences 函数加载源语言句子，返回一个包含所有源句子的列表。

python 复制代码

    translated_sentences = []

初始化一个空列表 translated_sentences，用于存储翻译的句子。

python 复制代码

    for src in src_sentences:

遍历源语言句子列表。

python 复制代码

        translated = translate_sentence(src, model, dataset, terminology, device)

调用 translate_sentence 函数翻译当前的源句子，返回翻译后的句子。

python 复制代码

        translated_sentences.append(translated)

将翻译后的句子添加到 translated_sentences 列表中。

python 复制代码

    text = '\n'.join(translated_sentences)

将 translated_sentences 列表中的所有句子连接成一个字符串，每个句子之间用换行符分隔。

python 复制代码

    with open(save_dir, 'w', encoding='utf-8') as f:

打开一个文件，如果不存在则创建，'w' 表示写模式，encoding='utf-8' 表示使用 UTF-8 编码。

python 复制代码

        f.write(text)

将连接后的字符串写入文件。

主程序

python 复制代码

if __name__ == '__main__':

python 复制代码

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

检查是否有可用的CUDA设备（即GPU），如果有则使用GPU，否则使用CPU。

python 复制代码

    terminology = load_terminology_dictionary('../dataset/en-zh.dic')

加载术语词典，该字典可能包含特定领域的术语翻译。

python 复制代码

    dataset = TranslationDataset('../dataset/train.txt', terminology=terminology)

创建一个 TranslationDataset 对象，该对象用于加载和处理训练数据，并传递术语词典。

python 复制代码

    INPUT_DIM = len(dataset.en_vocab)
    OUTPUT_DIM = len(dataset.zh_vocab)
    ENC_EMB_DIM = 256
    DEC_EMB_DIM = 256
    HID_DIM = 512
    N_LAYERS = 2
    ENC_DROPOUT = 0.5
    DEC_DROPOUT = 0.5

定义模型参数，包括输入和输出维度、嵌入维度、隐藏层维度、层数和 dropout 率。

python 复制代码

    enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)
    dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)
    model = Seq2Seq(enc, dec, device).to(device)

初始化编码器和解码器，并将其组合成一个 Seq2Seq 模型。然后将模型移动到指定的设备（GPU或CPU）。

python 复制代码

    model.load_state_dict(torch.load('./translation_model_GRU.pth'))

加载预训练的模型参数。

python 复制代码

    save_dir = '../dataset/submit.txt'

定义保存翻译结果的文件路径。

python 复制代码

    inference(model, dataset, src_file="../dataset/test_en.txt", save_dir=save_dir, terminology=terminology, device=device)

调用 inference 函数进行模型推理，并将翻译结果保存到指定文件中。传递的参数包括模型、数据集、源语言文件路径、保存路径、术语词典和设备。

python 复制代码

    print(f"翻译完成！文件已保存到{save_dir}")

打印翻译完成的消息，并显示保存文件的路径。

datawhale - 基于术语词典干预的机器翻译挑战赛 （一）

文章目录

torchtext 库是干什么用的 ？

TranslationDataset 类

定义 Seq2Seq模型

Encoder

Decoder

Seq2Seq 类

load_terminology_dictionary 函数

示例用法

train 函数

主程序代码

模型评价

load_sentences 函数

translate_sentence 函数

evaluate_bleu 函数

主程序

测试集上进行推理

inference 函数

主程序

datawhale - 基于术语词典干预的机器翻译挑战赛（一）

torchtext 库是干什么用的？