一、数据预处理
- 句子内容 分类标签
- 数据分为训练集和测试集
代码:
import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import BertTokenizerFast
# 读取数据
data = pd.read_csv('sentences.csv')
# 分割数据集为训练集和测试集
train_texts, test_texts, train_labels, test_labels = train_test_split(data['text'], data['label'], test_size=0.2, random_state=42)
# 初始化BERT tokenizer
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
# 对文本进行编码
train_encodings = tokenizer(list(train_texts), truncation=True, padding=True)
test_encodings = tokenizer(list(test_texts), truncation=True, padding=True)
# 转换为PyTorch的数据集格式
class Dataset(torch.utils.data.Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item['labels'] = torch.tensor(self.labels[idx])
return item
def __len__(self):
return len(self.labels)
train_dataset = Dataset(train_encodings, train_labels)
test_dataset = Dataset(test_encodings, test_labels)
二、模型定义
定义基于BERT的分类模型。Transformers库允许我们轻松地在预训练的BERT模型上添加一个分类头。
新添加的分类头是一个全连接层(Dense Layer),其输入维度等于BERT模型的输出维度(对于BERT-base通常是768维),输出维度则等于你任务中的类别数量。例如,如果你正在进行二分类任务,输出维度就是2;如果是多分类任务,输出维度则与类别数量相匹配。
from transformers import BertForSequenceClassification
# 加载预训练的BERT模型并添加分类头
num_labels = len(data['label'].unique()) # 假设标签是离散的,计算类别数
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=num_labels)
三、训练模型
PyTorch的DataLoader来加载数据,并用Trainer类进行模型训练,这是Transformers提供的一个高级API。
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir='./results', # 输出目录
num_train_epochs=3, # 总的训练轮数
per_device_train_batch_size=16, # 每个GPU的训练批次大小
per_device_eval_batch_size=64, # 每个GPU的评估批次大小
warmup_steps=500, # 预热步数
weight_decay=0.01, # 权重衰减
logging_dir='./logs', # 日志目录
)
trainer = Trainer(
model=model, # 要训练的模型
args=training_args, # 训练参数
train_dataset=train_dataset, # 训练数据集
eval_dataset=test_dataset # 评估数据集
)
trainer.train()
四、评估模型
使用trainer.evaluate()来评估模型在测试集上的性能。
eval_results = trainer.evaluate()
print(f"Accuracy on Test Set: {eval_results['accuracy']*100:.2f}%")