216 lines
8.0 KiB
Python
216 lines
8.0 KiB
Python
"""Deterministic lexical baseline over a bounded, reviewed knowledge corpus."""
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import unicodedata
|
|
from collections import Counter
|
|
from functools import lru_cache
|
|
from time import perf_counter
|
|
from uuid import uuid4
|
|
|
|
import jieba
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import func, or_, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models import Document, KnowledgeItem, RetrievalRun, Source
|
|
from app.models.catalog import utc_now
|
|
from app.schemas.retrieval import RetrievalConfig, RetrievalHit, RetrievalRequest, RetrievalResponse
|
|
|
|
MAX_ITEMS = 5000
|
|
MAX_CHARACTERS = 2_000_000
|
|
SNIPPET_LENGTH = 240
|
|
STOPWORDS = frozenset("的 了 和 与 是 在 有 哪些 什么 请问 如何 吗 呢 我 你 想 了解".split())
|
|
CONFIG = RetrievalConfig(
|
|
tokenizer="jieba-0.42.1-search-hmm-off-nalati-v1",
|
|
scorer="bm25-weighted-text-v1",
|
|
k1=1.2,
|
|
b=0.75,
|
|
title_weight=2,
|
|
stopwords_version="zh-small-v1",
|
|
)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def tokenizer():
|
|
jieba.setLogLevel(logging.WARNING)
|
|
instance = jieba.Tokenizer()
|
|
instance.initialize()
|
|
instance.add_word("那拉提", freq=100000)
|
|
return instance
|
|
|
|
|
|
def normalized_with_offsets(text: str) -> tuple[str, list[int]]:
|
|
pieces, offsets = [], []
|
|
for position, character in enumerate(text):
|
|
piece = unicodedata.normalize("NFKC", character).casefold()
|
|
pieces.append(piece)
|
|
offsets.extend([position] * len(piece))
|
|
return "".join(pieces), offsets
|
|
|
|
|
|
def tokenize(text: str) -> list[str]:
|
|
normalized, _ = normalized_with_offsets(text)
|
|
return [
|
|
term
|
|
for term in tokenizer().cut_for_search(normalized, HMM=False)
|
|
if term not in STOPWORDS and any(character.isalnum() for character in term)
|
|
]
|
|
|
|
|
|
def bm25(counters: list[Counter], query_terms: list[str]) -> list[float]:
|
|
"""Positive log-IDF BM25; repeated query terms count once."""
|
|
count = len(counters)
|
|
lengths = [sum(counter.values()) for counter in counters]
|
|
average = sum(lengths) / count if count else 0
|
|
if not average:
|
|
return [0.0] * count
|
|
scores = [0.0] * count
|
|
for term in dict.fromkeys(query_terms):
|
|
frequency = sum(counter[term] > 0 for counter in counters)
|
|
if not frequency:
|
|
continue
|
|
idf = math.log(1 + (count - frequency + 0.5) / (frequency + 0.5))
|
|
for index, counter in enumerate(counters):
|
|
tf = counter[term]
|
|
if tf:
|
|
denominator = tf + CONFIG.k1 * (1 - CONFIG.b + CONFIG.b * lengths[index] / average)
|
|
scores[index] += idf * tf * (CONFIG.k1 + 1) / denominator
|
|
return scores
|
|
|
|
|
|
def excerpt(content: str, matched_terms: list[str]) -> tuple[str, int, int]:
|
|
normalized, offsets = normalized_with_offsets(content)
|
|
positions = [normalized.find(term) for term in matched_terms if term in normalized]
|
|
start = max(0, offsets[min(positions)] - 60) if positions else 0
|
|
end = min(len(content), start + SNIPPET_LENGTH)
|
|
return content[start:end], start, end
|
|
|
|
|
|
def load_corpus(db: Session, category: str | None) -> list[dict]:
|
|
statement = (
|
|
select(
|
|
KnowledgeItem,
|
|
Source.name.label("source_name"),
|
|
Source.version.label("source_version"),
|
|
Document.version.label("document_version"),
|
|
func.count().over().label("corpus_total"),
|
|
func.sum(func.length(KnowledgeItem.title) + func.length(KnowledgeItem.content))
|
|
.over()
|
|
.label("corpus_characters"),
|
|
)
|
|
.outerjoin(Source, KnowledgeItem.source_id == Source.id)
|
|
.outerjoin(Document, KnowledgeItem.document_id == Document.id)
|
|
.where(
|
|
KnowledgeItem.verification_status == "verified",
|
|
KnowledgeItem.is_deleted.is_(False),
|
|
KnowledgeItem.source_url.is_not(None),
|
|
func.length(func.trim(KnowledgeItem.source_url)) > 0,
|
|
KnowledgeItem.collected_at.is_not(None),
|
|
KnowledgeItem.reviewed_at.is_not(None),
|
|
KnowledgeItem.reviewer.is_not(None),
|
|
KnowledgeItem.verification_notes.is_not(None),
|
|
or_(KnowledgeItem.source_id.is_(None), Source.is_deleted.is_(False)),
|
|
or_(
|
|
KnowledgeItem.document_id.is_(None),
|
|
(Document.is_deleted.is_(False)) & (Document.source_id == KnowledgeItem.source_id),
|
|
),
|
|
or_(KnowledgeItem.source_id.is_(None), KnowledgeItem.source_url == Source.source_url),
|
|
)
|
|
.order_by(KnowledgeItem.id)
|
|
.limit(MAX_ITEMS + 1)
|
|
)
|
|
if category is not None:
|
|
statement = statement.where(KnowledgeItem.category == category)
|
|
corpus = []
|
|
# 一条 SQL 的快照同时给出语料与容量；流式读取避免先载入超限正文。
|
|
with db.execute(statement.execution_options(stream_results=True, max_row_buffer=1)) as rows:
|
|
for row in rows:
|
|
if row.corpus_total > MAX_ITEMS or row.corpus_characters > MAX_CHARACTERS:
|
|
raise HTTPException(
|
|
503, "已核验语料超过关键词基线容量，请缩小分类范围或建立检索索引"
|
|
)
|
|
item = row.KnowledgeItem
|
|
corpus.append(
|
|
{
|
|
"knowledge_id": item.id,
|
|
"knowledge_version": item.version,
|
|
"document_id": item.document_id,
|
|
"document_version": row.document_version,
|
|
"source_id": item.source_id,
|
|
"source_version": row.source_version,
|
|
"source_name": row.source_name,
|
|
"source_url": item.source_url,
|
|
"title": item.title,
|
|
"category": item.category,
|
|
"content": item.content,
|
|
}
|
|
)
|
|
return corpus
|
|
|
|
|
|
def retrieve(db: Session, payload: RetrievalRequest) -> RetrievalResponse:
|
|
started = perf_counter()
|
|
terms = list(dict.fromkeys(tokenize(payload.query)))
|
|
if not terms:
|
|
raise HTTPException(422, "问题中没有可检索词，请补充具体关键词")
|
|
corpus = load_corpus(db, payload.category)
|
|
fingerprint = hashlib.sha256(
|
|
json.dumps(corpus, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()
|
|
).hexdigest()
|
|
counters = [
|
|
Counter(tokenize(item["title"]) * CONFIG.title_weight + tokenize(item["content"]))
|
|
for item in corpus
|
|
]
|
|
scores = bm25(counters, terms)
|
|
ranked = sorted(
|
|
(index for index, score in enumerate(scores) if score > 0),
|
|
key=lambda index: (-scores[index], corpus[index]["knowledge_id"]),
|
|
)
|
|
results = []
|
|
for rank, index in enumerate(ranked[: payload.top_k], start=1):
|
|
item = corpus[index]
|
|
matched = [term for term in terms if counters[index][term]]
|
|
snippet, start, end = excerpt(item["content"], matched)
|
|
results.append(
|
|
RetrievalHit(
|
|
**{key: value for key, value in item.items() if key != "content"},
|
|
rank=rank,
|
|
score=round(scores[index], 8),
|
|
matched_terms=matched,
|
|
snippet=snippet,
|
|
snippet_start=start,
|
|
snippet_end=end,
|
|
)
|
|
)
|
|
response = RetrievalResponse(
|
|
run_id=uuid4(),
|
|
query=payload.query,
|
|
mode="keyword",
|
|
status="ok" if results else "no_evidence",
|
|
top_k=payload.top_k,
|
|
category=payload.category,
|
|
query_terms=terms,
|
|
corpus_size=len(corpus),
|
|
corpus_fingerprint=fingerprint,
|
|
total_hits=len(ranked),
|
|
results=results,
|
|
elapsed_ms=round((perf_counter() - started) * 1000, 3),
|
|
created_at=utc_now(),
|
|
config=CONFIG,
|
|
)
|
|
db.add(
|
|
RetrievalRun(
|
|
id=str(response.run_id),
|
|
query=response.query,
|
|
mode=response.mode,
|
|
corpus_fingerprint=fingerprint,
|
|
response=response.model_dump(mode="json"),
|
|
created_at=response.created_at,
|
|
)
|
|
)
|
|
db.commit()
|
|
return response
|