Improve reranker ordering, low-confidence threshold and memory use

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.
 

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:

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())]

tests/test_stoogle.py

@@ -8,14 +8,15 @@
 import sys
 import time
 from pathlib import Path
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
 from fastapi.testclient import TestClient
 
 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: