基于 BERT 的文本分类系统
需要一个基于历史标注数据的自动分类系统,将领导们的 issues 分配到三个预定义的类别中。这种场景非常适合使用自然语言处理(NLP)中的文本分类技术。
·
需要一个基于历史标注数据的自动分类系统,将领导们的 issues 分配到三个预定义的类别中。这种场景非常适合使用自然语言处理(NLP)中的文本分类技术。
解决方案思路
- 数据准备:收集历史标注的 issues 数据
- 特征工程:将文本描述转换为数值特征向量
- 模型训练:使用标注数据训练分类模型
- 模型部署:实现自动化分类流程
代码实现:基于 BERT 的文本分类系统
下面是一个使用 BERT-base-chinese 预训练模型的文本分类解决方案,它能够根据 issue 描述自动分配到三个预定义类别中:
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from tqdm import tqdm
import os
class IssueClassificationSystem:
"""领导issue自动分类系统"""
def __init__(self, num_classes=3, model_name="bert-base-chinese", max_length=128):
"""初始化分类系统"""
self.tokenizer = BertTokenizer.from_pretrained(model_name)
self.model = BertForSequenceClassification.from_pretrained(
model_name,
num_labels=num_classes,
output_attentions=False,
output_hidden_states=False
)
self.max_length = max_length
self.num_classes = num_classes
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
def prepare_data(self, data_path):
"""准备训练数据"""
# 读取CSV文件,假设文件包含两列:text(issue描述)和label(分类标签)
df = pd.read_csv(data_path)
# 划分训练集和验证集
train_df, val_df = train_test_split(df, test_size=0.2, random_state=42)
# 创建数据集
train_dataset = IssueDataset(train_df, self.tokenizer, self.max_length)
val_dataset = IssueDataset(val_df, self.tokenizer, self.max_length)
# 创建数据加载器
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=16, shuffle=False)
return train_dataloader, val_dataloader
def train(self, train_dataloader, val_dataloader, epochs=4, learning_rate=2e-5):
"""训练分类模型"""
optimizer = AdamW(self.model.parameters(), lr=learning_rate)
for epoch in range(epochs):
# 训练阶段
self.model.train()
total_train_loss = 0
for batch in tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{epochs} [Training]"):
# 将数据移至设备
b_input_ids = batch['input_ids'].to(self.device)
b_input_mask = batch['attention_mask'].to(self.device)
b_labels = batch['labels'].to(self.device)
# 清零梯度
self.model.zero_grad()
# 前向传播
outputs = self.model(
b_input_ids,
token_type_ids=None,
attention_mask=b_input_mask,
labels=b_labels
)
loss = outputs.loss
total_train_loss += loss.item()
# 反向传播
loss.backward()
optimizer.step()
# 验证阶段
self.model.eval()
total_val_accuracy = 0
total_val_loss = 0
for batch in tqdm(val_dataloader, desc=f"Epoch {epoch+1}/{epochs} [Validation]"):
b_input_ids = batch['input_ids'].to(self.device)
b_input_mask = batch['attention_mask'].to(self.device)
b_labels = batch['labels'].to(self.device)
with torch.no_grad():
outputs = self.model(
b_input_ids,
token_type_ids=None,
attention_mask=b_input_mask,
labels=b_labels
)
loss = outputs.loss
total_val_loss += loss.item()
# 计算准确率
logits = outputs.logits
logits = logits.detach().cpu().numpy()
label_ids = b_labels.to('cpu').numpy()
total_val_accuracy += self._flat_accuracy(logits, label_ids)
# 输出训练结果
avg_train_loss = total_train_loss / len(train_dataloader)
avg_val_loss = total_val_loss / len(val_dataloader)
avg_val_accuracy = total_val_accuracy / len(val_dataloader)
print(f"Epoch: {epoch+1}")
print(f" Train Loss: {avg_train_loss:.4f}")
print(f" Validation Loss: {avg_val_loss:.4f}")
print(f" Validation Accuracy: {avg_val_accuracy:.4f}")
def predict(self, text):
"""预测单个issue的分类"""
self.model.eval()
# 对文本进行分词和编码
encoded_dict = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
# 将数据移至设备
input_ids = encoded_dict['input_ids'].to(self.device)
attention_mask = encoded_dict['attention_mask'].to(self.device)
# 预测
with torch.no_grad():
outputs = self.model(input_ids, token_type_ids=None, attention_mask=attention_mask)
# 获取预测结果
logits = outputs.logits
logits = logits.detach().cpu().numpy()
predicted_class = np.argmax(logits, axis=1)[0]
return predicted_class
def save_model(self, save_dir):
"""保存模型"""
if not os.path.exists(save_dir):
os.makedirs(save_dir)
self.model.save_pretrained(save_dir)
self.tokenizer.save_pretrained(save_dir)
def load_model(self, save_dir):
"""加载模型"""
self.model = BertForSequenceClassification.from_pretrained(save_dir)
self.tokenizer = BertTokenizer.from_pretrained(save_dir)
self.model.to(self.device)
def _flat_accuracy(self, preds, labels):
"""计算准确率"""
pred_flat = np.argmax(preds, axis=1).flatten()
labels_flat = labels.flatten()
return np.sum(pred_flat == labels_flat) / len(labels_flat)
class IssueDataset(Dataset):
"""领导issue数据集"""
def __init__(self, df, tokenizer, max_length):
self.texts = df['text'].tolist()
self.labels = df['label'].tolist()
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
# 对文本进行分词和编码
encoding = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
return {
'input_ids': encoding['input_ids'].flatten(),
'attention_mask': encoding['attention_mask'].flatten(),
'labels': torch.tensor(label, dtype=torch.long)
}
# 使用示例
def demo():
# 创建示例数据
sample_data = {
'text': [
"团队成员绩效评估和晋升计划",
"季度财务报表审核和分析",
"新市场拓展策略讨论",
"产品研发进度跟踪",
"预算分配和成本控制",
"客户投诉处理流程优化",
"招聘计划和人才储备",
"投资回报率分析",
"合作伙伴关系维护",
"员工培训计划制定"
],
'label': [0, 1, 2, 0, 1, 2, 0, 1, 2, 0] # 假设0=人力资源, 1=财务, 2=业务发展
}
df = pd.DataFrame(sample_data)
df.to_csv("sample_issues.csv", index=False)
# 初始化分类系统
classifier = IssueClassificationSystem(num_classes=3)
# 准备数据
train_dataloader, val_dataloader = classifier.prepare_data("sample_issues.csv")
# 训练模型
classifier.train(train_dataloader, val_dataloader, epochs=2)
# 保存模型
classifier.save_model("issue_classifier_model")
# 加载模型(实际使用时只需加载已训练好的模型)
classifier.load_model("issue_classifier_model")
# 预测新issue
new_issue = "下季度的营销预算申请"
predicted_class = classifier.predict(new_issue)
# 假设类别映射: 0=人力资源, 1=财务, 2=业务发展
class_mapping = {0: "人力资源", 1: "财务", 2: "业务发展"}
print(f"\n新issue: {new_issue}")
print(f"预测分类: {class_mapping[predicted_class]}")
if __name__ == "__main__":
demo()
系统使用说明
-
数据格式:
- 准备一个 CSV 文件,包含两列:
text
(issue 描述)和label
(分类标签,0、1、2 对应三个类别)
- 准备一个 CSV 文件,包含两列:
-
模型训练:
- 运行
demo()
函数或按照示例代码训练模型 - 训练完成后,模型会保存在
issue_classifier_model
目录
- 运行
-
模型应用:
- 加载已训练模型:
classifier.load_model("issue_classifier_model")
- 对新 issue 进行分类:
predicted_class = classifier.predict("新issue描述")
- 加载已训练模型:
-
类别映射:
- 根据实际业务需求,将数字标签映射到具体类别名称(示例中:0 = 人力资源,1 = 财务,2 = 业务发展)
优化建议
- 增加训练数据:收集更多标注数据以提高模型准确率
- 调整超参数:尝试不同的学习率、批次大小和训练轮次
- 使用多标签分类:如果一个 issue 可能属于多个类别,可以改为多标签分类
- 引入主动学习:让模型对不确定的样本进行标记请求,减少人工标注工作量
- 结合规则引擎:在模型分类基础上,加入业务规则进行修正
这个系统可以帮助您的团队自动将领导的 issues 分配到三个预定义类别中,大幅提高工作效率和分类一致性。
更多推荐
所有评论(0)