支持向量机类实现

2023-11-18

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_modelload_model方法,用于保存和加载模型。在保存模型时,我们使用joblib库将模型的分类类别、TF-IDF向量化器和支持向量机分类器保存到文件中。在加载模型时,我们使用joblib库从文件中加载模型的分类类别、TF-IDF向量化器和支持向量机分类器。

在使用此代码时,需要将数据集放在traintest文件夹中,并按照类别分别存储。例如,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}

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

支持向量机类实现 的相关文章