TF-IDF算法的代码实践应用——关键词提取、文本分类、信息检索

回顾：TF-IDF算法详解与实践总结

上一篇文章我们深入剖析了TF-IDF的原理与细节，但实践才是检验真理的唯一标准！今天，我们将从“纸上谈兵”转向“实战演练”：通过纯Python手写实现与调用sklearn工具包两种方式，带你一步步完成TF-IDF在真实场景中的应用

• 关键词提取：如何从一篇长文中精准抓取核心词汇？
• 文本分类：如何让机器自动识别“AI”、“ML”、“DL”类别的文本？
• 信息检索：如何根据查询语句，从文档库中找到最相关的内容？

TF-IDF 基类实现

手搓版本

import math

class TFIDFBase:
    def __init__(self, corpus):
        self.corpus = corpus
        self.tokenized_corpus = [preprocess(doc) for doc in corpus]
        self.N = len(corpus)  # 文档总数
        self.df = self._compute_df()
        self.idf = self._compute_idf()

    def _compute_df(self):
        """计算文档频率 (Document Frequency)"""
        df = defaultdict(int)
        for tokens in self.tokenized_corpus:
            unique_tokens = set(tokens)
            for token in unique_tokens:
                df[token] += 1
        return df

    def _compute_idf(self):
        """计算逆文档频率 (IDF)"""
        idf = defaultdict(float)
        for token, count in self.df.items():
            idf[token] = math.log(self.N / (count + 1)) + 1  # 平滑处理
        return idf

    def compute_tf(self, tokens):
        """计算词频 (TF)"""
        tf = defaultdict(int)
        total = len(tokens)
        for token in tokens:
            tf[token] += 1 / total  # 归一化
        return tf

    def compute_tfidf(self, tf):
        """计算 TF-IDF 向量"""
        return {token: tf[token] * self.idf[token] for token in tf}

    def get_tfidf_vector(self):
        """生成整个语料库的 TF-IDF 向量"""
        tfidf_vectors = []
        for tokens in self.tokenized_corpus:
            tf = self.compute_tf(tokens)
            tfidf = self.compute_tfidf(tf)
            tfidf_vectors.append(tfidf)
        return tfidf_vectors

依赖版本

from sklearn.feature_extraction.text import TfidfVectorizer

vectorizer = TfidfVectorizer(
    tokenizer=preprocess,  # 指定自定义分词函数
    token_pattern=r"(?u)\b\w+\b",  # 匹配中文分词结果
    max_features=1000  # 控制词汇表大小
)

文本预处理

import jieba
import nltk
from nltk.corpus import stopwords


try:
    nltk.download('stopwords')
    stopwords_cn = stopwords.words('chinese')
    print("已下载 NLTK 中文停用词")
except:
    # 若 nltk 中文停用词缺失，手动补充常用中文停用词
    stopwords_cn = [
        '的', '了', '和', '是', '在', '中', '也', '都', '而', '就',
        '那', '这', '之', '为', '对', '与', '于', '上', '下', '前'
    ]


# 中文文本预处理（分词 + 去停用词）
def preprocess(text):
    tokens = jieba.cut(text)
    return [token for token in tokens if token not in stopwords_cn and token.strip()]

TD-IDF 关键词提取

手搓版本

提取关键词

class KeywordExtractor(TFIDFBase):
    def extract_keywords(self, top_k=5):
        """提取每个文档的 top-k 关键词"""
        tfidf_vectors = self.get_tfidf_vector()
        results = []
        for idx, tfidf in enumerate(tfidf_vectors):
            sorted_keywords = sorted(tfidf.items(), key=lambda x: -x[1])[:top_k]
            results.append([kw[0] for kw in sorted_keywords])
        return results

执行

