【Datawhale AI 夏令营】基于术语词典干预的机器翻译挑战赛——Task2baseline解读

嗨！大家好，这是一期Task2baseline解读的文章，拖更了好几天，希望对大家有用。

对于Task2教程里对相关的知识已经做了很详细的说明，大家可以去夏令营的教程里去学习这部分的知识，接下里是我对Task2的baseline一个理解

首先是第一部分，安装依赖包

!pip install torchtext
!pip install jieba
!pip install sacrebleu

这些命令用于安装所需的Python包

安装的包	作用
torchtext	用于文本数据处理的工具包，配合PyTorch使用
jieba	中文分词工具
sacrebleu	用于计算BLEU评分的工具

接下来是数据预处理

定义Tokenizer

en_tokenizer = get_tokenizer('spacy', language='en_core_web_trf')
zh_tokenizer = lambda x: list(jieba.cut(x))  # 使用jieba分词

这里定义了两个tokenizer（分词器），分别用于英文和中文，en_tokenizer 使用 spacy 进行英文分词，zh_tokenizer 使用 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

这个函数用于构建英文和中文的词汇表，词汇表包含一些特殊标记（如 <unk> 表示未知单词，<pad> 表示填充符，<bos> 表示句子开始，<eos> 表示句子结束），构建完成后，将默认索引设置为 <unk>，即任何不在词汇表中的单词都会映射到 <unk>。

第三步定义数据集

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

这个类继承自 torch.utils.data.Dataset，用于创建一个翻译数据集，每个样本由分词后的英文句子和中文句子构成，并将这些句子转换为对应的词汇索引。

第四部分定义数据加载器的 collate_fn

def collate_fn(batch):
    en_batch, zh_batch = [], []
    for en_item, zh_item in batch:
        if en_item and zh_item:  # 确保两个序列都不为空
            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([])
    
    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>'])
    
    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=1, collate_fn=collate_fn, drop_last=True)
​
    return train_loader, dev_loader, test_loader, en_vocab, zh_vocab

这个函数用于加载训练集、开发集和测试集数据，并创建相应的数据加载器（DataLoader）。

第六步构建模型

编码器

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):
        embedded = self.dropout(self.embedding(src))
        outputs, hidden = self.gru(embedded)
        return outputs, hidden

这个类定义了编码器（Encoder），它使用GRU网络对输入的源语言序列进行编码。

注意力机制

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):
        batch_size = encoder_outputs.shape[0]
        src_len = encoder_outputs.shape[1]
        
        hidden =
​
 hidden.unsqueeze(1).repeat(1, src_len, 1)
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
        attention = self.v(energy).squeeze(2)
        return F.softmax(attention, dim=1)

这个类定义了注意力机制（Attention），用于计算解码器隐藏状态与编码器输出之间的注意力权重。

解码器

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(emb_dim + hid_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
        self.fc_out = nn.Linear(emb_dim + hid_dim * 2, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, encoder_outputs):
        input = input.unsqueeze(1)
        embedded = self.dropout(self.embedding(input))
        
        a = self.attention(hidden[-1], encoder_outputs).unsqueeze(1)
        weighted = torch.bmm(a, encoder_outputs)
        
        rnn_input = torch.cat((embedded, weighted), dim=2)
        output, hidden = self.gru(rnn_input, hidden)
        
        embedded = embedded.squeeze(1)
        output = output.squeeze(1)
        weighted = weighted.squeeze(1)
        
        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim=1))
        return prediction, hidden

这个类定义了解码器（Decoder），它结合了注意力机制来生成目标语言序列。

Seq2Seq模型

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):
        trg_len = trg.shape[1]
        batch_size = trg.shape[0]
        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

这个类定义了Seq2Seq模型，结合了编码器和解码器，用于训练和预测。

第七部分训练和评估

训练函数

