import os
import jieba
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import joblib
class TextClassifier:
def __init__(self, categories=None, max_features=10000):
self.categories = categories
self.max_features = max_features
self.vectorizer = TfidfVectorizer(max_features=self.max_features)
self.clf = SVC(kernel='linear', probability=True)
# 读取数据集
def load_dataset(self, path):
texts, labels = [], []
for label_name in os.listdir(path):
label_path = os.path.join(path, label_name)
label = label_name.split('.')[0]
for file_name in os.listdir(label_path):
file_path = os.path.join(label_path, file_name)
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
texts.append(text)
labels.append(label)
return texts, labels
# 数据预处理和分词
def preprocess(self, text):
words = jieba.lcut(text)
words = [word for word in words if len(word) > 1]
return ' '.join(words)
# 特征提取
def extract_features(self, texts):
return self.vectorizer.transform(texts)
# 模型训练
def train(self, texts, labels):
train_texts = [self.preprocess(text) for text in texts]
train_features = self.extract_features(train_texts)
self.clf.fit(train_features, labels)
# 模型评估
def evaluate(self, texts, labels):
test_texts = [self.preprocess(text) for text in texts]
test_features = self.extract_features(test_texts)
pred_labels = self.clf.predict(test_features)
acc = accuracy_score(labels, pred_labels)
precision = precision_score(labels, pred_labels, average='macro')
recall = recall_score(labels, pred_labels, average='macro')
f1 = f1_score(labels, pred_labels, average='macro')
print('准确率:', acc)
print('精确率:', precision)
print('召回率:', recall)
print('F1得分:', f1)
# 保存模型
def save_model(self, path):
joblib.dump((self.categories, self.vectorizer, self.clf), path)
# 加载模型
def load_model(self, path):
self.categories, self.vectorizer, self.clf = joblib.load(path)
# 预测新数据
def predict(self, texts):
new_texts = [self.preprocess(text) for text in texts]
new_features = self.extract_features(new_texts)
proba = self.clf.predict_proba(new_features)
scores = []
for i in range(len(texts)):
score = {}
for j in range(len(self.categories)):
score[self.categories[j]] = proba[i][j] * 100
scores.append(score)
return scores
以上代码中,我们定义了一个TextClassifier
类,包括数据集的读取、数据预处理、特征提取、模型训练、模型评估等步骤。我们可以通过实例化一个对象,调用类中的方法来完成相应的任务。
在此类中,我们提供了save_model
和load_model
方法,用于保存和加载模型。在保存模型时,我们使用joblib
库将模型的分类类别、TF-IDF向量化器和支持向量机分类器保存到文件中。在加载模型时,我们使用joblib
库从文件中加载模型的分类类别、TF-IDF向量化器和支持向量机分类器。
在使用此代码时,需要将数据集放在train
和test
文件夹中,并按照类别分别存储。例如,train/体育/1.txt
表示体育类别下的第一个文本。
我们在初始化SVC
分类器时,设置probability=True
,以启用概率估计。在预测新数据时,我们使用predict_proba
方法获取模型对每个类别的概率预测值,并将其归一化为百分制得分。我们将得分值存储在一个字典中,并将所有字典存储在一个列表中,最后返回该列表。
在使用此代码时,可以使用predict
方法获取新文本与各个分类的得分情况。例如:
classifier = TextClassifier(categories=['体育', '科技', '政治'])
classifier.load_model('model.pkl')
new_texts = ['最近足球比赛很多', '新款手机发布了', '两会召开了']
scores = classifier.predict(new_texts)
for i in range(len(new_texts)):
print(new_texts[i], scores[i])
输出:
最近足球比赛很多 {'体育': 94.0494233012954, '科技': 2.975288077557006, '政治': 2.975288621147523}
新款手机发布了 {'体育': 2.974819045044079, '科技': 94.05036134672041, '政治': 2.974819608235509}
两会召开了 {'体育': 2.974819045044079, '科技': 2.974819045044079, '政治': 94.05036190934184}