(9)创建自定义数据集类CustomDataset,用于将文本数据集与BERT模型的输入格式相对应。通过使用PyTorch中的类Dataset,它可以被PyTorch的DataLoader加载,用于实现模型的训练工作。
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, df, tokenizer):
super().__init__()
self.tokenizer = tokenizer
self.inputs = df["Word"].values
self.labels = df["TagId"].values
self.len = len(self.labels)
def __getitem__(self, idx):
# 使用分词器对输入文本进行编码
encoded_dict = self.tokenizer(
self.inputs[idx],
is_split_into_words=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors='pt',
)
# 创建标记级别的标签
token_labels = self.__create_token_labels(self.labels[idx], encoded_dict.word_ids())
# 返回编码后的输入及其对应的标签
return encoded_dict['input_ids'][0], encoded_dict['token_type_ids'][0], encoded_dict['attention_mask'][0], token_labels
def __create_token_labels(self, labels, word_ids):
# 将标签映射到每个单词的位置,未知单词位置用-100表示
extended_labels = [-100 if word_id is None else labels[word_id] for word_id in word_ids]
return torch.tensor(extended_labels).to(torch.int64)
def __len__(self):
return self.len
对上述代码的具体说明如下所示:
- __init__ 方法:用于初始化数据集,接收一个DataFrame df 和一个分词器 tokenizer。
- __getitem__ 方法:用于获取数据集中的单个样本,返回编码后的输入和相应的标签。
- __create_token_labels:用于将整个序列的标签映射到每个单词的位置,使用-100表示未知单词的位置。
- __len__ 方法:用于返回数据集的长度。
这个数据集类CustomDataset的设计目的旨在与BERT模型的输入和标签格式相匹配,以便在训练过程中使用。
(10)创建函数collate_fn,用于处理数据批次,确保批次中的样本具有相同的长度。函数collate_fn能够计算批次中样本的最大长度,然后对每个样本的输入、token类型、注意力掩码和标签进行填充,以使它们具有相同的长度。在填充操作中,使用0填充输入和掩码,使用-100填充标签。最后,函数collate_fn将填充后的张量堆叠起来,并将其移到指定的设备上(GPU或CPU)。
# 定义一个用于处理数据批次的函数 collate_fn
def collate_fn(batch):
# 计算批次中样本的最大长度
max_len = max([len(sample[0]) for sample in batch])
# 初始化空列表,用于存储输入、token类型、注意力掩码和标签
ids = []
token_types = []
attention_masks = []
labels = []
# 遍历批次中的每个样本
for sample in batch:
# 计算需要填充的长度
pad_length = max_len - len(sample[0])
# 分别对输入、token类型、注意力掩码和标签进行填充,使用0填充输入和掩码,使用-100填充标签
ids.append(pad(sample[0], (0, pad_length), value=0))
token_types.append(pad(sample[1], (0, pad_length), value=0))
attention_masks.append(pad(sample[2], (0, pad_length), value=0))
labels.append(pad(sample[3], (0, pad_length), value=-100))
# 将填充后的张量堆叠起来,并将其移到指定的设备上(GPU或CPU)
return torch.stack(ids).to(device), torch.stack(token_types).to(device), torch.stack(attention_masks).to(device), torch.stack(labels).to(device)
(11)下面的代码首先创建了一个自定义数据集的实例 dataset,并使用 random_split 函数将其划分为训练集 train_dataset 和验证集 validation_dataset,其中80%用于训练,20%用于验证。最后,打印了训练集和验证集的样本数量。
# 创建自定义数据集实例
dataset = CustomDataset(df, tokenizer)
# 划分训练集和验证集
train_size = int(0.8 * len(dataset))
validation_size = len(dataset) - train_size
train_dataset, validation_dataset = random_split(dataset, [train_size, validation_size])
#