基于 Python 实现文档搜索引擎:手把手写好每一行核心代码
·
下面我将手把手实现一个基于TF-IDF算法的文档搜索引擎核心代码,逐步讲解每个关键模块:
1. 文档预处理模块
import jieba
import re
def preprocess_text(text):
"""文本预处理:清洗、分词、去除停用词"""
# 清洗特殊字符
cleaned = re.sub(r'[^\w\s]', '', text)
# 中文分词
words = jieba.lcut(cleaned)
# 加载停用词表 (需准备stopwords.txt文件)
with open('stopwords.txt', encoding='utf-8') as f:
stopwords = set(f.read().splitlines())
# 过滤停用词和单字
return [word for word in words if word not in stopwords and len(word) > 1]
2. 倒排索引构建模块
from collections import defaultdict
import os
def build_inverted_index(doc_dir):
"""构建倒排索引:词项 -> (文档ID, 词频)"""
inverted_index = defaultdict(list)
doc_lengths = {}
for filename in os.listdir(doc_dir):
doc_id = filename
with open(os.path.join(doc_dir, filename), encoding='utf-8') as f:
text = f.read()
tokens = preprocess_text(text)
doc_lengths[doc_id] = len(tokens) # 记录文档长度
# 计算词频
term_freq = defaultdict(int)
for token in tokens:
term_freq[token] += 1
# 更新倒排索引
for term, freq in term_freq.items():
inverted_index[term].append((doc_id, freq))
return inverted_index, doc_lengths
3. TF-IDF计算模块
import math
def compute_tfidf(inverted_index, doc_lengths):
"""计算TF-IDF权重矩阵"""
# 文档总数
N = len(doc_lengths)
tfidf_index = {}
for term, postings in inverted_index.items():
# 计算逆文档频率 (IDF)
df = len(postings) # 包含该词的文档数
idf = math.log(N / (df + 1)) # 平滑处理
# 计算每个文档的TF-IDF
for doc_id, freq in postings:
# 词频标准化 (TF)
tf = freq / doc_lengths[doc_id]
tfidf = tf * idf
# 存储结果
if doc_id not in tfidf_index:
tfidf_index[doc_id] = {}
tfidf_index[doc_id][term] = tfidf
return tfidf_index
4. 查询处理模块
def process_query(query, inverted_index, doc_lengths, N):
"""处理用户查询并计算相关度"""
# 预处理查询词
query_terms = preprocess_text(query)
# 计算查询向量 (使用词频)
query_vector = {}
for term in query_terms:
query_vector[term] = query_vector.get(term, 0) + 1
# 计算文档得分
scores = defaultdict(float)
for term in query_terms:
if term not in inverted_index:
continue
# 计算查询词的IDF
df = len(inverted_index[term])
idf = math.log(N / (df + 1))
# 更新查询向量权重
query_vector[term] *= idf
# 计算文档相关度
for doc_id, freq in inverted_index[term]:
tf = freq / doc_lengths[doc_id]
scores[doc_id] += query_vector[term] * tf * idf
return scores
5. 主搜索引擎类
class DocumentSearchEngine:
def __init__(self, data_path):
self.index, self.doc_lengths = build_inverted_index(data_path)
self.N = len(self.doc_lengths)
self.tfidf_index = compute_tfidf(self.index, self.doc_lengths)
def search(self, query, top_k=10):
"""执行搜索并返回top_k结果"""
scores = process_query(
query,
self.index,
self.doc_lengths,
self.N
)
# 按得分排序
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
return ranked[:top_k]
使用示例
# 初始化引擎 (假设文档存储在./docs目录)
engine = DocumentSearchEngine('./docs')
# 执行搜索
results = engine.search("人工智能发展现状", top_k=5)
# 打印结果
print("搜索排名:")
for rank, (doc_id, score) in enumerate(results, 1):
print(f"{rank}. {doc_id} (相关度: {score:.4f})")
关键算法说明
-
TF-IDF权重:
- 词频 $tf(t,d) = \frac{f_{t,d}}{\sum_{t' \in d} f_{t',d}}$
- 逆文档频率 $idf(t) = \log \frac{N}{df_t + 1}$
- 最终权重 $w_{t,d} = tf(t,d) \times idf(t)$
-
余弦相似度: 文档与查询的相关度通过向量夹角余弦计算: $$\text{similarity} = \frac{\vec{d} \cdot \vec{q}}{|\vec{d}| |\vec{q}|}$$
此实现包含完整的索引构建、查询处理和排序功能,可根据需求扩展加入PageRank等权重因子。
更多推荐
所有评论(0)