From daa6629879405fdb90bcb62689d627ab5d8a3ebe Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 13:37:52 +0300 Subject: [PATCH] Add the federated ClapNQ retrieval dataset Measures whether cross-database fusion reaches what a query needs, scored on retrieval alone so no model or judge sits between the fusion and the number. The corpus is MTRAG ClapNQ partitioned by article title, whole titles to a collection, so an article's passages never split and a query's gold stays concentrated in one collection, which is the condition a per-collection depth quota punishes. collection_of keys on sha256 rather than hash(), which is salted per process: the partition is never stored, and scoring recomputes it in a different process than the one that ingested. The 148 titles holding a gold passage carry 10,723 passages between them, so a budget near that floor leaves no cross-topic distractors and inflates recall. The default is 40,000 and the build reports the gold/distractor split, warning when there are none. build_databases opens each collection by configured name with a scope of one, since populate_db writes to a single database. The operator entry point emits the config for the partition it just built, so a config cannot search a differently-partitioned build. Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc --- CHANGELOG.md | 3 + evaluations/configs/mtrag_federated.yaml | 56 ++++ evaluations/evaluations/datasets/__init__.py | 2 + .../evaluations/datasets/mtrag_federated.py | 280 ++++++++++++++++++ evaluations/tests/test_mtrag_federated.py | 195 ++++++++++++ 5 files changed, 536 insertions(+) create mode 100644 evaluations/configs/mtrag_federated.yaml create mode 100644 evaluations/evaluations/datasets/mtrag_federated.py create mode 100644 evaluations/tests/test_mtrag_federated.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5971c43b..80c95fb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,9 @@ - `dot` removed from `search.vector_index_metric`; switch to `cosine` or `l2` and rerun `create-index`. +### Added + +- `mtrag_federated` evaluation dataset and its reference config `evaluations/configs/mtrag_federated.yaml`: MTRAG ClapNQ partitioned by article title into `n` collections, scored on retrieval only with Recall@5/@10, nDCG@5 and MAP. `python -m evaluations.datasets.mtrag_federated --config REF --n N --out PATH` builds the partition and emits the config that searches it. ### Fixed diff --git a/evaluations/configs/mtrag_federated.yaml b/evaluations/configs/mtrag_federated.yaml new file mode 100644 index 00000000..1902974c --- /dev/null +++ b/evaluations/configs/mtrag_federated.yaml @@ -0,0 +1,56 @@ +# Reference config for `mtrag_federated`: IBM MTRAG ClapNQ, partitioned by +# article title into four collections, scored on retrieval only. +# +# Build the partition and emit the config that searches exactly it: +# uv run python -m evaluations.datasets.mtrag_federated \ +# --config configs/mtrag_federated.yaml --n 4 --out ~/configs/fed-n4.yaml +# evaluations run mtrag_federated --config ~/configs/fed-n4.yaml \ +# --skip-db --skip-qa +# +# The databases below are the canonical n=4 partition at seed 20260831. Sweep +# configs for other collection counts live outside the repo, because a config +# here must be named after a registered dataset. +# +# No reranking block on purpose: this eval measures the reciprocal-rank fusion +# path, where retrieval depth per collection is `limit // n`. Adding a reranker +# is the comparison arm, not the baseline. +# base_url uses the `vllm` host serving each model over an OpenAI-compatible API. + +environment: development + +storage: + auto_vacuum: false + +lancedb: + # Declaration order is load-bearing: fusion resolves equal ranks to this + # order, so permuting these four keys is an arm. + databases: + clapnq_0: ${HOME}/.local/share/haiku.rag/evaluations/dbs/mtrag_federated_s20260831_n4_0.lancedb + clapnq_1: ${HOME}/.local/share/haiku.rag/evaluations/dbs/mtrag_federated_s20260831_n4_1.lancedb + clapnq_2: ${HOME}/.local/share/haiku.rag/evaluations/dbs/mtrag_federated_s20260831_n4_2.lancedb + clapnq_3: ${HOME}/.local/share/haiku.rag/evaluations/dbs/mtrag_federated_s20260831_n4_3.lancedb + +embeddings: + model: + provider: openai + name: qwen3-embedding-4b + vector_dim: 2560 + base_url: http://vllm:11431/v1 + +search: + # Matches the spec's retrieval_limit and the product default. + limit: 5 + +evaluations: + judge: + provider: openai + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 + temperature: 0.6 + max_tokens: 16384 + extra_body: + top_p: 0.95 + top_k: 20 + min_p: 0 + chat_template_kwargs: + reasoning_effort: low diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index 4a968b5f..61aa4bc8 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/evaluations/evaluations/datasets/__init__.py @@ -8,6 +8,7 @@ from .mtrag import ( MTRAG_CLAPNQ_REWRITE_SPEC, MTRAG_CLAPNQ_SPEC, ) +from .mtrag_federated import MTRAG_FEDERATED_SPEC from .open_rag_bench import ( ORB_MULTIMODAL_NEMOTRON_SPEC, ORB_MULTIMODAL_SPEC, @@ -24,6 +25,7 @@ DATASETS: dict[str, DatasetSpec] = { MTRAG_CLAPNQ_REWRITE_SPEC, MTRAG_CLAPNQ_LIVE_SPEC, MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC, + MTRAG_FEDERATED_SPEC, ORB_TEXT_SPEC, ORB_MULTIMODAL_SPEC, ORB_MULTIMODAL_NEMOTRON_SPEC, diff --git a/evaluations/evaluations/datasets/mtrag_federated.py b/evaluations/evaluations/datasets/mtrag_federated.py new file mode 100644 index 00000000..c1e77dc7 --- /dev/null +++ b/evaluations/evaluations/datasets/mtrag_federated.py @@ -0,0 +1,280 @@ +import argparse +import asyncio +import hashlib +import random +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import yaml +from datasets import Dataset + +from evaluations.config import DatasetSpec +from evaluations.datasets.mtrag import ( + _load_qrels, + build_mtrag_case, + load_clapnq_corpus, + load_clapnq_retrieval, + map_mtrag_document, + map_mtrag_retrieval, +) +from evaluations.evaluators import ( + CitationMAPEvaluator, + MAPEvaluator, + NDCGEvaluator, + RecallEvaluator, +) +from haiku.rag.config import load_yaml_config +from haiku.rag.config.models import AppConfig +from haiku.rag.utils import get_default_data_dir + +COLLECTION_PREFIX = "clapnq" +DEFAULT_SEED = 20260831 +# Whole titles are kept, and the 148 titles holding a gold passage carry 10,723 +# passages between them, so that is the floor. A budget near it leaves no +# cross-topic distractors at all and inflates recall; this default leaves about +# 29,000, a gold-title share near a quarter. +GOLD_TITLE_FLOOR = 10_723 +DEFAULT_BUDGET = 40_000 +INGEST_BATCH_SIZE = 512 + + +def collection_of(title: str, n: int, seed: int = DEFAULT_SEED) -> int: + """Which of `n` collections holds a title's passages. + + Keyed on the title, so an article's passages never split. sha256 rather than + `hash()`, which is salted per process: the partition is never stored, and + scoring recomputes it in a different process than the one that ingested. + """ + if n < 1: + raise ValueError("a partition needs at least one collection") + digest = hashlib.sha256(f"{seed}/{title}".encode()).digest() + return int.from_bytes(digest[:8], "big") % n + + +def collection_names(n: int) -> tuple[str, ...]: + """The collection names in declaration order. + + Order is load-bearing: fusion resolves equal ranks to configured order, so + permuting these names is an arm rather than a cosmetic change. + """ + return tuple(f"{COLLECTION_PREFIX}_{index}" for index in range(n)) + + +def database_paths(n: int, seed: int = DEFAULT_SEED) -> dict[str, str]: + """Where each collection's database lives. + + The partition is in the filename, so a build at one `(n, seed)` can never + overwrite another's databases or be searched by the wrong config. + """ + root = get_default_data_dir() / "evaluations" / "dbs" + return { + name: str(root / f"mtrag_federated_s{seed}_n{n}_{index}.lancedb") + for index, name in enumerate(collection_names(n)) + } + + +def sample_records( + records: Sequence[Mapping[str, Any]], + gold_ids: Iterable[str], + budget: int = DEFAULT_BUDGET, + seed: int = DEFAULT_SEED, +) -> list[Mapping[str, Any]]: + """A fixed sub-corpus: every gold passage, plus seeded distractor titles. + + Whole titles are kept or dropped together. Gold is mandatory, so a budget + below the gold floor yields the gold titles alone rather than an incomplete + corpus that would score as missing retrievals. + """ + by_title: dict[str, list[Mapping[str, Any]]] = {} + title_of: dict[str, str] = {} + for row in records: + by_title.setdefault(row["title"], []).append(row) + title_of[row["_id"]] = row["title"] + + wanted = set(gold_ids) + missing = sorted(wanted - set(title_of)) + if missing: + raise ValueError( + f"{len(missing)} gold passages do not resolve to the corpus, " + f"first few: {missing[:3]}" + ) + + gold_titles = {title_of[passage_id] for passage_id in wanted} + kept = {title for title in by_title if title in gold_titles} + total = sum(len(by_title[title]) for title in kept) + + distractors = [title for title in by_title if title not in gold_titles] + random.Random(seed).shuffle(distractors) + for title in distractors: + size = len(by_title[title]) + if total + size > budget: + continue + kept.add(title) + total += size + + return [row for row in records if row["title"] in kept] + + +def partition_records( + records: Sequence[Mapping[str, Any]], + n: int, + seed: int = DEFAULT_SEED, +) -> dict[str, list[Mapping[str, Any]]]: + """Route every record to its collection, naming all `n` even when empty.""" + names = collection_names(n) + grouped: dict[str, list[Mapping[str, Any]]] = {name: [] for name in names} + for row in records: + grouped[names[collection_of(row["title"], n, seed)]].append(row) + return grouped + + +def pool_composition( + records: Sequence[Mapping[str, Any]], gold_ids: Iterable[str] +) -> tuple[int, int]: + """Passages in gold-bearing titles, and passages in distractor titles. + + A pool with no distractors scores as an easy retrieval task and says + nothing, so the build reports this rather than leaving it to be inferred + from the budget. + """ + wanted = set(gold_ids) + gold_titles = {row["title"] for row in records if row["_id"] in wanted} + gold_side = sum(1 for row in records if row["title"] in gold_titles) + return gold_side, len(records) - gold_side + + +def gold_passage_ids() -> set[str]: + """Every corpus id the ClapNQ qrels reference.""" + return {passage_id for ids in _load_qrels().values() for passage_id in ids} + + +def load_pool( + budget: int = DEFAULT_BUDGET, seed: int = DEFAULT_SEED +) -> list[Mapping[str, Any]]: + records = [dict(row) for row in load_clapnq_corpus()] + return sample_records(records, gold_passage_ids(), budget, seed) + + +def _unused_document_loader() -> Dataset: + raise RuntimeError( + "the federated corpus is built by build_databases(); run with --skip-db" + ) + + +def emitted_config(reference: Path, n: int, seed: int = DEFAULT_SEED) -> dict[str, Any]: + """The reference config with this partition's databases placed in it.""" + settings = load_yaml_config(reference) + lancedb = dict(settings.get("lancedb") or {}) + lancedb.pop("uri", None) + lancedb["databases"] = database_paths(n, seed) + settings["lancedb"] = lancedb + return settings + + +async def build_databases( + config: AppConfig, + n: int, + seed: int = DEFAULT_SEED, + budget: int = DEFAULT_BUDGET, +) -> dict[str, int]: + """Ingest the partition into one database per collection. + + Each member is opened by configured name with a scope of one, which is what + makes `create=True` legal on a client whose config places several. + """ + from haiku.rag.client import HaikuRAG + + from evaluations.population import _ingest_batched + + names = collection_names(n) + configured = set(config.lancedb.databases or {}) + missing = sorted(set(names) - configured) + if missing: + raise ValueError( + f"lancedb.databases must place every collection; missing {missing}" + ) + + pool = load_pool(budget, seed) + gold_side, distractors = pool_composition(pool, gold_passage_ids()) + print( + f"pool: {len(pool)} passages, {gold_side} in gold-bearing titles, " + f"{distractors} distractors" + ) + if not distractors: + print( + " WARNING: no distractor titles, so every passage belongs to an " + f"answer-bearing article; raise --budget above {GOLD_TITLE_FLOOR}" + ) + grouped = partition_records(pool, n, seed) + written: dict[str, int] = {} + for name in names: + async with HaikuRAG(config=config, sources=[name], create=True) as client: + await _ingest_batched( + client, MTRAG_FEDERATED_SPEC, grouped[name], INGEST_BATCH_SIZE + ) + written[name] = len(grouped[name]) + return written + + +MTRAG_FEDERATED_SPEC = DatasetSpec( + key="mtrag_federated", + # Never read: the run searches the configured set. Present because the spec + # requires one, and pointed at the first collection so a stray --db is + # obviously wrong rather than silently plausible. + db_filename="mtrag_federated_unused.lancedb", + document_loader=_unused_document_loader, + document_mapper=map_mtrag_document, + # The QA phase is not wired yet, and the loader is what run_qa_benchmark + # reaches first, so a forgotten --skip-qa fails loudly there. The builder is + # the one the generation tasks will need when QA arrives. + qa_loader=_unused_document_loader, + qa_case_builder=build_mtrag_case, + retrieval_loader=lambda: load_clapnq_retrieval("lastturn"), + retrieval_mapper=map_mtrag_retrieval, + retrieval_evaluators=[ + RecallEvaluator(5), + RecallEvaluator(10), + NDCGEvaluator(5), + MAPEvaluator(), + ], + citation_evaluator=CitationMAPEvaluator(), + # The product default, which is where the fusion depth quota bites. + retrieval_limit=5, + ingest_batch_size=INGEST_BATCH_SIZE, +) + + +async def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Build the federated ClapNQ partition and emit the config that " + "searches exactly it." + ) + ) + parser.add_argument("--config", type=Path, required=True, help="reference config") + parser.add_argument("--n", type=int, required=True, help="collection count") + parser.add_argument("--seed", type=int, default=DEFAULT_SEED) + parser.add_argument("--budget", type=int, default=DEFAULT_BUDGET) + parser.add_argument( + "--out", + type=Path, + required=True, + help="where to write the emitted config for this partition", + ) + args = parser.parse_args() + + settings = emitted_config(args.config, args.n, args.seed) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(yaml.safe_dump(settings, sort_keys=False)) + print(f"wrote {args.out}") + + # Reload from disk, so the config that builds is the file that will search. + config = AppConfig.model_validate(load_yaml_config(args.out)) + written = await build_databases(config, args.n, args.seed, args.budget) + for name, count in written.items(): + print(f"{name}: {count} passages") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/evaluations/tests/test_mtrag_federated.py b/evaluations/tests/test_mtrag_federated.py new file mode 100644 index 00000000..fd93357d --- /dev/null +++ b/evaluations/tests/test_mtrag_federated.py @@ -0,0 +1,195 @@ +import os +import subprocess +import sys + +import pytest + +from evaluations.datasets import DATASETS +from evaluations.datasets.mtrag_federated import ( + DEFAULT_BUDGET, + GOLD_TITLE_FLOOR, + MTRAG_FEDERATED_SPEC, + collection_names, + collection_of, + partition_records, + pool_composition, + sample_records, +) + + +def record(passage_id: str, title: str) -> dict[str, str]: + return {"_id": passage_id, "title": title, "text": f"text of {passage_id}"} + + +def corpus(titles: dict[str, int]) -> list[dict[str, str]]: + """One record per passage, `titles` mapping a title to its passage count.""" + return [ + record(f"{title}_{index}", title) + for title, count in titles.items() + for index in range(count) + ] + + +class TestCollectionOf: + def test_assigns_within_range(self) -> None: + for n in (2, 4, 8): + assigned = {collection_of(f"title {i}", n) for i in range(200)} + assert assigned <= set(range(n)) + + def test_uses_every_collection(self) -> None: + """A partition that leaves a collection empty is not a partition.""" + for n in (2, 4, 8): + assigned = {collection_of(f"title {i}", n) for i in range(200)} + assert assigned == set(range(n)) + + def test_is_stable_across_processes(self) -> None: + """Salted `hash()` would make a build unreproducible between runs. + + The partition is never stored, so scoring recomputes it in a different + process than the one that ingested. + """ + code = ( + "from evaluations.datasets.mtrag_federated import collection_of;" + "print([collection_of(f'title {i}', 8) for i in range(12)])" + ) + runs = { + subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + check=True, + env={**os.environ, "PYTHONHASHSEED": seed}, + ).stdout.strip() + for seed in ("0", "1", "12345") + } + assert len(runs) == 1, f"assignment varies with PYTHONHASHSEED: {runs}" + + def test_seed_changes_the_assignment(self) -> None: + titles = [f"title {i}" for i in range(200)] + one = [collection_of(t, 8, seed=1) for t in titles] + two = [collection_of(t, 8, seed=2) for t in titles] + assert one != two + + def test_rejects_a_collection_count_below_one(self) -> None: + with pytest.raises(ValueError, match="at least one collection"): + collection_of("title", 0) + + +class TestCollectionNames: + def test_names_one_per_collection(self) -> None: + assert collection_names(3) == ( + "clapnq_0", + "clapnq_1", + "clapnq_2", + ) + + def test_declaration_order_is_the_name_order(self) -> None: + """Fusion resolves ties to configured order, so the order is load-bearing.""" + names = collection_names(4) + assert list(names) == sorted(names, key=lambda name: int(name.split("_")[1])) + + +class TestPartitionRecords: + def test_keeps_every_record(self) -> None: + records = corpus({"a": 3, "b": 2, "c": 4}) + grouped = partition_records(records, 2) + assert sum(len(rows) for rows in grouped.values()) == len(records) + + def test_never_splits_a_title(self) -> None: + """A title is the atom: its passages must share a collection, or a + query's gold spreads for reasons the partition never intended.""" + records = corpus({f"title {i}": 5 for i in range(40)}) + grouped = partition_records(records, 4) + holders: dict[str, set[str]] = {} + for name, rows in grouped.items(): + for row in rows: + holders.setdefault(row["title"], set()).add(name) + split = {title: names for title, names in holders.items() if len(names) > 1} + assert not split, f"titles split across collections: {split}" + + def test_names_every_collection_even_when_one_is_empty(self) -> None: + """The config declares n databases, so the build must create n.""" + records = corpus({"only": 2}) + grouped = partition_records(records, 4) + assert set(grouped) == set(collection_names(4)) + + +class TestSampleRecords: + def test_keeps_every_gold_passage(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + gold = {"title 3_1", "title 17_4", "title 42_9"} + sampled = sample_records(records, gold, budget=60) + assert gold <= {row["_id"] for row in sampled} + + def test_keeps_whole_titles_holding_gold(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + sampled = sample_records(records, {"title 3_1"}, budget=0) + assert sorted(row["_id"] for row in sampled) == sorted( + f"title 3_{i}" for i in range(10) + ) + + def test_respects_the_budget(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + sampled = sample_records(records, {"title 3_1"}, budget=100) + assert len(sampled) <= 100 + + def test_budget_below_the_gold_floor_still_keeps_gold(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + gold = {f"title {i}_0" for i in range(20)} + sampled = sample_records(records, gold, budget=5) + assert len(sampled) == 200 + assert gold <= {row["_id"] for row in sampled} + + def test_is_stable_for_a_seed(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + first = sample_records(records, {"title 0_0"}, budget=100, seed=7) + second = sample_records(records, {"title 0_0"}, budget=100, seed=7) + assert [row["_id"] for row in first] == [row["_id"] for row in second] + + def test_seed_changes_the_distractors(self) -> None: + records = corpus({f"title {i}": 10 for i in range(50)}) + first = sample_records(records, {"title 0_0"}, budget=100, seed=7) + second = sample_records(records, {"title 0_0"}, budget=100, seed=8) + assert {row["_id"] for row in first} != {row["_id"] for row in second} + + def test_rejects_gold_the_corpus_does_not_hold(self) -> None: + records = corpus({"a": 2}) + with pytest.raises(ValueError, match="do not resolve"): + sample_records(records, {"missing"}, budget=10) + + +class TestSpec: + def test_registers_under_its_key(self) -> None: + assert DATASETS[MTRAG_FEDERATED_SPEC.key] is MTRAG_FEDERATED_SPEC + + def test_opts_out_of_the_shared_population(self) -> None: + """The databases are built by build_databases, not populate_db.""" + with pytest.raises(RuntimeError, match="build_databases"): + MTRAG_FEDERATED_SPEC.document_loader() + + def test_retrieval_limit_matches_the_product_default(self) -> None: + """5 is config's search.limit, the setting the depth quota bites at.""" + assert MTRAG_FEDERATED_SPEC.retrieval_limit == 5 + + def test_scores_retrieval_without_a_judge(self) -> None: + assert MTRAG_FEDERATED_SPEC.retrieval_evaluators + assert MTRAG_FEDERATED_SPEC.retrieval_loader is not None + assert MTRAG_FEDERATED_SPEC.retrieval_mapper is not None + + +class TestPoolComposition: + def test_separates_gold_bearing_titles_from_distractors(self) -> None: + records = corpus({"answers": 4, "filler": 6}) + gold_side, distractors = pool_composition(records, {"answers_2"}) + assert (gold_side, distractors) == (4, 6) + + def test_reports_no_distractors_when_the_budget_is_at_the_floor(self) -> None: + """A pool of only answer-bearing articles scores as an easy task and + says nothing, so the build has to be able to see it.""" + records = corpus({f"title {i}": 10 for i in range(5)}) + gold = {f"title {i}_0" for i in range(5)} + sampled = sample_records(records, gold, budget=1) + assert pool_composition(sampled, gold) == (50, 0) + + def test_the_default_budget_clears_the_gold_floor(self) -> None: + assert DEFAULT_BUDGET > GOLD_TITLE_FLOOR