youpaishidifu
1441278082
TF-IDF 余弦相似度单独使用对短中文查询过于严格（3-5 个词 vs 数百个词的长文档）。 改为混合评分：0.4×TF-IDF + 0.6×关键词命中率，让短查询只要命中关键词就能返回结果。 同时降低 auto_reply_min_score 从 0.35→0.30，适配混合评分的数值范围。
187 lines
6.4 KiB
Python
187 lines
6.4 KiB
Python
"""本地知识库 — 加载 Markdown/TXT 文件，jieba 分词 + TF-IDF 搜索"""
|
||
|
||
import os
|
||
import re
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import List, Tuple, Optional
|
||
|
||
import jieba
|
||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
from sklearn.metrics.pairwise import cosine_similarity
|
||
import numpy as np
|
||
|
||
|
||
class KnowledgeBase:
|
||
"""知识库引擎"""
|
||
|
||
def __init__(self, root_path: str = ""):
|
||
self.root_path: str = root_path
|
||
self.docs: List[dict] = [] # [{path, title, content, segments}]
|
||
self._vectorizer = TfidfVectorizer(tokenizer=self._tokenize, lowercase=False)
|
||
self._tfidf_matrix = None
|
||
self._ready = False
|
||
|
||
@staticmethod
|
||
def _tokenize(text: str) -> List[str]:
|
||
"""jieba 分词"""
|
||
return list(jieba.cut(text))
|
||
|
||
@staticmethod
|
||
def _parse_qa_segments(content: str) -> list:
|
||
"""解析 Q&A 格式文件为独立问答段。
|
||
格式：问题[N]：xxx\\n答：xxx
|
||
返回 [{"question": str, "answer": str, "raw": str}]
|
||
"""
|
||
segments = []
|
||
cur_q = []
|
||
cur_a = []
|
||
state = None # None, 'q', 'a'
|
||
|
||
def flush():
|
||
if not cur_q and not cur_a:
|
||
return
|
||
q = '\n'.join(cur_q).strip()
|
||
a = '\n'.join(cur_a).strip()
|
||
if q or a:
|
||
segments.append({
|
||
"question": q,
|
||
"answer": a,
|
||
"raw": f"问题：{q}\n答：{a}",
|
||
})
|
||
cur_q.clear()
|
||
cur_a.clear()
|
||
|
||
for line in content.split('\n'):
|
||
qm = re.match(r'^问题\d*\s*[：:]\s*(.*)', line)
|
||
am = re.match(r'^答\s*[：:]\s*(.*)', line)
|
||
if qm:
|
||
flush()
|
||
state = 'q'
|
||
cur_q.append(qm.group(1))
|
||
elif am:
|
||
state = 'a'
|
||
cur_a.append(am.group(1))
|
||
elif state == 'q':
|
||
cur_q.append(line)
|
||
elif state == 'a':
|
||
cur_a.append(line)
|
||
# else: 文件头空白/注释，忽略
|
||
|
||
flush()
|
||
return segments
|
||
|
||
def set_path(self, path: str):
|
||
self.root_path = path
|
||
self.reload()
|
||
|
||
def reload(self):
|
||
"""重新加载知识库文件"""
|
||
self.docs = []
|
||
if not self.root_path or not os.path.isdir(self.root_path):
|
||
self._ready = False
|
||
return
|
||
|
||
root = Path(self.root_path)
|
||
extensions = (".md", ".txt", ".markdown")
|
||
for fp in sorted(root.rglob("*")):
|
||
if fp.suffix.lower() in extensions and fp.is_file():
|
||
try:
|
||
content = fp.read_text("utf-8", errors="ignore")
|
||
title = fp.stem
|
||
# 解析 Q&A 段
|
||
segments = self._parse_qa_segments(content)
|
||
if segments:
|
||
for seg in segments:
|
||
self.docs.append({
|
||
"path": str(fp),
|
||
"title": title,
|
||
"content": seg["raw"],
|
||
"answer": seg["answer"],
|
||
"question": seg["question"],
|
||
})
|
||
else:
|
||
# 非 Q&A 格式，整文件作为一段
|
||
self.docs.append({
|
||
"path": str(fp),
|
||
"title": title,
|
||
"content": content,
|
||
})
|
||
except Exception:
|
||
continue
|
||
|
||
if not self.docs:
|
||
self._ready = False
|
||
return
|
||
|
||
# 构建 TF-IDF
|
||
corpus = [d["content"] for d in self.docs]
|
||
try:
|
||
self._tfidf_matrix = self._vectorizer.fit_transform(corpus)
|
||
self._ready = True
|
||
except Exception:
|
||
self._ready = False
|
||
|
||
@property
|
||
def is_ready(self) -> bool:
|
||
return self._ready and len(self.docs) > 0
|
||
|
||
def search(self, query: str, top_k: int = 3, min_score: float = 0.05) -> List[Tuple[dict, float]]:
|
||
"""搜索知识库，返回 (doc, score) 列表。
|
||
|
||
混合评分策略：
|
||
1) TF-IDF 余弦相似度（权重 0.4）— 适合长查询，对短查询偏严
|
||
2) 关键词命中率（权重 0.6）— 适合短查询，对关键词命中敏感
|
||
最终得分 = 0.4 * tfidf_score + 0.6 * keyword_hit_ratio
|
||
|
||
这种混合方式让短查询（3-5 个词）只要命中关键词就能返回结果，
|
||
同时长查询仍能利用 TF-IDF 的语义区分能力。
|
||
"""
|
||
if not self.is_ready or not query.strip():
|
||
return []
|
||
|
||
# 短文本：至少 2 个字且至少有 1 个有效词才搜索
|
||
tokens = [t for t in jieba.cut(query.strip()) if len(t) >= 2]
|
||
if len(query.strip()) < 2 or len(tokens) == 0:
|
||
return []
|
||
|
||
query_stripped = query.strip()
|
||
query_vec = self._vectorizer.transform([query_stripped])
|
||
tfidf_scores = cosine_similarity(query_vec, self._tfidf_matrix)[0]
|
||
|
||
# 关键词命中率：查询 tokens 中有多少个出现在文档内容中
|
||
query_tokens_set = set(tokens)
|
||
keyword_ratios = []
|
||
for doc in self.docs:
|
||
content = doc["content"]
|
||
hit = sum(1 for t in query_tokens_set if t in content)
|
||
keyword_ratios.append(hit / max(len(query_tokens_set), 1))
|
||
keyword_ratios = np.array(keyword_ratios)
|
||
|
||
# 混合评分
|
||
combined_scores = 0.4 * tfidf_scores + 0.6 * keyword_ratios
|
||
|
||
top_indices = np.argsort(combined_scores)[::-1][:top_k]
|
||
|
||
results = []
|
||
for idx in top_indices:
|
||
score = float(combined_scores[idx])
|
||
if score < min_score:
|
||
continue
|
||
results.append((self.docs[idx], score))
|
||
if results:
|
||
return results
|
||
|
||
# Fallback: 子串匹配（查询包含在文档标题或内容中）
|
||
q = query_stripped
|
||
if len(q) >= 4:
|
||
for doc in self.docs:
|
||
if q in doc["content"] or q in doc["title"]:
|
||
logger = logging.getLogger("msgbot.kb")
|
||
logger.debug(f"子串匹配: \"{q}\" → {doc['title']}")
|
||
return [(doc, min_score)]
|
||
return results
|
||
|
||
def get_doc_count(self) -> int:
|
||
return len(self.docs)
|