-
配置环境
-
torchtext :是一个用于自然语言处理(NLP)任务的库,它提供了丰富的功能,包括数据预处理、词汇构建、序列化和批处理等,特别适合于文本分类、情感分析、机器翻译等任务
-
jieba
:是一个中文分词库,用于将中文文本切分成有意义的词语 -
sacrebleu:用于评估机器翻译质量的工具,主要通过计算BLEU(Bilingual Evaluation Understudy)得分来衡量生成文本与参考译文之间的相似度
!pip install torchtext !pip install jieba !pip install sacrebleu
-
spacy:是一个强大的自然语言处理库,支持70+语言的分词与训练
这里,需要安装 spacy 用于英文的 tokenizer(分词,就是将句子、段落、文章这种长文本,分解为以字词为单位的数据结构,方便后续的处理分析工作),不同环境的安装请参考:https://spacy.io/usage
2.数据预处理
机器翻译任务的预处理是确保模型能够有效学习源语言到目标语言映射的关键步骤。预处理阶段通常包括多个步骤,旨在清理、标准化和转换数据,使之适合模型训练。以下是机器翻译任务预处理中常见的几个处理步骤:
-
清洗和规范化数据
-
去除无关信息:删除HTML标签、特殊字符、非文本内容等,确保文本的纯净性(本赛题的训练集中出现了非常多的脏数据,如“Joey. (掌声) (掌声) 乔伊”、“Thank you. (马嘶声) 谢谢你们”等这种声音词)
-
统一格式:转换所有文本为小写,确保一致性;标准化日期、数字等格式。
-
分句和分段:将长文本分割成句子或段落,便于处理和训练。
-
-
分词
-
分词:将句子分解成单词或词素(构成单词的基本组成部分,一个词素可以是一个完整的单词,也可以是单词的一部分,但每一个词素都至少携带一部分语义或语法信息),这是NLP中最基本的步骤之一。我们这里使用了使用
jieba
对中文进行分词,使用spaCy
对英文进行分词。
-
-
构建词汇表和词向量
-
词汇表构建:从训练数据中收集所有出现过的词汇,构建词汇表,并为每个词分配一个唯一的索引。
-
词向量:使用预训练的词向量或自己训练词向量,将词汇表中的词映射到高维空间中的向量,以捕捉语义信息(当前大模型领域训练的 embedding 模型就是用来完成此任务的)。
-
-
序列截断和填充
-
序列截断:限制输入序列的长度,过长的序列可能增加计算成本,同时也可能包含冗余信息。
-
序列填充:将所有序列填充至相同的长度,便于批量处理。通常使用
<PAD>
标记填充。
-
-
添加特殊标记
-
序列开始和结束标记:在序列两端添加
<SOS>
(Sequence Start)和<EOS>
(Sequence End)标记,帮助模型识别序列的起始和结束。 -
未知词标记:为不在词汇表中的词添加
<UNK>
(Unknown)标记,使模型能够处理未见过的词汇。
-
-
数据增强
-
随机替换或删除词:在训练数据中随机替换或删除一些词,增强模型的鲁棒性。
-
同义词替换:使用同义词替换原文中的词,增加训练数据的多样性。
-
-
数据分割
-
划分数据集:将数据划分为训练集、验证集和测试集,分别用于模型训练、参数调整和最终性能评估(该赛题中已划分好,不需要自己进行划分)
-
3.模型训练
说到神经机器翻译就不得不提编码器-解码器模型,或编码器-解码器框架(EncoderDecoder Paradigm)。本质上,编码器解码器模型是描述输入输出之间关系的一种方式。编码器解码器这个概念在日常生活中并不少见。
例如,在电视系统上为了便于视频的传播,会使用各种编码器将视频编码成数字信号,在客户端,相应的解码器组件会把收到的数字信号解码为视频。另外一个更贴近生活的例子是电话,它通过对声波和电信号进行相互转换,达到传递声音的目的。
这种“先编码,再解码”的思想被应用到密码学、信息论等多个领域。不难看出,机器翻译问题也完美的贴合编码器解码器结构的特点。可以将源语言编码为类似信息传输中的数字信号,然后利用解码器对其进行转换,生成目标语言。
在源语言句子的表示形式确定之后,需要设计相应的编码器和解码器结构。在当今主流的神经机器翻译系统中,编码器由词嵌入层和中间网络层组成:
-
当输入一串单词序列时,词嵌入层(embedding)会将每个单词映射到多维实数表示空间,这个过程也被称为词嵌入。
-
之后中间层会对词嵌入向量进行更深层的抽象,得到输入单词序列的中间表示。中间层的实现方式有很多,比如:循环神经网络、卷积神经网络、自注意力机制等都是模型常用的结构。
解码器的结构基本上和编码器是一致的,在基于循环神经网络的翻译模型中,解码器只比编码器多了输出层,用于输出每个目标语言位置的单词生成概率,而在基于自注意力机制的翻译模型中,除了输出层,解码器还比编码器多一个编码解码注意力子层,用于帮助模型更好地利用源语言信息。
以 基于循环神经网络的机器翻译模型为例,说明其整体结构。其中,左侧为编码器部分,源语言单词按照其在文本序列中的先后顺序被依次送入到循环神经网络(RNN)当中。在每个时间步 t 中,模型依据送入的源语言单词$$x_{t} $$对应修改并维护其模型内部的隐状态 $$h_{t}$$,这个隐状态编码了输入的源语言序列前 t 个时刻的所有必要信息。按照这种方式当 m 个输入全部被送入到编码器之后,所对应的 $$h_{m}$$可以认为包含了源语言序列的所有信息。
右半部分是 RNN 解码器部分,它接收编码器输出的编码源语言句子信息的向量 $$h_{m}$$作为初始隐状态 $$s_{0}$$。由于 RNN 的循环过程在每个时间步都要求一个输入单词,为了启动解码过程,一般会使用一个保留的特殊符号 “[Start]” 作为翻译开始的标记送入到 RNN 解码器当中并解码出目标语言序列的第一个单词 $$z_{1}$$。接下来,$$z_{1}$$ 会作为下一个时刻的输入被送入到循环神经网络当中,并按照不断迭代产生后续的预测。由于目标语言序列的长度无法被提前预知,因此使用另一个保留符号 “[Stop]” 作为预测结束的标志。当某一个时刻 t 预测出的目标语言单词为 zt =“[Stop]” 时,解码过程动态地停止。在上述过程当中,主要涉及到两步运算,第一步是 RNN 接收前一时刻隐状态 $$s_{t-1}$$ 并依据当前时刻输入 $$z_{t-1}$$(目标语言单词 $$z_{t-1}$$ 对应的语义嵌入)对隐状态进行维护并生成$$s_{t}$$的运算过程,第二步是依据当前时刻隐状态生成目标语言单词的过程:
其中 U,W,V 是可学习的参数。U,W 负责维护循环状态,而 V 负责将当前时刻状态转换到词表大小的概率分布
$$P \in R^{vocab_size} $$,从中可以采样得到目标语言单词 $$z_{t}$$。
通过循环网络对源语言文本进行编码,并生成目标语言翻译结果的过程十分简单。然而,它仅仅使用一个定长的向量 $$h_{m}$$ 编码整个源语言序列。这对于较短的源语言文本没有什么问题,但随着文本序列长度的逐渐加长,单一的一个向量 hm 可能不足以承载源语言序列当中的所有信息。
蓝色的线代表上述简单循环神经网络性能随源语言文本长度的变化趋势。当文本长度在 20 个单词以内时,单一向量能够承载源语言文本中的必要信息。随着文本序列的进一步增加,翻译性能的评价指标 BLEU 的值就开始出现明显地下降。因此,这就启发我们使用更加有效地机制从编码器向解码器传递源语言信息,这就是接下来要讲到的注意力机制。
引入注意力机制的循环机器翻译架构与基于简单循环网络的机器翻译模型大体结构相似,均采用循环神经网络作为编码器与解码器的实现。关键的不同点在于注意力机制的引入使得不再需要把原始文本中的所有必要信息压缩到一个向量当中。引入注意力机制的循环神经网络机器翻译架构如图所示:
-
翻译质量评价
人们在使用机器翻译系统时需要评估系统输出结果的质量。这个过程也被称作机器翻译译文质量评价,简称为译文质量评价(Quality Evaluation of Translation)。在机器翻译的发展进程中,译文质量评价有着非常重要的作用。不论在系统研发的反复迭代中,还是在诸多的机器翻译应用场景中,都存在大量的译文质量评价环节。从某种意义上说,没有译文质量评价,机器翻译也不会发展成今天的样子。比如,本世纪初研究人员提出了译文质量自动评价方法 BLEU(Bilingual Evaluation Understudy)(Task 1知识文档已详细介绍过)。该方法使得机器翻译系统的评价变得自动、快速、便捷,而且评价过程可以重复。正是由于 BLEU 等自动评价方法的提出,机器翻译研究人员可以在更短的时间内得到译文质量的评价结果,加速系统研发的进程。
传统观点把翻译分为“信”、“达”、“雅”三个层次,而忠诚度体现的是一种“信”的思想,而流畅度体现的是一种“达”的思想。不过“雅”在机器翻译评价中还不是一个常用的标准,而且机器翻译还没有达到“雅”的水平,是未来所追求的目标。给定评价标准,译文质量评价有很多实现方式,下图给出了机器翻译译文评价方法的逻辑关系图:
-
人工评价。当需要对系统进行准确的评估时,往往采用人工评价。比如,对于机器翻译的一些互联网应用,在系统上线前都会采用人工评价对机器翻译系统性能进行测试。当然,这种方法的时间和人力成本是最高的。
-
有参考答案的自动评价。由于机器翻译系统研发过程中需要频繁地对系统性能进行评价,这时可以让人标注一些正确的译文,之后把这些译文作为参考答案与机器翻译系统输出的结果进行比对。这种自动评价的结果获取成本低,可以多次重复,而且可以用于对系统结果的快速反馈,指导系统优化的方向。
-
无参考答案的自动评价。在很多应用场景中,在系统输出译文时,使用者希望提前知道译文的质量,即使这时并没有可比对的参考答案。这样,系统使用者可以根据这个对质量的“估计”结果有选择地使用机器翻译译文。严格意义上说,这并不是一个传统的译文质量评价方法,而是一种对译文置信度和可能性的估计。
完整代码:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_
from torchtext.data.metrics import bleu_score
from torch.utils.data import Dataset, DataLoader
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from typing import List, Tuple
import jieba
import random
from torch.nn.utils.rnn import pad_sequence
import sacrebleu
import time
import math
!pip install -U pip setuptools wheel -i https://pypi.tuna.tsinghua.edu.cn/simple
!pip install -U 'spacy[cuda12x]' -i https://pypi.tuna.tsinghua.edu.cn/simple
!pip install ../dataset/en_core_web_trf-3.7.3-py3-none-any.whl
# !python -m spacy download en_core_web_sm
# 定义tokenizer
en_tokenizer = get_tokenizer('spacy', language='en_core_web_trf')
zh_tokenizer = lambda x: list(jieba.cut(x)) # 使用jieba分词
# 读取数据函数
def read_data(file_path: str) -> List[str]:
with open(file_path, 'r', encoding='utf-8') as f:
return [line.strip() for line in f]
# 数据预处理函数
def preprocess_data(en_data: List[str], zh_data: List[str]) -> List[Tuple[List[str], List[str]]]:
processed_data = []
for en, zh in zip(en_data, zh_data):
en_tokens = en_tokenizer(en.lower())[:MAX_LENGTH]
zh_tokens = zh_tokenizer(zh)[:MAX_LENGTH]
if en_tokens and zh_tokens: # 确保两个序列都不为空
processed_data.append((en_tokens, zh_tokens))
return processed_data
# 构建词汇表
def build_vocab(data: List[Tuple[List[str], List[str]]]):
en_vocab = build_vocab_from_iterator(
(en for en, _ in data),
specials=['<unk>', '<pad>', '<bos>', '<eos>']
)
zh_vocab = build_vocab_from_iterator(
(zh for _, zh in data),
specials=['<unk>', '<pad>', '<bos>', '<eos>']
)
en_vocab.set_default_index(en_vocab['<unk>'])
zh_vocab.set_default_index(zh_vocab['<unk>'])
return en_vocab, zh_vocab
class TranslationDataset(Dataset):
def __init__(self, data: List[Tuple[List[str], List[str]]], en_vocab, zh_vocab):
self.data = data
self.en_vocab = en_vocab
self.zh_vocab = zh_vocab
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
en, zh = self.data[idx]
en_indices = [self.en_vocab['<bos>']] + [self.en_vocab[token] for token in en] + [self.en_vocab['<eos>']]
zh_indices = [self.zh_vocab['<bos>']] + [self.zh_vocab[token] for token in zh] + [self.zh_vocab['<eos>']]
return en_indices, zh_indices
def collate_fn(batch):
en_batch, zh_batch = [], []
for en_item, zh_item in batch:
if en_item and zh_item: # 确保两个序列都不为空
# print("都不为空")
en_batch.append(torch.tensor(en_item))
zh_batch.append(torch.tensor(zh_item))
else:
print("存在为空")
if not en_batch or not zh_batch: # 如果整个批次为空,返回空张量
return torch.tensor([]), torch.tensor([])
# src_sequences = [item[0] for item in batch]
# trg_sequences = [item[1] for item in batch]
en_batch = nn.utils.rnn.pad_sequence(en_batch, batch_first=True, padding_value=en_vocab['<pad>'])
zh_batch = nn.utils.rnn.pad_sequence(zh_batch, batch_first=True, padding_value=zh_vocab['<pad>'])
# en_batch = pad_sequence(en_batch, batch_first=True, padding_value=en_vocab['<pad>'])
# zh_batch = pad_sequence(zh_batch, batch_first=True, padding_value=zh_vocab['<pad>'])
return en_batch, zh_batch
# 数据加载函数
def load_data(train_path: str, dev_en_path: str, dev_zh_path: str, test_en_path: str):
# 读取训练数据
train_data = read_data(train_path)
train_en, train_zh = zip(*(line.split('\t') for line in train_data))
# 读取开发集和测试集
dev_en = read_data(dev_en_path)
dev_zh = read_data(dev_zh_path)
test_en = read_data(test_en_path)
# 预处理数据
train_processed = preprocess_data(train_en, train_zh)
dev_processed = preprocess_data(dev_en, dev_zh)
test_processed = [(en_tokenizer(en.lower())[:MAX_LENGTH], []) for en in test_en if en.strip()]
# 构建词汇表
global en_vocab, zh_vocab
en_vocab, zh_vocab = build_vocab(train_processed)
# 创建数据集
train_dataset = TranslationDataset(train_processed, en_vocab, zh_vocab)
dev_dataset = TranslationDataset(dev_processed, en_vocab, zh_vocab)
test_dataset = TranslationDataset(test_processed, en_vocab, zh_vocab)
from torch.utils.data import Subset
# 假设你有10000个样本,你只想用前1000个样本进行测试
indices = list(range(N))
train_dataset = Subset(train_dataset, indices)
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn, drop_last=True)
dev_loader = DataLoader(dev_dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn, drop_last=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn, drop_last=True)
return train_loader, dev_loader, test_loader, en_vocab, zh_vocab
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
super().__init__()
self.hid_dim = hid_dim
self.n_layers = n_layers
self.embedding = nn.Embedding(input_dim, emb_dim)
self.gru = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
# src = [batch size, src len]
embedded = self.dropout(self.embedding(src))
# embedded = [batch size, src len, emb dim]
outputs, hidden = self.gru(embedded)
# outputs = [batch size, src len, hid dim * n directions]
# hidden = [n layers * n directions, batch size, hid dim]
return outputs, hidden
class Attention(nn.Module):
def __init__(self, hid_dim):
super().__init__()
self.attn = nn.Linear(hid_dim * 2, hid_dim)
self.v = nn.Linear(hid_dim, 1, bias=False)
def forward(self, hidden, encoder_outputs):
# hidden = [1, batch size, hid dim]
# encoder_outputs = [batch size, src len, hid dim]
batch_size = encoder_outputs.shape[0]
src_len = encoder_outputs.shape[1]
hidden = hidden.repeat(src_len, 1, 1).transpose(0, 1)
# hidden = [batch size, src len, hid dim]
energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
# energy = [batch size, src len, hid dim]
attention = self.v(energy).squeeze(2)
# attention = [batch size, src len]
return F.softmax(attention, dim=1)
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout, attention):
super().__init__()
self.output_dim = output_dim
self.hid_dim = hid_dim
self.n_layers = n_layers
self.attention = attention
self.embedding = nn.Embedding(output_dim, emb_dim)
self.gru = nn.GRU(hid_dim + emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
self.fc_out = nn.Linear(hid_dim * 2 + emb_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, input, hidden, encoder_outputs):
# input = [batch size, 1]
# hidden = [n layers, batch size, hid dim]
# encoder_outputs = [batch size, src len, hid dim]
input = input.unsqueeze(1)
embedded = self.dropout(self.embedding(input))
# embedded = [batch size, 1, emb dim]
a = self.attention(hidden[-1:], encoder_outputs)
# a = [batch size, src len]
a = a.unsqueeze(1)
# a = [batch size, 1, src len]
weighted = torch.bmm(a, encoder_outputs)
# weighted = [batch size, 1, hid dim]
rnn_input = torch.cat((embedded, weighted), dim=2)
# rnn_input = [batch size, 1, emb dim + hid dim]
output, hidden = self.gru(rnn_input, hidden)
# output = [batch size, 1, hid dim]
# hidden = [n layers, batch size, hid dim]
embedded = embedded.squeeze(1)
output = output.squeeze(1)
weighted = weighted.squeeze(1)
prediction = self.fc_out(torch.cat((output, weighted, embedded), dim=1))
# prediction = [batch size, output dim]
return prediction, hidden
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 = [batch size, src len]
# trg = [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)
encoder_outputs, hidden = self.encoder(src)
input = trg[:, 0]
for t in range(1, trg_len):
output, hidden = self.decoder(input, hidden, encoder_outputs)
outputs[:, t] = output
teacher_force = random.random() < teacher_forcing_ratio
top1 = output.argmax(1)
input = trg[:, t] if teacher_force else top1
return outputs
# 初始化模型
def initialize_model(input_dim, output_dim, emb_dim, hid_dim, n_layers, dropout, device):
attn = Attention(hid_dim)
enc = Encoder(input_dim, emb_dim, hid_dim, n_layers, dropout)
dec = Decoder(output_dim, emb_dim, hid_dim, n_layers, dropout, attn)
model = Seq2Seq(enc, dec, device).to(device)
return model
# 运行时间
def epoch_time(start_time, end_time):
elapsed_time = end_time - start_time
elapsed_mins = int(elapsed_time / 60)
elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
return elapsed_mins, elapsed_secs
def train(model, iterator, optimizer, criterion, clip):
model.train()
epoch_loss = 0
for i, batch in enumerate(iterator):
#print(f"Training batch {i}")
src, trg = batch
#print(f"Source shape before: {src.shape}, Target shape before: {trg.shape}")
if src.numel() == 0 or trg.numel() == 0:
#print("Empty batch detected, skipping...")
continue # 跳过空的批次
src, trg = src.to(DEVICE), trg.to(DEVICE)
optimizer.zero_grad()
output = model(src, trg)
output_dim = output.shape[-1]
output = output[:, 1:].contiguous().view(-1, output_dim)
trg = trg[:, 1:].contiguous().view(-1)
loss = criterion(output, trg)
loss.backward()
clip_grad_norm_(model.parameters(), clip)
optimizer.step()
epoch_loss += loss.item()
print(f"Average loss for this epoch: {epoch_loss / len(iterator)}")
return epoch_loss / len(iterator)
def evaluate(model, iterator, criterion):
model.eval()
epoch_loss = 0
with torch.no_grad():
for i, batch in enumerate(iterator):
#print(f"Evaluating batch {i}")
src, trg = batch
if src.numel() == 0 or trg.numel() == 0:
continue # 跳过空批次
src, trg = src.to(DEVICE), trg.to(DEVICE)
output = model(src, trg, 0) # 关闭 teacher forcing
output_dim = output.shape[-1]
output = output[:, 1:].contiguous().view(-1, output_dim)
trg = trg[:, 1:].contiguous().view(-1)
loss = criterion(output, trg)
epoch_loss += loss.item()
return epoch_loss / len(iterator)
# 翻译函数
def translate_sentence(src_indexes, src_vocab, tgt_vocab, model, device, max_length=50):
model.eval()
src_tensor = src_indexes.unsqueeze(0).to(device) # 添加批次维度
# with torch.no_grad():
# encoder_outputs = model.encoder(model.positional_encoding(model.src_embedding(src_tensor) * math.sqrt(model.d_model)))
trg_indexes = [tgt_vocab['<bos>']]
for i in range(max_length):
trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)
# print("src_tensor:",src_tensor)
# print("trg_tensor:",trg_tensor)
with torch.no_grad():
output = model(src_tensor, trg_tensor)
pred_token = output.argmax(2)[:, -1].item()
trg_indexes.append(pred_token)
if pred_token == tgt_vocab['<eos>']:
break
trg_tokens = [tgt_vocab.get_itos()[i] for i in trg_indexes]
return trg_tokens[1:-1] # 移除<bos>和<eos>标记
def calculate_bleu(dev_loader, src_vocab, tgt_vocab, model, device):
model.eval()
translations = []
references = []
with torch.no_grad():
for src, tgt in dev_loader:
src = src.to(device)
for sentence in src:
translated = translate_sentence(sentence, src_vocab, tgt_vocab, model, device)
translations.append(' '.join(translated))
for reference in tgt:
ref_tokens = [tgt_vocab.get_itos()[idx] for idx in reference if idx not in [tgt_vocab['<bos>'], tgt_vocab['<eos>'], tgt_vocab['<pad>']]]
references.append([' '.join(ref_tokens)])
bleu = sacrebleu.corpus_bleu(translations, references)
return bleu.score
# 主训练循环
def train_model(model, train_iterator, valid_iterator, optimizer, criterion, N_EPOCHS = 10, CLIP = 1, save_path = '../model/best-model.pt'):
best_valid_loss = float('inf')
for epoch in range(N_EPOCHS):
start_time = time.time()
#print(f"Starting Epoch {epoch + 1}")
train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
valid_loss = evaluate(model, valid_iterator, criterion)
end_time = time.time()
epoch_mins, epoch_secs = epoch_time(start_time, end_time)
if valid_loss < best_valid_loss:
best_valid_loss = valid_loss
torch.save(model.state_dict(), save_path)
print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
print(f'\t Val. Loss: {valid_loss:.3f} | Val. PPL: {math.exp(valid_loss):7.3f}')
# 定义常量
MAX_LENGTH = 100 # 最大句子长度
BATCH_SIZE = 32
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
N = 1000 # 采样训练集的数量,最大148363
train_path = '../dataset/train.txt'
dev_en_path = '../dataset/dev_en.txt'
dev_zh_path = '../dataset/dev_zh.txt'
test_en_path = '../dataset/test_en.txt'
train_loader, dev_loader, test_loader, en_vocab, zh_vocab = load_data(
train_path, dev_en_path, dev_zh_path, test_en_path
)
print(f"英语词汇表大小: {len(en_vocab)}")
print(f"中文词汇表大小: {len(zh_vocab)}")
print(f"训练集大小: {len(train_loader.dataset)}")
print(f"开发集大小: {len(dev_loader.dataset)}")
print(f"测试集大小: {len(test_loader.dataset)}")
# 主函数
if __name__ == '__main__':
N_EPOCHS = 5
CLIP=1
# 模型参数
INPUT_DIM = len(en_vocab)
OUTPUT_DIM = len(zh_vocab)
EMB_DIM = 128
HID_DIM = 256
N_LAYERS = 2
DROPOUT = 0.5
# 初始化模型
model = initialize_model(INPUT_DIM, OUTPUT_DIM, EMB_DIM, HID_DIM, N_LAYERS, DROPOUT, DEVICE)
print(f'The model has {sum(p.numel() for p in model.parameters() if p.requires_grad):,} trainable parameters')
# 定义损失函数
criterion = nn.CrossEntropyLoss(ignore_index=zh_vocab['<pad>'])
# 初始化优化器
optimizer = optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
# 训练模型
save_path = '../model/best-model.pt'
train_model(model, train_loader, dev_loader, optimizer, criterion, N_EPOCHS, CLIP, save_path = save_path)
print(f"训练完成!模型已保存到:{save_path}")
save_dir = '../results/submit_task2.txt'
with open(save_dir, 'w') as f:
translated_sentences = []
for batch in test_loader: # 遍历所有数据
src, _ = batch
src = src.to(DEVICE)
translated = translate_sentence(src[0], en_vocab, zh_vocab, model, DEVICE, max_length=50) # 翻译结果,max_length生成翻译的最大长度
#print(translated)
results = "".join(translated)
f.write(results + '\n') # 将结果写入文件
print(f"翻译完成,结果已保存到{save_dir}")