lean-explore 0.3.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

lean_explore/search/scoring.py

@@ -0,0 +1,156 @@
+"""Score normalization and fusion algorithms for search ranking.
+
+This module provides utilities for combining multiple retrieval signals into
+a unified ranking using techniques like Reciprocal Rank Fusion (RRF) and
+weighted score fusion.
+"""
+
+import difflib
+import math
+
+EPSILON = 1e-9
+
+
+def normalize_scores(scores: list[float]) -> list[float]:
+    """Min-max normalize scores to [0, 1] range.
+
+    Args:
+        scores: List of raw scores.
+
+    Returns:
+        List of normalized scores.
+    """
+    if not scores:
+        return []
+
+    min_score = min(scores)
+    max_score = max(scores)
+    score_range = max_score - min_score
+
+    if score_range < EPSILON:
+        if max_score > EPSILON:
+            return [1.0] * len(scores)
+        return [0.0] * len(scores)
+
+    return [(s - min_score) / score_range for s in scores]
+
+
+def normalize_dependency_counts(counts: list[int]) -> list[float]:
+    """Log-scale normalization for dependency counts.
+
+    Uses log(1 + count) / log(1 + max_count) to compress the range
+    and give more credit to items with moderate dependency counts.
+
+    Args:
+        counts: List of dependency counts.
+
+    Returns:
+        List of normalized scores in [0, 1] range.
+    """
+    if not counts:
+        return []
+
+    max_count = max(counts)
+    if max_count == 0:
+        return [0.0] * len(counts)
+
+    log_max = math.log(1 + max_count)
+    return [math.log(1 + c) / log_max for c in counts]
+
+
+def compute_ranks(scores: list[float]) -> list[int]:
+    """Compute ranks for a list of scores (1-indexed, higher score = lower rank).
+
+    Candidates with score 0 get rank len(scores)+1 (worst possible).
+
+    Args:
+        scores: List of raw scores.
+
+    Returns:
+        List of ranks (1 = best).
+    """
+    n = len(scores)
+    indexed = [(i, s) for i, s in enumerate(scores)]
+    indexed.sort(key=lambda x: x[1], reverse=True)
+
+    ranks = [0] * n
+    for rank, (idx, score) in enumerate(indexed, 1):
+        if score > 0:
+            ranks[idx] = rank
+        else:
+            ranks[idx] = n + 1
+
+    return ranks
+
+
+def reciprocal_rank_fusion(rank_lists: list[list[int]], k: int = 0) -> list[float]:
+    """Compute RRF scores from multiple rank lists.
+
+    RRF(d) = sum(1 / (k + rank_i(d)) for each signal i)
+
+    Args:
+        rank_lists: List of rank lists, one per signal.
+        k: Constant to prevent top rank from dominating. Default 0 means 1/rank.
+
+    Returns:
+        List of RRF scores for each candidate.
+    """
+    n = len(rank_lists[0])
+    rrf_scores = []
+
+    for i in range(n):
+        score = sum(1.0 / (k + ranks[i]) for ranks in rank_lists)
+        rrf_scores.append(score)
+
+    return rrf_scores
+
+
+def weighted_score_fusion(
+    score_lists: list[list[float]],
+    weights: list[float],
+) -> list[float]:
+    """Combine multiple score lists using weighted normalized scores.
+
+    Each score list is normalized to [0, 1] using min-max scaling,
+    then combined with the given weights.
+
+    Args:
+        score_lists: List of score lists, one per signal.
+        weights: Weight for each signal (should sum to 1.0 for interpretability).
+
+    Returns:
+        List of combined scores for each candidate.
+    """
+    if not score_lists:
+        return []
+
+    n = len(score_lists[0])
+    if n == 0:
+        return []
+
+    normalized_lists = [normalize_scores(scores) for scores in score_lists]
+
+    combined = []
+    for i in range(n):
+        score = sum(w * normalized_lists[j][i] for j, w in enumerate(weights))
+        combined.append(score)
+
+    return combined
+
+
+def fuzzy_name_score(query: str, name: str) -> float:
+    """Compute fuzzy match score between query and declaration name.
+
+    Normalizes both strings (dots/underscores -> spaces) and uses
+    SequenceMatcher ratio for character-level similarity.
+
+    Args:
+        query: Search query string.
+        name: Declaration name to match against.
+
+    Returns:
+        Similarity score between 0 and 1.
+    """
+    normalized_query = query.lower().replace(".", " ").replace("_", " ")
+    normalized_name = name.lower().replace(".", " ").replace("_", " ")
+    return difflib.SequenceMatcher(None, normalized_query, normalized_name).ratio()

lean_explore/search/service.py