corpus = [
    "自然语言处理是人工智能领域的一个重要分支",
    "信息检索技术帮助我们从大量文档中找到相关内容",
    "BM25算法是信息检索中常用的经典算法",
    "中文分词是中文自然语言处理的基础步骤",
    "搜索引擎使用各种算法来提高搜索结果的相关性",
    "TF-IDF和BM25都是基于统计的检索模型",
    "深度学习在自然语言处理中取得了显著进展",
    "倒排索引是信息检索系统的核心技术之一",
    "查询扩展可以提高信息检索的召回率",
    "准确率和召回率是评价信息检索系统的重要指标"
]

extractor = KeywordExtractor(corpus)
keywords = extractor.extract_keywords(top_k=3)
for i, kws in enumerate(keywords):
    print(f"Document {i + 1} Keywords: {kws}")

结果

Document 1 Keywords: ['人工智能', '领域', '一个']
Document 2 Keywords: ['技术', '文档', '找到']
Document 3 Keywords: ['算法', '常用', '经典']
Document 4 Keywords: ['中文', '分词', '基础']
Document 5 Keywords: ['搜索引擎', '搜索', '相关性']
Document 6 Keywords: ['TF', '-', 'IDF']
Document 7 Keywords: ['深度', '学习', '进展']
Document 8 Keywords: ['倒排', '索引', '核心技术']
Document 9 Keywords: ['查询', '扩展', '提高']
Document 10 Keywords: ['准确率', '评价', '指标']

依赖版本

执行

# 拟合语料库并生成TF-IDF矩阵
tfidf_matrix = vectorizer.fit_transform(corpus)

# 提取关键词（每个文档的Top N关键词）
for i, text in enumerate(corpus):
    feature_index = tfidf_matrix[i, :].nonzero()[1]  # 非零元素的索引
    tfidf_scores = zip(feature_index, [tfidf_matrix[i, x] for x in feature_index])  # (词索引, 权重)
    sorted_tfidf = sorted(tfidf_scores, key=lambda x: x[1], reverse=True)  # 按权重排序
    top_n = 3  # 提取每个文档的前3个关键词
    keywords = [vectorizer.get_feature_names_out()[idx] for idx, score in sorted_tfidf[:top_n]]
    print(f"文档 {i + 1} 的关键词: {keywords}")

结果

文档 1 的关键词: ['人工智能', '领域', '一个']
文档 2 的关键词: ['技术', '文档', '找到']
文档 3 的关键词: ['算法', '常用', '经典']
文档 4 的关键词: ['中文', '分词', '基础']
文档 5 的关键词: ['搜索引擎', '搜索', '相关性']
文档 6 的关键词: ['tf', '-', 'idf']
文档 7 的关键词: ['深度', '学习', '进展']
文档 8 的关键词: ['倒排', '索引', '核心技术']
文档 9 的关键词: ['查询', '扩展', '提高']
文档 10 的关键词: ['准确率', '评价', '指标']

TF-IDF 文本分类

设置的分类语料以及标签

corpus = [
    # AI类
    "人工智能正在改变医疗、金融和教育等多个行业",
    "自然语言处理技术使计算机能够理解和生成人类语言",
    "计算机视觉是人工智能领域的重要研究方向之一",
    "智能语音助手通过人工智能实现了语音识别和对话功能",
    "人工智能伦理问题成为学术界和工业界讨论的焦点",
    "自动驾驶技术依赖人工智能实现环境感知和决策",
    "人工智能在游戏开发中用于创建智能NPC和动态剧情",
    "工业机器人通过人工智能优化生产流程和质量检测",
    "人工智能与物联网结合推动了智能家居的发展",
    "人工智能辅助诊断系统能快速分析医学影像并提供建议",

    # ML类
    "机器学习算法通过数据训练模型来预测未来趋势",
    "监督学习需要标注数据来训练分类和回归模型",
    "随机森林是一种集成学习方法，常用于分类任务",
    "特征工程是机器学习中提升模型性能的关键步骤",
    "支持向量机（SVM）适用于高维数据的分类问题",
    "交叉验证技术用于评估机器学习模型的泛化能力",
    "梯度下降法是优化机器学习模型参数的常用方法",
    "K近邻算法基于数据点的距离进行分类或回归",
    "贝叶斯分类器利用概率统计实现高效的文本分类",
    "机器学习中的过拟合问题可通过正则化技术缓解",

    # DL类
    "深度学习通过多层神经网络模拟人脑的信息处理方式",
    "卷积神经网络（CNN）在图像识别任务中表现出色",
    "循环神经网络（RNN）擅长处理时序数据和自然语言",
    "生成对抗网络（GAN）可用于生成逼真的图像和文本",
    "Transformer模型推动了自然语言处理领域的突破",
    "深度学习需要大量数据和计算资源进行模型训练",
    "注意力机制提升了深度学习模型对关键信息的捕捉能力",
    "迁移学习在深度学习中广泛应用于小样本场景",
    "深度强化学习被用于训练机器人完成复杂任务",
    "神经网络的参数调优对深度学习模型性能至关重要"
]

