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

实战2：利用Python与AI模型实现文本分类

news 2025/5/18 6:18:16

本文将带你使用Python和机器学习算法实现一个简单的文本分类应用。通过该项目，你将学习如何预处理数据、训练分类模型，并使用现有的AI模型（如BERT）来提升分类效果。

一、项目背景

文本分类是自然语言处理（NLP）中的一个常见任务，它的应用场景广泛，如垃圾邮件分类、情感分析、新闻分类等。在本篇中，我们将通过Python实现一个文本分类器，利用BERT模型进行预训练，提升模型的准确性。

二、核心思路

数据收集与预处理：首先，我们需要收集文本数据，并对其进行必要的清洗和预处理。
模型选择：使用BERT模型进行文本表示，利用预训练模型来增强文本特征。
训练与评估：训练分类模型，评估其性能，并进行适当的优化。

三、开发环境准备

需要安装以下Python库：

pip install transformers torch sklearn pandas numpy

transformers：用于加载BERT等预训练模型。
torch：PyTorch深度学习框架，用于训练和评估模型。
sklearn：用于机器学习中常见的评估指标，如精度、召回率等。
pandas：用于数据处理和加载。

四、数据准备与预处理

假设我们使用一个公开的新闻数据集，数据格式如下：

文本内容	标签
经济下行压力大	经济类
疫苗研发取得突破	健康类
公司发布新产品	商业类
...	...

1. 数据清洗

首先，我们将加载数据并进行基础清洗，如去除停用词和标点符号。

import pandas as pd
import re
from sklearn.model_selection import train_test_split# 加载数据
df = pd.read_csv("news_data.csv")# 简单清洗文本：去除标点和多余空格
def clean_text(text):text = re.sub(r"[^A-Za-z0-9]+", " ", text)  # 只保留字母和数字text = text.lower().strip()  # 小写化并去除首尾空格return textdf['cleaned_text'] = df['文本内容'].apply(clean_text)# 划分训练集与测试集
X_train, X_test, y_train, y_test = train_test_split(df['cleaned_text'], df['标签'], test_size=0.2, random_state=42)

2. 将文本转换为向量

BERT模型将文本转化为向量表示。我们使用transformers库加载BERT模型。

from transformers import BertTokenizer# 加载预训练的BERT Tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')def tokenize_data(texts):return tokenizer(texts.tolist(), padding=True, truncation=True, return_tensors="pt")train_data = tokenize_data(X_train)
test_data = tokenize_data(X_test)

五、构建分类模型

1. 加载BERT模型

我们将使用BERTForSequenceClassification，这是Hugging Face提供的适合文本分类任务的模型。

from transformers import BertForSequenceClassification
import torch
from torch.utils.data import DataLoader, TensorDataset# 加载BERT模型
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=len(df['标签'].unique()))# 将数据转化为TensorDataset
train_dataset = TensorDataset(train_data['input_ids'], train_data['attention_mask'], torch.tensor(y_train.values))
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)

2. 模型训练

from transformers import AdamW
from torch.optim.lr_scheduler import StepLR# 定义优化器和学习率调度器
optimizer = AdamW(model.parameters(), lr=2e-5)
scheduler = StepLR(optimizer, step_size=1, gamma=0.1)# 训练循环
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)epochs = 3
for epoch in range(epochs):model.train()for batch in train_dataloader:input_ids, attention_mask, labels = [b.to(device) for b in batch]optimizer.zero_grad()# 前向传播outputs = model(input_ids, attention_mask=attention_mask, labels=labels)loss = outputs.lossloss.backward()optimizer.step()scheduler.step()print(f"Epoch {epoch+1}/{epochs} - Loss: {loss.item()}")

六、模型评估

训练完成后，我们在测试集上评估模型的表现：

from sklearn.metrics import classification_report# 模型评估
model.eval()
with torch.no_grad():outputs = model(test_data['input_ids'].to(device), attention_mask=test_data['attention_mask'].to(device))predictions = torch.argmax(outputs.logits, dim=-1)# 打印评估报告
print(classification_report(y_test, predictions.cpu().numpy()))