@@ -0,0 +1,68 @@
+"""Service layer for search operations."""
+
+import time
+
+from lean_explore.models import SearchResponse, SearchResult
+from lean_explore.search.engine import SearchEngine
+
+
+class Service:
+    """Service wrapper for search operations.
+
+    Provides a clean interface for searching and retrieving declarations.
+    """
+
+    def __init__(self, engine: SearchEngine | None = None):
+        """Initialize the search service.
+
+        Args:
+            engine: SearchEngine instance. Defaults to new engine.
+        """
+        self.engine = engine or SearchEngine()
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 20,
+        rerank_top: int | None = 50,
+        packages: list[str] | None = None,
+    ) -> SearchResponse:
+        """Search for Lean declarations.
+
+        Args:
+            query: Search query string.
+            limit: Maximum number of results to return.
+            rerank_top: Number of candidates to rerank with cross-encoder.
+            packages: Filter results to specific packages (e.g., ["Mathlib"]).
+
+        Returns:
+            SearchResponse containing results and metadata.
+        """
+        start_time = time.time()
+
+        results = await self.engine.search(
+            query=query,
+            limit=limit,
+            rerank_top=rerank_top,
+            packages=packages,
+        )
+
+        processing_time_ms = int((time.time() - start_time) * 1000)
+
+        return SearchResponse(
+            query=query,
+            results=results,
+            count=len(results),
+            processing_time_ms=processing_time_ms,
+        )
+
+    async def get_by_id(self, declaration_id: int) -> SearchResult | None:
+        """Retrieve a declaration by ID.
+
+        Args:
+            declaration_id: The declaration ID.
+
+        Returns:
+            SearchResult if found, None otherwise.
+        """
+        return await self.engine.get_by_id(declaration_id)

lean_explore/search/tokenization.py

@@ -0,0 +1,71 @@
+"""Text tokenization utilities for search indexing and querying.
+
+This module provides tokenization strategies for Lean declaration names,
+supporting both spaced tokenization (splits on dots, underscores, camelCase)
+and raw tokenization (preserves structure for exact matching).
+"""
+
+import re
+
+
+def tokenize_spaced(text: str) -> list[str]:
+    """Tokenize text with spacing on dots, underscores, and camelCase.
+
+    Args:
+        text: Input text to tokenize.
+
+    Returns:
+        List of lowercase word tokens.
+    """
+    if not text:
+        return []
+    # Replace dots and underscores with spaces
+    text = text.replace(".", " ").replace("_", " ")
+    # Split camelCase: insert space before uppercase letters
+    text = re.sub(r"([a-z])([A-Z])", r"\1 \2", text)
+    return re.findall(r"\w+", text.lower())
+
+
+def tokenize_raw(text: str) -> list[str]:
+    """Tokenize text as single token (preserves dots).
+
+    Args:
+        text: Input text to tokenize.
+
+    Returns:
+        List with the full text as a single lowercase token.
+    """
+    if not text:
+        return []
+    return [text.lower()]
+
+
+def tokenize_words(text: str) -> list[str]:
+    """Simple word tokenization for natural language text.
+
+    Splits on whitespace and punctuation, returns lowercase tokens.
+
+    Args:
+        text: Input text to tokenize.
+
+    Returns:
+        List of lowercase word tokens.
+    """
+    if not text:
+        return []
+    return [w.lower() for w in re.findall(r"\w+", text)]
+
+
+def is_autogenerated(name: str) -> bool:
+    """Check if a declaration name is auto-generated by Lean.
+
+    Auto-generated declarations include:
+    - .mk constructors (e.g., Nat.mk)
+
+    Args:
+        name: Fully qualified declaration name.
+
+    Returns:
+        True if the declaration is auto-generated.
+    """
+    return name.endswith(".mk")

lean_explore/util/__init__.py

@@ -0,0 +1,28 @@
+"""Shared utilities for lean_explore.
+
+Imports are lazy to avoid loading torch when not needed.
+"""
+
+
+def __getattr__(name: str):
+    """Lazy import attributes to avoid loading torch unnecessarily."""
+    if name == "EmbeddingClient":
+        from lean_explore.util.embedding_client import EmbeddingClient
+
+        return EmbeddingClient
+    if name == "RerankerClient":
+        from lean_explore.util.reranker_client import RerankerClient
+
+        return RerankerClient
+    if name == "OpenRouterClient":
+        from lean_explore.util.openrouter_client import OpenRouterClient
+
+        return OpenRouterClient
+    if name == "setup_logging":
+        from lean_explore.util.logging import setup_logging
+
+        return setup_logging
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["EmbeddingClient", "RerankerClient", "OpenRouterClient", "setup_logging"]

lean_explore/util/embedding_client.py