labels = ['AI'] * 10 + ['ML'] * 10 + ['DL'] * 10  # 生成对应标签

手搓版本

class TextClassifier(TFIDFBase):
    def __init__(self, corpus, labels):
        super().__init__(corpus)
        self.labels = labels
        self.class_vectors = self._build_class_vectors()

    def _build_class_vectors(self):
        """按类别聚合 TF-IDF 向量（简单平均）"""
        class_vectors = defaultdict(list)
        tfidf_vectors = self.get_tfidf_vector()
        for label, tfidf in zip(self.labels, tfidf_vectors):
            class_vectors[label].append(tfidf)
        # 计算类别向量（平均）
        avg_class_vectors = {}
        for label, vectors in class_vectors.items():
            avg_vector = defaultdict(float)
            count = len(vectors)
            for vec in vectors:
                for token, val in vec.items():
                    avg_vector[token] += val / count
            avg_class_vectors[label] = avg_vector
        return avg_class_vectors

    def classify(self, text):
        """对新文本分类（余弦相似度）"""
        tokens = preprocess(text)
        tf = self.compute_tf(tokens)
        tfidf = self.compute_tfidf(tf)

        # 计算余弦相似度
        def cosine_sim(vec1, vec2):
            common_tokens = set(vec1.keys()) & set(vec2.keys())
            dot = sum(vec1[t] * vec2[t] for t in common_tokens)
            norm1 = math.sqrt(sum(v ** 2 for v in vec1.values()))
            norm2 = math.sqrt(sum(v ** 2 for v in vec2.values()))
            return dot / (norm1 * norm2 + 1e-8)

        max_sim = -1
        predicted_label = None
        for label, class_vec in self.class_vectors.items():
            sim = cosine_sim(tfidf, class_vec)
            if sim > max_sim:
                max_sim = sim
                predicted_label = label
        return predicted_label

执行

new_text = "强化学习是人工智能的前沿领域"


classifier = TextClassifier(corpus, labels)
print(f"Classified Label: {classifier.classify(new_text)}")

结果

Classified Label: AI

依赖版本

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report

# 使用TF - IDF向量化器
vectorizer = TfidfVectorizer(
    tokenizer=preprocess,  # 指定自定义分词函数
    token_pattern=r"(?u)\b\w+\b",  # 匹配中文分词结果（可忽略）
    max_features=1000  # 控制词汇表大小
)

执行

# 将文本转换为TF-IDF特征向量
X = vectorizer.fit_transform(corpus)
y = np.array(labels)

# 划分训练集和测试集（由于样本太少，此处仅演示流程）
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练分类模型（以逻辑回归为例）
model = LogisticRegression()
model.fit(X_train, y_train)

# 预测与评估
# 使用训练好的TF-IDF向量化器转换新文本
X_new = vectorizer.transform([new_text])
y_pred = model.predict(X_new)
print("分类标签:", y_pred)

结果

分类标签: ['AI']

TF-IDF 信息检索

手搓版本

