From 79f169b03634e269f8f22abe9f2f216f7bc4087e Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 2 Sep 2026 10:09:15 +0800 Subject: [PATCH] feat(retrieval): switch map-unit BM25 to word-level tokens Replace character-level CJK tokenization with the shared jieba path and bump map-unit index format to v2 so query hashes match stored tokens. Add --tokens-only/--skip-current-format for lightweight reindex after tokenizer changes without rebuilding manifests or namespace snapshots. Co-authored-by: Cursor --- apps/api/scripts/backfill_map_unit_indexes.py | 132 +++++++++++++++--- .../test_retrieval_map_unit_index_contract.py | 2 +- .../services/retrieval/map_unit_index.py | 4 +- .../services/retrieval/nav/knowhere_hybrid.py | 30 ++-- .../services/retrieval/nav/nav_knowhere.py | 4 +- .../retrieval/search/map_unit_discovery.py | 9 +- .../tests/test_knowhere_hybrid_tokenize.py | 47 +++++++ 7 files changed, 187 insertions(+), 41 deletions(-) create mode 100644 packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index ebf1a04d..43a9b5c4 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -7,9 +7,20 @@ deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. +Re-run after map-unit tokenizer / ``MAP_UNIT_INDEX_FORMAT_VERSION`` changes +(e.g. v1 character-level → v2 word-level, or loading a domain userdict) so +query hashes match stored tokens. Prefer:: + + python /app/scripts/backfill_map_unit_indexes.py --apply --tokens-only + python /app/scripts/backfill_map_unit_indexes.py --apply --tokens-only --skip-current-format + +``--tokens-only`` rebuilds map-unit tokens/index only (skips serving manifest +and namespace MAP snapshot). Use full ``--apply`` only when those artifacts +are missing. + Use ``--check`` after backfill to verify whether query-time snapshot fallbacks (manifest_merge / table_scan) would still fire, and whether -map-unit indexes are complete for scoring. +map-unit indexes are complete for scoring at the current format version. """ # ruff: noqa: E402 @@ -57,7 +68,10 @@ def _bootstrap_python_path() -> None: RetrievalNamespaceMapSnapshot, RetrievalServingRevisionManifest, ) -from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.map_unit_index import ( + MAP_UNIT_INDEX_FORMAT_VERSION, + replace_document_map_units, +) from shared.services.retrieval.namespace_map_snapshot import ( patch_namespace_map_snapshot, ) @@ -92,6 +106,23 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument( "--document-id", default="", help="Limit the backfill to one document." ) + parser.add_argument( + "--tokens-only", + action="store_true", + help=( + "With --apply: rebuild only document_map_units / tokens / index " + "(skip serving manifest, namespace MAP snapshot, and generation bump). " + "Use after tokenizer / format_version changes when manifests are already ready." + ), + ) + parser.add_argument( + "--skip-current-format", + action="store_true", + help=( + "With --apply: skip revisions whose map-unit index already has " + f"format_version={MAP_UNIT_INDEX_FORMAT_VERSION}." + ), + ) return parser @@ -196,6 +227,7 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback DocumentMapUnitIndex.unit_count, DocumentMapUnitIndex.average_idf_path, DocumentMapUnitIndex.average_idf_content, + DocumentMapUnitIndex.format_version, ).where( DocumentMapUnitIndex.document_id.in_( [document_id_value for document_id_value, _ in revisions] @@ -208,8 +240,16 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback int(unit_count or 0), float(average_idf_path or 0.0), float(average_idf_content or 0.0), + int(format_version or 0), ) - for document_id_value, job_result_id, unit_count, average_idf_path, average_idf_content in index_rows + for ( + document_id_value, + job_result_id, + unit_count, + average_idf_path, + average_idf_content, + format_version, + ) in index_rows } missing_map_index = 0 suspicious_zero_idf = 0 @@ -218,7 +258,10 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback if stats is None: missing_map_index += 1 continue - unit_count, average_idf_path, average_idf_content = stats + unit_count, average_idf_path, average_idf_content, format_version = stats + if format_version != MAP_UNIT_INDEX_FORMAT_VERSION: + missing_map_index += 1 + continue if ( unit_count > 0 and average_idf_path == 0.0 @@ -295,26 +338,36 @@ def print_fallback_check(reports: list[NamespaceFallbackReport]) -> int: return 1 if failed else 0 -def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: +def backfill_map_unit_indexes( + *, + apply: bool, + document_id: str = "", + tokens_only: bool = False, + skip_current_format: bool = False, +) -> int: documents = _load_documents(document_id) if not apply: for document in documents: print( f"would backfill document={document.document_id} revision={document.current_job_result_id}" + + (" tokens_only" if tokens_only else "") ) return len(documents) session_factory = get_sync_session_factory() + backfilled = 0 + skipped_current = 0 for document in documents: job_result_id = document.current_job_result_id if not job_result_id: continue with session_factory() as db: - lock_namespace_generation( - db, - user_id=document.user_id, - namespace=document.namespace, - ) + if not tokens_only: + lock_namespace_generation( + db, + user_id=document.user_id, + namespace=document.namespace, + ) locked_document = db.execute( select(Document) .where(Document.document_id == document.document_id) @@ -333,6 +386,26 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: f"revision={job_result_id}" ) continue + if skip_current_format: + existing = db.execute( + select(DocumentMapUnitIndex.format_version) + .where( + DocumentMapUnitIndex.document_id + == locked_document.document_id + ) + .where(DocumentMapUnitIndex.job_result_id == job_result_id) + ).scalar_one_or_none() + if ( + existing is not None + and int(existing) == MAP_UNIT_INDEX_FORMAT_VERSION + ): + db.rollback() + skipped_current += 1 + print( + f"skipped current-format document={document.document_id} " + f"revision={job_result_id} format_version={existing}" + ) + continue scope = DocumentPublicationScope( user_id=locked_document.user_id, namespace=locked_document.namespace, @@ -341,18 +414,25 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: source_file_name=str(locked_document.source_file_name or ""), ) replace_document_map_units(db, scope=scope) - manifest_payload = persist_revision_serving_state(db, scope=scope) - patch_namespace_map_snapshot( - db, scope=scope, manifest_payload=manifest_payload - ) - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) + if not tokens_only: + manifest_payload = persist_revision_serving_state(db, scope=scope) + patch_namespace_map_snapshot( + db, scope=scope, manifest_payload=manifest_payload + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) db.commit() - print(f"backfilled document={document.document_id} revision={job_result_id}") - return len(documents) + backfilled += 1 + mode = "tokens_only" if tokens_only else "full" + print( + f"backfilled document={document.document_id} revision={job_result_id} mode={mode}" + ) + if skip_current_format: + print(f"skipped_current_format={skipped_current}") + return backfilled def main() -> None: @@ -360,11 +440,19 @@ def main() -> None: if arguments.check: if arguments.apply: raise SystemExit("use either --check or --apply, not both") + if arguments.tokens_only or arguments.skip_current_format: + raise SystemExit("--tokens-only/--skip-current-format require --apply") reports = check_fallback_readiness(document_id=str(arguments.document_id)) raise SystemExit(print_fallback_check(reports)) + if (arguments.tokens_only or arguments.skip_current_format) and not arguments.apply: + raise SystemExit("--tokens-only/--skip-current-format require --apply") + count = backfill_map_unit_indexes( - apply=bool(arguments.apply), document_id=str(arguments.document_id) + apply=bool(arguments.apply), + document_id=str(arguments.document_id), + tokens_only=bool(arguments.tokens_only), + skip_current_format=bool(arguments.skip_current_format), ) action = "backfilled" if arguments.apply else "found" print(f"{action} revisions={count}") diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index b4b047d8..e71dd6b0 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -56,7 +56,7 @@ def __init__(self) -> None: def execute(self, statement: str, parameters: object = None) -> None: executions.append((statement, parameters)) if "document_map_unit_indexes" in statement: - self.rows = [(document_id, job_result_id, 1, 1, 0.0, 0.0)] + self.rows = [(document_id, job_result_id, 2, 1, 0.0, 0.0)] elif "FROM document_map_units AS units" in statement: self.rows = [("unit-frequency", document_id, "chunk-frequency", "section-frequency", 1, 1)] elif "FROM document_map_unit_tokens" in statement: diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 4fff8206..ee18434d 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -23,11 +23,11 @@ UnitRow, ) from shared.services.retrieval.nav.nav_map_scores import build_score_units +from shared.services.retrieval.nav.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION from shared.services.retrieval.nav.persisted_score_load import average_idf_from_unit_dfs from shared.services.retrieval.publication_models import DocumentPublicationScope - -MAP_UNIT_INDEX_FORMAT_VERSION = 1 +__all__ = ["MAP_UNIT_INDEX_FORMAT_VERSION", "replace_document_map_units"] def replace_document_map_units( diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 2c4bbb48..6e5a95f8 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -8,28 +8,34 @@ from __future__ import annotations -import os -import re import math +import os from dataclasses import dataclass from typing import Dict, List, Mapping, Optional, Sequence, Tuple +from shared.utils.text_utils import tokenize_for_retrieval as _tokenize_word_level + RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 +# Persisted map-unit token format. Bump when tokenizer semantics change. +# v1: character-level CJK regex. v2: word-level jieba (text_utils). +MAP_UNIT_INDEX_FORMAT_VERSION = 2 def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: - tokens = re.findall(r"[a-z0-9_]+|[\u4e00-\u9fff]", str(text or "").lower()) - if not dedupe: - return [t for t in tokens if t] - seen: set[str] = set() - out: List[str] = [] - for t in tokens: - if t and t not in seen: - seen.add(t) - out.append(t) - return out + """Word-level retrieval tokens (jieba + English), aligned with chunk publication. + + Uses the same knobs as ``search.lexical_text``: no stopword filtering and + ``min_token_length=2`` so map-unit indexes match ``document_chunks`` search + text. Character-level regex tokenization was removed. + """ + return _tokenize_word_level( + text, + stopwords=[], + dedupe=dedupe, + min_token_length=2, + ) def tokenize_query_for_ranker(query: str) -> List[str]: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index bd4e24f7..3dadd57e 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -42,6 +42,7 @@ from .nav_address import NavLevel from .nav_hierarchy import NodeMeta from .knowhere_hybrid import ( + MAP_UNIT_INDEX_FORMAT_VERSION, PersistedScoreCorpus, PersistedScoreUnit, tokenize_query_for_ranker, @@ -51,7 +52,6 @@ # Knowhere sentinel path for the virtual document container (not a collectable leaf). ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" -_MAP_UNIT_INDEX_FORMAT_VERSION = 1 _MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") _logger = logging.getLogger(__name__) @@ -353,7 +353,7 @@ def load_persisted_score_corpus( len(index_rows), ) if len(index_rows) != len(revisions) or any( - len(row) < 6 or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION + len(row) < 6 or int(row[2]) != MAP_UNIT_INDEX_FORMAT_VERSION for row in index_rows ): return None diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 95436458..4e7d1a4a 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -32,6 +32,7 @@ normalize_chunk_type, ) from shared.services.retrieval.nav.knowhere_hybrid import ( + MAP_UNIT_INDEX_FORMAT_VERSION, PersistedScoreCorpus, PersistedScoreUnit, score_persisted_corpus_many, @@ -282,7 +283,7 @@ async def map_unit_discovery( ) + """ SELECT indexes.average_idf_path, indexes.average_idf_content, - indexes.unit_count + indexes.unit_count, indexes.format_version FROM document_map_unit_indexes AS indexes JOIN ( SELECT DISTINCT document_id, job_result_id FROM scoped_units @@ -293,9 +294,13 @@ async def map_unit_discovery( ), params, ) + # Only count indexes written with the current tokenizer (v2 = word-level). + # Stale char-level (v1) rows look complete by unit_count but cannot match + # word-level query hashes — treat them as missing so readiness fails closed. index_parts = [ (float(path_idf or 0.0), float(content_idf or 0.0), int(unit_count or 0)) - for path_idf, content_idf, unit_count in index_result.all() + for path_idf, content_idf, unit_count, format_version in index_result.all() + if int(format_version or 0) == MAP_UNIT_INDEX_FORMAT_VERSION ] expected_revisions = { (str(row["document_id"]), str(row["job_result_id"])) for row in unit_rows diff --git a/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py b/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py new file mode 100644 index 00000000..2d4b96e1 --- /dev/null +++ b/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py @@ -0,0 +1,47 @@ +"""Word-level map-unit tokenization (replaces character-level CJK regex).""" + +from __future__ import annotations + +from shared.services.retrieval.nav.knowhere_hybrid import ( + MAP_UNIT_INDEX_FORMAT_VERSION, + build_content_search_text, + build_path_search_text, + tokenize_for_retrieval, + tokenize_query_for_ranker, +) +from shared.utils.text_utils import tokenize_for_retrieval as tokenize_publication + + +def test_map_unit_index_format_is_word_level_v2() -> None: + assert MAP_UNIT_INDEX_FORMAT_VERSION == 2 + + +def test_query_uses_chinese_word_tokens_not_characters() -> None: + tokens = tokenize_query_for_ranker("冠心病的诊断标准") + assert "冠心病" in tokens + assert "诊断" in tokens + assert "标准" in tokens + assert "冠" not in tokens + assert "诊" not in tokens + + +def test_index_text_matches_publication_tokenizer() -> None: + text = "冠心病患者的诊断标准与心肌炎鉴别" + map_tokens = tokenize_for_retrieval(text, dedupe=False) + pub_tokens = tokenize_publication( + text, stopwords=[], dedupe=False, min_token_length=2 + ) + assert map_tokens == pub_tokens + assert "冠心病" in map_tokens + assert "心肌炎" in map_tokens + + +def test_build_search_text_joins_word_tokens() -> None: + content = build_content_search_text("冠心病诊断标准") + path = build_path_search_text(section_path="指南 / 冠心病 / 诊断") + assert "冠心病" in content.split() + assert "诊断" in content.split() + assert "冠心病" in path.split() + # Single CJK characters must not appear as standalone tokens. + assert "冠" not in content.split() + assert "冠" not in path.split()