基于BertModel()的中文文本分类模型实现
发布时间:2024-01-06 17:48:17
BERT(Bidirectional Encoder Representations from Transformers)是一种预训练语言模型,可以用于各种自然语言处理任务,包括文本分类。在这个示例中,我们将使用基于BERT模型的中文文本分类。
首先,我们需要安装所需的库。我们将使用transformers库来加载BERT模型,使用torch库来构建和训练模型。
!pip install transformers torch
接下来,我们将导入所需的库并加载BERT模型和tokenizer。
import torch from transformers import BertModel, BertTokenizer # 加载BERT模型和tokenizer model_name = 'bert-base-chinese' model = BertModel.from_pretrained(model_name) tokenizer = BertTokenizer.from_pretrained(model_name)
我们将使用一个中文的情感分析任务作为示例。我们将使用情感分析数据集,并将其分为训练集和测试集。我们的目标是根据给定的中文文本预测情感类别,例如正面或负面情感。
import pandas as pd
from sklearn.model_selection import train_test_split
# 加载情感分析数据集
data = pd.read_csv('sentiment_data.csv')
# 划分训练集和测试集
train_data, test_data = train_test_split(data, test_size=0.2, random_state=42)
现在,我们将使用BERT tokenizer将中文文本转换为模型可接受的输入格式。BERT模型需要输入文本的标记化表示形式,即每个词汇都会映射到一个唯一的标记ID。此外,BERT模型还需要输入每个句子的“attention mask”,以指示哪些标记是实际的词汇,哪些是填充标记。我们将使用tokenizer.encode_plus方法来进行标记化和生成注意力蒙版。
# 限制输入文本的最大长度
max_length = 128
def tokenize_text(text):
tokens = tokenizer.encode_plus(
text,
max_length=max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
return tokens['input_ids'], tokens['attention_mask']
# 对训练集进行标记化和生成注意力蒙版
train_tokens = train_data['text'].apply(tokenize_text)
train_input_ids = torch.cat([x[0] for x in train_tokens])
train_attention_mask = torch.cat([x[1] for x in train_tokens])
train_labels = torch.tensor(train_data['label'])
# 对测试集进行标记化和生成注意力蒙版
test_tokens = test_data['text'].apply(tokenize_text)
test_input_ids = torch.cat([x[0] for x in test_tokens])
test_attention_mask = torch.cat([x[1] for x in test_tokens])
test_labels = torch.tensor(test_data['label'])
现在准备好了训练集和测试集的输入。我们可以定义一个简单的分类模型来使用BERT模型进行文本分类。我们将用BERT模型对输入的文本进行编码,并将其传入一个线性层进行分类。
import torch.nn as nn
class BERTClassifier(nn.Module):
def __init__(self, bert_model, num_labels):
super(BERTClassifier, self).__init__()
self.bert = bert_model
self.linear = nn.Linear(self.bert.config.hidden_size, num_labels)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs[1] # 取CLS标记的输出向量,用于分类
logits = self.linear(pooled_output)
return logits
# 定义模型并移动到GPU(如果可用)
num_labels = len(data['label'].unique())
model = BERTClassifier(model, num_labels)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
接下来,我们需要定义训练模型所需的损失函数和优化器。
import torch.optim as optim # 定义损失函数 criterion = nn.CrossEntropyLoss() # 定义优化器 optimizer = optim.Adam(model.parameters(), lr=2e-5)
我们可以开始训练模型了。
import time
# 训练参数
epochs = 10
batch_size = 32
# 训练模型
start_time = time.time()
model.train()
for epoch in range(epochs):
running_loss = 0.0
# 批量训练
for i in range(0, len(train_input_ids), batch_size):
input_ids = train_input_ids[i:i+batch_size].to(device)
attention_mask = train_attention_mask[i:i+batch_size].to(device)
labels = train_labels[i:i+batch_size].to(device)
# 梯度置零
optimizer.zero_grad()
# 前向传播
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
loss = criterion(outputs, labels)
# 反向传播和优化
loss.backward()
optimizer.step()
running_loss += loss.item()
# 打印训练损失
print(f'Epoch {epoch + 1} - Training loss: {running_loss / len(train_input_ids):.4f}')
end_time = time.time()
total_time = end_time - start_time
print(f'Training took {total_time} seconds')
训练完成后,我们可以在测试集上评估模型的性能。
model.eval() # 将模型切换为评估模式
correct = 0
for i in range(0, len(test_input_ids), batch_size):
input_ids = test_input_ids[i:i+batch_size].to(device)
attention_mask = test_attention_mask[i:i+batch_size].to(device)
labels = test_labels[i:i+batch_size].to(device)
with torch.no_grad():
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
_, predicted = torch.max(outputs, dim=1)
correct += (predicted == labels).sum().item()
accuracy = correct / len(test_input_ids)
print(f'Test accuracy: {accuracy:.4f}')
这就是使用BERT模型的中文文本分类模型的实现示例。我们可以根据需要调整模型的超参数,例如批次大小、学习率和训练的时期来优化模型的性能。
