Task1 baseline学习记录 Datawhale AI 夏令营

#Datawhale AI 夏令营

        baseline共分为三个模块：数据预处理与模型的训练；在开发集对模型性能评估；在测试集上进行推理。

一、数据预处理与模型的训练

1.导入PyTorch及相关库

torchtext 用于文本处理，random 用于随机化，time 用于计时。

!pip install torchtext
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

2. 数据预处理

定义数据集类

# 定义数据集类
# 修改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))
        
        self.terminology = terminology
        
        # 创建词汇表，注意这里需要确保术语词典中的词也被包含在词汇表中
        self.en_tokenizer = get_tokenizer('basic_english')
        self.zh_tokenizer = list  # 使用字符级分词
        
        en_vocab = Counter(self.terminology.keys())  # 确保术语在词汇表中
        zh_vocab = Counter()
        
        for en, zh in self.data:
            en_vocab.update(self.en_tokenizer(en))
            zh_vocab.update(self.zh_tokenizer(zh))
        
        # 添加术语到词汇表
        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_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 __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

        这个类实现了自定义数据集，用于加载和处理平行语料库。TranslationDataset 读取数据文件，将英文和中文句子对加载到内存中，并使用分词器进行分词。还创建了词汇表，确保术语词典中的词包含在内。

数据批处理函数

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

用于将不同长度的句子填充为相同长度，以便于批处理。

3. 模型设计

编码器

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

编码器将输入句子编码为隐藏状态，使用GRU层处理输入序列。

解码器

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):
        embedded = self.dropout(self.embedding(input))
        output, hidden = self.rnn(embedded, hidden)
        prediction = self.fc_out(output.squeeze(1))
        return prediction, hidden

解码器将编码器的隐藏状态作为输入，生成目标句子的预测。

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):
        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

Seq2Seq模型将编码器和解码器结合起来，实现序列到序列的转换。

4. 训练函数

def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    for i, (src, trg) in enumerate(iterator):
        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()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        epoch_loss += loss.item()
    return epoch_loss / len(iterator)

用于训练模型，通过前向传播、计算损失、反向传播和参数更新来训练模型。

5. 主函数

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')
    terminology = load_terminology_dictionary('../dataset/en-zh.dic')

    # 加载数据
    dataset = TranslationDataset('../dataset/train.txt', terminology=terminology)
    N = 1000  # 选择数据集的前1000个样本进行训练
    subset_indices = list(range(N))
    subset_dataset = Subset(dataset, subset_indices)
    train_loader = DataLoader(subset_dataset, batch_size=32, shuffle=True, collate_fn=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)

    # 定义优化器和损失函数
    optimizer = optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss(ignore_index=dataset.zh_word2idx['<pad>'])

    # 训练模型
    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:

定义模型训练参数并将所有函数统合调用

二、在开发集评估模型性能

1.导入必要库

import torch
from sacrebleu.metrics import BLEU
from typing import List

导入了PyTorch库、sacrebleu用于计算BLEU分数的库，以及Python的类型提示库。

2.定义辅助函数

加载句子函数

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]

这个函数用于从文件中加载句子，每行一个句子，并返回一个句子列表。

翻译句子函数

def translate_sentence(sentence: str, model: Seq2Seq, dataset: TranslationDataset, terminology, device: torch.device, max_length: int = 50):
    model.eval()
    tokens = dataset.en_tokenizer(sentence)
    tensor = torch.LongTensor([dataset.en_word2idx.get(token, dataset.en_word2idx['<sos>']) for token in tokens]).unsqueeze(0).to(device)  # [1, seq_len]
    
    with torch.no_grad():
        _, hidden = model.encoder(tensor)

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

    for _ in range(max_length):
        output, hidden = model.decoder(input_token, hidden)
        top_token = output.argmax(1)
        translated_token = dataset.zh_vocab[top_token.item()]
        
        if translated_token == '<eos>':
            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)
        input_token = top_token.unsqueeze(1)  # [1, 1]

    return ''.join(translated_tokens)

这个函数用于翻译一个句子，并使用术语词典中的术语。具体步骤如下：

