b17786b57b
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
186 lines
5.7 KiB
Python
186 lines
5.7 KiB
Python
"""
|
||
BM25 倒排索引搜索引擎
|
||
用于图片 OCR 内容的全文搜索
|
||
参考 Meilisearch / Whoosh 设计
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import math
|
||
import logging
|
||
from collections import Counter, defaultdict
|
||
from typing import List, Dict, Optional, Tuple
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 尝试导入 jieba，失败则用简单的字符分割
|
||
try:
|
||
import jieba
|
||
jieba.setLogLevel(logging.WARNING)
|
||
def tokenize(text: str) -> List[str]:
|
||
"""中文分词"""
|
||
return [w for w in jieba.cut_for_search(text) if len(w.strip()) > 1]
|
||
except ImportError:
|
||
logger.warning("jieba 未安装，使用简单分词（pip install jieba 获得更好效果）")
|
||
def tokenize(text: str) -> List[str]:
|
||
"""简单分词：按空格和标点分割"""
|
||
return [w for w in re.split(r'[\s,.\-;:!?"\'/\\|()\[\]{}<>，。；：！？""''、（）【】]+', text) if len(w.strip()) > 1]
|
||
|
||
|
||
class BM25Index:
|
||
"""
|
||
BM25 倒排索引
|
||
- 支持增量添加文档
|
||
- 支持删除文档
|
||
- 支持中文分词（jieba）
|
||
- BM25 评分排序
|
||
"""
|
||
|
||
def __init__(self, k1: float = 1.5, b: float = 0.75):
|
||
"""
|
||
Args:
|
||
k1: 词频饱和参数（默认 1.5）
|
||
b: 文档长度归一化参数（默认 0.75）
|
||
"""
|
||
self.k1 = k1
|
||
self.b = b
|
||
|
||
# 倒排索引：term -> {doc_id: tf}
|
||
self.inverted_index: Dict[str, Dict[int, int]] = defaultdict(lambda: defaultdict(int))
|
||
|
||
# 文档信息
|
||
self.doc_lengths: Dict[int, int] = {} # doc_id -> 文档长度（词数）
|
||
self.doc_count: int = 0
|
||
self.avg_doc_length: float = 0.0
|
||
|
||
# 文档存储（存完整数据）
|
||
self.documents: Dict[int, dict] = {}
|
||
|
||
# IDF 缓存
|
||
self._idf_cache: Dict[str, float] = {}
|
||
|
||
def add_document(self, doc_id: int, text: str, metadata: Optional[dict] = None):
|
||
"""
|
||
添加文档到索引
|
||
|
||
Args:
|
||
doc_id: 文档 ID
|
||
text: 待索引的文本内容
|
||
metadata: 附加元数据（如 file_path, tags 等）
|
||
"""
|
||
# 如果已存在，先删除
|
||
if doc_id in self.documents:
|
||
self.remove_document(doc_id)
|
||
|
||
tokens = tokenize(text)
|
||
self.doc_lengths[doc_id] = len(tokens)
|
||
self.documents[doc_id] = {
|
||
"text": text,
|
||
"tokens": tokens,
|
||
"metadata": metadata or {},
|
||
}
|
||
|
||
# 构建倒排索引
|
||
tf = Counter(tokens)
|
||
for term, count in tf.items():
|
||
self.inverted_index[term][doc_id] = count
|
||
|
||
self.doc_count += 1
|
||
self._update_avg_doc_length()
|
||
self._idf_cache.clear() # 文档数变了，IDF 需要重算
|
||
|
||
def remove_document(self, doc_id: int):
|
||
"""从索引中删除文档"""
|
||
if doc_id not in self.documents:
|
||
return
|
||
|
||
# 从倒排索引中移除
|
||
tokens = self.documents[doc_id].get("tokens", [])
|
||
for term in set(tokens):
|
||
if term in self.inverted_index and doc_id in self.inverted_index[term]:
|
||
del self.inverted_index[term][doc_id]
|
||
if not self.inverted_index[term]:
|
||
del self.inverted_index[term]
|
||
|
||
del self.documents[doc_id]
|
||
del self.doc_lengths[doc_id]
|
||
self.doc_count -= 1
|
||
self._update_avg_doc_length()
|
||
self._idf_cache.clear()
|
||
|
||
def _update_avg_doc_length(self):
|
||
"""更新平均文档长度"""
|
||
if self.doc_count > 0:
|
||
self.avg_doc_length = sum(self.doc_lengths.values()) / self.doc_count
|
||
else:
|
||
self.avg_doc_length = 0.0
|
||
|
||
def _idf(self, term: str) -> float:
|
||
"""
|
||
计算逆文档频率 IDF
|
||
IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
|
||
其中 N 是文档总数，df(t) 是包含 term 的文档数
|
||
"""
|
||
if term in self._idf_cache:
|
||
return self._idf_cache[term]
|
||
|
||
df = len(self.inverted_index.get(term, {}))
|
||
idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)
|
||
|
||
self._idf_cache[term] = idf
|
||
return idf
|
||
|
||
def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float]]:
|
||
"""
|
||
BM25 搜索
|
||
|
||
Args:
|
||
query: 搜索查询（会自动分词）
|
||
top_k: 返回前 K 个结果
|
||
|
||
Returns:
|
||
[(doc_id, score), ...] 按 BM25 分数降序排列
|
||
"""
|
||
if self.doc_count == 0 or self.avg_doc_length == 0:
|
||
return []
|
||
|
||
query_tokens = tokenize(query)
|
||
if not query_tokens:
|
||
return []
|
||
|
||
scores: Dict[int, float] = defaultdict(float)
|
||
|
||
for term in query_tokens:
|
||
idf = self._idf(term)
|
||
postings = self.inverted_index.get(term, {})
|
||
|
||
for doc_id, tf in postings.items():
|
||
dl = self.doc_lengths[doc_id]
|
||
# BM25 公式
|
||
numerator = tf * (self.k1 + 1)
|
||
denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_doc_length)
|
||
scores[doc_id] += idf * numerator / denominator
|
||
|
||
# 按分数降序排序
|
||
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
|
||
return ranked[:top_k]
|
||
|
||
def get_document(self, doc_id: int) -> Optional[dict]:
|
||
"""获取文档完整数据"""
|
||
return self.documents.get(doc_id)
|
||
|
||
def clear(self):
|
||
"""清空索引"""
|
||
self.inverted_index.clear()
|
||
self.doc_lengths.clear()
|
||
self.documents.clear()
|
||
self._idf_cache.clear()
|
||
self.doc_count = 0
|
||
self.avg_doc_length = 0.0
|
||
|
||
@property
|
||
def size(self) -> int:
|
||
"""索引中的文档数量"""
|
||
return self.doc_count
|