Harshavard21/FinRAG

1"""2src/vectorstore/bm25_index.py3==============================4BM25 sparse index for keyword-based retrieval.5 6Why BM25 alongside dense vectors?7    Dense vectors capture SEMANTIC meaning ("revenue growth" ≈ "sales increase")8    BM25 captures EXACT TERM matching ("₹4,234 crore" or "NPA 2.3%")9 10Financial documents are full of exact numbers and acronyms that11dense retrieval may miss. BM25 ensures we always catch exact matches.12 13This index is combined with Qdrant dense search via RRF in hybrid_retriever.py.14 15Storage:16    - BM25 index serialized to ./data/bm25_index/bm25.pkl17    - Document list serialized to ./data/bm25_index/docs.json18    - Both files must be regenerated when documents change19"""20 21from __future__ import annotations22 23import json24import pickle25import re26from dataclasses import dataclass27from pathlib import Path28from typing import Optional29 30from rank_bm25 import BM25Okapi31 32from config.settings import settings33from src.utils.logger import logger34from src.chunking.hierarchical_chunker import Chunk35 36 37@dataclass38class BM25Result:39    """A single result from BM25 keyword search."""40    chunk_id: str41    content: str42    bm25_score: float43    company: str44    ticker: str45    source_file: str46    fiscal_year: str47    page_number: int48    section: str49    subsection: str50    content_type: str51    chunk_level: str52    parent_chunk_id: Optional[str] = None53 54 55class BM25Index:56    """57    BM25 sparse index over all indexed chunks.58 59    Build once → persist to disk → reload on startup.60    Supports company/year filtering (post-filter after search).61    """62 63    INDEX_FILE = "bm25.pkl"64    DOCS_FILE = "docs.json"65 66    def __init__(self):67        self.index_dir = settings.bm25_index_path68        self.index_dir.mkdir(parents=True, exist_ok=True)69        self._bm25: Optional[BM25Okapi] = None70        self._docs: list[dict] = []   # stores chunk metadata + content71 72        # Try to load existing index73        if self._index_exists():74            self._load()75 76    # ---------------------------------------------------------------- #77    # Build78    # ---------------------------------------------------------------- #79 80    def build(self, chunks: list[Chunk]) -> None:81        """82        Build BM25 index from a list of chunks.83        Indexes the same chunks that go into Qdrant (child + atomic).84 85        Args:86            chunks: List of Chunk objects to index87        """88        logger.info(f"Building BM25 index over {len(chunks)} chunks...")89 90        self._docs = []91        tokenized_corpus = []92 93        for chunk in chunks:94            tokens = self._tokenize(chunk.content)95            tokenized_corpus.append(tokens)96 97            # Store chunk metadata for result reconstruction98            self._docs.append({99                "chunk_id": chunk.chunk_id,100                "content": chunk.content,101                "company": chunk.company,102                "ticker": chunk.ticker,103                "source_file": chunk.source_file,104                "fiscal_year": chunk.fiscal_year,105                "page_number": chunk.page_number,106                "section": chunk.section,107                "subsection": chunk.subsection,108                "content_type": chunk.content_type,109                "chunk_level": chunk.chunk_level,110                "parent_chunk_id": chunk.parent_chunk_id or "",111            })112 113        self._bm25 = BM25Okapi(tokenized_corpus)114        self._save()115        logger.info(f"BM25 index built and saved | {len(chunks)} documents")116 117    # ---------------------------------------------------------------- #118    # Search119    # ---------------------------------------------------------------- #120 121    def search(122        self,123        query: str,124        top_k: int = 30,125        company_filter: Optional[str | list[str]] = None,126        fiscal_year_filter: Optional[str] = None,127    ) -> list[BM25Result]:128        """129        BM25 keyword search with optional post-filtering.130 131        Args:132            query: Search query string133            top_k: Number of results to return134            company_filter: Filter by company name or list of companies135            fiscal_year_filter: Filter by fiscal year136 137        Returns:138            List of BM25Result sorted by score descending139        """140        if self._bm25 is None:141            logger.warning("BM25 index not loaded — call build() first")142            return []143 144        query_tokens = self._tokenize(query)145        scores = self._bm25.get_scores(query_tokens)146 147        # Pair scores with doc metadata and sort148        scored_docs = sorted(149            enumerate(scores),150            key=lambda x: x[1],151            reverse=True,152        )153 154        results = []155        for idx, score in scored_docs:156            if score <= 0:157                continue158            if len(results) >= top_k * 3:   # over-retrieve for post-filtering159                break160 161            doc = self._docs[idx]162 163            # Post-filter by company164            if company_filter:165                if isinstance(company_filter, list):166                    if doc["company"] not in company_filter:167                        continue168                else:169                    if doc["company"] != company_filter:170                        continue171 172            # Post-filter by fiscal year173            if fiscal_year_filter and doc["fiscal_year"] != fiscal_year_filter:174                continue175 176            results.append(BM25Result(177                chunk_id=doc["chunk_id"],178                content=doc["content"],179                bm25_score=float(score),180                company=doc["company"],181                ticker=doc["ticker"],182                source_file=doc["source_file"],183                fiscal_year=doc["fiscal_year"],184                page_number=doc["page_number"],185                section=doc["section"],186                subsection=doc["subsection"],187                content_type=doc["content_type"],188                chunk_level=doc["chunk_level"],189                parent_chunk_id=doc["parent_chunk_id"] or None,190            ))191 192            if len(results) >= top_k:193                break194 195        return results196 197    # ---------------------------------------------------------------- #198    # Tokenizer199    # ---------------------------------------------------------------- #200 201    def _tokenize(self, text: str) -> list[str]:202        """203        Tokenize text for BM25.204 205        Financial document specific:206        - Preserve numbers (4234, 2.3%) as tokens207        - Preserve ₹ symbol208        - Lowercase everything209        - Remove punctuation except . % ₹ (financial symbols)210        - Preserve common financial acronyms as single tokens211        """212        text = text.lower()213 214        # Keep numbers with decimals and % intact215        text = re.sub(r'[^\w\s₹.%,/-]', ' ', text)216 217        # Split on whitespace218        tokens = text.split()219 220        # Remove very short tokens (< 2 chars) except numbers221        tokens = [t for t in tokens if len(t) >= 2 or t.isdigit()]222 223        # Remove pure punctuation tokens224        tokens = [t for t in tokens if re.search(r'[a-zA-Z0-9]', t)]225 226        return tokens227 228    # ---------------------------------------------------------------- #229    # Persistence230    # ---------------------------------------------------------------- #231 232    def _index_exists(self) -> bool:233        return (234            (self.index_dir / self.INDEX_FILE).exists() and235            (self.index_dir / self.DOCS_FILE).exists()236        )237 238    def _save(self) -> None:239        """Persist index to disk."""240        with open(self.index_dir / self.INDEX_FILE, "wb") as f:241            pickle.dump(self._bm25, f)242        with open(self.index_dir / self.DOCS_FILE, "w", encoding="utf-8") as f:243            json.dump(self._docs, f, ensure_ascii=False)244        logger.debug("BM25 index saved to disk")245 246    def _load(self) -> None:247        """Load index from disk."""248        try:249            with open(self.index_dir / self.INDEX_FILE, "rb") as f:250                self._bm25 = pickle.load(f)251            with open(self.index_dir / self.DOCS_FILE, "r", encoding="utf-8") as f:252                self._docs = json.load(f)253            logger.info(f"BM25 index loaded | {len(self._docs)} documents")254        except Exception as e:255            logger.warning(f"Failed to load BM25 index: {e} — will rebuild on next ingest")256            self._bm25 = None257            self._docs = []258