1. 模型进入评估模式。
2. 对输入句子进行分词并转换为张量。
3. 使用编码器生成隐藏状态。
4. 使用解码器逐步生成翻译结果。
5. 在每一步中，如果生成的词在术语词典中，则替换为术语词典中的词。
6. 返回生成的翻译结果。

3.评估BLEU分数函数

def evaluate_bleu(model: Seq2Seq, dataset: TranslationDataset, src_file: str, ref_file: str, terminology, device: torch.device):
    model.eval()
    src_sentences = load_sentences(src_file)
    ref_sentences = load_sentences(ref_file)
    
    translated_sentences = []
    for src in src_sentences:
        translated = translate_sentence(src, model, dataset, terminology, device)
        translated_sentences.append(translated)
    
    bleu = BLEU()
    score = bleu.corpus_score(translated_sentences, [ref_sentences])
    
    return score

这个函数用于计算BLEU分数。具体步骤如下：

1. 模型进入评估模式。
2. 加载源句子和参考句子。
3. 使用translate_sentence函数逐个翻译源句子，并保存翻译结果。
4. 使用BLEU类计算并返回BLEU分数。

4.主函数

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)

    # 定义模型参数
    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)

    # 加载训练好的模型
    model.load_state_dict(torch.load('./translation_model_GRU.pth'))

    # 评估BLEU分数
    bleu_score = evaluate_bleu(model, dataset, '../dataset/dev_en.txt', '../dataset/dev_zh.txt', terminology, device)
    print(f'BLEU-4 score: {bleu_score.score:.2f}')

主函数执行以下步骤：

1. 确定使用GPU或CPU。
2. 加载术语词典。
3. 创建数据集实例。
4. 定义模型参数。
5. 初始化模型。
6. 加载训练好的模型参数。
7. 使用开发集评估模型的BLEU分数并打印结果。

       通过加载训练好的翻译模型，并结合术语词典，在开发集上进行翻译，并计算BLEU分数来评估模型性能。

三、在测试集上进行推理（翻译）

        该代码加载术语词典和训练好的翻译模型，使用测试集上的源句子进行翻译并将翻译结果保存到指定文件中，在主函数中进行一系列初始化操作并调用推理函数，最终输出翻译结果。

1.定义推理函数

def inference(model: Seq2Seq, dataset: TranslationDataset, src_file: str, save_dir: str, terminology, device: torch.device):
    model.eval()
    src_sentences = load_sentences(src_file)
    
    translated_sentences = []
    for src in src_sentences:
        translated = translate_sentence(src, model, dataset, terminology, device)
        translated_sentences.append(translated)

    # 将列表元素连接成一个字符串，每个元素后换行
    text = '\n'.join(translated_sentences)

    # 打开一个文件，如果不存在则创建，'w'表示写模式
    with open(save_dir, 'w', encoding='utf-8') as f:
        # 将字符串写入文件
        f.write(text)

这个函数的功能是在测试集上进行翻译，并将翻译结果保存到文件中。具体步骤如下：

1. 将模型设置为评估模式。
2. 从文件中加载源句子。
3. 使用translate_sentence函数对每个句子进行翻译，并将翻译结果保存到列表中。
4. 将翻译结果列表转换为字符串，每个翻译结果用换行符分隔。
5. 将转换后的字符串写入指定文件。

2.主函数

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)

    # 定义模型参数
    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)

    # 加载训练好的模型
    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)
    print(f"翻译完成！文件已保存到{save_dir}")

主函数执行以下步骤：

1. 确定使用GPU或CPU。
2. 加载术语词典。
3. 创建数据集实例并传递术语词典。
4. 定义模型参数。
5. 初始化模型。
6. 加载训练好的模型参数。
7. 指定保存翻译结果的文件路径。
8. 调用inference函数在测试集上进行翻译，并将结果保存到指定文件。
9. 打印消息，告知翻译完成并且结果已保存。

        通过这一系列操作，可以将训练好的翻译模型应用于测试集，并将翻译结果保存下来以供进一步使用或分析。