Spaces:
Running on Zero
Running on Zero
| """BM25 index storage and retrieval.""" | |
| import hashlib | |
| import json | |
| import logging | |
| import pickle | |
| import re | |
| from pathlib import Path | |
| from typing import Any | |
| from langchain_core.documents import Document | |
| from rank_bm25 import BM25Okapi | |
| logger = logging.getLogger(__name__) | |
| class BM25Store: | |
| """Storage and retrieval for BM25 index. | |
| This class manages a BM25 index for lexical/keyword search. | |
| It handles tokenization, index building, persistence, and retrieval. | |
| Attributes: | |
| persist_path: Directory to store the index | |
| tokenizer: Tokenization method ('simple' or 'nltk') | |
| """ | |
| def __init__( | |
| self, | |
| persist_path: str = "./bm25_index", | |
| tokenizer: str = "simple", | |
| ): | |
| """Initialize BM25 store. | |
| Args: | |
| persist_path: Directory to store the index files | |
| tokenizer: Tokenization method ('simple' or 'nltk') | |
| """ | |
| self.persist_path = Path(persist_path) | |
| self.tokenizer_type = tokenizer | |
| self.bm25: BM25Okapi | None = None | |
| self.documents: list[Document] = [] | |
| self.tokenized_corpus: list[list[str]] = [] | |
| self._doc_hash: str | None = None | |
| def _simple_tokenize(self, text: str) -> list[str]: | |
| """Simple whitespace and punctuation-based tokenization. | |
| Args: | |
| text: Text to tokenize | |
| Returns: | |
| List of tokens | |
| """ | |
| # Lowercase and split on non-alphanumeric characters | |
| text = text.lower() | |
| tokens = re.findall(r"\b\w+\b", text) | |
| return tokens | |
| def _tokenize(self, text: str) -> list[str]: | |
| """Tokenize text using configured tokenizer. | |
| Args: | |
| text: Text to tokenize | |
| Returns: | |
| List of tokens | |
| """ | |
| if self.tokenizer_type == "nltk": | |
| try: | |
| import nltk | |
| # Ensure punkt tokenizer is available | |
| try: | |
| nltk.data.find("tokenizers/punkt") | |
| except LookupError: | |
| nltk.download("punkt", quiet=True) | |
| tokens = nltk.word_tokenize(text.lower()) | |
| return [t for t in tokens if t.isalnum()] | |
| except ImportError: | |
| logger.warning("NLTK not installed, falling back to simple tokenizer") | |
| return self._simple_tokenize(text) | |
| else: | |
| return self._simple_tokenize(text) | |
| def _compute_hash(self, documents: list[Document]) -> str: | |
| """Compute hash of document contents for change detection. | |
| Args: | |
| documents: List of documents | |
| Returns: | |
| Hash string | |
| """ | |
| content = "".join(doc.page_content for doc in documents) | |
| return hashlib.sha256(content.encode()).hexdigest()[:16] | |
| def build_index(self, documents: list[Document]) -> None: | |
| """Build BM25 index from documents. | |
| Args: | |
| documents: List of Document objects to index | |
| """ | |
| if not documents: | |
| raise ValueError("No documents provided for indexing") | |
| logger.info(f"Building BM25 index with {len(documents)} documents...") | |
| # Store documents | |
| self.documents = documents | |
| # Tokenize all documents | |
| self.tokenized_corpus = [self._tokenize(doc.page_content) for doc in documents] | |
| # Build BM25 index | |
| self.bm25 = BM25Okapi(self.tokenized_corpus) | |
| # Compute hash for change detection | |
| self._doc_hash = self._compute_hash(documents) | |
| logger.info(f"BM25 index built successfully with {len(documents)} documents") | |
| def save(self) -> None: | |
| """Save BM25 index to disk.""" | |
| if self.bm25 is None: | |
| raise ValueError("No index to save. Build index first.") | |
| # Create directory if needed | |
| self.persist_path.mkdir(parents=True, exist_ok=True) | |
| # Save BM25 index | |
| index_path = self.persist_path / "bm25_index.pkl" | |
| with open(index_path, "wb") as f: | |
| pickle.dump(self.bm25, f) | |
| # Save documents | |
| docs_path = self.persist_path / "documents.pkl" | |
| with open(docs_path, "wb") as f: | |
| pickle.dump(self.documents, f) | |
| # Save tokenized corpus (for rebuilding if needed) | |
| corpus_path = self.persist_path / "tokenized_corpus.pkl" | |
| with open(corpus_path, "wb") as f: | |
| pickle.dump(self.tokenized_corpus, f) | |
| # Save metadata | |
| metadata = { | |
| "doc_hash": self._doc_hash, | |
| "num_documents": len(self.documents), | |
| "tokenizer": self.tokenizer_type, | |
| } | |
| metadata_path = self.persist_path / "metadata.json" | |
| with open(metadata_path, "w") as f: | |
| json.dump(metadata, f, indent=2) | |
| logger.info(f"BM25 index saved to {self.persist_path}") | |
| def load(self) -> bool: | |
| """Load BM25 index from disk. | |
| Returns: | |
| True if loaded successfully, False otherwise | |
| """ | |
| index_path = self.persist_path / "bm25_index.pkl" | |
| docs_path = self.persist_path / "documents.pkl" | |
| corpus_path = self.persist_path / "tokenized_corpus.pkl" | |
| metadata_path = self.persist_path / "metadata.json" | |
| # Check if all required files exist | |
| if not all(p.exists() for p in [index_path, docs_path, corpus_path, metadata_path]): | |
| logger.warning(f"BM25 index not found at {self.persist_path}") | |
| return False | |
| try: | |
| # Load BM25 index | |
| with open(index_path, "rb") as f: | |
| self.bm25 = pickle.load(f) | |
| # Load documents | |
| with open(docs_path, "rb") as f: | |
| self.documents = pickle.load(f) | |
| # Load tokenized corpus | |
| with open(corpus_path, "rb") as f: | |
| self.tokenized_corpus = pickle.load(f) | |
| # Load metadata | |
| with open(metadata_path) as f: | |
| metadata = json.load(f) | |
| self._doc_hash = metadata.get("doc_hash") | |
| logger.info( | |
| f"BM25 index loaded from {self.persist_path} ({len(self.documents)} documents)" | |
| ) | |
| return True | |
| except Exception as e: | |
| logger.error(f"Error loading BM25 index: {e}") | |
| return False | |
| def search(self, query: str, k: int = 4) -> list[Document]: | |
| """Search the BM25 index. | |
| Args: | |
| query: Search query | |
| k: Number of documents to return | |
| Returns: | |
| List of relevant documents | |
| """ | |
| if self.bm25 is None: | |
| raise ValueError("BM25 index not built or loaded") | |
| # Tokenize query | |
| tokenized_query = self._tokenize(query) | |
| if not tokenized_query: | |
| logger.warning("Query tokenization produced no tokens") | |
| return [] | |
| # Get BM25 scores | |
| scores = self.bm25.get_scores(tokenized_query) | |
| # Get top k indices | |
| top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k] | |
| # Return documents | |
| results = [self.documents[i] for i in top_indices if scores[i] > 0] | |
| logger.debug(f"BM25 search returned {len(results)} documents for query: {query[:50]}...") | |
| return results | |
| def search_with_scores(self, query: str, k: int = 4) -> list[tuple[Document, float]]: | |
| """Search the BM25 index and return scores. | |
| Args: | |
| query: Search query | |
| k: Number of documents to return | |
| Returns: | |
| List of (document, score) tuples | |
| """ | |
| if self.bm25 is None: | |
| raise ValueError("BM25 index not built or loaded") | |
| # Tokenize query | |
| tokenized_query = self._tokenize(query) | |
| if not tokenized_query: | |
| return [] | |
| # Get BM25 scores | |
| scores = self.bm25.get_scores(tokenized_query) | |
| # Get top k indices with scores | |
| top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k] | |
| # Return documents with scores | |
| return [(self.documents[i], scores[i]) for i in top_indices if scores[i] > 0] | |
| def is_built(self) -> bool: | |
| """Check if index is built. | |
| Returns: | |
| True if index is ready | |
| """ | |
| return self.bm25 is not None and len(self.documents) > 0 | |
| def get_stats(self) -> dict[str, Any]: | |
| """Get index statistics. | |
| Returns: | |
| Dictionary with index statistics | |
| """ | |
| return { | |
| "num_documents": len(self.documents), | |
| "tokenizer": self.tokenizer_type, | |
| "persist_path": str(self.persist_path), | |
| "doc_hash": self._doc_hash, | |
| "is_built": self.is_built(), | |
| } | |
| def delete(self) -> None: | |
| """Delete the persisted index.""" | |
| import shutil | |
| if self.persist_path.exists(): | |
| shutil.rmtree(self.persist_path) | |
| logger.info(f"Deleted BM25 index at {self.persist_path}") | |
| self.bm25 = None | |
| self.documents = [] | |
| self.tokenized_corpus = [] | |
| self._doc_hash = None | |