Make TREC CAsT evaluation tractable and publish Gemini baseline

CHANGELOG.md

@@ -17,6 +17,31 @@ All notable changes to this project are documented here.
   checksum, selected conversation IDs, model, graph, package, and git provenance.
 - Cache repeated query embeddings in the eval provider and reuse one bounded
   shortest-path search across all targets for a graph source.
+- Build embedding kNN graphs with blockwise NumPy top-k selection instead of
+  Python-sorting every corpus pair, keeping large judged corpora tractable.
+- Bound graph shortest-path caching by source and bypass graph searches when
+  `transition_weight=0`, preventing pointwise evaluation from materializing an
+  all-pairs distance cache.
+- Accelerate `KNNGraph` shortest paths with SciPy's compiled sparse-graph
+  implementation and compact dense-distance cache when available while
+  retaining the standard-library fallback.
+- Vectorize evaluation retrieval scores with NumPy and retain a bounded query
+  score cache, removing per-dimension Python loops from full-corpus baselines.
+- Checkpoint Gemini rewrites after each successful request and retry transient
+  rate-limit/server failures with bounded exponential backoff.
+- Prefetch and pace Gemini rewrites before local retrieval evaluation so
+  low-request-rate API keys can resume without repeating graph computation.
+- Pin the CAsT baseline to stable, high-volume `gemini-3.1-flash-lite` after
+  sustained capacity errors from the 3.5 Flash and 3 Flash preview models.
+- Cap individual Gemini HTTP attempts at 30 seconds and extend bounded retries
+  for occasional capacity stalls during long prefetch runs.
+- Omit the resolved-query oracle when a benchmark does not provide resolved
+  turns instead of silently evaluating raw queries under an oracle label.
+- Publish the CAsT 2019 judged-passage result: map-matched β=1.0 lifts follow-up
+  nDCG@3 by `+0.027` (`[+0.008, +0.049]`), while Gemini Flash-Lite rewriting
+  lifts it by `+0.191` (`[+0.078, +0.296]`).
+- Correct the development `build` dependency floor to the available 1.5.0
+  release.
 - Publish the reproducible TopiOCQA n=25 MiniLM/kNN micro-corpus result:
   map-matched β=1.0 lifts follow-up nDCG@3 by `+0.084` with paired 95% CI
   `[+0.046, +0.128]`; two runs produced byte-identical reports.
@@ -47,7 +72,8 @@ All notable changes to this project are documented here.
 - Fix the TREC CAsT 2019 loader: use the correct ir-datasets id
   `trec-cast/v1/2019/judged`, load real passage text from the collection
   `docs_store()` (previously the doc id was used as the text), read
-  `raw_utterance` / `manual_rewritten_utterance`, and populate resolved queries.
+  `raw_utterance` / `manual_rewritten_utterance`, populate resolved queries,
+  install TREC CAR support, and stop retrying a failed docstore build per passage.
 - Add an optional `SentenceTransformerEmbedder` (extra: `[st]`) and an
   `--embedder {hash,sentence-transformers}` CLI flag so eval runs can use real
   semantic embeddings instead of the deterministic hash fixture. Both embedders

README.md

