脚本如何实现文件内容朴素贝叶斯

本文目录导读：

1. 完整的Python实现
2. 使用示例
3. 高级功能扩展
4. 批量处理多个文件
5. 关键技术点总结

进行分类的实现。

完整的Python实现

import os
import re
import math
from collections import defaultdict, Counter
import pickle
class NaiveBayesClassifier:
    def __init__(self, alpha=1.0):
        """
        初始化朴素贝叶斯分类器
        alpha: 拉普拉斯平滑参数
        """
        self.alpha = alpha
        self.class_probabilities = {}  # P(类别)
        self.word_probabilities = {}   # P(单词|类别)
        self.classes = []
        self.vocabulary = set()
    def _tokenize(self, text):
        """文本预处理和分词"""
        # 转为小写
        text = text.lower()
        # 去除标点符号和数字
        text = re.sub(r'[^\w\s]', ' ', text)
        text = re.sub(r'\d+', ' ', text)
        # 分词
        tokens = text.split()
        # 去除停用词（可以根据需要扩充）
        stop_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'by', 'with', 'from'}
        tokens = [t for t in tokens if t not in stop_words and len(t) > 2]
        return tokens
    def _read_file(self, filepath):
        """读取文件内容"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            # 尝试其他编码
            try:
                with open(filepath, 'r', encoding='latin-1') as f:
                    return f.read()
            except Exception as e:
                print(f"Error reading file {filepath}: {e}")
                return ""
    def train_from_directory(self, directory_path):
        """
        从目录结构训练模型
        目录结构: 
            train/
                category1/
                    file1.txt
                    file2.txt
                category2/
                    file3.txt
                    ...
        """
        # 获取所有类别
        self.classes = [d for d in os.listdir(directory_path) 
                       if os.path.isdir(os.path.join(directory_path, d))]
        # 统计每个类别的文档数量和单词数量
        class_doc_count = defaultdict(int)
        class_word_count = defaultdict(Counter)
        total_docs = 0
        # 遍历所有文件
        for category in self.classes:
            category_path = os.path.join(directory_path, category)
            for filename in os.listdir(category_path):
                filepath = os.path.join(category_path, filename)
                if os.path.isfile(filepath):
                    content = self._read_file(filepath)
                    if content:
                        tokens = self._tokenize(content)
                        class_doc_count[category] += 1
                        class_word_count[category].update(tokens)
                        self.vocabulary.update(tokens)
                        total_docs += 1
        # 计算类先验概率 P(类别)
        self.class_probabilities = {
            c: count / total_docs 
            for c, count in class_doc_count.items()
        }
        # 计算条件概率 P(单词|类别) 使用拉普拉斯平滑
        vocab_size = len(self.vocabulary)
        self.word_probabilities = {}
        for category in self.classes:
            word_counts = class_word_count[category]
            total_words = sum(word_counts.values())
            self.word_probabilities[category] = {}
            for word in self.vocabulary:
                # 使用拉普拉斯平滑
                word_prob = (word_counts.get(word, 0) + self.alpha) / (total_words + self.alpha * vocab_size)
                self.word_probabilities[category][word] = word_prob
        print(f"训练完成！共 {total_docs} 个文档，{vocab_size} 个唯一词汇，{len(self.classes)} 个类别")
    def predict(self, text_or_filepath):
        """
        预测文本或文件的类别
        返回: (预测类别, 各类别概率)
        """
        # 如果是文件路径，读取文件内容
        if os.path.isfile(text_or_filepath):
            text = self._read_file(text_or_filepath)
        else:
            text = text_or_filepath
        # 分词
        tokens = self._tokenize(text)
        if not tokens:
            return None, {}
        # 计算每个类别的对数概率
        log_probabilities = {}
        for category in self.classes:
            # 初始化为对数类先验概率
            log_prob = math.log(self.class_probabilities[category])
            # 加上每个词的对数条件概率
            for token in tokens:
                if token in self.vocabulary:
                    log_prob += math.log(self.word_probabilities[category].get(token, self.alpha))
            log_probabilities[category] = log_prob
        # 计算实际概率（用于返回）
        probabilities = self._softmax(log_probabilities)
        # 选择概率最大的类别
        predicted_class = max(probabilities, key=probabilities.get)
        return predicted_class, probabilities
    def _softmax(self, log_probs):
        """将对数概率转换为实际概率"""
        # 找到最大值避免数值溢出
        max_log = max(log_probs.values())
        exp_probs = {c: math.exp(log_p - max_log) for c, log_p in log_probs.items()}
        total = sum(exp_probs.values())
        return {c: p / total for c, p in exp_probs.items()}
    def save_model(self, filepath):
        """保存模型"""
        model = {
            'class_probabilities': self.class_probabilities,
            'word_probabilities': self.word_probabilities,
            'classes': self.classes,
            'vocabulary': self.vocabulary,
            'alpha': self.alpha
        }
        with open(filepath, 'wb') as f:
            pickle.dump(model, f)
        print(f"模型已保存到 {filepath}")
    def load_model(self, filepath):
        """加载模型"""
        with open(filepath, 'rb') as f:
            model = pickle.load(f)
        self.class_probabilities = model['class_probabilities']
        self.word_probabilities = model['word_probabilities']
        self.classes = model['classes']
        self.vocabulary = model['vocabulary']
        self.alpha = model['alpha']
        print(f"模型已从 {filepath} 加载")
    def evaluate(self, test_dir):
        """评估模型性能"""
        correct = 0
        total = 0
        confusion_matrix = defaultdict(lambda: defaultdict(int))
        for actual_class in self.classes:
            class_path = os.path.join(test_dir, actual_class)
            if not os.path.isdir(class_path):
                continue
            for filename in os.listdir(class_path):
                filepath = os.path.join(class_path, filename)
                if os.path.isfile(filepath):
                    predicted, _ = self.predict(filepath)
                    if predicted == actual_class:
                        correct += 1
                    confusion_matrix[actual_class][predicted] += 1
                    total += 1
        accuracy = correct / total if total > 0 else 0
        print(f"准确率: {accuracy:.2%} ({correct}/{total})")
        # 打印混淆矩阵
        print("\n混淆矩阵:")
        print("真实\\预测", end="\t")
        for c in self.classes:
            print(c[:8], end="\t")
        print()
        for true_class in self.classes:
            print(true_class[:8], end="\t")
            for pred_class in self.classes:
                print(confusion_matrix[true_class][pred_class], end="\t")
            print()
        return accuracy, confusion_matrix

使用示例

# 1. 创建数据目录结构
"""
train/
    spam/
        email1.txt
        email2.txt
    ham/
        email3.txt
        email4.txt
