当前位置：首页 > news >正文

NLP:微调BERT进行文本分类

news 来源：原创 2024/9/22 1:01:21

本篇博客的重点在于BERT的使用。
transformers包版本：4.44.2

1. 微调BERT进行文本分类

这里我们使用stanford大学的SST2数据集来演示BERT模型的微调过程。SST-2数据集(Stanford Sentiment Treebank 2)是一个用于情感分类的经典数据集，常用于自然语言处理(NLP)领域的情感分析任务。

第1步: 下载数据。其代码如下：

import pandas as pd
from transformers import BertTokenizer
from datasets import DatasetDict, Dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArgumentssplits = {'train': 'data/train-00000-of-00001.parquet', 'validation': 'data/validation-00000-of-00001.parquet', 'test': 'data/test-00000-of-00001.parquet'}
train = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["train"])
validation = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["validation"])
test = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["test"])
dataset = DatasetDict({'train': Dataset.from_pandas(train), 'validation': Dataset.from_pandas(validation), 'test': Dataset.from_pandas(test)})

要注意一下，这里并没有使用datasets包从hugging face上直接下载数据集的方式来获取数据，这是因为使用load_datesets方法获取数据时仍然会提示：NotImplementedError: Loading a dataset cached in a LocalFileSystem is not supported。样例数据集如下：
在这里插入图片描述

第2步: 构造训练集、验证集和测试集

SST2数据集中训练集(train)共67349条，验证集(validation)共872条，而测试集(test)共1821条。因为训练集数目较大微调会比较耗时，所以从这三个数据集分别抽取出了1000条、200条、200条进行后续的任务。具体代码如下：

dataset['train'] = dataset['train'].shuffle(seed=42).select(range(1000))
dataset['validation'] = dataset['validation'].shuffle(seed=42).select(range(200))
dataset['test'] = dataset['test'].shuffle(seed=42).select(range(200))
print(dataset)

其输出结果如下：

Dataset({features: ['idx', 'sentence', 'label'],num_rows: 1000
})
Dataset({features: ['idx', 'sentence', 'label'],num_rows: 200
})
Dataset({features: ['idx', 'sentence', 'label'],num_rows: 200
})

第3步：从bert中提取嵌入

训练集、验证集及测试集生成后，接着需要将这些语料全都转化成embedding向量。具体代码如下：

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def tokenize_function(examples):return tokenizer(examples['sentence'], padding='max_length', truncation=True)
dataset =dataset.map(tokenize_function, batched=True)
dataset=dataset.remove_columns(['sentence',"idx"])
dataset=dataset.rename_column("label","labels")
dataset.set_format("torch")
train_dataset=dataset['train']
eval_dataset=dataset['validation']
test_dataset=dataset['test']

第4步:模型训练。 具体代码如下：

model=BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
training_args = TrainingArguments(output_dir='results',         per_device_train_batch_size=8,  per_device_eval_batch_size=8,      num_train_epochs=1,
)
trainer = Trainer(model=model,                         args=training_args,                 train_dataset=train_dataset,        eval_dataset=eval_dataset,
)
trainer.train()
eval_result=trainer.evaluate(eval_dataset)
print(eval_result)
test_dataset=test_dataset.remove_columns(["labels"])
pred_result = trainer.predict(test_dataset=test_dataset)
print(pred_result)

关于上述代码，有以下几点需要说明：

训练模型的选择: tranformers库中有多个分类模型，其中BertForSequenceClassification类适用于序列分类任务，比如情感分析和文本分类；而BertForTokenClassification类适用于token级的分类任务，比如命名实体识别。
TrainingArguments方法中的主要参数及其作用如下表所示：

参数名	作用
output_dir	指定模型和训练日志保存的记录;
num_train_epochs	设置训练的周期数(即遍历整个训练数据集的次数，指的是整个训练集将被遍历多少次以进行训练);
per_device_train_batch_size	设置每个设备(如GPU)上的训练批次大小，训练批次是指在一次训练迭代中，模型同时处理的数据样本数量;bert微调时会将整个训练数据集分成多个批次逐步进行训练；
per_device_eval_batch_size	设置每个设备上的评估批次大小;
logging_dir	指定训练日志的保存目录;
evaluation_strategy	设置评估策略。可以是 ‘no’(不评估)、‘steps’(每隔一定步数评估)或 ‘epoch’(每个周期评估);
save_total_limit	设置保存模型检查点的总数限制，超过限制的检查点会被删除；
fp16	启用半精度浮点数(FP16)训练，以减少显存使用并加速训练(需要支持 FP16 的硬件);
max_steps	最大迭代次数