436 lines
22 KiB
Python
436 lines
22 KiB
Python
"""Indexing and hybrid retrieval with page-type aware candidate filtering.
|
|
|
|
Chinese keyword recall does not use PostgreSQL ``to_tsvector('simple', …)``. That
|
|
configuration has no Chinese segmenter, so a whole run of Chinese characters becomes a
|
|
single token; measured on the course reports, the query ``文献调研`` does not match the
|
|
page that contains ``第2章文献调研与技术现状``. Recall therefore uses character bigrams
|
|
plus whole ASCII/number tokens, which needs no extension and behaves the same on
|
|
PostgreSQL and SQLite.
|
|
|
|
Semantic and keyword scores are fused by rank (reciprocal rank fusion), never by adding
|
|
raw numbers: a cosine similarity and a bigram overlap are not on a common scale.
|
|
"""
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
import re
|
|
|
|
from sqlalchemy import case, cast, delete, func, literal, or_, select
|
|
|
|
from .config import get_settings
|
|
from .models import KnowledgeChunk, KnowledgeDocument, KnowledgePage
|
|
from .pages import (BLOCK_BODY, BLOCK_TOC, classify_document, is_toc_entry,
|
|
normalize_heading, searchable_pages)
|
|
|
|
INDEX_VERSION = 3
|
|
|
|
#: Fusion constant for reciprocal rank fusion; the usual value from the IR literature.
|
|
RRF_K = 60
|
|
#: Candidate pools are gathered independently so both retrievers get a chance.
|
|
SEMANTIC_POOL = 60
|
|
KEYWORD_POOL = 60
|
|
#: A chunk must be at least this long to count as substantive content evidence.
|
|
MIN_CHUNK_CHARS = 60
|
|
#: At most this many chunks from one page in the first retrieval pass.
|
|
DEFAULT_PAGE_LIMIT = 2
|
|
#: Soft cap on chunks from one section once a reliable section path exists.
|
|
DEFAULT_SECTION_CAP = 3
|
|
#: Maximum chunks handed to the model for one criterion, across all targets.
|
|
SCORE_EVIDENCE_LIMIT = 14
|
|
|
|
_ASCII_TOKEN = re.compile(r"[A-Za-z][A-Za-z0-9_+\-.]*|\d+(?:\.\d+)*")
|
|
_CJK_RUN = re.compile(r"[㐀-䶿一-鿿]+")
|
|
|
|
|
|
def query_terms(query: str) -> tuple[list[str], list[str]]:
|
|
"""Split a query into (ascii tokens lowercased, CJK character bigrams).
|
|
|
|
Single-character CJK runs are dropped: they carry too little signal and would make
|
|
every chunk with a common character look relevant.
|
|
"""
|
|
value = str(query or "")
|
|
ascii_tokens = [token.lower() for token in _ASCII_TOKEN.findall(value)]
|
|
bigrams: list[str] = []
|
|
seen = set()
|
|
for run in _CJK_RUN.findall(value):
|
|
for index in range(len(run) - 1):
|
|
bigram = run[index:index + 2]
|
|
if bigram not in seen:
|
|
seen.add(bigram)
|
|
bigrams.append(bigram)
|
|
return sorted(set(ascii_tokens)), bigrams
|
|
|
|
|
|
def _keyword_query(db, filters, query: str, limit: int):
|
|
"""Keyword recall ranked by bigram/token overlap over the already-filtered rows.
|
|
|
|
PostgreSQL's ``to_tsvector('simple', …)`` cannot segment Chinese, so it is not used
|
|
for recall. The overlap is computed with ``strpos``/``substr``/``generate_series``
|
|
so the same expression works on every dialect SQLAlchemy targets here.
|
|
"""
|
|
ascii_tokens, bigrams = query_terms(query)
|
|
if not ascii_tokens and not bigrams:
|
|
return []
|
|
total = len(ascii_tokens) + len(bigrams)
|
|
score = None
|
|
for token in ascii_tokens:
|
|
hit = case((func.strpos(func.lower(KnowledgeChunk.text), literal(token)) > 0, 1.0), else_=0.0)
|
|
score = hit if score is None else score + hit
|
|
if bigrams:
|
|
series = select(func.generate_series(1, func.greatest(func.length(KnowledgeChunk.text) - 1, 0)).label("i")).subquery("series")
|
|
matches = (select(func.count()).select_from(series)
|
|
.where(func.substr(KnowledgeChunk.text, series.c.i, 2).in_(bigrams)).scalar_subquery())
|
|
score = matches if score is None else score + matches
|
|
score = (score / float(total)).label("keyword")
|
|
rows = db.execute(select(KnowledgeChunk.id, score).where(*filters)
|
|
.order_by(score.desc(), KnowledgeChunk.id.asc()).limit(limit)).all()
|
|
return [row[0] for row in rows if float(row[1] or 0) > 0]
|
|
|
|
|
|
def split_pages(pages: list[dict], *, size=650, overlap=100):
|
|
"""Split pages into searchable chunks carrying page type and section path."""
|
|
for page in pages:
|
|
page_type = page.get("page_type") or "body"
|
|
chapter = page.get("chapter", "")
|
|
section = page.get("section", "")
|
|
blocks = page.get("blocks")
|
|
if not blocks:
|
|
blocks = [{"kind": BLOCK_BODY, "text": str(page.get("text") or "").strip()}]
|
|
blocks.extend({"kind": BLOCK_BODY, "text": "表格: " + str(table.get("rows", table))}
|
|
for table in page.get("tables", []))
|
|
for block in blocks:
|
|
if block["kind"] == BLOCK_TOC:
|
|
# A contents page is navigation, so its text is not body content.
|
|
yield {"page": page["page_number"], "page_type": "toc", "block_type": BLOCK_TOC,
|
|
"chapter": chapter, "section": section, "text": block["text"]}
|
|
continue
|
|
text_block = block["text"].strip()
|
|
if not text_block:
|
|
continue
|
|
for table in ([block["tables"]] if block.get("tables") else []):
|
|
for table_item in table:
|
|
text_block += "\n表格: " + str(table_item.get("rows", table_item))
|
|
# Table rows of a page are appended once, on the page's first body chunk.
|
|
start = 0
|
|
while start < len(text_block):
|
|
yield {"page": page["page_number"], "page_type": page_type, "block_type": BLOCK_BODY,
|
|
"chapter": chapter, "section": section, "text": text_block[start:start + size]}
|
|
if start + size >= len(text_block):
|
|
break
|
|
start += size - overlap
|
|
|
|
|
|
def build_chunks(pages: list[dict]) -> list[dict]:
|
|
"""Chunks for one document, classifying pages first when they are not yet typed."""
|
|
prepared = pages if pages and "page_type" in (pages[0] or {}) else searchable_pages(pages)
|
|
return list(split_pages(prepared))
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def local_encoder(model_path: str):
|
|
path = Path(model_path)
|
|
if not path.is_dir() or not any(path.iterdir()):
|
|
raise RuntimeError(f"离线嵌入模型不存在: {path}；请挂载本地模型包")
|
|
from sentence_transformers import SentenceTransformer
|
|
return SentenceTransformer(str(path), local_files_only=True, trust_remote_code=False)
|
|
|
|
|
|
def embed(texts: list[str], *, model_path: str | None = None):
|
|
if not texts:
|
|
return []
|
|
model = local_encoder(model_path or str(get_settings().embedding_model_path))
|
|
vectors = model.encode(texts, normalize_embeddings=True)
|
|
output = []
|
|
for vector in vectors:
|
|
values = vector.tolist()
|
|
if len(values) > 1024:
|
|
raise RuntimeError("嵌入维度超过pgvector列容量1024")
|
|
output.append(values + [0.0] * (1024 - len(values)))
|
|
return output
|
|
|
|
|
|
def ingest(db, *, course_version_id: int, project_id: int | None, submission_file_id: int | None,
|
|
title: str, category: str, source_hash: str, pages: list[dict], embedder=embed):
|
|
existing = db.scalar(select(KnowledgeDocument).where(KnowledgeDocument.course_version_id == course_version_id,
|
|
KnowledgeDocument.project_id == project_id, KnowledgeDocument.submission_file_id == submission_file_id,
|
|
KnowledgeDocument.title == title))
|
|
if existing and existing.source_hash == source_hash and existing.index_version == INDEX_VERSION and existing.status == "ready":
|
|
return existing
|
|
prepared = searchable_pages(pages)
|
|
chunks = list(split_pages(prepared))
|
|
vectors = embedder([chunk["text"] for chunk in chunks])
|
|
if len(vectors) != len(chunks):
|
|
raise RuntimeError("嵌入向量数量与知识片段数量不一致")
|
|
if existing:
|
|
db.execute(delete(KnowledgePage).where(KnowledgePage.document_id == existing.id))
|
|
db.execute(delete(KnowledgeChunk).where(KnowledgeChunk.document_id == existing.id))
|
|
doc = existing
|
|
doc.document_type, doc.source_hash = category, source_hash
|
|
else:
|
|
doc = KnowledgeDocument(course_version_id=course_version_id, project_id=project_id, submission_file_id=submission_file_id,
|
|
title=title, document_type=category, source_hash=source_hash)
|
|
db.add(doc)
|
|
doc.index_version, doc.status = INDEX_VERSION, "indexing"
|
|
db.flush()
|
|
for source, prepared_page in zip(pages, prepared):
|
|
classification = prepared_page["classification"]
|
|
db.add(KnowledgePage(
|
|
document_id=doc.id, page=source["page_number"], text=page_text(source),
|
|
image_path=relative_image_path(source.get("rendered_image_path")),
|
|
page_type=classification["page_type"], mixed=bool(classification["mixed"]),
|
|
section=classification["section"], chapter=classification["chapter"],
|
|
confidence=classification["confidence"],
|
|
classification={**classification["signals"], "classifier_version": classification["classifier_version"],
|
|
"blocks": [{"kind": block["kind"], "chars": len(block["text"]),
|
|
"substantive": block.get("substantive", False)}
|
|
for block in classification["blocks"]]}))
|
|
for chunk, vector in zip(chunks, vectors):
|
|
db.add(KnowledgeChunk(document_id=doc.id, course_version_id=course_version_id, project_id=project_id,
|
|
document_type=category, page=chunk["page"], section=chunk["section"],
|
|
chapter=chunk["chapter"], page_type=chunk["page_type"],
|
|
block_type=chunk["block_type"], text=chunk["text"], embedding=vector))
|
|
doc.status = "ready"
|
|
db.commit()
|
|
return doc
|
|
|
|
|
|
def _scope_filter(project_id: int | None, include_public: bool, public_only: bool):
|
|
if public_only:
|
|
return KnowledgeChunk.project_id.is_(None)
|
|
if include_public:
|
|
return or_(KnowledgeChunk.project_id == project_id, KnowledgeChunk.project_id.is_(None))
|
|
return KnowledgeChunk.project_id == project_id
|
|
|
|
|
|
def _base_filters(*, course_version_id: int, project_id: int | None, category, categories,
|
|
include_public: bool, public_only: bool, allowed_page_types, exclude_document_ids):
|
|
filters = [KnowledgeChunk.course_version_id == course_version_id,
|
|
_scope_filter(project_id, include_public, public_only)]
|
|
if category:
|
|
filters.append(KnowledgeChunk.document_type == category)
|
|
if categories:
|
|
filters.append(KnowledgeChunk.document_type.in_(list(categories)))
|
|
if allowed_page_types is not None:
|
|
filters.append(KnowledgeChunk.page_type.in_(list(allowed_page_types)))
|
|
if exclude_document_ids:
|
|
filters.append(KnowledgeChunk.document_id.not_in(list(exclude_document_ids)))
|
|
return filters
|
|
|
|
|
|
def _keyword_query(db, filters, query: str, limit: int):
|
|
"""Keyword recall ranked by bigram/token overlap over the already-filtered rows.
|
|
|
|
PostgreSQL's ``to_tsvector('simple', …)`` cannot segment Chinese, so it is not used
|
|
for recall. The overlap is computed with ``strpos``/``substr`` so the same expression
|
|
works on every dialect SQLAlchemy targets here.
|
|
"""
|
|
ascii_tokens, bigrams = query_terms(query)
|
|
if not ascii_tokens and not bigrams:
|
|
return []
|
|
total = len(ascii_tokens) + len(bigrams)
|
|
score = None
|
|
for token in ascii_tokens:
|
|
term = literal(token)
|
|
hit = case((func.strpos(func.lower(KnowledgeChunk.text), term) > 0, 1.0), else_=0.0)
|
|
score = hit if score is None else score + hit
|
|
if bigrams:
|
|
start = func.generate_series(1, func.greatest(func.length(KnowledgeChunk.text) - 1, 0))
|
|
series = select(start.label("i")).subquery("series")
|
|
matches = (select(func.count()).select_from(series)
|
|
.where(func.substr(KnowledgeChunk.text, series.c.i, 2).in_(bigrams))
|
|
.scalar_subquery())
|
|
score = matches if score is None else score + matches
|
|
score = (score / float(total)).label("keyword")
|
|
rows = db.execute(select(KnowledgeChunk.id, score).where(*filters)
|
|
.order_by(score.desc(), KnowledgeChunk.id.asc()).limit(limit)).all()
|
|
return [row[0] for row in rows if float(row[1] or 0) > 0]
|
|
|
|
|
|
def keyword_overlap(text_value: str, ascii_tokens: list[str], bigrams: list[str]) -> float:
|
|
"""Bigram/token overlap used by the non-PostgreSQL path and by tests."""
|
|
value = str(text_value or "")
|
|
lower = value.lower()
|
|
hits = sum(1 for token in ascii_tokens if token in lower)
|
|
total = len(ascii_tokens)
|
|
hits += sum(1 for bigram in bigrams if bigram in value)
|
|
total += len(bigrams)
|
|
return hits / total if total else 0.0
|
|
|
|
|
|
def _semantic_query(db, filters, vector, limit: int):
|
|
if db.bind.dialect.name == "postgresql":
|
|
from pgvector.sqlalchemy import Vector
|
|
distance = cast(KnowledgeChunk.embedding, Vector(1024)).cosine_distance(vector)
|
|
rows = db.execute(select(KnowledgeChunk.id).where(*filters)
|
|
.order_by(distance.asc(), KnowledgeChunk.id.asc()).limit(limit)).all()
|
|
return [row[0] for row in rows]
|
|
rows = db.execute(select(KnowledgeChunk.id, KnowledgeChunk.embedding).where(*filters)).all()
|
|
scored = [(chunk_id, sum(a * b for a, b in zip(embedding or [], vector)))
|
|
for chunk_id, embedding in rows]
|
|
scored.sort(key=lambda pair: (-pair[1], pair[0]))
|
|
return [chunk_id for chunk_id, _ in scored[:limit]]
|
|
|
|
|
|
def _fuse(db, ids: list[int], semantic_ids: list[int], keyword_ids: list[int]):
|
|
"""Reciprocal rank fusion over the two ranked lists, keeping discovery order stable."""
|
|
scores: dict[int, float] = {chunk_id: 0.0 for chunk_id in ids}
|
|
for ranked in (semantic_ids, keyword_ids):
|
|
for rank, chunk_id in enumerate(ranked):
|
|
if chunk_id in scores:
|
|
scores[chunk_id] += 1.0 / (RRF_K + rank + 1)
|
|
return scores
|
|
|
|
|
|
def _row_dict(chunk):
|
|
return {"chunk_id": chunk.id, "document_id": chunk.document_id, "project_id": chunk.project_id,
|
|
"document_type": chunk.document_type, "page": chunk.page, "section": chunk.section,
|
|
"chapter": chunk.chapter, "page_type": chunk.page_type, "block_type": chunk.block_type,
|
|
"text": chunk.text}
|
|
|
|
|
|
def normalize_hit(hit: dict) -> dict:
|
|
"""Fill the fields the rest of the pipeline reads, so one shape flows through.
|
|
|
|
An alternate retriever or a test double may return only the essentials. Filling the
|
|
defaults here keeps every consumer from re-deriving them; the authoritative page type
|
|
is still re-read from storage when a score is produced.
|
|
"""
|
|
return {"chunk_id": hit.get("chunk_id"), "document_id": hit.get("document_id"),
|
|
"project_id": hit.get("project_id"), "document_type": hit.get("document_type"),
|
|
"page": hit.get("page"), "text": hit.get("text") or "",
|
|
"section": hit.get("section") or "", "chapter": hit.get("chapter") or "",
|
|
"page_type": hit.get("page_type") or "unknown", "block_type": hit.get("block_type") or "legacy",
|
|
"rank": float(hit.get("rank") or 0.0)}
|
|
|
|
|
|
def retrieve(db, *, course_version_id: int, project_id: int, query: str, category: str | None = None,
|
|
categories=None, limit=10, embedder=embed, include_public=True, public_only=False,
|
|
allowed_page_types=None, exclude_document_ids=(), substantive_only=False,
|
|
semantic_pool: int | None = None, keyword_pool: int | None = None):
|
|
"""Filtered hybrid retrieval, then rank fusion over the two independent recalls.
|
|
|
|
Page filtering happens inside the query, so it always precedes the candidate limit:
|
|
a contents page can never displace a body page simply by scoring well.
|
|
"""
|
|
filters = _base_filters(course_version_id=course_version_id, project_id=project_id, category=category,
|
|
categories=categories, include_public=include_public, public_only=public_only,
|
|
allowed_page_types=allowed_page_types, exclude_document_ids=exclude_document_ids)
|
|
if substantive_only:
|
|
filters.append(KnowledgeChunk.block_type != BLOCK_TOC)
|
|
vector = embedder([query])[0]
|
|
semantic_ids = _semantic_query(db, filters, vector, semantic_pool or SEMANTIC_POOL)
|
|
keyword_ids = _keyword_query(db, filters, query, keyword_pool or KEYWORD_POOL)
|
|
ids = list(dict.fromkeys(semantic_ids + keyword_ids))
|
|
if not ids:
|
|
return []
|
|
scores = _fuse(db, ids, semantic_ids, keyword_ids)
|
|
chunks = db.scalars(select(KnowledgeChunk).where(KnowledgeChunk.id.in_(ids))).all()
|
|
usable = []
|
|
for chunk in chunks:
|
|
if substantive_only and (chunk.block_type == BLOCK_TOC
|
|
or (chunk.page_type == "body" and is_toc_entry(chunk.text)
|
|
and len(chunk.text.strip()) < MIN_CHUNK_CHARS)):
|
|
continue
|
|
usable.append(chunk)
|
|
usable.sort(key=lambda chunk: (-scores.get(chunk.id, 0.0), chunk.id))
|
|
return [{**_row_dict(chunk), "rank": float(scores.get(chunk.id, 0.0))} for chunk in usable[:limit]]
|
|
|
|
|
|
def _page_diverse(chunks: list[dict], limit: int, *, page_limit: int, section_cap: int | None):
|
|
"""Spread the first pass over pages, then allow a bounded number per section."""
|
|
picked: list[dict] = []
|
|
per_page: dict[tuple, int] = {}
|
|
per_section: dict[str, int] = {}
|
|
remaining = []
|
|
for chunk in chunks:
|
|
page_key = (chunk.get("document_id"), chunk.get("page"))
|
|
section = str(chunk.get("section") or "")
|
|
if len(picked) < limit and per_page.get(page_key, 0) < page_limit:
|
|
if section_cap and section and per_section.get(section, 0) >= section_cap:
|
|
remaining.append(chunk)
|
|
continue
|
|
picked.append(chunk)
|
|
per_page[page_key] = per_page.get(page_key, 0) + 1
|
|
if section:
|
|
per_section[section] = per_section.get(section, 0) + 1
|
|
else:
|
|
remaining.append(chunk)
|
|
for chunk in remaining:
|
|
if len(picked) >= limit:
|
|
break
|
|
picked.append(chunk)
|
|
return picked[:limit]
|
|
|
|
|
|
def retrieve_targets(db, *, course_version_id: int, project_id: int, targets: list[dict], category=None,
|
|
categories=None, limit=SCORE_EVIDENCE_LIMIT, embedder=embed, include_public=False,
|
|
page_limit=DEFAULT_PAGE_LIMIT, section_cap=DEFAULT_SECTION_CAP,
|
|
per_target_limit: int | None = None, navigator=None, retrieve_fn=retrieve):
|
|
"""Search once per required-evidence target and merge with page-level diversity.
|
|
|
|
Every target is searched separately with its own query and its own allowed page
|
|
types, so a contents page permitted for a structural target can never flood a
|
|
content target's evidence. A target counts as unmet when its *own* search found
|
|
nothing admissible, which is what decides a teacher review.
|
|
"""
|
|
merged: list[dict] = []
|
|
seen: set[int] = set()
|
|
per_target: dict[str, list[dict]] = {}
|
|
for target in targets:
|
|
allowed = target.get("allowed_page_types")
|
|
raw_hits = retrieve_fn(db, course_version_id=course_version_id, project_id=project_id, category=category,
|
|
categories=categories, query=target.get("query") or target.get("label") or "",
|
|
limit=per_target_limit or max(limit, 8), embedder=embed,
|
|
include_public=include_public, allowed_page_types=allowed,
|
|
substantive_only=target.get("use") != "structure",
|
|
semantic_pool=SEMANTIC_POOL, keyword_pool=KEYWORD_POOL)
|
|
hits = [normalize_hit(hit) for hit in (raw_hits or [])]
|
|
if target.get("use") == "content" and not hits and navigator is not None:
|
|
# The contents list may only be used to navigate: its headings are matched
|
|
# against headings that really occur in body text, and those sections and
|
|
# their neighbouring pages are searched. The quotation still lands on body.
|
|
hits = [dict(hit, via="toc_navigation") for hit in (navigator(target) or [])]
|
|
per_target[target["target_id"]] = hits
|
|
for hit in hits:
|
|
if hit["chunk_id"] not in seen:
|
|
seen.add(hit["chunk_id"])
|
|
merged.append({**hit, "target_id": target["target_id"], "target_label": target["label"],
|
|
"target_use": target["use"]})
|
|
merged.sort(key=lambda chunk: -chunk["rank"])
|
|
selected = _page_diverse(merged, limit, page_limit=page_limit, section_cap=section_cap)
|
|
uncovered = [target for target in targets if not per_target.get(target["target_id"])]
|
|
return {"evidence": selected, "per_target": per_target, "uncovered_targets": uncovered}
|
|
|
|
|
|
def body_heading_index(db, *, course_version_id: int, project_id: int, category: str | None = None):
|
|
"""Sections present in real body text, used to navigate from a contents page."""
|
|
filters = [KnowledgeChunk.course_version_id == course_version_id, KnowledgeChunk.project_id == project_id,
|
|
KnowledgeChunk.page_type == "body"]
|
|
if category:
|
|
filters.append(KnowledgeChunk.document_type == category)
|
|
rows = db.execute(select(KnowledgeChunk.document_id, KnowledgeChunk.page, KnowledgeChunk.section,
|
|
KnowledgeChunk.chapter).where(*filters).distinct()).all()
|
|
index: dict[str, list[dict]] = {}
|
|
for document_id, page, section, chapter in rows:
|
|
for value in (section, chapter):
|
|
key = normalize_heading(value or "")
|
|
if key:
|
|
index.setdefault(key, []).append({"document_id": document_id, "page": page,
|
|
"section": section, "chapter": chapter})
|
|
return index
|
|
|
|
|
|
def page_text(source: dict) -> str:
|
|
blocks = [str(source.get("text") or "").strip()]
|
|
blocks.extend("表格: " + str(table.get("rows", table)) for table in source.get("tables", []))
|
|
return "\n".join(block for block in blocks if block)
|
|
|
|
|
|
def relative_image_path(value: str | None) -> str | None:
|
|
if not value:
|
|
return None
|
|
root = get_settings().processed_path.resolve()
|
|
path = Path(value).resolve()
|
|
if not path.is_relative_to(root):
|
|
raise ValueError("页面图像路径不在处理目录内")
|
|
return str(path.relative_to(root))
|