diff --git a/README.md b/README.md index 2bca0b3..010761e 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ Stoogle follows the retrieval design of [LeanExplore](https://github.com/justinc 1. **Descriptions.** When an index is built, an LLM writes a short textbook-style description of each declaration. Declarations are described in dependency order, so the description of a lemma can draw on those of the concepts it uses. 2. **Retrieval.** Keyword search (BM25) over names and over descriptions and signatures is combined with semantic search using [Qwen3-Embedding-0.6B](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B). -3. **Reranking.** [Qwen3-Reranker-0.6B](https://huggingface.co/Qwen/Qwen3-Reranker-0.6B) scores the top 20 candidates against the query. Its score is reported as `relevance`; when no candidate reaches 0.5, the response sets `low_confidence`, meaning StatLean may not contain the result. +3. **Reranking.** [Qwen3-Reranker-0.6B](https://huggingface.co/Qwen/Qwen3-Reranker-0.6B) scores the top 20 candidates against the query, and they are reordered by their reranker rank combined with their retrieval rank. Its score is reported as `relevance`; when no candidate reaches 0.9, the response sets `low_confidence`, meaning StatLean may not contain the result. Exact Lean names always rank first. Reranking is on by default when a GPU is available and takes over ten seconds per query on a CPU; `--rerank on|off` overrides the default. diff --git a/src/stoogle/index.py b/src/stoogle/index.py index 564743b..b615efb 100644 --- a/src/stoogle/index.py +++ b/src/stoogle/index.py @@ -158,11 +158,12 @@ def build_index(corpus: Path, output: Path, mode="hybrid", model=DEFAULT_MODEL, "query_prefix": embedder.query_prefix if embedder else None, "max_seq_length": embedder.model.max_seq_length if embedder else None, "reranker": reranker or None, "reranker_revision": reranker_revision if reranker else None, - "rerank_top": 20, "low_confidence_below": 0.5, + "rerank_top": 20, "low_confidence_below": 0.9, "corpus_sha256": manifest["corpus_sha256"], "embedding_cache": embedding_stats, "ranking": ("weighted RRF(k=60): name BM25 (1) + prose/signature BM25 (1) + dense cosine (2); " - "optional cross-encoder rerank of the top candidates; exact names first"), + "optional cross-encoder rerank of the top candidates, blended with their retrieval rank; " + "exact names first"), }, indent=2) + "\n") staging.rename(output) except BaseException: diff --git a/src/stoogle/search.py b/src/stoogle/search.py index 6d370f7..a02cf1c 100644 --- a/src/stoogle/search.py +++ b/src/stoogle/search.py @@ -31,6 +31,9 @@ } RRF_K = 60 WEIGHTS = {"name": 1.0, "text": 1.0, "dense": 2.0} +# Reranked order: 1 / (BLEND_K + reranker rank) + RETRIEVAL_WEIGHT / (BLEND_K + retrieval rank). +BLEND_K = 1 +RETRIEVAL_WEIGHT = 0.5 def expand_query(query): @@ -42,6 +45,15 @@ def expand_query(query): return " ".join([query, *extra]) +def blend_rerank(scores): + """Order candidates, given in retrieval order, by reranker rank combined with retrieval rank.""" + n = len(scores) + rerank_rank = np.empty(n) + rerank_rank[np.argsort(-np.asarray(scores), kind="stable")] = np.arange(1, n + 1) + score = 1 / (BLEND_K + rerank_rank) + RETRIEVAL_WEIGHT / (BLEND_K + np.arange(1, n + 1)) + return np.argsort(-score, kind="stable") + + class Reranker: """Qwen3-Reranker scored as P(yes) under its documented judgement prompt.""" TASK = "Given a description of a statistics theorem or definition, judge whether the Lean declaration states it" @@ -70,7 +82,7 @@ def score(self, query: str, documents: list[str], batch_size=10) -> np.ndarray: for d in documents[start:start + batch_size]] batch = self.tokenizer(prompts, padding=True, truncation=True, max_length=1024, return_tensors="pt").to(self.device) - logits = self.model(**batch).logits[:, -1, :] + logits = self.model(**batch, logits_to_keep=1).logits[:, -1, :] pair = self.torch.stack([logits[:, self.no], logits[:, self.yes]], 1).float() scores += pair.softmax(1)[:, 1].tolist() return np.asarray(scores) @@ -175,8 +187,7 @@ def _search(self, query, limit, module, kind): top = order[:max(limit, self.config.get("rerank_top", 20))] scores = self.reranker.score(query, [self._rerank_texts[i] for i in top]) relevance = {i: float(s) for i, s in zip(top, scores)} - top = [top[j] for j in np.argsort(-scores, kind="stable")] - order = top + order[len(top):] + order = [top[j] for j in blend_rerank(scores)] + order[len(top):] folded = query.casefold() exact = [i for i in order if folded in (self.records[i]["name"].casefold(), self.records[i]["name"].rsplit(".", 1)[-1].casefold())] diff --git a/tests/test_stoogle.py b/tests/test_stoogle.py index 7f9bc3c..c00a2e9 100644 --- a/tests/test_stoogle.py +++ b/tests/test_stoogle.py @@ -8,6 +8,7 @@ import sys import time from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -15,7 +16,7 @@ from stoogle.corpus import RenderedCode, export_corpus, plain, read_corpus, source_identity from stoogle.index import build_index -from stoogle.search import Search, expand_query +from stoogle.search import Reranker, Search, blend_rerank, expand_query SAMPLE = Path(__file__).parent / "fixtures" / "parametric-family" SCORE = "AsymptoticStatistics.Score.score_mean_zero" @@ -300,9 +301,43 @@ def blob(text): class FakeReranker: """Scores documents by whether they mention the query's first word.""" + def __init__(self, match=0.95): + self.match = match + def score(self, query, documents): word = query.split()[0].casefold() - return np.asarray([0.9 if word in d.casefold() else 0.1 for d in documents]) + return np.asarray([self.match if word in d.casefold() else 0.05 for d in documents]) + + +def test_blend_rerank_combines_reranker_and_retrieval_ranks(): + # The reranker ranks candidate 15 first and the retrieval leader second: the leader stays first. + scores = np.zeros(20) + scores[15], scores[0] = 0.99, 0.98 + assert blend_rerank(scores)[:2].tolist() == [0, 15] + # The reranker's first choice moves above a retrieval leader it ranks third. + assert blend_rerank(np.asarray([0.2, 0.9, 0.5])).tolist() == [1, 0, 2] + + +def test_reranker_scores_only_the_last_position(): + torch = pytest.importorskip("torch") + requested = [] + + class Batch(dict): + def to(self, device): + return self + + class Model: + def __call__(self, input_ids, logits_to_keep=0): + requested.append(logits_to_keep) + logits = torch.zeros(input_ids.shape[0], logits_to_keep or input_ids.shape[1], 3) + logits[:, -1, 2], logits[:, -1, 1] = 3.0, 1.0 + return SimpleNamespace(logits=logits) + + reranker = Reranker.__new__(Reranker) + reranker.torch, reranker.device, reranker.model, reranker.yes, reranker.no = torch, "cpu", Model(), 2, 1 + reranker.tokenizer = lambda prompts, **kwargs: Batch(input_ids=torch.zeros(len(prompts), 8, dtype=torch.long)) + assert reranker.score("query", ["a", "b", "c"], batch_size=2) == pytest.approx([1 / (1 + np.exp(-2.0))] * 3) + assert requested == [1, 1] def test_rerank_reorders_and_flags_low_confidence(index): @@ -315,9 +350,12 @@ async def scenario(): result = await search.search("fisherInformation of the model", limit=5) assert result["reranked"] and not result["low_confidence"] assert "fisherinformation" in result["results"][0]["name"].casefold() - assert result["results"][0]["relevance"] == 0.9 + assert result["results"][0]["relevance"] == 0.95 missing = await search.search("zzzz score", limit=5) - assert missing["low_confidence"] and missing["confidence"] == 0.1 + assert missing["low_confidence"] and missing["confidence"] == 0.05 + search._reranker = FakeReranker(match=0.7) + weak = await search.search("fisherInformation of the model", limit=5) + assert weak["low_confidence"] and weak["confidence"] == 0.7 exact = await search.search("score_mean_zero", limit=3) assert exact["results"][0]["name"] == SCORE # exact names stay pinned above reranking finally: