diff --git a/reme2/component/base_component.py b/reme2/component/base_component.py index 4cab9f58..7b7c122e 100644 --- a/reme2/component/base_component.py +++ b/reme2/component/base_component.py @@ -1,7 +1,7 @@ """Base class for components.""" import asyncio -from abc import ABC, abstractmethod +from abc import ABC from typing import TYPE_CHECKING from ..enumeration import ComponentEnum diff --git a/reme2/component/file_store/bm25_lite.py b/reme2/component/file_store/bm25_lite.py new file mode 100644 index 00000000..c81a206f --- /dev/null +++ b/reme2/component/file_store/bm25_lite.py @@ -0,0 +1,273 @@ +"""Lightweight BM25 search engine with persistent index support. + +BM25Lite implements the Okapi BM25 ranking algorithm for text retrieval. +It uses an inverted index for efficient document lookup and supports +incremental updates, persistence via pickle, and automatic vocab compaction. +""" + +import math +import pickle +from collections import Counter +from pathlib import Path +from typing import TypedDict + +from ..base_component import BaseComponent +from ..tokenizer import BaseTokenizer +from ...enumeration import ComponentEnum + + +class DocMeta(TypedDict): + """Document metadata stored in the index. + + Attributes: + len: Number of tokens in the document. + token_ids: Set of unique token IDs present in the document. + """ + len: int + token_ids: set[int] + + +class BM25Lite(BaseComponent): + """Lightweight BM25 search engine with file-based persistence. + + BM25 (Best Matching 25) is a probabilistic ranking function that scores + documents based on term frequency and document length normalization. + + Attributes: + k1: Term frequency saturation parameter (default: 1.5). + b: Document length normalization parameter (default: 0.75). + vocab: Token to token ID mapping. + inverted_index: Token ID to {doc_id: term_frequency} mapping. + doc_meta: Document ID to metadata mapping. + + Args: + index_dir: Directory to store index files. + k1: BM25 term frequency saturation parameter. + b: BM25 document length normalization parameter. + tokenizer: Name of tokenizer component to use. + """ + + def __init__( + self, + index_dir: str | Path, + k1: float = 1.5, + b: float = 0.75, + tokenizer: str = "default", + **kwargs): + super().__init__(**kwargs) + self.k1 = k1 + self.b = b + self.index_dir = Path(index_dir) + self.index_dir.mkdir(parents=True, exist_ok=True) + self.vocab: dict[str, int] = {} + self.inverted_index: dict[int, dict[str, int]] = {} + self.doc_meta: dict[str, DocMeta] = {} + self.total_len = 0 + + self._idf_cache: dict[int, float] = {} + self.tokenizer_name = tokenizer + self._tokenizer: BaseTokenizer | None = None + + def clear(self): + """Reset the index to empty state.""" + self.vocab = {} + self.inverted_index = {} + self.doc_meta = {} + self.total_len = 0 + self._idf_cache = {} + + @property + def n_docs(self) -> int: + """Number of indexed documents.""" + return len(self.doc_meta) + + @property + def avg_len(self) -> float: + """Average document length in tokens.""" + return self.total_len / self.n_docs if self.n_docs > 0 else 0.0 + + @property + def index_file(self) -> Path: + """Path to the index pickle file based on tokenizer name.""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not initialized. Call start() first.") + name = type(self._tokenizer).__name__.replace("Tokenizer", "").lower() + return self.index_dir / f"index_{name}.pkl" + + async def _start(self) -> None: + """Initialize tokenizer and load existing index if available.""" + if self.app_context is not None: + tokenizer_dict: dict = self.app_context.components.get(ComponentEnum.TOKENIZER, {}) + if self.tokenizer_name in tokenizer_dict: + self._tokenizer = tokenizer_dict[self.tokenizer_name] + + if self._tokenizer is None: + from ..tokenizer import RegexTokenizer + self._tokenizer = RegexTokenizer(filter_stopwords=False) + + if self._tokenizer is not None: + await self._tokenizer.start() + + if self.index_file.exists(): + await self.load() + self.logger.info(f"Loaded index from {self.index_dir}") + + async def _close(self) -> None: + """Save index and cleanup tokenizer on shutdown.""" + if self.inverted_index: + await self.dump() + self.logger.info(f"Saved index to {self.index_dir}") + + if self._tokenizer is not None: + await self._tokenizer.close() + + def _tokenize(self, text: str) -> list[str]: + if self._tokenizer is None: + raise RuntimeError("Tokenizer not initialized. Call start() first.") + return self._tokenizer.tokenize([text])[0] + + def _tokens_to_ids(self, tokens: list[str]) -> list[int]: + ids = [] + for token in tokens: + token = token.strip() + if token: + ids.append(self.vocab.setdefault(token, len(self.vocab))) + return ids + + def add_docs(self, docs_dict: dict[str, str]): + """Index or update multiple documents. + + Args: + docs_dict: Mapping of document ID to document content. + """ + for doc_id, content in docs_dict.items(): + if doc_id in self.doc_meta: + self._remove_doc(doc_id) + + tokens = self._tokenize(content) + if not tokens: + continue + + token_ids = self._tokens_to_ids(tokens) + token_counts = Counter(token_ids) + + for tid, tf in token_counts.items(): + self.inverted_index.setdefault(tid, {})[doc_id] = tf + + self.doc_meta[doc_id] = {"len": len(token_ids), "token_ids": set(token_counts)} + self.total_len += len(token_ids) + + self._idf_cache = {} + + def reindex(self): + """Rebuild vocab to remove unused tokens and compact token IDs.""" + # Collect all token IDs still in use + used_token_ids: set[int] = set() + for tid in self.inverted_index: + used_token_ids.add(tid) + + if not used_token_ids: + self.clear() + return + + # Build new vocab with compact IDs + old_to_new: dict[int, int] = {} + new_vocab: dict[str, int] = {} + for token, old_tid in self.vocab.items(): + if old_tid in used_token_ids: + new_tid = len(new_vocab) + new_vocab[token] = new_tid + old_to_new[old_tid] = new_tid + + # Rebuild inverted_index with new token IDs + new_inverted_index: dict[int, dict[str, int]] = {} + for old_tid, postings in self.inverted_index.items(): + new_tid = old_to_new[old_tid] + new_inverted_index[new_tid] = postings + + # Update doc_meta token_ids + for doc_id, meta in self.doc_meta.items(): + meta["token_ids"] = {old_to_new[old_tid] for old_tid in meta["token_ids"] if old_tid in old_to_new} + + self.vocab = new_vocab + self.inverted_index = new_inverted_index + self._idf_cache = {} + + def _remove_doc(self, doc_id: str): + if doc_id not in self.doc_meta: + return + meta = self.doc_meta[doc_id] + self.total_len -= meta["len"] + for tid in meta["token_ids"]: + if tid in self.inverted_index: + self.inverted_index[tid].pop(doc_id, None) + if not self.inverted_index[tid]: + del self.inverted_index[tid] + del self.doc_meta[doc_id] + + def _get_idf(self, token_id: int) -> float: + if token_id in self._idf_cache: + return self._idf_cache[token_id] + df = len(self.inverted_index.get(token_id, {})) + self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0 + return self._idf_cache[token_id] + + def retrieve(self, query: str, k: int = 3) -> dict[str, float]: + """Search for documents matching the query. + + Args: + query: Search query string. + k: Maximum number of results to return. + + Returns: + Dict mapping document IDs to BM25 scores, sorted by score descending. + """ + query_ids = [self.vocab[t] for t in self._tokenize(query) if t in self.vocab] + if not query_ids or self.n_docs == 0: + return {} + + scores: dict[str, float] = {} + avg_len = self.avg_len + + for tid in query_ids: + if tid not in self.inverted_index: + continue + idf = self._get_idf(tid) + for doc_id, tf in self.inverted_index[tid].items(): + doc_len = self.doc_meta[doc_id]["len"] + tf_score = tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * doc_len / avg_len)) + scores[doc_id] = scores.get(doc_id, 0.0) + idf * tf_score + + return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:k]) if scores else {} + + async def dump(self): + """Persist index to disk via pickle.""" + with open(self.index_file, "wb") as f: + pickle.dump( + { + "vocab": self.vocab, + "inverted_index": self.inverted_index, + "doc_meta": self.doc_meta, + "total_len": self.total_len, + "k1": self.k1, + "b": self.b, + }, + f, + ) + + async def load(self): + """Load index from disk. Clears index on failure.""" + try: + with open(self.index_file, "rb") as f: + data = pickle.load(f) + self.vocab = data["vocab"] + self.inverted_index = data["inverted_index"] + self.doc_meta = data["doc_meta"] + self.total_len = data.get("total_len", 0) + self.k1 = data.get("k1", 1.5) + self.b = data.get("b", 0.75) + self._idf_cache = {} + except Exception as e: + self.logger.exception(f"Failed to load index: {e}") + self.index_file.unlink(missing_ok=True) + self.clear() diff --git a/test/reme2/test_bm25_lite.py b/test/reme2/test_bm25_lite.py new file mode 100644 index 00000000..a0921b05 --- /dev/null +++ b/test/reme2/test_bm25_lite.py @@ -0,0 +1,556 @@ +"""Tests for BM25Lite search engine.""" + +import asyncio +import tempfile +from pathlib import Path + +import sys + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from reme2.component.file_store.bm25_lite import BM25Lite + + +async def create_bm25(index_dir: Path, k1: float = 1.5, b: float = 0.75) -> BM25Lite: + """Create and start a BM25Lite instance.""" + bm25 = BM25Lite(index_dir=index_dir, k1=k1, b=b) + await bm25.start() + return bm25 + + +def test_basic_init(): + """Test BM25Lite initialization.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = BM25Lite(index_dir=tmpdir) + assert bm25.k1 == 1.5 + assert bm25.b == 0.75 + assert bm25.vocab == {} + assert bm25.inverted_index == {} + assert bm25.doc_meta == {} + assert bm25.n_docs == 0 + assert bm25.avg_len == 0.0 + print("✓ test_basic_init passed") + + asyncio.run(run()) + + +def test_start_with_tokenizer(): + """Test BM25Lite starts and initializes tokenizer.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + assert bm25._tokenizer is not None + assert bm25.is_started + + await bm25.close() + assert not bm25.is_started + print("✓ test_start_with_tokenizer passed") + + asyncio.run(run()) + + +def test_add_single_doc(): + """Test adding a single document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({"doc1": "hello world"}) + + assert bm25.n_docs == 1 + assert bm25.total_len > 0 + assert "doc1" in bm25.doc_meta + + await bm25.close() + print("✓ test_add_single_doc passed") + + asyncio.run(run()) + + +def test_add_multiple_docs(): + """Test adding multiple documents.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "hello world", + "doc2": "hello python", + "doc3": "world python", + } + bm25.add_docs(docs) + + assert bm25.n_docs == 3 + assert len(bm25.vocab) > 0 + assert len(bm25.inverted_index) > 0 + + await bm25.close() + print("✓ test_add_multiple_docs passed") + + asyncio.run(run()) + + +def test_retrieve_basic(): + """Test basic retrieval functionality.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "python programming language", + "doc2": "java programming language", + "doc3": "python data analysis", + } + bm25.add_docs(docs) + + results = bm25.retrieve("python", k=3) + assert len(results) <= 3 + assert "doc1" in results or "doc3" in results + + await bm25.close() + print("✓ test_retrieve_basic passed") + + asyncio.run(run()) + + +def test_retrieve_with_limit(): + """Test retrieval with result limit.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + f"doc{i}": f"python programming {i}" for i in range(10) + } + bm25.add_docs(docs) + + results = bm25.retrieve("python", k=3) + assert len(results) == 3 + + results = bm25.retrieve("python", k=5) + assert len(results) == 5 + + await bm25.close() + print("✓ test_retrieve_with_limit passed") + + asyncio.run(run()) + + +def test_retrieve_empty_query(): + """Test retrieval with empty or unknown query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = {"doc1": "hello world"} + bm25.add_docs(docs) + + results = bm25.retrieve("", k=3) + assert results == {} + + results = bm25.retrieve("unknownxyz", k=3) + assert results == {} + + await bm25.close() + print("✓ test_retrieve_empty_query passed") + + asyncio.run(run()) + + +def test_retrieve_empty_index(): + """Test retrieval from empty index.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + results = bm25.retrieve("python", k=3) + assert results == {} + + await bm25.close() + print("✓ test_retrieve_empty_index passed") + + asyncio.run(run()) + + +def test_update_doc(): + """Test updating an existing document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({"doc1": "hello world python"}) + old_len = bm25.total_len + + bm25.add_docs({"doc1": "java"}) + assert bm25.n_docs == 1 + assert bm25.total_len != old_len + + results = bm25.retrieve("java", k=1) + assert "doc1" in results + + await bm25.close() + print("✓ test_update_doc passed") + + asyncio.run(run()) + + +def test_remove_doc(): + """Test removing a document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "hello world", + "doc2": "hello python", + } + bm25.add_docs(docs) + assert bm25.n_docs == 2 + + bm25._remove_doc("doc1") + assert bm25.n_docs == 1 + assert "doc1" not in bm25.doc_meta + + results = bm25.retrieve("hello", k=2) + assert "doc1" not in results + assert "doc2" in results + + await bm25.close() + print("✓ test_remove_doc passed") + + asyncio.run(run()) + + +def test_remove_nonexistent_doc(): + """Test removing a nonexistent document (should be no-op).""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({"doc1": "hello world"}) + bm25._remove_doc("nonexistent") + assert bm25.n_docs == 1 + + await bm25.close() + print("✓ test_remove_nonexistent_doc passed") + + asyncio.run(run()) + + +def test_clear(): + """Test clearing the index.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({ + "doc1": "hello world", + "doc2": "hello python", + }) + assert bm25.n_docs == 2 + + bm25.clear() + assert bm25.n_docs == 0 + assert bm25.vocab == {} + assert bm25.inverted_index == {} + assert bm25.doc_meta == {} + assert bm25.total_len == 0 + assert bm25._idf_cache == {} + + await bm25.close() + print("✓ test_clear passed") + + asyncio.run(run()) + + +def test_reindex(): + """Test reindex functionality to compact vocab.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({"doc1": "hello world"}) + bm25._remove_doc("doc1") + + assert bm25.n_docs == 0 + assert len(bm25.vocab) > 0 + + bm25.reindex() + assert bm25.vocab == {} + assert bm25.inverted_index == {} + + await bm25.close() + print("✓ test_reindex passed") + + asyncio.run(run()) + + +def test_reindex_with_docs(): + """Test reindex with remaining documents.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({ + "doc1": "hello world", + "doc2": "hello python", + }) + + old_vocab = bm25.vocab.copy() + bm25._remove_doc("doc1") + + bm25.reindex() + + assert bm25.n_docs == 1 + assert "doc2" in bm25.doc_meta + assert len(bm25.vocab) < len(old_vocab) + + results = bm25.retrieve("hello", k=1) + assert "doc2" in results + + await bm25.close() + print("✓ test_reindex_with_docs passed") + + asyncio.run(run()) + + +def test_persistence(): + """Test dump and load persistence.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + tmpdir_path = Path(tmpdir) + + bm25 = await create_bm25(tmpdir_path) + docs = { + "doc1": "hello world", + "doc2": "hello python", + "doc3": "programming language", + } + bm25.add_docs(docs) + + old_vocab = bm25.vocab.copy() + old_doc_meta = {k: dict(v) for k, v in bm25.doc_meta.items()} + + await bm25.dump() + await bm25.close() + + bm25_new = await create_bm25(tmpdir_path) + + assert bm25_new.vocab == old_vocab + assert bm25_new.n_docs == 3 + for doc_id, meta in old_doc_meta.items(): + assert doc_id in bm25_new.doc_meta + + results = bm25_new.retrieve("hello", k=2) + assert "doc1" in results or "doc2" in results + + await bm25_new.close() + print("✓ test_persistence passed") + + asyncio.run(run()) + + +def test_custom_params(): + """Test custom k1 and b parameters.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir), k1=2.0, b=0.5) + + assert bm25.k1 == 2.0 + assert bm25.b == 0.5 + + bm25.add_docs({"doc1": "test document"}) + results = bm25.retrieve("test", k=1) + assert "doc1" in results + + await bm25.close() + print("✓ test_custom_params passed") + + asyncio.run(run()) + + +def test_chinese_text(): + """Test with Chinese text.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "我爱北京天安门", + "doc2": "北京是中国的首都", + "doc3": "上海的天气很好", + } + bm25.add_docs(docs) + + results = bm25.retrieve("北京", k=2) + assert len(results) <= 2 + assert "doc1" in results or "doc2" in results + + await bm25.close() + print("✓ test_chinese_text passed") + + asyncio.run(run()) + + +def test_mixed_chinese_english(): + """Test with mixed Chinese and English text.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "Python 是一种编程语言", + "doc2": "Java 编程语言", + "doc3": "Python 数据分析", + } + bm25.add_docs(docs) + + results = bm25.retrieve("Python", k=3) + assert len(results) > 0 + + results = bm25.retrieve("编程", k=2) + assert len(results) > 0 + + await bm25.close() + print("✓ test_mixed_chinese_english passed") + + asyncio.run(run()) + + +def test_idf_cache(): + """Test IDF cache functionality.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({ + "doc1": "hello world", + "doc2": "hello python", + }) + + token = "hello" + if token in bm25.vocab: + tid = bm25.vocab[token] + idf1 = bm25._get_idf(tid) + assert tid in bm25._idf_cache + idf2 = bm25._get_idf(tid) + assert idf1 == idf2 + + await bm25.close() + print("✓ test_idf_cache passed") + + asyncio.run(run()) + + +def test_avg_len(): + """Test average document length calculation.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + assert bm25.avg_len == 0.0 + + bm25.add_docs({"doc1": "hello world python"}) + assert bm25.avg_len > 0 + + bm25.add_docs({"doc2": "test"}) + new_avg = bm25.avg_len + assert new_avg > 0 + + await bm25.close() + print("✓ test_avg_len passed") + + asyncio.run(run()) + + +def test_score_ordering(): + """Test that results are ordered by score descending.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + docs = { + "doc1": "python python python", + "doc2": "python python", + "doc3": "python", + } + bm25.add_docs(docs) + + results = bm25.retrieve("python", k=3) + scores = list(results.values()) + + for i in range(len(scores) - 1): + assert scores[i] >= scores[i + 1] + + await bm25.close() + print("✓ test_score_ordering passed") + + asyncio.run(run()) + + +def test_empty_doc(): + """Test adding empty document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir: + bm25 = await create_bm25(Path(tmpdir)) + + bm25.add_docs({"doc1": ""}) + assert bm25.n_docs == 0 + + bm25.add_docs({"doc2": " "}) + assert bm25.n_docs == 0 + + await bm25.close() + print("✓ test_empty_doc passed") + + asyncio.run(run()) + + +if __name__ == "__main__": + print("\n=== BM25Lite Tests ===") + test_basic_init() + test_start_with_tokenizer() + test_add_single_doc() + test_add_multiple_docs() + test_retrieve_basic() + test_retrieve_with_limit() + test_retrieve_empty_query() + test_retrieve_empty_index() + test_update_doc() + test_remove_doc() + test_remove_nonexistent_doc() + test_clear() + test_reindex() + test_reindex_with_docs() + test_persistence() + test_custom_params() + test_chinese_text() + test_mixed_chinese_english() + test_idf_cache() + test_avg_len() + test_score_ordering() + test_empty_doc() + print("\n所有测试通过!") \ No newline at end of file