torchtext
是一个用于处理文本数据的库,它是 PyTorch 生态系统的一部分。这个库主要用于简化文本数据的预处理和加载过程,使得在深度学习模型中使用文本数据变得更加容易。以下是 torchtext
库的一些主要功能:
数据加载:torchtext
提供了方便的 API 来加载和处理各种文本数据集,包括常见的数据集格式如 CSV、TSV 等。
文本预处理:它包含了一系列的文本预处理工具,如分词、词干提取、词性标注等,这些工具可以帮助你将原始文本转换为模型可以理解的格式。
词汇表管理:torchtext
可以自动构建词汇表,并将文本转换为数值表示(如词嵌入),这对于深度学习模型来说是必不可少的。
数据迭代器:它提供了数据迭代器,可以方便地将数据分批加载到模型中进行训练,支持多线程和多进程加载,提高了数据加载的效率。
与 PyTorch 集成:torchtext
与 PyTorch 深度集成,可以直接将处理好的数据输入到 PyTorch 模型中,简化了整个深度学习流程。
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
列表中。 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
:使用字符级分词,即将中文句子拆分成单个字符。 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)
:打印中文词汇表的值。 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))
:更新中文词汇表,统计每个字符的出现频率。 # 添加术语到词汇表
self.en_vocab = ['' , '' , '' ] + list(self.terminology.keys()) + [word for word, _ in en_vocab.most_common(10000)]
self.zh_vocab = ['' , '' , '' ] + [word for word, _ in zh_vocab.most_common(10000)]
self.en_vocab = ['', '', ''] + list(self.terminology.keys()) + [word for word, _ in en_vocab.most_common(10000)]
:构建英文词汇表,包含特殊标记(
,
,
)、术语词典中的词以及出现频率最高的10000个词。self.zh_vocab = ['', '', ''] + [word for word, _ in zh_vocab.most_common(10000)]
:构建中文词汇表,包含特殊标记(
,
,
)以及出现频率最高的10000个字符。 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)}
:构建中文词到索引的映射字典。 def __len__(self):
return len(self.data)
def __len__(self):
:定义数据集的长度方法。return len(self.data)
:返回数据集的长度。 def __getitem__(self, idx):
en, zh = self.data[idx]
en_tensor = torch.tensor([self.en_word2idx.get(word, self.en_word2idx['' ]) for word in self.en_tokenizer(en)] + [self.en_word2idx['' ]])
zh_tensor = torch.tensor([self.zh_word2idx.get(word, self.zh_word2idx['' ]) for word in self.zh_tokenizer(zh)] + [self.zh_word2idx['' ]])
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['']) for word in self.en_tokenizer(en)] + [self.en_word2idx['']])
:将英文句子转换为索引张量,并在末尾添加
标记。zh_tensor = torch.tensor([self.zh_word2idx.get(word, self.zh_word2idx['']) for word in self.zh_tokenizer(zh)] + [self.zh_word2idx['']])
:将中文句子转换为索引张量,并在末尾添加
标记。return en_tensor, zh_tensor
:返回转换后的英文和中文张量。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
:返回填充后的英文和中文批量数据。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]
。embedded
,形状为[batch_size, src_len, emb_dim]
。outputs
和hidden
,形状分别为[batch_size, src_len, hid_dim]
和[n_layers, batch_size, hid_dim]
。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]
。embedded
,形状为[batch_size, 1, emb_dim]
。output
和hidden
,形状分别为[batch_size, 1, hid_dim]
和[n_layers, batch_size, hid_dim]
。prediction
,形状为[batch_size, output_dim]
。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
决定是否使用教师强制(即使用目标序列的下一个词作为输入,还是使用解码器的预测结果)。函数定义:
def load_terminology_dictionary(dict_file):
dict_file
: 这是一个字符串参数,表示包含术语词典的文件路径。初始化术语词典:
terminology = {}
terminology
,用于存储从文件中读取的术语。打开文件并读取内容:
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
字典中。返回术语词典:
return terminology
假设有一个文件 terminology.txt
,内容如下:
apple 苹果
banana 香蕉
orange 橙子
调用 load_terminology_dictionary
函数:
terminology = load_terminology_dictionary('terminology.txt')
print(terminology)
输出结果将是:
{'apple': '苹果', 'banana': '香蕉', 'orange': '橙子'}
def train(model, iterator, optimizer, criterion, clip):
train
的函数,该函数接受五个参数:model
(模型)、iterator
(数据迭代器)、optimizer
(优化器)、criterion
(损失函数)和 clip
(梯度裁剪的阈值)。 model.train()
epoch_loss = 0
epoch_loss
变量为 0,用于累积整个 epoch 的损失。 for i, (src, trg) in enumerate(iterator):
enumerate
遍历数据迭代器 iterator
,每次迭代获取一个批次的数据 (src, trg)
,其中 src
是源序列,trg
是目标序列。 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)
(batch_size * (sequence_length - 1), output_dim)
。 trg = trg[:, 1:].contiguous().view(-1)
(batch_size * (sequence_length - 1))
。 loss = criterion(output, trg)
criterion
。 loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
clip
指定。 optimizer.step()
epoch_loss += loss.item()
epoch_loss
中。 return epoch_loss / len(iterator)
if __name__ == '__main__':
start_time = time.time() # 开始计时
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
terminology = load_terminology_dictionary('../dataset/en-zh.dic')
dataset = TranslationDataset('../dataset/train.txt', terminology=terminology)
TranslationDataset
对象,该对象用于加载和处理训练数据。 N = 1000 # int(len(dataset) * 1) # 或者你可以设置为数据集大小的一定比例,如 int(len(dataset) * 0.1)
subset_indices = list(range(N))
subset_dataset = Subset(dataset, subset_indices)
subset_indices
是一个包含前1000个索引的列表,subset_dataset
是原始数据集的一个子集。 train_loader = DataLoader(subset_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
DataLoader
对象,用于批量加载数据。batch_size
设置为32,shuffle
设置为 True
以打乱数据顺序,collate_fn
是一个自定义函数,用于处理批次数据的拼接。 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
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)。 optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=dataset.zh_word2idx['' ])
。 N_EPOCHS = 10
CLIP = 1
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}')
torch.save(model.state_dict(), './translation_model_GRU.pth')
translation_model_GRU.pth
。 end_time = time.time() # 结束计时
elapsed_time_minute = (end_time - start_time)/60
print(f"Total running time: {elapsed_time_minute:.2f} minutes")
import torch
from sacrebleu.metrics import BLEU
from typing import List
# 假设已经定义了TranslationDataset, Encoder, Decoder, Seq2Seq类
torch
用于深度学习,sacrebleu.metrics.BLEU
用于评估翻译质量,typing.List
用于类型注解。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
是文件路径,函数返回一个包含所有句子的列表。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。 model.eval()
tokens = dataset.en_tokenizer(sentence)
dataset
中的 en_tokenizer
对输入句子进行分词,得到一个 token 列表。 tensor = torch.LongTensor([dataset.en_word2idx.get(token, dataset.en_word2idx['' ]) for token in tokens]).unsqueeze(0).to(device) # [1, seq_len]
的索引。然后将张量增加一个维度并移动到指定的设备。 with torch.no_grad():
_, hidden = model.encoder(tensor)
hidden
。 translated_tokens = []
translated_tokens
,用于存储翻译的 token。 input_token = torch.LongTensor([[dataset.zh_word2idx['' ]]]).to(device) # [1, 1]
的索引,并将其转换为张量并移动到指定的设备。 for _ in range(max_length):
max_length
次。 output, hidden = model.decoder(input_token, hidden)
output
和新的隐藏状态 hidden
。 top_token = output.argmax(1)
translated_token = dataset.zh_vocab[top_token.item()]
if translated_token == '' :
break
,则结束循环。 if translated_token in terminology.values():
for en_term, ch_term in terminology.items():
if translated_token == ch_term:
translated_token = en_term
break
translated_tokens.append(translated_token)
translated_tokens
列表中。 input_token = top_token.unsqueeze(1) # [1, 1]
return ''.join(translated_tokens)
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)。 model.eval()
src_sentences = load_sentences(src_file)
load_sentences
函数加载源语言句子,返回一个包含所有源句子的列表。 ref_sentences = load_sentences(ref_file)
load_sentences
函数加载参考翻译句子,返回一个包含所有参考翻译句子的列表。 translated_sentences = []
translated_sentences
,用于存储翻译的句子。 for src in src_sentences:
translated = translate_sentence(src, model, dataset, terminology, device)
translate_sentence
函数翻译当前的源句子,返回翻译后的句子。 translated_sentences.append(translated)
translated_sentences
列表中。 bleu = BLEU()
BLEU
对象,用于计算 BLEU 分数。 score = bleu.corpus_score(translated_sentences, [ref_sentences])
corpus_score
方法计算翻译句子的 BLEU 分数。translated_sentences
是模型生成的翻译句子列表,[ref_sentences]
是参考翻译句子列表。 return score
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
terminology = load_terminology_dictionary('../dataset/en-zh.dic')
dataset = TranslationDataset('../dataset/train.txt', terminology)
TranslationDataset
对象,该对象用于加载和处理训练数据,并传递术语词典。 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
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)。 model.load_state_dict(torch.load('./translation_model_GRU.pth'))
bleu_score = evaluate_bleu(model, dataset, '../dataset/dev_en.txt', '../dataset/dev_zh.txt', terminology=terminology, device=device)
evaluate_bleu
函数评估模型的 BLEU 分数。传递的参数包括模型、数据集、源语言文件路径、参考翻译文件路径、术语词典和设备。 print(f'BLEU-4 score: {bleu_score.score:.2f}')
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)。 model.eval()
src_sentences = load_sentences(src_file)
load_sentences
函数加载源语言句子,返回一个包含所有源句子的列表。 translated_sentences = []
translated_sentences
,用于存储翻译的句子。 for src in src_sentences:
translated = translate_sentence(src, model, dataset, terminology, device)
translate_sentence
函数翻译当前的源句子,返回翻译后的句子。 translated_sentences.append(translated)
translated_sentences
列表中。 text = '\n'.join(translated_sentences)
translated_sentences
列表中的所有句子连接成一个字符串,每个句子之间用换行符分隔。 with open(save_dir, 'w', encoding='utf-8') as f:
'w'
表示写模式,encoding='utf-8'
表示使用 UTF-8 编码。 f.write(text)
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
terminology = load_terminology_dictionary('../dataset/en-zh.dic')
dataset = TranslationDataset('../dataset/train.txt', terminology=terminology)
TranslationDataset
对象,该对象用于加载和处理训练数据,并传递术语词典。 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
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)。 model.load_state_dict(torch.load('./translation_model_GRU.pth'))
save_dir = '../dataset/submit.txt'
inference(model, dataset, src_file="../dataset/test_en.txt", save_dir=save_dir, terminology=terminology, device=device)
inference
函数进行模型推理,并将翻译结果保存到指定文件中。传递的参数包括模型、数据集、源语言文件路径、保存路径、术语词典和设备。 print(f"翻译完成!文件已保存到{save_dir}")