206 lines
7.2 KiB
Python
206 lines
7.2 KiB
Python
"""Bounded chunk BM25 and equal-weight reciprocal rank fusion on one publication."""
|
|
|
|
from collections import Counter
|
|
from time import perf_counter
|
|
from uuid import uuid4
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.catalog import utc_now
|
|
from app.models.indexing import Chunk
|
|
from app.models.retrieval import RetrievalRun
|
|
from app.schemas.hybrid import (
|
|
ChunkKeywordRequest,
|
|
ChunkKeywordResponse,
|
|
ChunkRetrievalConfig,
|
|
HybridConfig,
|
|
HybridHit,
|
|
HybridRequest,
|
|
HybridResponse,
|
|
RankedCandidate,
|
|
)
|
|
from app.schemas.vectors import SpaceRead, VectorHit
|
|
from app.services.indexing import locked_state
|
|
from app.services.retrieval import CONFIG, bm25, tokenize
|
|
from app.services.vectors import (
|
|
checked_chunks,
|
|
current_text,
|
|
encode,
|
|
rank_vectors,
|
|
retrieval_scope,
|
|
)
|
|
|
|
RRF_K = 60
|
|
CHUNK_LEXICAL_CONFIG = CONFIG.model_copy(update={"scorer": "bm25-chunk-body-v1", "title_weight": 0})
|
|
|
|
|
|
def chunk_key(chunk):
|
|
return chunk.knowledge_id, chunk.ordinal, chunk.id
|
|
|
|
|
|
def chunks_in_scope(db, text_version, category):
|
|
statement = select(Chunk).where(Chunk.index_version_id == text_version.id)
|
|
if category is not None:
|
|
statement = statement.where(Chunk.provenance["category"].as_string() == category)
|
|
return list(db.scalars(statement.order_by(Chunk.knowledge_id, Chunk.ordinal, Chunk.id)))
|
|
|
|
|
|
def lexical_ranking(chunks, terms):
|
|
counters = {chunk.id: Counter(tokenize(chunk.text)) for chunk in chunks}
|
|
scores = bm25(list(counters.values()), terms)
|
|
ranked = sorted(
|
|
((chunk, score) for chunk, score in zip(chunks, scores) if score > 0),
|
|
key=lambda pair: (-pair[1], chunk_key(pair[0])),
|
|
)
|
|
return ranked, counters
|
|
|
|
|
|
def candidate_list(rows):
|
|
return [
|
|
RankedCandidate(chunk_id=chunk.id, rank=rank, score=score)
|
|
for rank, (chunk, score) in enumerate(rows, 1)
|
|
]
|
|
|
|
|
|
def fuse(lexical, vector):
|
|
"""Ranks start at one; an absent candidate contributes zero, never a padded rank."""
|
|
scores = {}
|
|
for leg in (lexical, vector):
|
|
for candidate in leg:
|
|
key = str(candidate.chunk_id)
|
|
scores[key] = scores.get(key, 0.0) + 1.0 / (RRF_K + candidate.rank)
|
|
return scores
|
|
|
|
|
|
def hit_fields(chunk, terms, counters):
|
|
return dict(
|
|
**chunk.provenance,
|
|
chunk_id=chunk.id,
|
|
chunk_ordinal=chunk.ordinal,
|
|
snippet=chunk.text,
|
|
snippet_start=chunk.start_offset,
|
|
snippet_end=chunk.end_offset,
|
|
matched_terms=[term for term in terms if counters[chunk.id][term]],
|
|
)
|
|
|
|
|
|
def save_run(db, response):
|
|
db.add(
|
|
RetrievalRun(
|
|
id=str(response.run_id),
|
|
query=response.query,
|
|
mode=response.mode,
|
|
corpus_fingerprint=response.corpus_fingerprint,
|
|
response=response.model_dump(mode="json"),
|
|
created_at=response.created_at,
|
|
)
|
|
)
|
|
db.commit()
|
|
return response
|
|
|
|
|
|
def response_fields(payload, state, text_version, chunks, terms, total, started):
|
|
return dict(
|
|
run_id=uuid4(),
|
|
query=payload.query,
|
|
status="ok" if total else "no_evidence",
|
|
top_k=payload.top_k,
|
|
category=payload.category,
|
|
query_terms=terms,
|
|
text_index_version_id=text_version.id,
|
|
text_generation=state.generation,
|
|
catalog_revision=state.catalog_revision,
|
|
corpus_fingerprint=text_version.corpus_fingerprint,
|
|
corpus_size=len(chunks),
|
|
total_hits=total,
|
|
elapsed_ms=round((perf_counter() - started) * 1000, 3),
|
|
created_at=utc_now(),
|
|
)
|
|
|
|
|
|
def retrieve_chunks(db: Session, payload: ChunkKeywordRequest):
|
|
started = perf_counter()
|
|
terms = list(dict.fromkeys(tokenize(payload.query)))
|
|
if not terms:
|
|
raise HTTPException(422, "问题中没有可检索词，请补充具体关键词")
|
|
state = locked_state(db)
|
|
text_version, corpus = current_text(db, state)
|
|
checked_chunks(db, text_version, corpus)
|
|
chunks = chunks_in_scope(db, text_version, payload.category)
|
|
ranked, counters = lexical_ranking(chunks, terms)
|
|
hits = [
|
|
VectorHit(**hit_fields(chunk, terms, counters), rank=rank, score=score)
|
|
for rank, (chunk, score) in enumerate(ranked[: payload.top_k], 1)
|
|
]
|
|
return save_run(
|
|
db,
|
|
ChunkKeywordResponse(
|
|
**response_fields(payload, state, text_version, chunks, terms, len(ranked), started),
|
|
results=hits,
|
|
config=ChunkRetrievalConfig(lexical=CHUNK_LEXICAL_CONFIG, chunking=text_version.config),
|
|
),
|
|
)
|
|
|
|
|
|
def retrieve_hybrid(db: Session, payload: HybridRequest, provider):
|
|
started = perf_counter()
|
|
terms = list(dict.fromkeys(tokenize(payload.query)))
|
|
*_, token = retrieval_scope(db, provider)
|
|
db.commit()
|
|
query_vector = encode(provider, [payload.query], query=True)[0]
|
|
state, vector_state, text_version, version, space, current = retrieval_scope(db, provider)
|
|
if token != current:
|
|
raise HTTPException(409, "查询编码期间发布版本或审核依据已变化，请重试")
|
|
# Hold the shared catalog/publication locks until both rankings and the run are saved.
|
|
chunks = chunks_in_scope(db, text_version, payload.category)
|
|
ranked, counters = lexical_ranking(chunks, terms)
|
|
lexical = candidate_list(ranked[: payload.candidate_k])
|
|
vector_total, vector_rows = rank_vectors(
|
|
db, version, space, payload.category, payload.candidate_k, query_vector
|
|
)
|
|
vector = candidate_list(
|
|
[(chunk, max(-1.0, min(1.0, 1.0 - float(distance)))) for chunk, distance in vector_rows]
|
|
)
|
|
scores = fuse(lexical, vector)
|
|
by_id = {chunk.id: chunk for chunk in chunks}
|
|
ordered = sorted(scores, key=lambda key: (-scores[key], chunk_key(by_id[key])))
|
|
lexical_by_id = {str(c.chunk_id): c for c in lexical}
|
|
vector_by_id = {str(c.chunk_id): c for c in vector}
|
|
hits = []
|
|
for rank, key in enumerate(ordered[: payload.top_k], 1):
|
|
left, right = lexical_by_id.get(key), vector_by_id.get(key)
|
|
hits.append(
|
|
HybridHit(
|
|
**hit_fields(by_id[key], terms, counters),
|
|
rank=rank,
|
|
score=scores[key],
|
|
lexical_rank=left.rank if left else None,
|
|
lexical_score=left.score if left else None,
|
|
vector_rank=right.rank if right else None,
|
|
vector_score=right.score if right else None,
|
|
lexical_contribution=1.0 / (RRF_K + left.rank) if left else 0.0,
|
|
vector_contribution=1.0 / (RRF_K + right.rank) if right else 0.0,
|
|
)
|
|
)
|
|
return save_run(
|
|
db,
|
|
HybridResponse(
|
|
**response_fields(payload, state, text_version, chunks, terms, len(scores), started),
|
|
vector_version_id=version.id,
|
|
generation=vector_state.generation,
|
|
lexical_total_hits=len(ranked),
|
|
vector_total_hits=vector_total,
|
|
lexical_candidates=lexical,
|
|
vector_candidates=vector,
|
|
results=hits,
|
|
config=HybridConfig(
|
|
lexical=CHUNK_LEXICAL_CONFIG,
|
|
chunking=text_version.config,
|
|
candidate_k=payload.candidate_k,
|
|
space=SpaceRead.model_validate(space),
|
|
),
|
|
),
|
|
)
|