test/
    spam/
        email5.txt
    ham/
        email6.txt
"""
# 批量创建测试文件
import os
def create_sample_data():
    """创建示例数据集"""
    # 创建目录
    for category in ['spam', 'ham']:
        os.makedirs(f'train/{category}', exist_ok=True)
        os.makedirs(f'test/{category}', exist_ok=True)
    # 创建示例文件
    spam_samples = [
        "Buy cheap medicine now!!! Limited offer",
        "You have won a free prize! Click here",
        "Make money fast, work from home",
        "Urgent! Your account needs verification"
    ]
    ham_samples = [
        "Meeting scheduled for tomorrow at 2pm",
        "Please review the attached document",
        "Thanks for your help with the project",
        "Can we reschedule our appointment?"
    ]
    # 写入训练数据
    for i, text in enumerate(spam_samples[:3]):
        with open(f'train/spam/spam_{i}.txt', 'w') as f:
            f.write(text)
    for i, text in enumerate(ham_samples[:3]):
        with open(f'train/ham/ham_{i}.txt', 'w') as f:
            f.write(text)
    # 写入测试数据
    with open('test/spam/test_spam.txt', 'w') as f:
        f.write(spam_samples[3])
    with open('test/ham/test_ham.txt', 'w') as f:
        f.write(ham_samples[3])
# 2. 使用分类器
def main():
    # 创建示例数据
    create_sample_data()
    # 初始化分类器
    classifier = NaiveBayesClassifier(alpha=1.0)
    # 训练模型
    classifier.train_from_directory('train')
    # 测试单个文件
    result, probs = classifier.predict('test/spam/test_spam.txt')
    print(f"\n预测结果: {result}")
    print("各类别概率:")
    for cat, prob in probs.items():
        print(f"  {cat}: {prob:.4f}")
    # 预测新文本
    new_text = "This is a normal message about work"
    result, probs = classifier.predict(new_text)
    print(f"\n新文本预测: {result}")
    print(f"概率: {probs}")
    # 评估模型
    classifier.evaluate('test')
    # 保存模型
    classifier.save_model('nb_model.pkl')
    # 加载模型
    new_classifier = NaiveBayesClassifier()
    new_classifier.load_model('nb_model.pkl')
    # 使用加载的模型预测
    result, probs = new_classifier.predict("Limited offer! Don't miss out!")
    print(f"\n加载模型后预测: {result}")
    # 清理示例文件
    import shutil
    for dir in ['train', 'test']:
        shutil.rmtree(dir)
    if os.path.exists('nb_model.pkl'):
        os.remove('nb_model.pkl')
if __name__ == "__main__":
    main()

高级功能扩展

class AdvancedNaiveBayes(NaiveBayesClassifier):
    """增强版朴素贝叶斯"""
    def __init__(self, alpha=1.0, use_bigram=False, stop_words_file=None):
        super().__init__(alpha)
        self.use_bigram = use_bigram
        self.stop_words = self._load_stop_words(stop_words_file) if stop_words_file else set()
    def _load_stop_words(self, filepath):
        """从文件加载停用词"""
        with open(filepath, 'r', encoding='utf-8') as f:
            return set(f.read().splitlines())
    def _tokenize(self, text):
        """增强的分词功能，支持n-gram"""
        # 基础处理
        tokens = super()._tokenize(text)
        # 去除停用词
        tokens = [t for t in tokens if t not in self.stop_words]
        # 可选：生成bigram
        if self.use_bigram and len(tokens) > 1:
            bigrams = ['_'.join(tokens[i:i+2]) for i in range(len(tokens)-1)]
            tokens.extend(bigrams)
        return tokens
    def feature_selection(self, min_df=2, max_df=0.8):
        """特征选择：基于文档频率"""
        # 计算每个词的出现文档数
        word_doc_count = Counter()
        total_docs = sum(self.class_probabilities.values())
        for category in self.word_probabilities:
            for word, prob in self.word_probabilities[category].items():
                if prob > 0:
                    word_doc_count[word] += 1
        # 过滤词汇
        filtered_vocab = set()
        for word, count in word_doc_count.items():
            doc_freq = count / total_docs
            if min_df <= count and doc_freq <= max_df:
                filtered_vocab.add(word)
        # 更新词汇表
        self.vocabulary = filtered_vocab
        # 重新计算概率
        for category in self.word_probabilities:
            self.word_probabilities[category] = {
                word: prob 
                for word, prob in self.word_probabilities[category].items()
                if word in filtered_vocab
            }
        print(f"特征选择后词汇量: {len(self.vocabulary)}")

批量处理多个文件

def batch_predict_files(classifier, file_list, output_file=None):
    """批量预测多个文件"""
    results = []
    for filepath in file_list:
        if os.path.isfile(filepath):
            predicted_class, probabilities = classifier.predict(filepath)
            results.append({
                'file': filepath,
                'predicted': predicted_class,
                'confidence': max(probabilities.values()) if probabilities else 0,
                'probabilities': probabilities
            })
    # 可选：保存结果到文件
    if output_file:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("文件\t预测类别\t置信度\n")
            for result in results:
                f.write(f"{result['file']}\t{result['predicted']}\t{result['confidence']:.4f}\n")
    return results
# 使用示例
# results = batch_predict_files(classifier, ['file1.txt', 'file2.txt', 'file3.txt'], 'predictions.txt')

关键技术点总结

1. 拉普拉斯平滑: 避免零概率问题
2. 对数概率: 防止数值下溢
3. 文本预处理: 分词、停用词过滤
4. 模型持久化: 保存和加载训练好的模型
5. 性能评估: 准确率、混淆矩阵

这个实现可以根据具体需求进行调整,比如词干提取、TF-IDF特征、增量学习等改进。