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

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

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

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

TF-IDF 基类实现

手搓版本

import mathclass TFIDFBase:def __init__(self, corpus):self.corpus = corpusself.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] += 1return dfdef _compute_idf(self):"""计算逆文档频率 (IDF)"""idf = defaultdict(float)for token, count in self.df.items():idf[token] = math.log(self.N / (count + 1)) + 1  # 平滑处理return idfdef compute_tf(self, tokens):"""计算词频 (TF)"""tf = defaultdict(int)total = len(tokens)for token in tokens:tf[token] += 1 / total  # 归一化return tfdef 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 TfidfVectorizervectorizer = TfidfVectorizer(tokenizer=preprocess,  # 指定自定义分词函数token_pattern=r"(?u)\b\w+\b",  # 匹配中文分词结果max_features=1000  # 控制词汇表大小
)

文本预处理

import jieba
import nltk
from nltk.corpus import stopwordstry: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 = labelsself.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 / countavg_class_vectors[label] = avg_vectorreturn avg_class_vectorsdef 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 = -1predicted_label = Nonefor label, class_vec in self.class_vectors.items():sim = cosine_sim(tfidf, class_vec)if sim > max_sim:max_sim = simpredicted_label = labelreturn 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]] = valreturn matrixdef 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 #关键词抽取 #文本分类 #信息检索 # 文本匹配