feat: add reproducible keyword retrieval baseline
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user