@@ -0,0 +1,92 @@
+"""Embedding generation client using sentence transformers."""
+
+import asyncio
+import logging
+
+import torch
+from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingResponse(BaseModel):
+    """Response from embedding generation."""
+
+    texts: list[str]
+    """Original input texts."""
+
+    embeddings: list[list[float]]
+    """List of embeddings (one per input text)."""
+
+    model: str
+    """Model name used for generation."""
+
+
+class EmbeddingClient:
+    """Client for generating text embeddings."""
+
+    def __init__(
+        self, model_name: str, device: str | None = None, max_length: int | None = None
+    ):
+        """Initialize the embedding client.
+
+        Args:
+            model_name: Name of the sentence transformer model
+            device: Device to use ("cuda", "mps", "cpu"). Auto-detects if None.
+            max_length: Maximum sequence length for tokenization. If None, uses
+                model default. Lower values reduce memory usage.
+        """
+        self.model_name = model_name
+        self.device = device or self._select_device()
+        self.max_length = max_length
+        logger.info(f"Loading embedding model {model_name} on {self.device}")
+        self.model = SentenceTransformer(model_name, device=self.device)
+
+        # Set max sequence length if specified
+        if max_length is not None:
+            self.model.max_seq_length = max_length
+            logger.info(f"Set max sequence length to {max_length}")
+
+    def _select_device(self) -> str:
+        """Select best available device."""
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    async def embed(
+        self, texts: list[str], is_query: bool = False
+    ) -> EmbeddingResponse:
+        """Generate embeddings for a list of texts.
+
+        Args:
+            texts: List of text strings to embed
+            is_query: If True, encode as search queries using the model's query
+                prompt. If False (default), encode as documents without prompt.
+                Qwen3-Embedding models are asymmetric and perform better when
+                queries use prompt_name="query".
+
+        Returns:
+            EmbeddingResponse with texts, embeddings, and model info
+        """
+        loop = asyncio.get_event_loop()
+
+        def _encode():
+            # Use query prompt for search queries, no prompt for documents
+            encode_kwargs = {
+                "show_progress_bar": False,
+                "convert_to_numpy": True,
+                "batch_size": 256,  # Larger batches for GPU utilization
+            }
+            if is_query:
+                encode_kwargs["prompt_name"] = "query"
+            return self.model.encode(texts, **encode_kwargs)
+
+        embeddings = await loop.run_in_executor(None, _encode)
+        return EmbeddingResponse(
+            texts=texts,
+            embeddings=[emb.tolist() for emb in embeddings],
+            model=self.model_name,
+        )

lean_explore/util/logging.py

@@ -0,0 +1,22 @@
+"""Logging utilities for lean_explore."""
+
+import logging
+import sys
+
+
+def setup_logging(verbose: bool = False) -> None:
+    """Configure logging for lean_explore applications.
+
+    Args:
+        verbose: If True, set level to DEBUG; otherwise INFO.
+    """
+    level = logging.DEBUG if verbose else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+    # Suppress noisy third-party loggers
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("httpcore").setLevel(logging.WARNING)

lean_explore/util/openrouter_client.py

@@ -0,0 +1,63 @@
+"""OpenRouter API wrapper using OpenAI SDK types.
+
+Provides a client class that wraps OpenRouter's API and returns OpenAI SDK-compatible
+types for easy integration.
+"""
+
+from __future__ import annotations
+
+import os
+
+from openai import AsyncOpenAI
+from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+
+class OpenRouterClient:
+    """Client for interacting with OpenRouter API using OpenAI SDK types."""
+
+    def __init__(self):
+        """Initialize OpenRouter client.
+
+        Reads API key from OPENROUTER_API_KEY environment variable.
+        """
+        api_key = os.getenv("OPENROUTER_API_KEY")
+        if not api_key:
+            raise ValueError("OPENROUTER_API_KEY environment variable not set")
+
+        self.client = AsyncOpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            api_key=api_key,
+        )
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+    )
+    async def generate(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs,
+    ) -> ChatCompletion:
+        """Generate a chat completion using OpenRouter.
+
+        Args:
+            model: Model name (e.g., "anthropic/claude-3.5-sonnet")
+            messages: List of message dicts with "role" and "content"
+            temperature: Sampling temperature
+            max_tokens: Maximum tokens to generate
+            **kwargs: Additional parameters to pass to the API
+
+        Returns:
+            ChatCompletion object from OpenAI SDK
+        """
+        return await self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )

lean_explore/util/reranker_client.py

