retriever.py

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import List

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
_CHUNK_SIZE = 400
_CHUNK_OVERLAP = 80


@dataclass
class Chunk:
    text: str
    source: str
    index: int


@dataclass
class HybridRetriever:
    chunks: List[Chunk] = field(default_factory=list)
    _bm25: BM25Okapi | None = None
    _embedder: SentenceTransformer | None = None
    _dense_matrix: np.ndarray | None = None

    # ------------------------------------------------------------------
    def _tokenize(self, text: str) -> List[str]:
        return re.findall(r"\w+", text.lower())

    def _split(self, text: str, source: str) -> List[Chunk]:
        words = text.split()
        results: List[Chunk] = []
        start = 0
        idx = 0
        while start < len(words):
            chunk_words = words[start : start + _CHUNK_SIZE]
            results.append(Chunk(" ".join(chunk_words), source, idx))
            start += _CHUNK_SIZE - _CHUNK_OVERLAP
            idx += 1
        return results

    # ------------------------------------------------------------------
    def add_documents(self, texts: List[tuple[str, str]]) -> None:
        """texts: list of (content, filename)"""
        self.chunks = []
        for content, name in texts:
            self.chunks.extend(self._split(content, name))
        self._build_index()

    def _build_index(self) -> None:
        corpus = [self._tokenize(c.text) for c in self.chunks]
        self._bm25 = BM25Okapi(corpus)
        if self._embedder is None:
            self._embedder = SentenceTransformer(_EMBED_MODEL)
        self._dense_matrix = self._embedder.encode(
            [c.text for c in self.chunks], show_progress_bar=False, normalize_embeddings=True
        )

    # ------------------------------------------------------------------
    def retrieve(self, query: str, top_k: int = 5, alpha: float = 0.5) -> List[Chunk]:
        """Hybrid BM25 + dense retrieval with linear interpolation."""
        if not self.chunks:
            return []

        tokens = self._tokenize(query)
        bm25_scores = np.array(self._bm25.get_scores(tokens))
        bm25_scores = (bm25_scores - bm25_scores.min()) / (bm25_scores.max() - bm25_scores.min() + 1e-9)

        q_emb = self._embedder.encode([query], normalize_embeddings=True)[0]
        dense_scores = self._dense_matrix @ q_emb
        dense_scores = (dense_scores - dense_scores.min()) / (dense_scores.max() - dense_scores.min() + 1e-9)

        hybrid = alpha * bm25_scores + (1 - alpha) * dense_scores
        top_idx = np.argsort(hybrid)[::-1][:top_k]
        return [self.chunks[i] for i in top_idx]

    @property
    def ready(self) -> bool:
        return bool(self.chunks) and self._bm25 is not None