diff --git a/quantmind/rag/document.py b/quantmind/rag/document.py index 4f5c3b7..3c6ed2a 100644 --- a/quantmind/rag/document.py +++ b/quantmind/rag/document.py @@ -16,6 +16,18 @@ ParsedPage, ) +_CJK_RANGES = "぀-ヿ㐀-䶿一-鿿豈-﫿" +_CJK_AWARE_TOKEN_PATTERN = rf"(?u)[^\W{_CJK_RANGES}]{{2,}}|[{_CJK_RANGES}]" +"""Tokenize CJK scripts per character, leaving other scripts untouched. + +The retriever default ``(?u)\\b\\w\\w+\\b`` assumes whitespace-delimited words. +Japanese and Chinese are not, so a whole phrase becomes a single token and a +query token can never equal a document token: every chunk scores ``0.0`` while +the retriever still returns ``top_k`` hits. Splitting only CJK characters keeps +the two-or-more-character rule for Latin, Cyrillic and other alphabets, so their +scores and ranking stay identical. +""" + @dataclass(frozen=True) class SentenceSplitterConfig: @@ -154,6 +166,7 @@ def retrieve_parsed_document( retriever = BM25Retriever.from_defaults( nodes=nodes, similarity_top_k=min(top_k, len(nodes)), + token_pattern=_CJK_AWARE_TOKEN_PATTERN, ) results = retriever.retrieve(query) return tuple( diff --git a/tests/rag/test_document.py b/tests/rag/test_document.py index f80dd51..7ed95ae 100644 --- a/tests/rag/test_document.py +++ b/tests/rag/test_document.py @@ -5,6 +5,7 @@ from quantmind.preprocess.format import parse_pdf from quantmind.rag import ( + ParsedChunk, SentenceSplitterConfig, chunk_parsed_document, retrieve_parsed_document, @@ -59,3 +60,63 @@ async def test_retrieval_rejects_invalid_query_arguments(self): retrieve_parsed_document(chunks, " ") with self.assertRaisesRegex(ValueError, "top_k"): retrieve_parsed_document(chunks, "fixture", top_k=0) + + +class ScriptAwareRetrievalTests(unittest.TestCase): + """Retrieval over scripts that are not whitespace-delimited.""" + + @staticmethod + def _chunks(texts: tuple[str, ...]) -> tuple[ParsedChunk, ...]: + return tuple( + ParsedChunk( + chunk_id=f"chunk-{index}", + text=text, + source_hash="0" * 64, + page_number=index + 1, + start_char=0, + end_char=len(text), + block_boxes=(), + screenshot_path=None, + image_paths=(), + ) + for index, text in enumerate(texts) + ) + + def test_japanese_query_ranks_the_matching_chunk_first(self): + chunks = self._chunks( + ( + "国際標準化に関する動向と規格策定プロセスの概観", + "人材育成の方針と研修体系の整備について", + ) + ) + + hits = retrieve_parsed_document(chunks, "国際標準化", top_k=2) + + self.assertEqual(hits[0].chunk.chunk_id, "chunk-0") + self.assertGreater(hits[0].score, 0.0) + + def test_chinese_query_ranks_the_matching_chunk_first(self): + chunks = self._chunks( + ( + "货币政策与利率变动对债券市场的影响", + "上市公司季度财报披露时间表", + ) + ) + + hits = retrieve_parsed_document(chunks, "利率变动", top_k=2) + + self.assertEqual(hits[0].chunk.chunk_id, "chunk-0") + self.assertGreater(hits[0].score, 0.0) + + def test_latin_query_ranks_the_matching_chunk_first(self): + chunks = self._chunks( + ( + "knowledge extraction architecture and retrieval evaluation", + "deterministic preprocessing keeps provenance replayable", + ) + ) + + hits = retrieve_parsed_document(chunks, "retrieval evaluation", top_k=2) + + self.assertEqual(hits[0].chunk.chunk_id, "chunk-0") + self.assertGreater(hits[0].score, 0.0)