@@ -0,0 +1,189 @@
+"""Reranker client using Qwen3-Reranker for query-document scoring."""
+
+import asyncio
+import logging
+
+import torch
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_INSTRUCTION = "Find relevant Lean 4 math declarations"
+
+
+class RerankerResponse(BaseModel):
+    """Response from reranking operation."""
+
+    query: str
+    """The original query."""
+
+    scores: list[float]
+    """Relevance scores for each document (same order as input)."""
+
+    model: str
+    """Model name used for reranking."""
+
+
+class RerankerClient:
+    """Client for reranking query-document pairs using Qwen3-Reranker."""
+
+    def __init__(
+        self,
+        model_name: str = "Qwen/Qwen3-Reranker-0.6B",
+        device: str | None = None,
+        max_length: int = 512,
+        instruction: str = DEFAULT_INSTRUCTION,
+    ):
+        """Initialize the reranker client.
+
+        Args:
+            model_name: Name of the reranker model from HuggingFace.
+            device: Device to use ("cuda", "mps", "cpu"). Auto-detects if None.
+            max_length: Maximum sequence length for tokenization.
+            instruction: Task instruction prepended to each query-document pair.
+        """
+        self.model_name = model_name
+        self.device = device or self._select_device()
+        self.max_length = max_length
+        self.instruction = instruction
+
+        logger.info(f"Loading reranker model {model_name} on {self.device}")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, padding_side="left", trust_remote_code=True
+        )
+
+        # Use float32 on CPU (faster than float16 which gets emulated)
+        # Use float16 on GPU for memory efficiency
+        dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=dtype, trust_remote_code=True
+        ).to(self.device)
+        self.model.eval()
+
+        # Get token IDs for true/false classification
+        self._token_true_id = self.tokenizer.convert_tokens_to_ids("true")
+        self._token_false_id = self.tokenizer.convert_tokens_to_ids("false")
+
+        logger.info("Reranker model loaded successfully")
+
+    def _select_device(self) -> str:
+        """Select best available device."""
+        if torch.cuda.is_available():
+            return "cuda"
+        return "cpu"
+
+    def _format_pair(self, query: str, document: str) -> str:
+        """Format a query-document pair with instruction.
+
+        Args:
+            query: The search query.
+            document: The document text to score.
+
+        Returns:
+            Formatted string for the reranker model.
+        """
+        return (
+            f"<Instruct>: {self.instruction}\n"
+            f"<Query>: {query}\n"
+            f"<Document>: {document}"
+        )
+
+    @torch.no_grad()
+    def _compute_scores_sync(self, pairs: list[str]) -> list[float]:
+        """Compute relevance scores for formatted pairs synchronously.
+
+        Args:
+            pairs: List of formatted query-document strings.
+
+        Returns:
+            List of relevance scores in [0, 1].
+        """
+        inputs = self.tokenizer(
+            pairs,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        ).to(self.device)
+
+        # Get logits for last token
+        outputs = self.model(**inputs)
+        logits = outputs.logits[:, -1, :]
+
+        # Extract true/false logits
+        true_logits = logits[:, self._token_true_id]
+        false_logits = logits[:, self._token_false_id]
+
+        # Compute probability of "true" using softmax
+        stacked = torch.stack([false_logits, true_logits], dim=1)
+        log_probs = torch.nn.functional.log_softmax(stacked, dim=1)
+        scores = log_probs[:, 1].exp()
+
+        return scores.cpu().tolist()
+
+    def rerank_sync(
+        self,
+        query: str,
+        documents: list[str],
+    ) -> RerankerResponse:
+        """Rerank documents synchronously (faster for small batches).
+
+        Args:
+            query: The search query.
+            documents: List of document texts to rerank.
+
+        Returns:
+            RerankerResponse with scores for each document.
+        """
+        if not documents:
+            return RerankerResponse(query=query, scores=[], model=self.model_name)
+
+        pairs = [self._format_pair(query, doc) for doc in documents]
+        scores = self._compute_scores_sync(pairs)
+        return RerankerResponse(query=query, scores=scores, model=self.model_name)
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        batch_size: int | None = None,
+    ) -> RerankerResponse:
+        """Rerank documents by relevance to query.
+
+        Args:
+            query: The search query.
+            documents: List of document texts to rerank.
+            batch_size: Number of pairs to process at once.
+
+        Returns:
+            RerankerResponse with scores for each document.
+        """
+        if not documents:
+            return RerankerResponse(query=query, scores=[], model=self.model_name)
+
+        # Default batch size: 16 on GPU (fits 8GB VRAM), 32 on CPU
+        if batch_size is None:
+            batch_size = 16 if self.device == "cuda" else 32
+
+        # For small batches, run synchronously to avoid executor overhead
+        if len(documents) <= batch_size:
+            return self.rerank_sync(query, documents)
+
+        # Format all pairs
+        pairs = [self._format_pair(query, doc) for doc in documents]
+
+        # Process in batches
+        loop = asyncio.get_event_loop()
+        all_scores: list[float] = []
+
+        for i in range(0, len(pairs), batch_size):
+            batch = pairs[i : i + batch_size]
+            batch_scores = await loop.run_in_executor(
+                None, self._compute_scores_sync, batch
+            )
+            all_scores.extend(batch_scores)
+
+        return RerankerResponse(query=query, scores=all_scores, model=self.model_name)