docsif 0.1.0__py3-none-any.whl

sif/search/rerank.py

@@ -0,0 +1,390 @@
+"""Result reranking using cross-encoders and LLM-based rerankers."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from sif.core.models import SearchResult
+from sif.utils.logging import get_logger, is_quiet, suppress_output
+
+
+if TYPE_CHECKING:
+    from sentence_transformers import CrossEncoder  # noqa: F401
+
+logger = get_logger(__name__)
+
+
+def _sort_and_build_results(
+    results: list[SearchResult],
+    scores: list[float],
+    top_k: int,
+) -> list[SearchResult]:
+    """Sort results by score descending and rebuild SearchResult objects."""
+    scored = list(zip(results, scores))
+    scored.sort(key=lambda x: x[1], reverse=True)
+
+    reranked: list[SearchResult] = []
+    for rank, (result, score) in enumerate(scored[:top_k], 1):
+        new_scores = dict(result.scores) if result.scores else {}
+        new_scores["reranker_score"] = score
+
+        reranked.append(
+            SearchResult(
+                document_id=result.document_id,
+                title=result.title,
+                path=result.path,
+                collection_name=result.collection_name,
+                score=score,
+                content=result.content,
+                highlights=result.highlights,
+                rank=rank,
+                scores=new_scores,
+                snippet=result.snippet,
+                context_description=result.context_description,
+            )
+        )
+    return reranked
+
+
+class LlamaCppReranker:
+    """Rerank search results using a GGUF cross-encoder via llama-cpp-python."""
+
+    def __init__(
+        self,
+        model_path: str | None = None,
+        model_name: str = "",
+        n_ctx: int = 512,
+        n_threads: int | None = None,
+        batch_size: int = 32,
+    ) -> None:
+        self._model_path = model_path
+        self._model_name = model_name
+        self._n_ctx = n_ctx
+        self._n_threads = n_threads
+        self._batch_size = batch_size
+        self._model = None
+
+    def load(self) -> None:
+        """Load the GGUF model."""
+        if self._model is not None:
+            return
+        from llama_cpp import Llama  # noqa: PLC0415
+
+        path = self._model_path or self._model_name
+        if not path:
+            raise ValueError("No model_path or model_name provided for LlamaCppReranker")
+        logger.info(f"Loading GGUF reranker model: {path}")
+        self._model = Llama(
+            model_path=str(path),
+            n_ctx=self._n_ctx,
+            n_threads=self._n_threads,
+            embedding=True,
+            verbose=False,
+        )
+        logger.info("GGUF reranker model loaded successfully")
+
+    def rerank(
+        self,
+        query: str,
+        results: list[SearchResult],
+        top_k: int = 10,
+    ) -> list[SearchResult]:
+        """Rerank search results."""
+        if not results:
+            return []
+
+        if self._model is None:
+            self.load()
+
+        import numpy as np  # noqa: PLC0415
+
+        # Embed query once
+        query_emb = self._model.embed(query)
+        query_vec = np.array(query_emb)
+        query_norm = np.linalg.norm(query_vec)
+
+        # Score each result by embedding similarity
+        scored = []
+        for result in results:
+            text = result.content or ""
+            if not text and result.highlights:
+                text = result.highlights[0]
+            if result.title:
+                text = f"{result.title}\n\n{text}" if text else result.title
+            doc_emb = self._model.embed(text)
+            doc_vec = np.array(doc_emb)
+            doc_norm = np.linalg.norm(doc_vec)
+            sim = np.dot(query_vec, doc_vec) / (query_norm * doc_norm + 1e-8)
+            scored.append((result, float(sim)))
+
+        # Normalize to [0, 1] via sigmoid
+        normalized = [1.0 / (1.0 + math.exp(-s)) for _, s in scored]
+
+        return _sort_and_build_results([r for r, _ in scored], normalized, top_k)
+
+
+class CrossEncoderReranker:
+    """Rerank search results using sentence-transformers CrossEncoder."""
+
+    def __init__(
+        self,
+        model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
+        model_path: str | None = None,
+        device: str | None = None,
+        batch_size: int = 32,
+        cache_dir: str | None = None,
+    ) -> None:
+        self._model_name = model_name
+        self._model_path = model_path
+        self._device = device
+        self._batch_size = batch_size
+        self._cache_dir = cache_dir
+        self._model = None
+
+    def load(self) -> None:
+        """Load the cross-encoder model."""
+        if self._model is not None:
+            return
+        from sentence_transformers import CrossEncoder  # noqa: PLC0415
+
+        model_id = self._model_path or self._model_name
+
+        # If model_id looks like a ModelScope ID (contains "/" and not a local path),
+        # download from ModelScope first
+        local_path = model_id
+        if "/" in model_id and not Path(model_id).exists():
+            try:
+                from sif.models.download import ModelDownloader  # noqa: PLC0415
+
+                downloader = ModelDownloader(self._cache_dir)
+                downloaded = downloader.download(model_id)
+                # snapshot_download may return a path with subdirectories
+                subdirs = [d for d in downloaded.iterdir() if d.is_dir()]
+                local_path = str(subdirs[0]) if subdirs else str(downloaded)
+                logger.info(f"Loading ModelScope reranker model: {model_id}")
+            except Exception:
+                logger.info(f"Loading ST reranker model: {model_id}")
+        else:
+            logger.info(f"Loading ST reranker model: {model_id}")
+
+        if is_quiet():
+            with suppress_output():
+                self._model = CrossEncoder(local_path, device=self._device, max_length=512)
+        else:
+            self._model = CrossEncoder(local_path, device=self._device, max_length=512)
+        logger.info("Reranker model loaded successfully")
+
+    def rerank(
+        self,
+        query: str,
+        results: list[SearchResult],
+        top_k: int = 10,
+    ) -> list[SearchResult]:
+        """Rerank search results."""
+        if not results:
+            return []
+
+        if self._model is None:
+            self.load()
+
+        # Build query-document pairs
+        pairs: list[tuple[str, str]] = []
+        for result in results:
+            text = result.content or ""
+            if not text and result.highlights:
+                text = result.highlights[0]
+            if result.title:
+                text = f"{result.title}\n\n{text}" if text else result.title
+            pairs.append((query, text))
+
+        # Score all pairs
+        raw_scores = self._model.predict(pairs, batch_size=self._batch_size)  # type: ignore[union-attr]
+
+        # Normalize scores to [0, 1] using sigmoid
+        normalized = [1.0 / (1.0 + math.exp(-float(s))) for s in raw_scores]
+
+        return _sort_and_build_results(results, normalized, top_k)
+
+
+class Qwen3Reranker:
+    """Reranker for Qwen3-Reranker models using AutoModelForCausalLM.
+
+    Qwen3-Reranker models are Causal LMs that predict "yes"/"no" tokens
+    to judge query-document relevance. See:
+    https://huggingface.co/Qwen/Qwen3-Reranker-0.6B
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        model_name: str = "Qwen/Qwen3-Reranker-0.6B",
+        model_path: str | None = None,
+        device: str | None = None,
+        batch_size: int = 8,
+        cache_dir: str | None = None,
+        max_length: int = 8192,
+    ) -> None:
+        self._model_name = model_name
+        self._model_path = model_path
+        self._device = device
+        self._batch_size = batch_size
+        self._cache_dir = cache_dir
+        self._max_length = max_length
+        self._model = None
+        self._tokenizer = None
+        self._token_true_id: int | None = None
+        self._token_false_id: int | None = None
+
+    def load(self) -> None:
+        """Load the Qwen3 reranker model."""
+        if self._model is not None:
+            return
+
+        import torch  # noqa: PLC0415
+        from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415
+
+        model_id = self._model_path or self._model_name
+        local_path = model_id
+
+        # Download from ModelScope if needed
+        if "/" in model_id and not Path(model_id).exists():
+            from sif.models.download import ModelDownloader  # noqa: PLC0415
+
+            downloader = ModelDownloader(self._cache_dir)
+            downloaded = downloader.download(model_id)
+            subdirs = [d for d in downloaded.iterdir() if d.is_dir()]
+            local_path = str(subdirs[0]) if subdirs else str(downloaded)
+
+        logger.info(f"Loading Qwen3 reranker model: {model_id}")
+
+        if is_quiet():
+            with suppress_output():
+                self._tokenizer = AutoTokenizer.from_pretrained(local_path, padding_side="left")
+                self._model = AutoModelForCausalLM.from_pretrained(
+                    local_path,
+                    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+                ).eval()
+
+                if self._device:
+                    self._model = self._model.to(self._device)
+                else:
+                    device = "cuda" if torch.cuda.is_available() else "cpu"
+                    self._model = self._model.to(device)
+
+                self._token_true_id = self._tokenizer.convert_tokens_to_ids("yes")
+                self._token_false_id = self._tokenizer.convert_tokens_to_ids("no")
+        else:
+            self._tokenizer = AutoTokenizer.from_pretrained(local_path, padding_side="left")
+            self._model = AutoModelForCausalLM.from_pretrained(
+                local_path,
+                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            ).eval()
+
+            if self._device:
+                self._model = self._model.to(self._device)
+            else:
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                self._model = self._model.to(device)
+
+            self._token_true_id = self._tokenizer.convert_tokens_to_ids("yes")
+            self._token_false_id = self._tokenizer.convert_tokens_to_ids("no")
+        logger.info("Qwen3 reranker model loaded successfully")
+
+    @staticmethod
+    def _format_prompt(instruction: str | None, query: str, doc: str) -> str:
+        """Format query-document pair into Qwen3 reranker prompt."""
+        if instruction is None:
+            instruction = (
+                "Given a web search query, retrieve relevant passages that answer the query"
+            )
+        return f"<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}"
+
+    def rerank(
+        self,
+        query: str,
+        results: list[SearchResult],
+        top_k: int = 10,
+    ) -> list[SearchResult]:
+        """Rerank search results using Qwen3 Causal LM."""
+        if not results:
+            return []
+
+        if self._model is None:
+            self.load()
+
+        import torch  # noqa: PLC0415
+
+        # Build prompts
+        texts: list[str] = []
+        for result in results:
+            text = result.content or ""
+            if not text and result.highlights:
+                text = result.highlights[0]
+            if result.title:
+                text = f"{result.title}\n\n{text}" if text else result.title
+            texts.append(self._format_prompt(None, query, text))
+
+        # Score in batches
+        all_scores: list[float] = []
+        for i in range(0, len(texts), self._batch_size):
+            batch = texts[i : i + self._batch_size]
+            inputs = self._tokenizer(
+                batch,
+                padding=True,
+                truncation=True,
+                max_length=self._max_length,
+                return_tensors="pt",
+            )
+            inputs = {k: v.to(self._model.device) for k, v in inputs.items()}  # type: ignore[union-attr]
+
+            with torch.no_grad():
+                outputs = self._model(**inputs)  # type: ignore[operator]
+
+            # Last token logits -> yes/no probabilities
+            logits = outputs.logits[:, -1, :]  # type: ignore[union-attr]
+            true_scores = logits[:, self._token_true_id]
+            false_scores = logits[:, self._token_false_id]
+            batch_logits = torch.stack([false_scores, true_scores], dim=1)
+            batch_probs = torch.nn.functional.softmax(batch_logits, dim=1)
+            scores = batch_probs[:, 1].tolist()
+            all_scores.extend(scores)
+
+        return _sort_and_build_results(results, all_scores, top_k)
+
+
+def create_reranker(
+    settings,
+) -> LlamaCppReranker | CrossEncoderReranker | Qwen3Reranker:
+    """Factory: create a reranker based on settings."""
+    model_type = getattr(settings, "reranker_model_type", "gguf")
+    model_name = getattr(settings, "reranker_model_name", "")
+    cache_dir = str(settings.get_cache_dir()) if hasattr(settings, "get_cache_dir") else None
+
+    # Auto-detect Qwen3 reranker models
+    if "qwen3-reranker" in model_name.lower() or "qwen3_reranker" in model_name.lower():
+        return Qwen3Reranker(
+            model_name=settings.reranker_model_name,
+            model_path=str(settings.reranker_model_path) if settings.reranker_model_path else None,
+            batch_size=settings.reranker_batch_size,
+            cache_dir=cache_dir,
+        )
+
+    if model_type == "gguf":
+        return LlamaCppReranker(
+            model_path=str(settings.reranker_model_path) if settings.reranker_model_path else None,
+            model_name=settings.reranker_model_name,
+            batch_size=settings.reranker_batch_size,
+        )
+    if model_type == "sentence_transformers":
+        return CrossEncoderReranker(
+            model_name=settings.reranker_model_name,
+            model_path=str(settings.reranker_model_path) if settings.reranker_model_path else None,
+            batch_size=settings.reranker_batch_size,
+            cache_dir=cache_dir,
+        )
+    raise ValueError(f"Unknown reranker model_type: {model_type}")
+
+
+# Backwards-compatible alias
+Reranker = LlamaCppReranker

sif/search/rrf.py

@@ -0,0 +1,186 @@
+"""Reciprocal Rank Fusion (RRF) implementation."""
+
+from __future__ import annotations
+
+from sif.core.models import SearchResult
+
+
+class RRFFusion:
+    """Reciprocal Rank Fusion for combining search results."""
+
+    def __init__(self, k: int = 60) -> None:
+        """Initialize RRF with k parameter.
+
+        Args:
+            k: Constant to prevent division by zero and dampen
+               the impact of low rankings. Default is 60.
+        """
+        self.k = k
+
+    def fuse(
+        self,
+        results_lists: list[list[SearchResult]],
+        limit: int = 10,
+    ) -> list[SearchResult]:
+        """Fuse multiple ranked lists using RRF.
+
+        Args:
+            results_lists: List of ranked result lists
+            limit: Maximum number of results to return
+
+        Returns:
+            Fused and re-ranked list of results
+        """
+        # Map document_id to scores
+        scores: dict[str, tuple[float, SearchResult]] = {}
+
+        for list_idx, results in enumerate(results_lists):
+            for rank, result in enumerate(results, 1):
+                doc_id = result.document_id
+
+                # Calculate RRF score: 1 / (k + rank)
+                rrf_score = 1.0 / (self.k + rank)
+
+                # Build scores dict
+                doc_scores: dict[str, float | None] = {}
+                if result.scores:
+                    doc_scores.update(result.scores)
+
+                # Track original score from this list
+                score_key = "bm25_score" if list_idx == 0 else "vector_score"
+                if score_key not in doc_scores:
+                    doc_scores[score_key] = result.score
+
+                if doc_id in scores:
+                    existing_score, existing_result = scores[doc_id]
+                    # Merge scores from existing result
+                    merged_scores: dict[str, float | None] = {}
+                    if existing_result.scores:
+                        merged_scores.update(existing_result.scores)
+                    merged_scores.update(doc_scores)
+                    scores[doc_id] = (existing_score + rrf_score, existing_result)
+                    # Update the stored result's scores
+                    existing_result.scores = merged_scores
+                else:
+                    new_result = SearchResult(
+                        document_id=result.document_id,
+                        title=result.title,
+                        path=result.path,
+                        collection_name=result.collection_name,
+                        score=rrf_score,
+                        content=result.content,
+                        highlights=result.highlights,
+                        rank=rank,
+                        scores=doc_scores,
+                        snippet=result.snippet,
+                        context_description=result.context_description,
+                    )
+                    scores[doc_id] = (rrf_score, new_result)
+
+        # Sort by RRF score (descending)
+        sorted_scores = sorted(scores.items(), key=lambda x: x[1][0], reverse=True)
+
+        # Build final results
+        fused_results = []
+        for rank, (_doc_id, (score, result)) in enumerate(sorted_scores[:limit], 1):
+            result.score = score
+            result.rank = rank
+            result.scores["rrf_score"] = score
+            fused_results.append(result)
+
+        return fused_results
+
+    def fuse_with_weights(
+        self,
+        results_lists: list[list[SearchResult]],
+        weights: list[float],
+        limit: int = 10,
+    ) -> list[SearchResult]:
+        """Fuse multiple ranked lists with weights.
+
+        Args:
+            results_lists: List of ranked result lists
+            weights: Weight for each result list (must sum to 1.0)
+            limit: Maximum number of results to return
+
+        Returns:
+            Fused and re-ranked list of results
+        """
+        if len(results_lists) != len(weights):
+            raise ValueError("Number of result lists must match number of weights")
+
+        # Normalize weights
+        total_weight = sum(weights)
+        normalized_weights = [w / total_weight for w in weights]
+
+        # Map document_id to weighted scores
+        scores: dict[str, tuple[float, SearchResult]] = {}
+
+        for list_idx, (results, weight) in enumerate(zip(results_lists, normalized_weights)):
+            for rank, result in enumerate(results, 1):
+                doc_id = result.document_id
+
+                # Calculate weighted RRF score
+                rrf_score = weight * (1.0 / (self.k + rank))
+
+                # Build scores dict
+                doc_scores: dict[str, float | None] = {}
+                if result.scores:
+                    doc_scores.update(result.scores)
+
+                # Track original score from this list
+                score_key = "bm25_score" if list_idx == 0 else "vector_score"
+                if score_key not in doc_scores:
+                    doc_scores[score_key] = result.score
+
+                if doc_id in scores:
+                    existing_score, existing_result = scores[doc_id]
+                    # Merge scores from existing result
+                    merged_scores: dict[str, float | None] = {}
+                    if existing_result.scores:
+                        merged_scores.update(existing_result.scores)
+                    merged_scores.update(doc_scores)
+                    scores[doc_id] = (existing_score + rrf_score, existing_result)
+                    # Update the stored result's scores
+                    existing_result.scores = merged_scores
+                else:
+                    new_result = SearchResult(
+                        document_id=result.document_id,
+                        title=result.title,
+                        path=result.path,
+                        collection_name=result.collection_name,
+                        score=rrf_score,
+                        content=result.content,
+                        highlights=result.highlights,
+                        rank=rank,
+                        scores=doc_scores,
+                        snippet=result.snippet,
+                        context_description=result.context_description,
+                    )
+                    scores[doc_id] = (rrf_score, new_result)
+
+        # Sort by weighted score (descending)
+        sorted_scores = sorted(scores.items(), key=lambda x: x[1][0], reverse=True)
+
+        # Build final results
+        fused_results = []
+        for rank, (_doc_id, (score, result)) in enumerate(sorted_scores[:limit], 1):
+            result.score = score
+            result.rank = rank
+            result.scores["rrf_score"] = score
+            fused_results.append(result)
+
+        return fused_results
+
+
+def compute_rrf_score(rank: int, k: int = 60) -> float:
+    """Compute RRF score for a single rank.
+
+    Args:
+        rank: The rank position (1-indexed)
+        k: RRF constant
+
+    Returns:
+        RRF score
+    """
+    return 1.0 / (k + rank)

sif/search/snippets.py

@@ -0,0 +1,130 @@
+"""Smart snippet extraction from document chunks."""
+
+from __future__ import annotations
+
+import re
+
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class SmartSnippetExtractor:
+    """Extract the most relevant snippet from a chunk based on query term frequency."""
+
+    def __init__(self, max_length: int = 300, context_radius: int = 80) -> None:
+        """Initialize snippet extractor.
+
+        Args:
+            max_length: Maximum snippet length in characters
+            context_radius: Context radius around best sentence (unused, kept for API compat)
+        """
+        self.max_length = max_length
+        self.context_radius = context_radius
+
+    def extract(self, text: str, query_terms: list[str]) -> str:
+        """Extract the best snippet from text matching query terms.
+
+        Args:
+            text: Full chunk or document text
+            query_terms: Lowercase query terms to score by
+
+        Returns:
+            Best snippet string (may be truncated with ellipsis)
+        """
+        if not text:
+            return ""
+
+        if not query_terms:
+            return self._fallback_snippet(text)
+
+        # Split into sentences
+        sentences = self._split_sentences(text)
+        if len(sentences) <= 1:
+            return self._truncate(text)
+
+        # Score each sentence by query term frequency
+        scored = []
+        for i, sentence in enumerate(sentences):
+            score = self._score_sentence(sentence, query_terms)
+            scored.append((i, score, sentence))
+
+        # Find the highest-scoring sentence
+        scored.sort(key=lambda x: x[1], reverse=True)
+        best_idx, best_score, _best_sentence = scored[0]
+
+        if best_score == 0:
+            return self._fallback_snippet(text)
+
+        # Build a window around the best sentence
+        return self._build_window(sentences, best_idx, query_terms)
+
+    def _split_sentences(self, text: str) -> list[str]:
+        """Split text into sentences."""
+        # Simple sentence splitting on .!? followed by space or end
+        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
+        return [s.strip() for s in sentences if s.strip()]
+
+    def _score_sentence(self, sentence: str, query_terms: list[str]) -> float:
+        """Score a sentence by weighted query term frequency."""
+        sentence_lower = sentence.lower()
+        score = 0.0
+        for term in query_terms:
+            count = sentence_lower.count(term.lower())
+            # Longer matches get higher weight
+            score += count * len(term)
+        return score
+
+    def _build_window(self, sentences: list[str], center_idx: int, _query_terms: list[str]) -> str:
+        """Build a contiguous window of sentences around the center."""
+        # Start with center sentence
+        window = [sentences[center_idx]]
+        total_len = len(window[0])
+
+        # Expand outward
+        left = center_idx - 1
+        right = center_idx + 1
+
+        while True:
+            added = False
+            # Try adding from left
+            if left >= 0:
+                candidate = sentences[left]
+                if total_len + len(candidate) + 1 <= self.max_length:
+                    window.insert(0, candidate)
+                    total_len += len(candidate) + 1
+                    left -= 1
+                    added = True
+
+            # Try adding from right
+            if right < len(sentences):
+                candidate = sentences[right]
+                if total_len + len(candidate) + 1 <= self.max_length:
+                    window.append(candidate)
+                    total_len += len(candidate) + 1
+                    right += 1
+                    added = True
+
+            if not added:
+                break
+
+        snippet = " ".join(window)
+
+        # Add ellipsis if truncated
+        if left >= 0:
+            snippet = "..." + snippet
+        if right < len(sentences):
+            snippet = snippet + "..."
+
+        return snippet.strip()
+
+    def _fallback_snippet(self, text: str) -> str:
+        """Return the beginning of text when no query terms match."""
+        return self._truncate(text)
+
+    def _truncate(self, text: str) -> str:
+        """Truncate text to max_length with ellipsis."""
+        if len(text) <= self.max_length:
+            return text.strip()
+        return text[: self.max_length].strip() + "..."