@@ -233,6 +233,25 @@ python -m pip install -e ".[eval,graph,st]"
 ./scripts/reproduce_topiocqa_n25.sh data/topiocqa_valid.jsonl
 ```
 
+The TREC CAsT 2019 judged-passage profile covers 20 topics, 173 turns, and
+21,726 judged passages. It uses the same MiniLM/kNN/full-ranking setup plus a
+checkpointed `gemini-3.1-flash-lite` rewrite baseline:
+
+| Slice | Method | nDCG@3 | Delta vs pointwise | Paired delta 95% CI |
+| --- | --- | ---: | ---: | ---: |
+| Follow-up | Pointwise | 0.325 | +0.000 | — |
+| Follow-up | Map-matched β=0.5 | 0.335 | +0.009 | [-0.006, +0.027] |
+| Follow-up | Map-matched β=1.0 | 0.352 | +0.027 | [+0.008, +0.049] |
+| Follow-up | Gemini rewrite | 0.516 | +0.191 | [+0.078, +0.296] |
+| Standalone | Map-matched β=1.0 | 0.216 | +0.015 | [+0.004, +0.027] |
+| Standalone | Gemini rewrite | 0.457 | +0.257 | [+0.187, +0.325] |
+
+These are judged-passage micro-corpus results, not full-corpus retrieval
+evidence. The ir-datasets CAsT query objects do not expose manual rewrites, so
+the profile does not report a resolved-query oracle. See the
+[`committed CAsT result`](results/cast2019_gemini_flash_lite_knn.md) for full
+provenance and limitations.
+
 ## Design choices and limits
 
 - Per-turn z-score normalization is the safe default; `center` and `none` are

docs/eval.md

@@ -58,7 +58,7 @@ best map-matched configuration against the β=0 pointwise baseline:
 | `history_concat` | Dense retrieval over concatenated query history |
 | `gemini_rewrite` | Gemini rewrites each turn into a standalone query before dense retrieval |
 | `maximal_marginal_relevance` | Per-turn MMR re-ranking (not map-matched retrieval) |
-| `resolved_oracle` | CAsT resolved utterances (upper bound) |
+| `resolved_oracle` | Caller-supplied resolved utterances (upper bound) |
 
 ## Benchmarks
 
@@ -116,13 +116,21 @@ Without the collection the loader degrades to using doc ids as passage text
 the follow-up-lift claim than TopiOCQA's topic switches.
 
 The script adds `gemini_rewrite` to the normal ablation grid. It sends each raw
-utterance and its prior user utterances to `gemini-3.5-flash` with minimal
+utterance and its prior user utterances to `gemini-3.1-flash-lite` with minimal
 thinking, retrieves with the returned standalone query, and compares it with both
-pointwise retrieval and CAsT's manual `resolved_oracle`. `GEMINI_API_KEY` is read
+pointwise retrieval and a `resolved_oracle` when every turn supplies a resolved
+query. The ir-datasets CAsT 2019 judged query objects currently expose raw
+utterances but not manual rewrites, so the reproduction profile omits the oracle
+instead of silently duplicating pointwise retrieval. `GEMINI_API_KEY` is read
 from the environment and is never written to reports. Reports record the Gemini
 model and prompt version. The baseline is opt-in because it makes one paid,
 networked model request per selected turn; `--conversation-limit` bounds those
 requests. Hosted-model output is not immutable across model revisions.
+Successful rewrites are checkpointed in `rewrites.json`, so rerunning the profile
+resumes after transient API failures instead of repeating completed requests.
+The reproduction script prefetches rewrites at a 13-second interval before
+starting retrieval evaluation, which also supports keys constrained to five
+requests per minute.
 
 ## Graph source and ranking mode
 

pyproject.toml

@@ -24,13 +24,14 @@ dependencies = []
 [project.optional-dependencies]
 graph = [
   "numpy>=2.2.6",
+  "scipy>=1.15.3",
 ]
 faiss = [
   "faiss-cpu>=1.14.3",
   "numpy>=2.2.6",
 ]
 eval = [
-  "ir-datasets>=0.5.11",
+  "ir-datasets[car]>=0.5.11",
   "numpy>=2.2.6",
 ]
 gemini = [
@@ -40,7 +41,7 @@ st = [
   "sentence-transformers>=3.0",
 ]
 dev = [
-  "build>=1.5.1",
+  "build>=1.5.0",
   "mypy>=1.16",
   "pytest>=8.4",
   "ruff>=0.12",

results/cast2019_gemini_flash_lite_knn.md

@@ -0,0 +1,51 @@
+# TREC CAsT 2019 Gemini Flash-Lite / MiniLM / kNN result
+
+This is a Tier B judged-passage micro-corpus result, not a full MS MARCO/TREC CAR
+retrieval result.
+
+- Dataset: `trec-cast/v1/2019/judged`
+- Conversations: 20 judged topics, 173 turns
+- Corpus: 21,726 unique judged passages
+- Embedder: `sentence-transformers/all-MiniLM-L6-v2`
+- Graph: 10-neighbor embedding kNN
+- Ranking: full, candidate limit 100
+- Rewriter: `gemini-3.1-flash-lite`
+- Rewrite prompt: `cast-standalone-v1`
+- Bootstrap: 1,000 conversation-level paired draws, seed 42
+- Profile: `cast2019-gemini-3.1-flash-lite-minilm-knn-full`
+- Git revision: `1d8ea0fdc7face208dfab53589972818cf94cf41`
+
+| Slice | Method | β | nDCG@3 | nDCG@3 95% CI | nDCG@5 | Recall@100 | Δ vs pointwise (95% CI) |
+| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| Follow-up | Pointwise | 0.0 | 0.325 | [0.232, 0.423] | 0.328 | 0.583 | +0.000 |
+| Standalone | Pointwise | 0.0 | 0.200 | [0.143, 0.256] | 0.217 | 0.343 | +0.000 |
+| Follow-up | Map-matched | 0.5 | 0.335 | [0.240, 0.432] | 0.340 | 0.583 | +0.009 [-0.006, +0.027] |
+| Standalone | Map-matched | 0.5 | 0.204 | [0.147, 0.258] | 0.220 | 0.343 | +0.003 [-0.003, +0.010] |
+| Follow-up | Map-matched | 1.0 | 0.352 | [0.263, 0.444] | 0.356 | 0.583 | +0.027 [+0.008, +0.049] |
+| Standalone | Map-matched | 1.0 | 0.216 | [0.157, 0.272] | 0.227 | 0.343 | +0.015 [+0.004, +0.027] |
+| Follow-up | History concat | — | 0.110 | [0.073, 0.160] | 0.112 | 0.328 | -0.215 [-0.325, -0.100] |
+| Standalone | History concat | — | 0.189 | [0.134, 0.256] | 0.197 | 0.410 | -0.011 [-0.077, +0.057] |
+| Follow-up | Gemini rewrite | — | 0.516 | [0.417, 0.598] | 0.507 | 0.700 | +0.191 [+0.078, +0.296] |
+| Standalone | Gemini rewrite | — | 0.457 | [0.366, 0.537] | 0.464 | 0.607 | +0.257 [+0.187, +0.325] |
+| Follow-up | MMR | — | 0.324 | [0.232, 0.423] | 0.327 | 0.583 | -0.001 [-0.004, +0.000] |
+| Standalone | MMR | — | 0.200 | [0.143, 0.256] | 0.217 | 0.343 | +0.000 [+0.000, +0.000] |
+
+## Interpretation
+
+Map-matched retrieval at β=1.0 improves follow-up nDCG@3 by 0.027, with its
+paired interval excluding zero, while also improving the standalone slice by
+0.015. The claim gate passes.
+
+Gemini rewriting is substantially stronger on both slices in this setup. It is
+an API-backed query transformation baseline rather than a trajectory decoder,
+and its hosted output can change across model revisions.
+
+## Limitations
+
+- The corpus contains only judged passages, so scores and recall do not estimate
+  full-corpus retrieval performance.
+- The CAsT query objects exposed by ir-datasets do not include manual rewrites;
+  a resolved-query oracle is therefore omitted.
+- Passage text comes from the locally built combined MS MARCO/TREC CAR docstore.
+- Gemini rewrites were checkpointed after each successful request; the API key
+  and billing data are not stored.

scripts/reproduce_cast2019_gemini.sh

@@ -2,9 +2,10 @@
 set -euo pipefail
 
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
-PROFILE="cast2019-gemini35flash-minilm-knn-full"
+GEMINI_MODEL="${MAPMATCHED_GEMINI_MODEL:-gemini-3.1-flash-lite}"
+MODEL_SLUG="${GEMINI_MODEL//\//-}"
+PROFILE="cast2019-$MODEL_SLUG-minilm-knn-full"
 REPORT_DIR="${MAPMATCHED_REPORT_DIR:-$ROOT_DIR/reports/$PROFILE}"
-GEMINI_MODEL="${MAPMATCHED_GEMINI_MODEL:-gemini-3.5-flash}"
 
 if [[ -z "${GEMINI_API_KEY:-}" ]]; then
   echo "GEMINI_API_KEY must be set" >&2
@@ -17,6 +18,15 @@ export CUDA_VISIBLE_DEVICES=""
 export PYTHONHASHSEED=0
 export TOKENIZERS_PARALLELISM=false
 
+python3 -m mapmatched.eval \
+  --benchmark cast2019 \
+  --conversation-limit 50 \
+  --include-gemini-rewrite \
+  --gemini-model "$GEMINI_MODEL" \
+  --gemini-rewrite-cache "$REPORT_DIR/rewrites.json" \
+  --gemini-min-request-interval 13 \
+  --gemini-prefetch-only
+
 python3 -m mapmatched.eval \
   --profile "$PROFILE" \
   --benchmark cast2019 \
@@ -33,6 +43,8 @@ python3 -m mapmatched.eval \
   --bootstrap-seed 42 \
   --include-gemini-rewrite \
   --gemini-model "$GEMINI_MODEL" \
+  --gemini-rewrite-cache "$REPORT_DIR/rewrites.json" \
+  --gemini-min-request-interval 13 \
   --output "$REPORT_DIR/report.json" \
   --markdown-output "$REPORT_DIR/report.md"
 

src/mapmatched/decoder.py

@@ -145,8 +145,13 @@ def _forward(
                     predecessor_score = cumulative_scores[turn_index - 1][predecessor_index]
                     if predecessor_score == float("-inf"):
                         continue
-                    graph_distance = graph.distance(predecessor.chunk_id, candidate.chunk_id)
-                    weighted_cost = transition_weight * graph_distance
+                    weighted_cost = 0.0
+                    if transition_weight != 0.0:
+                        graph_distance = graph.distance(
+                            predecessor.chunk_id,
+                            candidate.chunk_id,
+                        )
+                        weighted_cost = transition_weight * graph_distance
                     score = predecessor_score - weighted_cost
                     if not math.isfinite(score):
                         raise ValueError("decoder accumulation produced a nonfinite score")
@@ -223,7 +228,7 @@ def _build_path(
         cumulative_score = 0.0
         for turn_index, candidate_index in enumerate(candidate_indices):
             candidate = trellis[turn_index][candidate_index]
-            if turn_index == 0:
+            if turn_index == 0 or transition_weight == 0.0:
                 graph_distance = 0.0
             else:
                 previous_candidate = trellis[turn_index - 1][candidate_indices[turn_index - 1]]

src/mapmatched/eval/__main__.py

@@ -14,6 +14,7 @@
     GEMINI_REWRITE_PROMPT_VERSION,
     ConversationQueryRewriter,
     create_gemini_query_rewriter,
+    rewrite_conversation_queries,
 )
 from .embedder import DeterministicHashEmbedder, SentenceTransformerEmbedder
 from .loaders import load_cast2019_micro, load_synthetic_fixture, load_topiocqa_micro
@@ -104,6 +105,23 @@ def build_parser() -> argparse.ArgumentParser:
         default=DEFAULT_GEMINI_MODEL,
         help="Gemini model used by --include-gemini-rewrite.",
     )
+    parser.add_argument(
+        "--gemini-rewrite-cache",
+        type=Path,
+        default=None,
+        help="JSON checkpoint for completed Gemini rewrites.",
+    )
+    parser.add_argument(
+        "--gemini-min-request-interval",
+        type=float,
+        default=0.0,
+        help="Minimum seconds between Gemini requests.",
+    )
+    parser.add_argument(
+        "--gemini-prefetch-only",
+        action="store_true",
+        help="Checkpoint Gemini rewrites without running retrieval evaluation.",
+    )
     return parser
 
 
@@ -131,13 +149,23 @@ def main(argv: list[str] | None = None) -> int:
     args = parser.parse_args(argv)
     query_rewriter: ConversationQueryRewriter | None = None
     if args.include_gemini_rewrite:
-        query_rewriter = create_gemini_query_rewriter(model=args.gemini_model)
+        query_rewriter = create_gemini_query_rewriter(
+            model=args.gemini_model,
+            cache_path=args.gemini_rewrite_cache,
+            minimum_request_interval=args.gemini_min_request_interval,
+        )
     data_path = _effective_data_path(args.benchmark, args.data_path)
     conversations, passages = load_benchmark(
         args.benchmark,
         conversation_limit=args.conversation_limit,
         data_path=data_path,
     )
+    if args.gemini_prefetch_only:
+        if query_rewriter is None:
+            parser.error("--gemini-prefetch-only requires --include-gemini-rewrite")
+        for conversation in conversations:
+            rewrite_conversation_queries(conversation, query_rewriter)
+        return 0
     embedder: DeterministicHashEmbedder | SentenceTransformerEmbedder
     if args.embedder == "sentence-transformers":
         embedder = SentenceTransformerEmbedder(args.st_model)
@@ -170,6 +198,9 @@ def main(argv: list[str] | None = None) -> int:
         query_rewrite_prompt_version=GEMINI_REWRITE_PROMPT_VERSION
         if args.include_gemini_rewrite
         else None,
+        query_rewrite_cache_filename=args.gemini_rewrite_cache.name
+        if args.gemini_rewrite_cache is not None
+        else None,
     )
     report = run_ablation_grid(
         conversations=conversations,
@@ -178,7 +209,15 @@ def main(argv: list[str] | None = None) -> int:
         eval_config=eval_config,
         candidate_limit=args.candidate_limit,
         include_gemini_rewrite=args.include_gemini_rewrite,
-        include_resolved_oracle=args.include_resolved_oracle or args.benchmark == "cast2019",
+        include_resolved_oracle=args.include_resolved_oracle
+        or (
+            args.benchmark == "cast2019"
+            and all(
+                turn.resolved_query is not None
+                for conversation in conversations
+                for turn in conversation.turns
+            )
+        ),
         query_rewriter=query_rewriter,
     )
     args.output.write_text(render_json(report), encoding="utf-8")