From 3923b8a2bc4b63647c5469ad823c2dd675a37a98 Mon Sep 17 00:00:00 2001 From: Tingzhou Date: Mon, 19 Oct 2026 00:52:52 -0400 Subject: [PATCH] Improve reranker ordering, low-confidence threshold and memory use - Order reranked candidates by 1/(1 + reranker rank) + 0.5/(1 + retrieval rank) instead of by the reranker score alone. The reranker gives many candidates a probability close to 1, so ordering by it alone could push the best retrieval match out of the top results. - Write low_confidence_below: 0.9 instead of 0.5 into new indexes. Queries for results in StatLean score close to 1, so the higher threshold also flags queries for results StatLean does not contain that score between 0.5 and 0.9. - Compute reranker logits for the last position only. Full-sequence logits took about 3 GB per batch of ten 1024-token prompts in half precision and ran out of memory on 11 GB GPUs. --- README.md | 2 +- src/stoogle/index.py | 5 +++-- src/stoogle/search.py | 17 +++++++++++++--- tests/test_stoogle.py | 46 +++++++++++++++++++++++++++++++++++++++---- 4 files changed, 60 insertions(+), 10 deletions(-) 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: