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

文本分析之关键词提取（TF-IDF算法）

article 2024/11/17 23:56:42

键词提取是自然语言处理中的一个重要步骤，可以帮助我们理解文本的主要内容。TF-IDF（Term Frequency-Inverse Document Frequency）是一种常用的关键词提取方法，它基于词频和逆文档频率的概念来确定词语的重要性。

准备工作

首先，我们需要准备一些工具和库，包括 Pandas、jieba（结巴分词）、sklearn 等。

Pandas：用于数据处理。
jieba：用于中文分词。
sklearn：用于实现 TF-IDF 算法。

数据加载与预处理

我们将从多个文本文件中读取数据，并对其进行分词和停用词过滤。

import pandas as pd
import os
from sklearn.feature_extraction.text import TfidfVectorizer
import jieba
from sklearn.externals import joblib

# 加载用户词典
jieba.load_userdict(r"红楼梦词库.txt")

# 加载停用词
stopwords = set(pd.read_csv(r"StopwordsCN.txt", encoding='utf8')['stopword'].values.tolist())

# 文件路径和内容列表
filePaths = []
fileContents = []

# 遍历指定目录下的所有文件
for root, dirs, files in os.walk(r"D:\PyCharm\PyCharmProjects\python_class\机器学习\10.TF-IDF\项目3\红楼梦\分卷"):
    for name in files:
        filePath = os.path.join(root, name)
        filePaths.append(filePath)
        with open(filePath, 'r', encoding='utf-8') as f:
            fileContent = f.read()
            fileContents.append(fileContent)

# 将文件路径和内容添加到 DataFrame 中
corpos = pd.DataFrame({
    'filePath': filePaths,
    'fileContent': fileContents
})

# 分词并去除停用词
def segment_text(text):
    return ' '.join([seg for seg in jieba.cut(text) if seg not in stopwords])

# 应用分词函数
corpos['segmented'] = corpos['fileContent'].apply(segment_text)

# 将分词结果保存到文件
with open(r'D:\PyCharm\PyCharmProjects\python_class\机器学习\10.TF-IDF\项目3\红楼梦\分词后汇总.txt', 'w', encoding='utf-8') as f:
    for segmented_text in corpos['segmented']:
        f.write(segmented_text + '\n')

TF-IDF 算法应用

接下来，我们将使用 TF-IDF 算法来提取关键词。

# 创建 TF-IDF 矢量化器
vectorizer = TfidfVectorizer(use_idf=True, token_pattern=r"(?u)\b\w+\b")

# 训练 TF-IDF 模型
tfidf_matrix = vectorizer.fit_transform(corpos['segmented'])

# 获取词汇表
feature_names = vectorizer.get_feature_names_out()

# 输出 TF-IDF 矩阵
print(tfidf_matrix.toarray())

# 保存 TF-IDF 模型
joblib.dump(vectorizer, r"D:\PyCharm\PyCharmProjects\python_class\机器学习\10.TF-IDF\项目3\红楼梦\tfidf_model.pkl")

关键词提取

最后，我们可以从 TF-IDF 矩阵中提取出关键词。

def extract_keywords(tfidf_matrix, feature_names, top_n=10):
    # 获取每篇文档的 TF-IDF 值
    doc_ids = range(len(corpos))
    for i in doc_ids:
        # 获取每篇文档的 TF-IDF 值
        tfidf_scores = zip(feature_names, tfidf_matrix[i].toarray()[0])
        
        # 排序并获取前 N 个关键词
        sorted_tfidf_scores = sorted(tfidf_scores, key=lambda x: x[1], reverse=True)
        top_keywords = [term for term, score in sorted_tfidf_scores[:top_n]]
        
        yield top_keywords

# 打印每篇文档的前 10 个关键词
for keywords in extract_keywords(tfidf_matrix, feature_names):
    print(keywords)