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

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

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

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

TF-IDF 基类实现

手搓版本

python 复制代码

 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

依赖版本

python 复制代码

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

文本预处理

python 复制代码

 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 关键词提取

手搓版本

提取关键词

python 复制代码

 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

执行

scss 复制代码

 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}")

结果

less 复制代码

 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: ['准确率', '评价', '指标']

依赖版本

执行

ini 复制代码

 # 拟合语料库并生成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}")

结果

less 复制代码

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

TF-IDF 文本分类

设置的分类语料以及标签

ini 复制代码

 corpus = [
     # AI类
     "人工智能正在改变医疗、金融和教育等多个行业",
     "自然语言处理技术使计算机能够理解和生成人类语言",
     "计算机视觉是人工智能领域的重要研究方向之一",
     "智能语音助手通过人工智能实现了语音识别和对话功能",
     "人工智能伦理问题成为学术界和工业界讨论的焦点",
     "自动驾驶技术依赖人工智能实现环境感知和决策",
     "人工智能在游戏开发中用于创建智能NPC和动态剧情",
     "工业机器人通过人工智能优化生产流程和质量检测",
     "人工智能与物联网结合推动了智能家居的发展",
     "人工智能辅助诊断系统能快速分析医学影像并提供建议",
 
     # ML类
     "机器学习算法通过数据训练模型来预测未来趋势",
     "监督学习需要标注数据来训练分类和回归模型",
     "随机森林是一种集成学习方法，常用于分类任务",
     "特征工程是机器学习中提升模型性能的关键步骤",
     "支持向量机（SVM）适用于高维数据的分类问题",
     "交叉验证技术用于评估机器学习模型的泛化能力",
     "梯度下降法是优化机器学习模型参数的常用方法",
     "K近邻算法基于数据点的距离进行分类或回归",
     "贝叶斯分类器利用概率统计实现高效的文本分类",
     "机器学习中的过拟合问题可通过正则化技术缓解",
 
     # DL类
     "深度学习通过多层神经网络模拟人脑的信息处理方式",
     "卷积神经网络（CNN）在图像识别任务中表现出色",
     "循环神经网络（RNN）擅长处理时序数据和自然语言",
     "生成对抗网络（GAN）可用于生成逼真的图像和文本",
     "Transformer模型推动了自然语言处理领域的突破",
     "深度学习需要大量数据和计算资源进行模型训练",
     "注意力机制提升了深度学习模型对关键信息的捕捉能力",
     "迁移学习在深度学习中广泛应用于小样本场景",
     "深度强化学习被用于训练机器人完成复杂任务",
     "神经网络的参数调优对深度学习模型性能至关重要"
 ]
 
 labels = ['AI'] * 10 + ['ML'] * 10 + ['DL'] * 10  # 生成对应标签

手搓版本

python 复制代码

 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

执行

ini 复制代码

 new_text = "强化学习是人工智能的前沿领域"
 
 
 classifier = TextClassifier(corpus, labels)
 print(f"Classified Label: {classifier.classify(new_text)}")

结果

css 复制代码

 Classified Label: AI

依赖版本

python 复制代码

 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  # 控制词汇表大小
 )

执行

ini 复制代码

 # 将文本转换为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)

结果

less 复制代码

 分类标签: ['AI']

TF-IDF 信息检索

手搓版本

python 复制代码

 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]]

执行

ini 复制代码

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

结果

yaml 复制代码

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

依赖版本

python 复制代码

 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  # 控制词汇表大小
 )

执行

ini 复制代码

 # 将语料库转换为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})")

结果

yaml 复制代码

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

结语

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

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

关于我

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