class InformationRetriever(TFIDFBase):
    def __init__(self, corpus):
        super().__init__(corpus)
        self.tfidf_matrix = self._build_tfidf_matrix()

    def _build_tfidf_matrix(self):
        """构建 TF-IDF 矩阵（文档向量化）"""
        tfidf_vectors = self.get_tfidf_vector()
        # 构建全局词表
        self.vocab = {token: idx for idx, token in enumerate(set(
            token for vec in tfidf_vectors for token in vec
        ))}
        # 转换为 NumPy 矩阵
        matrix = np.zeros((self.N, len(self.vocab)))
        for i, vec in enumerate(tfidf_vectors):
            for token, val in vec.items():
                matrix[i, self.vocab[token]] = val
        return matrix

    def search(self, query, top_k=2):
        """根据查询语句检索最相关文档"""
        # 预处理查询
        tokens = preprocess(query)
        tf = self.compute_tf(tokens)
        tfidf = self.compute_tfidf(tf)

        # 构建查询向量
        query_vec = np.zeros(len(self.vocab))
        for token, val in tfidf.items():
            if token in self.vocab:
                query_vec[self.vocab[token]] = val

        # 计算余弦相似度
        sims = []
        for i in range(self.N):
            doc_vec = self.tfidf_matrix[i]
            sim = np.dot(query_vec, doc_vec) / (np.linalg.norm(query_vec) * np.linalg.norm(doc_vec) + 1e-8)
            sims.append((i, sim))
        # 返回 top-k 相关文档
        return [self.corpus[i] for i, _ in sorted(sims, key=lambda x: -x[1])[:top_k]]

执行

corpus = [
    "自然语言处理是人工智能的一个重要分支",
    "机器学习方法在自然语言处理中起着关键作用",
    "深度学习技术推动了自然语言处理的发展",
    "计算机视觉是人工智能领域的重要研究方向之一",
    "强化学习是人工智能的前沿领域"
]
# "信息检索"
query = "人工智能和机器学习的关系"
retriever = InformationRetriever(corpus)
results = retriever.search(query, top_k=3)
for idx, r in enumerate(results):
    print(f"Matched Document {idx}: {r}")

结果

Matched Document 0: 机器学习方法在自然语言处理中起着关键作用
Matched Document 1: 强化学习是人工智能的前沿领域
Matched Document 2: 自然语言处理是人工智能的一个重要分支

依赖版本

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 使用TF-IDF向量化器
vectorizer = TfidfVectorizer(
    tokenizer=preprocess,  # 自定义分词函数
    token_pattern=r"(?u)\b\w+\b",  # 对中文无效，但不影响
    max_features=1000  # 控制词汇表大小
)

执行

# 将语料库转换为TF-IDF特征向量
X_corpus = vectorizer.fit_transform(corpus)

# 示例查询
query = "人工智能和机器学习的关系"

# 预处理并转换查询
X_query = vectorizer.transform([query])  # 注意：必须用列表包裹

# 计算余弦相似度
similarities = cosine_similarity(X_query, X_corpus)

# 输出最相关的文档
top_n = 3  # 返回前3个最相关文档
top_indices = np.argsort(similarities[0])[::-1][:top_n]  # 降序排序取索引

print("查询:", query)
print("最相关的文档:")
for idx in top_indices:
    print(f"文档 {idx + 1}: {corpus[idx]} (相似度: {similarities[0][idx]:.4f})")

结果

查询: 人工智能和机器学习的关系
最相关的文档:
文档 2: 机器学习方法在自然语言处理中起着关键作用 (相似度: 0.4752)
文档 5: 强化学习是人工智能的前沿领域 (相似度: 0.3458)
文档 1: 自然语言处理是人工智能的一个重要分支 (相似度: 0.1913)

结语

通过本文的实践，我们不仅验证了TF-IDF算法的实用性，更展示了它在关键词提取、文本分类和信息检索中的多样化应用场景。

当然，TF-IDF并非万能钥匙：它无法捕捉词序、上下文语义，对动态语境和复杂任务的支持有限。但正如我们所见，它依然是许多现代算法（如BM25、深度学习模型）的重要基础。

TFIDF #关键词抽取 #文本分类 #信息检索 # 文本匹配