模型载入: 使用transformers库,加载了一个预训练的问答模型。
训练参数的设置: 为模型训练定制了多个参数,如学习率、批量大小和权重衰减等。
DualEncoder: 该模型的核心是一个双编码器结构,独立地计算问题和答案的向量表示,然后通过余弦相似度来计算它们之间的相关性。
数据预处理: 分词器用于将文本转化为数字序列,此外提供了文本到张量的转换功能。
数据加载: 使用pandas库从CSV文件中提取了所需的字段,然后将数据分为训练集和测试集。
模型训练: 在这一部分,定义了一个训练函数,该函数处理批量数据、计算损失、进行反向传播和优化。计算了每个批次的准确率,是衡量模型性能的指标。
训练开始: 初始化了两个GRU编码器,并将其传递给DualEncoder。为训练过程设置了损失函数和优化器。
结果保存: 保存了分词器,以便后续可以重复使用它。
from collections import Counter
import pickle
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd
from tqdm import tqdm
加载预训练模型,设定训练参数
from transformers import AutoModelForQuestionAnswering, TrainingArguments, Trainer
model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)
因为有这些随机初始化的参数,所以我们要在新的数据集上重新fine-tune我们的模型
# 训练设定参数
args = TrainingArguments(
f"test-squad",
evaluation_strategy = "epoch",
learning_rate=2e-5, #学习率
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
num_train_epochs=3, # 训练的轮次
weight_decay=0.01,
)
Dual Encoder即两个独立的Encoder,这里分别计算问题和回答的句向量,最后通过余弦相似度计算它们之间的关联程度。
# DualEncoder
class DualEncoder(nn.Module):
def __init__(self, encoder1, encoder2, type="cosine"):
super(DualEncoder, self).__init__()
self.encoder1 = encoder1
self.encoder2 = encoder2
if type != 'cosine':
# 训练一个简单的神经网络来计算相似度
self.linear = nn.Sequential(
# 拼接encoder1和encoder2的输出(向量),转换成100维的表示
nn.Linear(self.encoder1.hidden_size + self.encoder2.hidden_size, 100),
# 经过ReLU激活
nn.ReLU(),
# 再转换成一个数值,表示相似程度
nn.Linear(100, 1)
)
def forward(self, x, x_mask, y, y_mask):
x_rep = self.encoder1(x, x_mask)
y_rep = self.encoder2(y, y_mask)
return x_rep, y_rep
主要是去实现向前传播方法
然后采用GRU作为Encoder的实现,支持多种表征的获取,默认是平均每个时刻的输出。
GRUEncoder
class GRUEncoder(nn.Module):
def __init__(self, vocab_size, embed_size, hidden_size, dropout_p=0.1, avg_hidden=True, n_layers=1, bidirectional=True):
super(GRUEncoder, self).__init__()
self.hidden_size = hidden_size
self.embed = nn.Embedding(vocab_size, embed_size)
if bidirectional:
# 大小除以2,使得拼接两个方向后大小不变
hidden_size //= 2
# 这种生成句子表征的建议使用bidirectional=True
self.rnn = nn.GRU(embed_size, hidden_size, num_layers=n_layers, bidirectional=bidirectional,dropout=dropout_p)
self.dropout = nn.Dropout(dropout_p)
self.bidirectional = bidirectional
self.avg_hidden = avg_hidden
def forward(self, x, mask):
x_embed = self.embed(x) # 先得到嵌入表示
x_embed = self.dropout(x_embed) # 再经过dropout
seq_len = mask.sum(1) # 计算有效长度
# 压缩批次内填充数据
# 通过压缩填充加快训练效率,具体可参考文章: https://blog.csdn.net/yjw123456/article/details/118855324
x_embed = torch.nn.utils.rnn.pack_padded_sequence(
input=x_embed,
lengths=seq_len.cpu(),
batch_first=True,
enforce_sorted=False
)
output, hidden = self.rnn(x_embed)
# output (batch_size, seq_len, hidden_size)
# hidden (num_directions * num_layers, batch_size, hidden_size)
output, seq_len = torch.nn.utils.rnn.pad_packed_sequence(
sequence=output,
batch_first=True,
padding_value=0,
total_length=mask.shape[1]
)
if self.avg_hidden:
# 对RNN输出每个时刻的输出求均值
# mask.unsqueeze(2) 使维度个数和output一致
# hiden (batch_size, hidden_size)
hidden = torch.sum(output * mask.unsqueeze(2), 1) / torch.sum(mask, 1, keepdim=True)
else:
if self.bidirectional:
# 拼接两个方向上的输出
# hidden[-2,:,:] (batch_size, hidden_size / 2)
# hidden[-1,:,:] (batch_size, hidden_size / 2)
# hidden (batch_size, hidden_size)
hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]),dim=1)
else:
# 取出最顶层(若num_layers > 1)的hidden
hidden = hidden[-1,:,:]
# 需要保证各种情况下的hidden大小都是一致的
# 经过一层dropout
hidden = self.dropout(hidden)
return hidden
词典,处理NLP任务基本上都需要一个词典
构建分词器
class Tokenizer:
def __init__(self, vocab):
self.id2word = ["UNK"] + vocab # 保证未知词UNK的id为0
self.word2id = {w:i for i,w in enumerate(vocab)}
def text2id(self, text):
# 对中文简单的按字拆分
return [self.word2id.get(w, 0) for w in str(text)]
def id2text(self, ids):
return "".join([self.id2word[id] for id in ids])
def __len__(self):
return len(self.id2word)
def create_tokenizer(texts, vocab_size):
"""
创建分词器,输入文本列表和词典大小
"""
all_vocab = ""
for text in texts:
all_vocab += str(text)
vocab_count = Counter(all_vocab) # 按字拆分
# 最频繁的vocab_size个单词
vocab = vocab_count.most_common(vocab_size)
# (char, count) 从中取出char
vocab = [w[0] for w in vocab]
return Tokenizer(vocab)
def list2tensor(sents, tokenizer):
"""
将文本列表结合分词器转换为tensor
"""
res = []
mask = []
for sent in sents:
res.append(tokenizer.text2id(sent))
max_len = max([len(sent) for sent in res])
# 按最大长度进行填充
for i in range(len(res)):
_mask = np.zeros((1, max_len)) # 中的元素0表示填充词,1表示非填充
_mask[:,:len(res[i])] = 1 # 有效位元素置1
res[i] = np.expand_dims(np.array(res[i] + [0] * (max_len - len(res[i]))), 0) # 增加一个维度
mask.append(_mask)
res = np.concatenate(res, axis=0) # 按维度0进行拼接
mask = np.concatenate(mask, axis=0)
# 分别转换为long类型和float类型的tensor
res = torch.tensor(res).long()
mask = torch.tensor(mask).float()
return res, mask
加载数据集
# 数据集位置
file_path = '../dataset/nonghangzhidao_filter.csv'
df = pd.read_csv(file_path)[["title", "reply", "is_best"]] # 只需要这三个字段
df.head()
然后进行拆分数据集
np.random.seed(42) # 设定随机种子可以防止每次的训练/测试集数据不一样
# 拆分训练/测试集
def shuffle_and_split_data(data, test_ratio):
shuffled_indices = np.random.permutation(len(data))
test_set_size = int(len(data) * test_ratio)
test_indices = shuffled_indices[:test_set_size] # 前test_set_size作为测试集
train_indices = shuffled_indices[test_set_size:] # 剩下的作为训练集
return data.iloc[train_indices], data.iloc[test_indices]
# 20%的数据作为测试集
train_set, test_set = shuffle_and_split_data(df, 0.2)
print(len(train_set), len(test_set))
将文本转换为张量
texts = list(train_set["title"]) + list(train_set["reply"]) # 取出文本内容
编写训练代码
# 编写训练代码
def train(df, model, loss_fn, optimizer, device, tokenizer, loss_function, batch_size):
# 设成训练模型,让dropout生效
model.train()
df = df.sample(frac=1) # 每次训练时shuffle数据
# 分批处理
for i in range(0, df.shape[0], batch_size):
# 得到批次数据
batch_df = df.iloc[i:i+batch_size]
title = list(batch_df["title"])
reply = list(batch_df["reply"])
# 构建目标tensor(1或0)
target = torch.tensor(batch_df["is_best"].to_numpy()).float()
if loss_function == "cosine":
# 为了符合CosineEmbeddingLoss的要求,将0替换成-1
target[target == 0] = -1
x, x_mask = list2tensor(title, tokenizer)
y, y_mask = list2tensor(reply, tokenizer)
# 都切换到同一设备(cpu/gpu)
x, x_mask, y, y_mask, target = x.to(device), x_mask.to(device), y.to(device), y_mask.to(device), target.to(device)
# 计算x和y的表征
x_rep, y_rep = model(x, x_mask, y, y_mask)
# 根据需要使用不同的损失
if loss_function == "cosine":
loss = loss_fn(x_rep, y_rep, target)
else:
# 拼接x_pre和y_rep,并传入linear
logits = model.linear(torch.cat([x_rep, y_rep], 1))
loss = loss_fn(logits, target)
optimizer.zero_grad()
loss.backward()
# 梯度裁剪
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if loss_function == "cosine":
sim = F.cosine_similarity(x_rep, y_rep)
# 余弦相似度范围在[-1,1]之间,因此这里使用0作为分界
sim[sim < 0] = -1
sim[sim >= 0] = 1
else:
sim = model.linear(torch.cat([x_rep, y_rep], 1))
# sim = torch.sigmoid(logits) 可以不用sigmoid
sim[sim < 0] = 0
sim[sim >= 0] = 1
sim = sim.view(-1)
target = target.view(-1)
# 计算准确率
num_corrects = torch.sum(sim == target).item()
total_counts = target.shape[0]
print(f"accuracy:{num_corrects / total_counts}")
return num_corrects / total_counts
定义参数
# 定义参数
loss_function = "cosine"
batch_size = 64
output_dir = "./models"
num_epochs = 10
vocab_size = 5000
hidden_size = 300
embed_size = 600
开始训练
# 构建两个Encoder
title_encoder = GRUEncoder(len(tokenizer), embed_size, hidden_size)
reply_encoder = GRUEncoder(len(tokenizer), embed_size, hidden_size)
# 传入DualEncoder模型
model = DualEncoder(title_encoder, reply_encoder, type=loss_function)
# 设置特定的损失函数
if loss_function == "cosine":
loss_fn = nn.CosineEmbeddingLoss()
else:
loss_fn = nn.BCEWithLogitsLoss()
# Adam优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# 有GPU就用GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 分词器(词典)可以重复使用,为了不出问题,在推理时要使用和训练时同样的分词器(词典)
pickle.dump(tokenizer, open(os.path.join(output_dir,"tokenizer.pickle"), "wb"))
for epoch in tqdm(range(num_epochs)):
train(train_set, model, loss_fn, optimizer, device, tokenizer, loss_function, batch_size)