def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    
    for i, (src, trg) in enumerate(iterator):
        if len(src) == 0 or len(trg) == 0:
            print(f"跳过批次{i}")
            continue
        
        src, trg = src.to(model.device), trg.to(model.device)
        
        optimizer.zero_grad()
        
        output = model(src, trg)
        
        output_dim = output.shape[-1]
        output = output[:, 1:].reshape(-1, output_dim)
        trg = trg[:, 1:].reshape(-1)
        
        loss = criterion(output, trg)
        loss.backward()
        
        clip_grad_norm_(model.parameters(), clip)
        
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

这个函数用于训练模型，其中model.train()设置模型为训练模式，并对每个批次的数据进行前向传播、计算损失、反向传播并更新参数。

评估函数

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0
    
    with torch.no_grad():
        for i, (src, trg) in enumerate(iterator):
            if len(src) == 0 or len(trg) == 0:
                print(f"跳过批次{i}")
                continue
            
            src, trg = src.to(model.device), trg.to(model.device)
            
            output = model(src, trg, 0)  # 评估时不使用教师强制
            
            output_dim = output.shape[-1]
            output = output[:, 1:].reshape(-1, output_dim)
            trg = trg[:, 1:].reshape(-1)
            
            loss = criterion(output, trg)
            epoch_loss += loss.item()
            
    return epoch_loss / len(iterator)

这个函数用于评估模型，其中model.eval()设置模型为评估模式，并来计算模型在验证集上的损失。

第八部分翻译句子和计算BLEU评分

翻译函数

def translate_sentence(sentence, src_vocab, trg_vocab, model, device, max_length=50):
    model.eval()
    
    tokens = [token.lower() for token in en_tokenizer(sentence)]
    
    src_indexes = [src_vocab['<bos>']] + [src_vocab[token] for token in tokens] + [src_vocab['<eos>']]
    
    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)
    
    with torch.no_grad():
        encoder_outputs, hidden = model.encoder(src_tensor)
        
    trg_indexes = [trg_vocab['<bos>']]
    
    for i in range(max_length):
        trg_tensor = torch.LongTensor([trg_indexes[-1]]).to(device)
        
        with torch.no_grad():
            output, hidden = model.decoder(trg_tensor, hidden, encoder_outputs)
        
        pred_token = output.argmax(1).item()
        trg_indexes.append(pred_token)
        
        if pred_token == trg_vocab['<eos>']:
            break
    
    trg_tokens = [trg_vocab.get_itos()[i] for i in trg_indexes]
    
    return trg_tokens[1:-1]

用于翻译给定的英文句子，通过编码器和解码器生成目标语言（中文）的翻译结果。

计算BLEU评分

def calculate_bleu(data, src_vocab, trg_vocab, model, device, max_length=50):
    trgs = []
    pred_trgs = []
    
    for datum in data:
        src = datum[0]
        trg = datum[1]
        
        pred_trg = translate_sentence(src, src_vocab, trg_vocab, model, device, max_length)
        
        pred_trgs.append(pred_trg[:-1])
        trgs.append([trg])
    
    return bleu_score(pred_trgs, trgs)

计算模型在测试集上的BLEU评分，用于评估模型翻译质量。

最后一部分是训练和评估模型

N_EPOCHS = 10
CLIP = 1
​
best_valid_loss = float('inf')
​
for epoch in range(N_EPOCHS):
    start_time = time.time()
    
    train_loss = train(model, train_loader, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, dev_loader, 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(), 'tut2-model.pt')
    
    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}')

这个部分定义了训练过程，训练10个epoch，并在每个epoch后评估模型的性能，保存验证集损失最小的模型。

载入最佳模型并评估BLEU评分

model.load_state_dict(torch.load('tut2-model.pt'))
​
test_bleu = calculate_bleu(test_loader, en_vocab, zh_vocab, model, DEVICE)
​
print(f'BLEU score = {test_bleu * 100:.2f}')

这段代码载入最佳模型并在测试集上计算BLEU评分，评估模型的翻译质量。