diff --git a/CHANGELOG.md b/CHANGELOG.md index 32d53d5..afbf93c 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 7d09a1c..72a38c4 100644 --- a/README.md +++ b/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 diff --git a/docs/eval.md b/docs/eval.md index 7af2620..ddfb18f 100644 --- a/docs/eval.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 3124522..de5f5ae 100644 --- a/pyproject.toml +++ b/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", diff --git a/results/cast2019_gemini_flash_lite_knn.md b/results/cast2019_gemini_flash_lite_knn.md new file mode 100644 index 0000000..077a13e --- /dev/null +++ b/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. diff --git a/scripts/reproduce_cast2019_gemini.sh b/scripts/reproduce_cast2019_gemini.sh index 9e3d633..20c1304 100755 --- a/scripts/reproduce_cast2019_gemini.sh +++ b/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" diff --git a/src/mapmatched/decoder.py b/src/mapmatched/decoder.py index c158d75..2326c3e 100644 --- a/src/mapmatched/decoder.py +++ b/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]] diff --git a/src/mapmatched/eval/__main__.py b/src/mapmatched/eval/__main__.py index 4ae5699..226d153 100644 --- a/src/mapmatched/eval/__main__.py +++ b/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") diff --git a/src/mapmatched/eval/baselines/gemini.py b/src/mapmatched/eval/baselines/gemini.py index 20ddef8..5bfc5d5 100644 --- a/src/mapmatched/eval/baselines/gemini.py +++ b/src/mapmatched/eval/baselines/gemini.py @@ -1,13 +1,16 @@ from __future__ import annotations +import hashlib import importlib import json import os +import time from collections.abc import Callable, Sequence from dataclasses import dataclass +from pathlib import Path from typing import Literal -DEFAULT_GEMINI_MODEL = "gemini-3.5-flash" +DEFAULT_GEMINI_MODEL = "gemini-3.1-flash-lite" GEMINI_REWRITE_PROMPT_VERSION = "cast-standalone-v1" _SYSTEM_INSTRUCTION = """\ @@ -25,6 +28,10 @@ class GeminiDependencyUnavailableError(ImportError): class GeminiRewriteConfig: model: str = DEFAULT_GEMINI_MODEL thinking_level: Literal["minimal", "low", "medium", "high"] = "minimal" + maximum_attempts: int = 12 + initial_retry_delay: float = 5.0 + maximum_retry_delay: float = 60.0 + minimum_request_interval: float = 0.0 class GeminiQueryRewriter: @@ -34,9 +41,21 @@ def __init__( generate_content: Callable[..., object], config: GeminiRewriteConfig, client: object | None = None, + cache_path: Path | None = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, ) -> None: + if config.maximum_attempts <= 0: + raise ValueError("maximum_attempts must be greater than zero") + if config.minimum_request_interval < 0.0: + raise ValueError("minimum_request_interval must be nonnegative") self._client = client self._generate_content = generate_content + self._cache_path = cache_path + self._cache = self._load_cache(cache_path) + self._sleep = sleep + self._monotonic = monotonic + self._last_request_time: float | None = None self.config = config def rewrite(self, *, history: Sequence[str], query: str) -> str: @@ -47,21 +66,99 @@ def rewrite(self, *, history: Sequence[str], query: str) -> str: prompt = ( f"{_SYSTEM_INSTRUCTION}\n\nConversation:\n{json.dumps(request, ensure_ascii=False)}" ) - response = self._generate_content( - model=self.config.model, - contents=prompt, - config={"thinking_config": {"thinking_level": self.config.thinking_level}}, - ) + cache_key = hashlib.sha256( + json.dumps( + { + "model": self.config.model, + "prompt_version": GEMINI_REWRITE_PROMPT_VERSION, + "request": request, + }, + ensure_ascii=False, + sort_keys=True, + ).encode("utf-8") + ).hexdigest() + cached_query = self._cache.get(cache_key) + if cached_query is not None: + return cached_query + response = self._generate_with_retry(prompt) response_text = getattr(response, "text", None) if not isinstance(response_text, str) or not response_text.strip(): raise RuntimeError("Gemini returned no rewritten query") - return response_text.strip() + rewritten_query = response_text.strip() + self._cache[cache_key] = rewritten_query + self._save_cache() + return rewritten_query + + def _generate_with_retry(self, prompt: str) -> object: + for attempt in range(self.config.maximum_attempts): + try: + self._wait_for_request_interval() + return self._generate_content( + model=self.config.model, + contents=prompt, + config={ + "thinking_config": { + "thinking_level": self.config.thinking_level, + } + }, + ) + except Exception as error: + final_attempt = attempt + 1 == self.config.maximum_attempts + if final_attempt or not self._is_retryable(error): + raise + delay = min( + self.config.initial_retry_delay * (2**attempt), + self.config.maximum_retry_delay, + ) + self._sleep(delay) + raise RuntimeError("Gemini retry loop ended without a response") + + def _wait_for_request_interval(self) -> None: + now = self._monotonic() + if self._last_request_time is not None: + elapsed = now - self._last_request_time + delay = self.config.minimum_request_interval - elapsed + if delay > 0.0: + self._sleep(delay) + self._last_request_time = self._monotonic() + + @staticmethod + def _is_retryable(error: Exception) -> bool: + status_code = getattr(error, "status_code", None) + return isinstance(status_code, int) and (status_code == 429 or 500 <= status_code < 600) + + @staticmethod + def _load_cache(cache_path: Path | None) -> dict[str, str]: + if cache_path is None or not cache_path.exists(): + return {} + payload = json.loads(cache_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError("Gemini rewrite cache must contain a JSON object") + cache: dict[str, str] = {} + for key, value in payload.items(): + if not isinstance(key, str) or not isinstance(value, str): + raise ValueError("Gemini rewrite cache keys and values must be strings") + cache[key] = value + return cache + + def _save_cache(self) -> None: + if self._cache_path is None: + return + self._cache_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self._cache_path.with_suffix(f"{self._cache_path.suffix}.tmp") + temporary_path.write_text( + json.dumps(self._cache, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary_path.replace(self._cache_path) def create_gemini_query_rewriter( *, api_key: str | None = None, model: str = DEFAULT_GEMINI_MODEL, + cache_path: Path | None = None, + minimum_request_interval: float = 0.0, ) -> GeminiQueryRewriter: effective_api_key = api_key if api_key is not None else os.environ.get("GEMINI_API_KEY") if not effective_api_key: @@ -76,7 +173,10 @@ def create_gemini_query_rewriter( client_factory = getattr(genai, "Client", None) if not callable(client_factory): raise GeminiDependencyUnavailableError("google.genai.Client is unavailable") - client = client_factory(api_key=effective_api_key) + client = client_factory( + api_key=effective_api_key, + http_options={"timeout": 30_000}, + ) models = getattr(client, "models", None) generate_content = getattr(models, "generate_content", None) if not callable(generate_content): @@ -85,6 +185,10 @@ def create_gemini_query_rewriter( ) return GeminiQueryRewriter( generate_content=generate_content, - config=GeminiRewriteConfig(model=model), + config=GeminiRewriteConfig( + model=model, + minimum_request_interval=minimum_request_interval, + ), client=client, + cache_path=cache_path, ) diff --git a/src/mapmatched/eval/corpus.py b/src/mapmatched/eval/corpus.py index 08bc8ec..ccc0c49 100644 --- a/src/mapmatched/eval/corpus.py +++ b/src/mapmatched/eval/corpus.py @@ -1,13 +1,18 @@ from __future__ import annotations +import heapq +import importlib +from collections import OrderedDict from collections.abc import Sequence from mapmatched import KNNGraph, ScoredCandidate from mapmatched.graph import InMemoryCorpusGraph -from .embedder import PassageEmbedder, QueryEmbedder, dot_product +from .embedder import PassageEmbedder, QueryEmbedder from .types import Passage +_SCORE_CACHE_SIZE = 512 + class BruteForceProvider: def __init__( @@ -21,10 +26,24 @@ def __init__( raise ValueError("passage_ids must contain at least one ID") if len(ids) != len(passage_embeddings): raise ValueError("passage_ids and passage_embeddings length mismatch") + try: + numpy = importlib.import_module("numpy") + except ImportError as error: + raise ImportError( + "BruteForceProvider requires the 'eval' extra: " + "pip install 'map-matched-retrieval[eval]'" + ) from error + embedding_matrix = numpy.asarray(passage_embeddings, dtype="float64") + shape = getattr(embedding_matrix, "shape", None) + if not isinstance(shape, tuple) or len(shape) != 2: + raise ValueError("passage_embeddings must be a two-dimensional matrix") self._passage_ids = ids - self._passage_embeddings = tuple(tuple(values) for values in passage_embeddings) + self._embedding_matrix = embedding_matrix + self._embedding_dimension = shape[1] + self._numpy = numpy self._embed_query = embed_query self._query_embedding_cache: dict[str, tuple[float, ...]] = {} + self._score_cache: OrderedDict[str, tuple[float, ...]] = OrderedDict() @property def passage_count(self) -> int: @@ -37,21 +56,43 @@ def candidates(self, query: str, limit: int) -> list[ScoredCandidate]: if query_embedding is None: query_embedding = tuple(self._embed_query.embed_query(query)) self._query_embedding_cache[query] = query_embedding - scored = [ - ( - dot_product(query_embedding, passage_embedding), - passage_id, - ) - for passage_id, passage_embedding in zip( - self._passage_ids, - self._passage_embeddings, - strict=True, + if len(query_embedding) != self._embedding_dimension: + raise ValueError("query and passage embedding dimensions must match") + scores = self._score_cache.get(query) + if scores is None: + query_vector = self._numpy.asarray(query_embedding, dtype="float64") + raw_scores = self._embedding_matrix @ query_vector + tolist = getattr(raw_scores, "tolist", None) + if not callable(tolist): + raise TypeError("NumPy returned invalid retrieval scores") + score_values = tolist() + if not isinstance(score_values, list): + raise TypeError("NumPy returned invalid retrieval scores") + scores = tuple(float(score) for score in score_values) + self._score_cache[query] = scores + if len(self._score_cache) > _SCORE_CACHE_SIZE: + self._score_cache.popitem(last=False) + else: + self._score_cache.move_to_end(query) + effective_limit = min(limit, len(self._passage_ids)) + + def ranking_key(index: int) -> tuple[float, str]: + return -scores[index], self._passage_ids[index] + + if effective_limit == len(self._passage_ids): + ranked_indices = sorted(range(len(self._passage_ids)), key=ranking_key) + else: + ranked_indices = heapq.nsmallest( + effective_limit, + range(len(self._passage_ids)), + key=ranking_key, ) - ] - scored.sort(key=lambda item: (-item[0], item[1])) return [ - ScoredCandidate(chunk_id=passage_id, score=score) - for score, passage_id in scored[:limit] + ScoredCandidate( + chunk_id=self._passage_ids[index], + score=scores[index], + ) + for index in ranked_indices ] diff --git a/src/mapmatched/eval/loaders/cast2019.py b/src/mapmatched/eval/loaders/cast2019.py index 59cbd28..4e5deef 100644 --- a/src/mapmatched/eval/loaders/cast2019.py +++ b/src/mapmatched/eval/loaders/cast2019.py @@ -89,6 +89,7 @@ def load_cast2019_micro() -> tuple[tuple[EvalConversation, ...], tuple[Passage, try: doc = docs_store.get(doc_id) except Exception: + docs_store = None doc = None if doc is not None: text = _first_text_attr(doc, ("text", "body")) diff --git a/src/mapmatched/eval/runner.py b/src/mapmatched/eval/runner.py index 8e49e75..6679083 100644 --- a/src/mapmatched/eval/runner.py +++ b/src/mapmatched/eval/runner.py @@ -114,9 +114,10 @@ def run_method_on_conversation( rank_full_corpus(provider, rewritten_query) for rewritten_query in rewritten_queries ), tuple(None for _ in queries) if method.name == "resolved_oracle": + if any(turn.resolved_query is None for turn in conversation.turns): + raise ValueError("resolved_oracle requires a resolved query for every turn") oracle_queries = tuple( - turn.resolved_query if turn.resolved_query is not None else turn.query - for turn in conversation.turns + turn.resolved_query for turn in conversation.turns if turn.resolved_query is not None ) return tuple(rank_full_corpus(provider, query) for query in oracle_queries), tuple( None for _ in queries diff --git a/src/mapmatched/eval/types.py b/src/mapmatched/eval/types.py index a74bed3..c662303 100644 --- a/src/mapmatched/eval/types.py +++ b/src/mapmatched/eval/types.py @@ -132,6 +132,7 @@ class EvalConfig: query_rewrite_provider: str | None = None query_rewrite_model: str | None = None query_rewrite_prompt_version: str | None = None + query_rewrite_cache_filename: str | None = None def to_dict(self) -> dict[str, object]: return { @@ -157,6 +158,7 @@ def to_dict(self) -> dict[str, object]: "query_rewrite_provider": self.query_rewrite_provider, "query_rewrite_model": self.query_rewrite_model, "query_rewrite_prompt_version": self.query_rewrite_prompt_version, + "query_rewrite_cache_filename": self.query_rewrite_cache_filename, } diff --git a/src/mapmatched/graph.py b/src/mapmatched/graph.py index b422d8d..f0f303c 100644 --- a/src/mapmatched/graph.py +++ b/src/mapmatched/graph.py @@ -2,6 +2,7 @@ import heapq import math +from collections import OrderedDict from collections.abc import Iterable from dataclasses import dataclass from typing import Protocol @@ -37,13 +38,17 @@ def __init__( nodes: Iterable[str] = (), directed: bool = False, maximum_distance: float = 10.0, + distance_cache_size: int = 256, ) -> None: if not math.isfinite(maximum_distance) or maximum_distance <= 0.0: raise ValueError("maximum_distance must be finite and greater than zero") + if distance_cache_size <= 0: + raise ValueError("distance_cache_size must be greater than zero") self._maximum_distance = maximum_distance self._directed = directed + self._distance_cache_size = distance_cache_size self._adjacency: dict[str, dict[str, float]] = {} - self._distance_cache: dict[tuple[str, str], float] = {} + self._distance_cache: OrderedDict[str, dict[str, float]] = OrderedDict() for node in nodes: if not node: @@ -62,12 +67,14 @@ def from_edges( nodes: Iterable[str] = (), directed: bool = False, maximum_distance: float = 10.0, + distance_cache_size: int = 256, ) -> InMemoryCorpusGraph: return cls( (GraphEdge(source, target) for source, target in edges), nodes=nodes, directed=directed, maximum_distance=maximum_distance, + distance_cache_size=distance_cache_size, ) @property @@ -81,33 +88,39 @@ def _add_edge(self, source: str, target: str, distance: float) -> None: if current is None or distance < current: neighbors[target] = distance - def _cache_key(self, source: str, target: str) -> tuple[str, str]: - if self._directed or source <= target: - return source, target - return target, source - def distance(self, source_chunk_id: str, target_chunk_id: str) -> float: if not source_chunk_id or not target_chunk_id: raise ValueError("distance chunk IDs must not be empty") if source_chunk_id == target_chunk_id: return 0.0 - cache_key = self._cache_key(source_chunk_id, target_chunk_id) - cached = self._distance_cache.get(cache_key) - if cached is not None: - return cached - distances = self._bounded_distances(source_chunk_id, self._maximum_distance) - for chunk_id in self._adjacency: - distance = min( - distances.get(chunk_id, self._maximum_distance), - self._maximum_distance, - ) - self._distance_cache[self._cache_key(source_chunk_id, chunk_id)] = distance - if cache_key not in self._distance_cache: - self._distance_cache[cache_key] = min( + distances = self._cached_distances(source_chunk_id) + if distances is not None: + return min( distances.get(target_chunk_id, self._maximum_distance), self._maximum_distance, ) - return self._distance_cache[cache_key] + if not self._directed: + reverse_distances = self._cached_distances(target_chunk_id) + if reverse_distances is not None: + return min( + reverse_distances.get(source_chunk_id, self._maximum_distance), + self._maximum_distance, + ) + distances = self._bounded_distances(source_chunk_id, self._maximum_distance) + self._distance_cache[source_chunk_id] = distances + self._distance_cache.move_to_end(source_chunk_id) + if len(self._distance_cache) > self._distance_cache_size: + self._distance_cache.popitem(last=False) + return min( + distances.get(target_chunk_id, self._maximum_distance), + self._maximum_distance, + ) + + def _cached_distances(self, source: str) -> dict[str, float] | None: + distances = self._distance_cache.get(source) + if distances is not None: + self._distance_cache.move_to_end(source) + return distances def neighborhood(self, chunk_id: str, radius: float) -> tuple[str, ...]: if not chunk_id: diff --git a/src/mapmatched/knn.py b/src/mapmatched/knn.py index 7a00726..8128247 100644 --- a/src/mapmatched/knn.py +++ b/src/mapmatched/knn.py @@ -2,11 +2,14 @@ import importlib import math -from collections.abc import Sequence +from collections import OrderedDict +from collections.abc import Iterable, Sequence from types import ModuleType from .graph import GraphEdge, InMemoryCorpusGraph +_SIMILARITY_BLOCK_SIZE = 256 + class GraphDependencyUnavailableError(ImportError): pass @@ -21,6 +24,16 @@ def _load_numpy() -> ModuleType: ) from error +def _load_scipy_sparse() -> tuple[ModuleType, ModuleType] | None: + try: + return ( + importlib.import_module("scipy.sparse"), + importlib.import_module("scipy.sparse.csgraph"), + ) + except ImportError: + return None + + def _coerce_embedding_rows( embeddings: Sequence[Sequence[float]], expected_count: int, @@ -57,6 +70,127 @@ def _coerce_embedding_rows( class KNNGraph(InMemoryCorpusGraph): + def __init__( + self, + edges: Iterable[GraphEdge] = (), + *, + nodes: Iterable[str] = (), + directed: bool = False, + maximum_distance: float = 10.0, + distance_cache_size: int = 256, + ) -> None: + super().__init__( + edges, + nodes=nodes, + directed=directed, + maximum_distance=maximum_distance, + distance_cache_size=distance_cache_size, + ) + self._ordered_nodes = tuple(sorted(self._adjacency)) + self._node_indices = {chunk_id: index for index, chunk_id in enumerate(self._ordered_nodes)} + self._sparse_distance_cache: OrderedDict[str, object] = OrderedDict() + scipy_sparse = _load_scipy_sparse() + if scipy_sparse is None: + self._sparse_adjacency = None + self._scipy_csgraph = None + return + sparse, self._scipy_csgraph = scipy_sparse + rows: list[int] = [] + columns: list[int] = [] + distances: list[float] = [] + for source_id, neighbors in self._adjacency.items(): + source_index = self._node_indices[source_id] + for target_id, distance in neighbors.items(): + rows.append(source_index) + columns.append(self._node_indices[target_id]) + distances.append(distance) + csr_matrix = getattr(sparse, "csr_matrix", None) + if not callable(csr_matrix): + self._sparse_adjacency = None + self._scipy_csgraph = None + return + self._sparse_adjacency = csr_matrix( + (distances, (rows, columns)), + shape=(len(self._ordered_nodes), len(self._ordered_nodes)), + ) + + def distance(self, source_chunk_id: str, target_chunk_id: str) -> float: + if self._sparse_adjacency is None or self._scipy_csgraph is None: + return super().distance(source_chunk_id, target_chunk_id) + if not source_chunk_id or not target_chunk_id: + raise ValueError("distance chunk IDs must not be empty") + if source_chunk_id == target_chunk_id: + return 0.0 + target_index = self._node_indices.get(target_chunk_id) + if target_index is None: + return self._maximum_distance + distances = self._cached_sparse_distances(source_chunk_id) + if distances is not None: + return self._sparse_distance_value(distances, target_index) + if not self._directed: + reverse_distances = self._cached_sparse_distances(target_chunk_id) + source_index = self._node_indices.get(source_chunk_id) + if reverse_distances is not None and source_index is not None: + return self._sparse_distance_value(reverse_distances, source_index) + distances = self._compute_sparse_distances(source_chunk_id, self._maximum_distance) + if distances is None: + return self._maximum_distance + self._sparse_distance_cache[source_chunk_id] = distances + self._sparse_distance_cache.move_to_end(source_chunk_id) + if len(self._sparse_distance_cache) > self._distance_cache_size: + self._sparse_distance_cache.popitem(last=False) + return self._sparse_distance_value(distances, target_index) + + def _cached_sparse_distances(self, source: str) -> object | None: + distances = self._sparse_distance_cache.get(source) + if distances is not None: + self._sparse_distance_cache.move_to_end(source) + return distances + + def _sparse_distance_value(self, distances: object, target_index: int) -> float: + get_item = getattr(distances, "__getitem__", None) + if not callable(get_item): + return self._maximum_distance + distance = float(get_item(target_index)) + if not math.isfinite(distance): + return self._maximum_distance + return min(distance, self._maximum_distance) + + def _compute_sparse_distances(self, source: str, cutoff: float) -> object | None: + if self._sparse_adjacency is None or self._scipy_csgraph is None: + return None + source_index = self._node_indices.get(source) + if source_index is None: + return None + dijkstra = getattr(self._scipy_csgraph, "dijkstra", None) + if not callable(dijkstra): + return None + distances: object = dijkstra( + self._sparse_adjacency, + directed=self._directed, + indices=source_index, + limit=cutoff, + ) + return distances + + def _bounded_distances(self, source: str, cutoff: float) -> dict[str, float]: + if self._sparse_adjacency is None or self._scipy_csgraph is None: + return super()._bounded_distances(source, cutoff) + raw_distances = self._compute_sparse_distances(source, cutoff) + if raw_distances is None: + return {source: 0.0} + tolist = getattr(raw_distances, "tolist", None) + if not callable(tolist): + return super()._bounded_distances(source, cutoff) + values = tolist() + if not isinstance(values, list): + return super()._bounded_distances(source, cutoff) + return { + chunk_id: float(distance) + for chunk_id, distance in zip(self._ordered_nodes, values, strict=True) + if isinstance(distance, (int, float)) and math.isfinite(distance) and distance <= cutoff + } + @classmethod def from_embeddings( cls, @@ -88,50 +222,53 @@ def from_embeddings( "and no greater than maximum_distance" ) - normalized = _coerce_embedding_rows(embeddings, len(ids)) - + numpy = _load_numpy() + normalized_rows = _coerce_embedding_rows(embeddings, len(ids)) + normalized = numpy.asarray(normalized_rows, dtype="float64") + del normalized_rows edge_distances: dict[tuple[int, int], float] = {} effective_count = min(neighbor_count, max(0, len(ids) - 1)) - for source_index in range(len(ids)): - ranked_neighbors = sorted( - ( + for source_start in range(0, len(ids), _SIMILARITY_BLOCK_SIZE): + source_end = min(source_start + _SIMILARITY_BLOCK_SIZE, len(ids)) + similarities = normalized[source_start:source_end] @ normalized.T + numpy.clip(similarities, -1.0, 1.0, out=similarities) + for block_index, source_index in enumerate(range(source_start, source_end)): + if effective_count == 0: + continue + source_similarities = similarities[block_index] + source_similarities[source_index] = float("-inf") + partition = numpy.argpartition( + -source_similarities, + effective_count - 1, + )[:effective_count] + cutoff_similarity = float(numpy.min(source_similarities[partition])) + candidate_indices = numpy.flatnonzero( + source_similarities >= cutoff_similarity + ).tolist() + ranked_neighbors = sorted( ( - 1.0 - - max( - -1.0, - min( - 1.0, - sum( - source_value * target_value - for source_value, target_value in zip( - normalized[source_index], - normalized[target_index], - strict=True, - ) - ), - ), - ), - ids[target_index], - target_index, - ) - for target_index in range(len(ids)) - if target_index != source_index - ), - key=lambda item: (item[0], item[1]), - ) - for cosine_distance, _, target_index in ranked_neighbors[:effective_count]: - pair = ( - (source_index, target_index) - if source_index < target_index - else (target_index, source_index) + ( + 1.0 - float(source_similarities[target_index]), + ids[target_index], + target_index, + ) + for target_index in candidate_indices + ), + key=lambda item: (item[0], item[1]), ) - edge_distance = min( - max(cosine_distance, minimum_edge_distance), - maximum_distance, - ) - existing = edge_distances.get(pair) - if existing is None or edge_distance < existing: - edge_distances[pair] = edge_distance + for cosine_distance, _, target_index in ranked_neighbors[:effective_count]: + pair = ( + (source_index, target_index) + if source_index < target_index + else (target_index, source_index) + ) + edge_distance = min( + max(cosine_distance, minimum_edge_distance), + maximum_distance, + ) + existing = edge_distances.get(pair) + if existing is None or edge_distance < existing: + edge_distances[pair] = edge_distance edges = ( GraphEdge(ids[source_index], ids[target_index], distance) diff --git a/tests/test_decoder.py b/tests/test_decoder.py index a0bc0a4..85427f1 100644 --- a/tests/test_decoder.py +++ b/tests/test_decoder.py @@ -41,14 +41,25 @@ def test_hand_computed_viterbi_fixture() -> None: def test_zero_transition_weight_equals_pointwise_argmax() -> None: + class CountingDistanceGraph(InMemoryCorpusGraph): + def __init__(self) -> None: + super().__init__() + self.distance_count = 0 + + def distance(self, source_chunk_id: str, target_chunk_id: str) -> float: + self.distance_count += 1 + return super().distance(source_chunk_id, target_chunk_id) + + graph = CountingDistanceGraph() path = StandaloneDecoder().decode( hand_computed_trellis(), - graph=line_graph(), + graph=graph, emission_weight=1.0, transition_weight=0.0, ) assert path.chunk_ids == ("0", "2", "0") + assert graph.distance_count == 0 def test_fixed_lag_behavior() -> None: diff --git a/tests/test_eval_baselines.py b/tests/test_eval_baselines.py index fe02700..256991d 100644 --- a/tests/test_eval_baselines.py +++ b/tests/test_eval_baselines.py @@ -30,6 +30,25 @@ def embed_query(self, query: str) -> tuple[float, float]: assert embedder.query_count == 1 +def test_brute_force_provider_vectorized_ranking_is_deterministic() -> None: + class FixedEmbedder: + def embed_query(self, query: str) -> tuple[float, float]: + del query + return 1.0, 0.0 + + provider = BruteForceProvider( + ("second", "first", "middle"), + ((0.0, 1.0), (1.0, 0.0), (0.5, 0.5)), + FixedEmbedder(), + ) + + assert tuple(candidate.chunk_id for candidate in provider.candidates("query", limit=3)) == ( + "first", + "middle", + "second", + ) + + def test_history_concat_changes_query_sequence() -> None: _, passages = load_synthetic_fixture() embedder = DeterministicHashEmbedder() diff --git a/tests/test_eval_cast2019_loader.py b/tests/test_eval_cast2019_loader.py new file mode 100644 index 0000000..25cac6a --- /dev/null +++ b/tests/test_eval_cast2019_loader.py @@ -0,0 +1,51 @@ +from types import SimpleNamespace + +import pytest + +from mapmatched.eval.loaders import cast2019 + + +def test_cast_loader_stops_using_docstore_after_retrieval_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FailingDocStore: + def __init__(self) -> None: + self.get_count = 0 + + def get(self, doc_id: str) -> object: + del doc_id + self.get_count += 1 + raise ImportError("missing corpus dependency") + + docstore = FailingDocStore() + dataset = SimpleNamespace( + qrels_iter=lambda: iter( + ( + SimpleNamespace(query_id="1_1", doc_id="doc-1", relevance=2), + SimpleNamespace(query_id="1_1", doc_id="doc-2", relevance=1), + ) + ), + queries_iter=lambda: iter( + ( + SimpleNamespace( + topic_number=1, + turn_number=1, + query_id="1_1", + raw_utterance="What is renewable energy?", + manual_rewritten_utterance="renewable energy definition", + ), + ) + ), + docs_store=lambda: docstore, + ) + ir_datasets = SimpleNamespace(load=lambda dataset_id: dataset) + monkeypatch.setattr(cast2019, "_load_ir_datasets", lambda: ir_datasets) + + conversations, passages = cast2019.load_cast2019_micro() + + assert docstore.get_count == 1 + assert conversations[0].turns[0].resolved_query == "renewable energy definition" + assert tuple((passage.passage_id, passage.text) for passage in passages) == ( + ("doc-1", "doc-1"), + ("doc-2", "doc-2"), + ) diff --git a/tests/test_eval_gemini_rewrite.py b/tests/test_eval_gemini_rewrite.py index 9f6cf8c..ad4a99f 100644 --- a/tests/test_eval_gemini_rewrite.py +++ b/tests/test_eval_gemini_rewrite.py @@ -54,6 +54,103 @@ def test_create_gemini_query_rewriter_requires_api_key( create_gemini_query_rewriter() +def test_gemini_query_rewriter_retries_transient_errors() -> None: + attempts = 0 + delays: list[float] = [] + + class TransientError(Exception): + status_code = 503 + + class Response: + text = "standalone query" + + def generate_content(**request: object) -> object: + del request + nonlocal attempts + attempts += 1 + if attempts < 3: + raise TransientError + return Response() + + rewriter = GeminiQueryRewriter( + generate_content=generate_content, + config=GeminiRewriteConfig( + model="test-model", + initial_retry_delay=1.0, + ), + sleep=delays.append, + ) + + assert rewriter.rewrite(history=(), query="query") == "standalone query" + assert attempts == 3 + assert delays == [1.0, 2.0] + + +def test_gemini_query_rewriter_checkpoints_completed_rewrites(tmp_path: Path) -> None: + cache_path = tmp_path / "rewrites.json" + request_count = 0 + + class Response: + text = "standalone query" + + def generate_content(**request: object) -> object: + del request + nonlocal request_count + request_count += 1 + return Response() + + first = GeminiQueryRewriter( + generate_content=generate_content, + config=GeminiRewriteConfig(model="test-model"), + cache_path=cache_path, + ) + assert first.rewrite(history=("context",), query="follow up") == "standalone query" + second = GeminiQueryRewriter( + generate_content=generate_content, + config=GeminiRewriteConfig(model="test-model"), + cache_path=cache_path, + ) + + assert second.rewrite(history=("context",), query="follow up") == "standalone query" + assert request_count == 1 + assert json.loads(cache_path.read_text(encoding="utf-8")) + + +def test_gemini_query_rewriter_paces_requests() -> None: + delays: list[float] = [] + current_time = 0.0 + + class Response: + text = "standalone query" + + def generate_content(**request: object) -> object: + del request + return Response() + + def sleep(delay: float) -> None: + nonlocal current_time + delays.append(delay) + current_time += delay + + def monotonic() -> float: + return current_time + + rewriter = GeminiQueryRewriter( + generate_content=generate_content, + config=GeminiRewriteConfig( + model="test-model", + minimum_request_interval=13.0, + ), + sleep=sleep, + monotonic=monotonic, + ) + + rewriter.rewrite(history=(), query="first") + rewriter.rewrite(history=(), query="second") + + assert delays == [13.0] + + def test_gemini_rewrite_runs_end_to_end_with_paired_comparison() -> None: conversations, passages = load_synthetic_fixture() embedder = DeterministicHashEmbedder() @@ -93,6 +190,29 @@ def test_gemini_rewrite_runs_end_to_end_with_paired_comparison() -> None: assert all(slice_metrics.ndcg_at_3_delta_ci is not None for slice_metrics in comparison.slices) +def test_cli_prefetches_gemini_rewrites_without_evaluation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + conversations, _ = load_synthetic_fixture() + rewriter = RecordingQueryRewriter() + monkeypatch.setattr( + "mapmatched.eval.__main__.create_gemini_query_rewriter", + lambda *, model, cache_path, minimum_request_interval: rewriter, + ) + + exit_code = main( + [ + "--benchmark", + "synthetic", + "--include-gemini-rewrite", + "--gemini-prefetch-only", + ] + ) + + assert exit_code == 0 + assert len(rewriter.calls) == sum(len(conversation.turns) for conversation in conversations) + + def test_cli_records_gemini_rewrite_metadata( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -100,9 +220,10 @@ def test_cli_records_gemini_rewrite_metadata( rewriter = RecordingQueryRewriter() monkeypatch.setattr( "mapmatched.eval.__main__.create_gemini_query_rewriter", - lambda *, model: rewriter, + lambda *, model, cache_path, minimum_request_interval: rewriter, ) output_path = tmp_path / "report.json" + cache_path = tmp_path / "rewrites.json" exit_code = main( [ @@ -113,6 +234,8 @@ def test_cli_records_gemini_rewrite_metadata( "--include-gemini-rewrite", "--gemini-model", "test-model", + "--gemini-rewrite-cache", + str(cache_path), "--candidate-limit", "4", "--output", @@ -125,4 +248,5 @@ def test_cli_records_gemini_rewrite_metadata( assert payload["config"]["query_rewrite_provider"] == "gemini" assert payload["config"]["query_rewrite_model"] == "test-model" assert payload["config"]["query_rewrite_prompt_version"] == "cast-standalone-v1" + assert payload["config"]["query_rewrite_cache_filename"] == "rewrites.json" assert any(method["method_name"] == "gemini_rewrite" for method in payload["methods"]) diff --git a/tests/test_eval_synthetic_runner.py b/tests/test_eval_synthetic_runner.py index 631a56e..ccd6cb3 100644 --- a/tests/test_eval_synthetic_runner.py +++ b/tests/test_eval_synthetic_runner.py @@ -1,3 +1,5 @@ +import pytest + from mapmatched.eval import ( DeterministicHashEmbedder, EvalConfig, @@ -49,3 +51,26 @@ def test_synthetic_eval_runs_end_to_end() -> None: if slice_metrics.slice_name == "follow_up" ) assert follow_up.ndcg_at_3_delta_ci is not None + + +def test_resolved_oracle_requires_every_turn_to_be_resolved() -> None: + conversations, passages = load_synthetic_fixture() + embedder = DeterministicHashEmbedder() + + with pytest.raises(ValueError, match="resolved query for every turn"): + run_eval( + conversations=conversations, + passages=passages, + embedder=embedder, + methods=(MethodSpec(name="resolved_oracle"),), + config=MapMatchedMethodConfig(candidate_limit=4), + eval_config=EvalConfig( + benchmark="synthetic", + tier="synthetic", + embedder_name=embedder.name, + recall_k=10, + entropy_threshold=None, + standalone_tolerance=0.02, + follow_up_min_delta=0.0, + ), + ) diff --git a/tests/test_graph_and_scoring.py b/tests/test_graph_and_scoring.py index 78a13a3..0a0d44f 100644 --- a/tests/test_graph_and_scoring.py +++ b/tests/test_graph_and_scoring.py @@ -67,6 +67,32 @@ def _bounded_distances(self, source: str, cutoff: float) -> dict[str, float]: assert graph.search_count == 1 +def test_distance_source_cache_evicts_least_recently_used_search() -> None: + class CountingGraph(InMemoryCorpusGraph): + def __init__(self) -> None: + super().__init__( + [ + GraphEdge("a", "b", 1.0), + GraphEdge("b", "c", 1.0), + GraphEdge("c", "d", 1.0), + ], + directed=True, + distance_cache_size=2, + ) + self.search_count = 0 + + def _bounded_distances(self, source: str, cutoff: float) -> dict[str, float]: + self.search_count += 1 + return super()._bounded_distances(source, cutoff) + + graph = CountingGraph() + assert graph.distance("a", "d") == 3.0 + assert graph.distance("b", "d") == 2.0 + assert graph.distance("c", "d") == 1.0 + assert graph.distance("a", "d") == 3.0 + assert graph.search_count == 4 + + @pytest.mark.parametrize( ("method", "expected"), [ diff --git a/tests/test_knn_graph.py b/tests/test_knn_graph.py index 84a7dad..8fc1439 100644 --- a/tests/test_knn_graph.py +++ b/tests/test_knn_graph.py @@ -57,6 +57,44 @@ def test_tie_breaking_is_stable_across_input_order() -> None: assert first.distance(source, target) == second.distance(source, target) +def test_knn_graph_connects_neighbors_across_similarity_blocks() -> None: + chunk_ids = [f"chunk-{index:03d}" for index in range(258)] + embeddings = [[1.0, index / 1000.0] for index in range(258)] + + graph = KNNGraph.from_embeddings( + chunk_ids, + embeddings, + neighbor_count=1, + maximum_distance=3.0, + ) + + assert graph.distance("chunk-255", "chunk-256") < graph.maximum_distance + + +def test_sparse_and_standard_library_shortest_paths_match( + monkeypatch: pytest.MonkeyPatch, +) -> None: + chunk_ids = ["a", "b", "c", "d"] + embeddings = [[1.0, 0.0], [0.9, 0.1], [0.1, 0.9], [0.0, 1.0]] + accelerated = KNNGraph.from_embeddings( + chunk_ids, + embeddings, + neighbor_count=1, + ) + monkeypatch.setattr(knn_module, "_load_scipy_sparse", lambda: None) + standard_library = KNNGraph.from_embeddings( + chunk_ids, + embeddings, + neighbor_count=1, + ) + + for source in chunk_ids: + for target in chunk_ids: + assert accelerated.distance(source, target) == pytest.approx( + standard_library.distance(source, target) + ) + + @pytest.mark.parametrize( ("chunk_ids", "embeddings", "message"), [