diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cba3c29..1e5cb7f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Added - `evaluations run --retrieval-limit N`: candidates each database fetches during the retrieval benchmark, overriding the dataset's `retrieval_limit`. +- `mtrag_pooled` evaluation dataset and its reference config `evaluations/configs/mtrag_pooled.yaml`: all four MTRAG domains pooled and partitioned across `n` collections, `--alpha` interpolating between one domain per collection and a uniform shard. - `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_pooled.yaml b/evaluations/configs/mtrag_pooled.yaml new file mode 100644 index 00000000..63916a33 --- /dev/null +++ b/evaluations/configs/mtrag_pooled.yaml @@ -0,0 +1,49 @@ +# Reference config for `mtrag_pooled`: all four MTRAG domains (clapnq, cloud, +# fiqa, govt) pooled and partitioned across `n` collections, scored on retrieval +# only. +# +# This is the heterogeneous corpus. `mtrag_federated` partitions one domain by +# article title, which is round-robin fusion's friendliest case: no collection is +# ever off-topic for a query, so the guaranteed-slot waste is never exercised. +# Here a query belongs to one domain and the others are genuinely off-topic. +# +# Build the partition and emit the config that searches exactly it: +# uv run python -m evaluations.datasets.mtrag_federated \ +# --config configs/mtrag_pooled.yaml --pooled --n 4 --alpha 0 \ +# --out ~/configs/pooled-n4-a0.yaml +# evaluations run mtrag_pooled --config ~/configs/pooled-n4-a0.yaml \ +# --skip-db --skip-qa +# +# alpha 0 keeps each collection to one domain; alpha 1 shards titles across all +# of them, which is the degenerate sharding endpoint rather than a rival design. +# The operator emits the databases, so none are listed here. + +environment: development + +storage: + auto_vacuum: false + +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 61aa4bc8..1fd73a8f 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/evaluations/evaluations/datasets/__init__.py @@ -8,7 +8,7 @@ from .mtrag import ( MTRAG_CLAPNQ_REWRITE_SPEC, MTRAG_CLAPNQ_SPEC, ) -from .mtrag_federated import MTRAG_FEDERATED_SPEC +from .mtrag_federated import MTRAG_FEDERATED_SPEC, MTRAG_POOLED_SPEC from .open_rag_bench import ( ORB_MULTIMODAL_NEMOTRON_SPEC, ORB_MULTIMODAL_SPEC, @@ -26,6 +26,7 @@ DATASETS: dict[str, DatasetSpec] = { MTRAG_CLAPNQ_LIVE_SPEC, MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC, MTRAG_FEDERATED_SPEC, + MTRAG_POOLED_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 index c75a0c05..801af935 100644 --- a/evaluations/evaluations/datasets/mtrag_federated.py +++ b/evaluations/evaluations/datasets/mtrag_federated.py @@ -1,7 +1,9 @@ import argparse import asyncio import hashlib +import json import random +import zipfile from collections.abc import Iterable, Mapping, Sequence from pathlib import Path from typing import Any @@ -29,6 +31,8 @@ from haiku.rag.config.models import AppConfig from haiku.rag.utils import get_default_data_dir COLLECTION_PREFIX = "clapnq" +# The four corpora MTRAG ships, in upstream order. +DOMAINS = ("clapnq", "cloud", "fiqa", "govt") 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 @@ -40,17 +44,66 @@ INGEST_BATCH_SIZE = 512 FTS_INDEX_NAME = "content_fts_idx" -def collection_of(title: str, n: int, seed: int = DEFAULT_SEED) -> int: +def _hash_int(payload: str) -> int: + """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.""" + return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], "big") + + +def _unit(payload: str) -> float: + """A stable value in [0, 1) for probabilistic assignment.""" + return (_hash_int(payload) % 10**9) / 10**9 + + +def domain_files(domain: str, variant: str = "lastturn") -> tuple[str, str, str]: + """Corpus, qrels and query paths for one MTRAG domain.""" + return ( + f"corpora/passage_level/{domain}.jsonl.zip", + f"mtrag-human/retrieval_tasks/{domain}/qrels/dev.tsv", + f"mtrag-human/retrieval_tasks/{domain}/{domain}_{variant}.jsonl", + ) + + +def _domain_collection(title: str, domain: str, n: int, seed: int) -> int: + """The collection a title takes from its domain. + + Domains map onto collections proportionally: with more collections than + domains each domain is subdivided by title, with fewer, domains are grouped. + """ + index = DOMAINS.index(domain) + count = len(DOMAINS) + if n >= count: + per = n // count + return index * per + (_hash_int(f"{seed}/sub/{title}") % per) + return index * n // count + + +def collection_of( + title: str, + n: int, + seed: int = DEFAULT_SEED, + alpha: float = 0.0, + domain: str | None = None, +) -> 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. + Keyed on the title, so an article's passages never split. + + Without a domain the assignment is a pure hash — a topically arbitrary + grouping of whole articles, which is all the quota and order-bias arms need. + With one, `alpha` interpolates between the domain partition (0, each + collection one topic) and a uniform shard (1, domain ignored). Sharding is + the endpoint of the knob rather than a rival design. """ 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 + shard = _hash_int(f"{seed}/{title}") % n + if domain is None or alpha >= 1.0: + return shard + if alpha > 0.0 and _unit(f"{seed}/alpha/{title}") < alpha: + return shard + return _domain_collection(title, domain, n, seed) def collection_names(n: int) -> tuple[str, ...]: @@ -157,6 +210,115 @@ def load_pool( return sample_records(records, gold_passage_ids(), budget, seed) +POOLED_PREFIX = "dom" + + +class PassageIdCollision(AssertionError): + """Two domains claim the same passage id, so uri-keyed gold is ambiguous.""" + + +def pooled_collection_names(n: int) -> tuple[str, ...]: + """Names for the pooled partition, positional and distinct from the + single-domain set so the two never share database paths.""" + return tuple(f"{POOLED_PREFIX}_{index}" for index in range(n)) + + +def pooled_database_paths( + n: int, alpha: float, seed: int = DEFAULT_SEED +) -> dict[str, str]: + root = get_default_data_dir() / "evaluations" / "dbs" + tag = f"s{seed}_a{alpha:g}_n{n}" + return { + name: str(root / f"mtrag_pooled_{tag}_{index}.lancedb") + for index, name in enumerate(pooled_collection_names(n)) + } + + +def load_pooled_records() -> list[Mapping[str, Any]]: + """Every passage of all four domains, each tagged with the domain it came + from. Raises when two domains claim one passage id, since gold is uri-keyed. + """ + from evaluations.datasets.mtrag import _download + + records: list[Mapping[str, Any]] = [] + seen: dict[str, str] = {} + for domain in DOMAINS: + corpus_file, _, _ = domain_files(domain) + path = _download(corpus_file) + with zipfile.ZipFile(path) as archive: + with archive.open(archive.namelist()[0]) as handle: + for line in handle: + row = json.loads(line) + passage_id = row["_id"] + if passage_id in seen and seen[passage_id] != domain: + raise PassageIdCollision( + f"{passage_id} claimed by {seen[passage_id]} and {domain}" + ) + seen[passage_id] = domain + records.append( + { + "_id": passage_id, + "title": row["title"], + "text": row["text"], + "domain": domain, + } + ) + return records + + +def load_pooled_queries(variant: str = "lastturn") -> list[dict[str, Any]]: + """Retrieval queries from every domain, each with its gold passage uris.""" + from evaluations.datasets.mtrag import _download, _parse_qrels + + out: list[dict[str, Any]] = [] + for domain in DOMAINS: + _, qrels_file, query_file = domain_files(domain, variant) + qrels = _parse_qrels(_download(qrels_file).read_text().splitlines()) + for line in _download(query_file).read_text().splitlines(): + if not line.strip(): + continue + query = json.loads(line) + expected = qrels.get(query["_id"]) + if not expected: + continue + out.append( + { + "query_id": f"{domain}/{query['_id']}", + "question": query["text"], + "expected_uris": expected, + "domain": domain, + } + ) + return out + + +def pooled_gold_ids(variant: str = "lastturn") -> set[str]: + return { + uri for query in load_pooled_queries(variant) for uri in query["expected_uris"] + } + + +def load_pooled( + budget: int = DEFAULT_BUDGET, seed: int = DEFAULT_SEED +) -> list[Mapping[str, Any]]: + return sample_records(load_pooled_records(), pooled_gold_ids(), budget, seed) + + +def partition_pooled( + records: Sequence[Mapping[str, Any]], + n: int, + alpha: float, + seed: int = DEFAULT_SEED, +) -> dict[str, list[Mapping[str, Any]]]: + """Route pooled records to collections, honouring each record's domain.""" + names = pooled_collection_names(n) + grouped: dict[str, list[Mapping[str, Any]]] = {name: [] for name in names} + for row in records: + index = collection_of(row["title"], n, seed, alpha=alpha, domain=row["domain"]) + grouped[names[index]].append(row) + return grouped + + def _unused_document_loader() -> Dataset: raise RuntimeError( "the federated corpus is built by build_databases(); run with --skip-db" @@ -244,6 +406,48 @@ async def build_databases( return written +async def build_pooled_databases( + config: AppConfig, + n: int, + alpha: float, + seed: int = DEFAULT_SEED, + budget: int = DEFAULT_BUDGET, +) -> dict[str, int]: + """Ingest the four-domain pooled partition, one database per collection.""" + from haiku.rag.client import HaikuRAG + + from evaluations.population import _ingest_batched + + names = pooled_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_pooled(budget, seed) + by_domain: dict[str, int] = {} + for row in pool: + by_domain[row["domain"]] = by_domain.get(row["domain"], 0) + 1 + gold_side, distractors = pool_composition(pool, pooled_gold_ids()) + print( + f"pool: {len(pool)} passages, {gold_side} in gold-bearing titles, " + f"{distractors} distractors, by domain {by_domain}" + ) + grouped = partition_pooled(pool, n, alpha, 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_POOLED_SPEC, grouped[name], INGEST_BATCH_SIZE + ) + await client.store.vacuum(retention_seconds=0) + await assert_fts_covers_rows(client.store.chunks_table, name) + 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 @@ -272,6 +476,27 @@ MTRAG_FEDERATED_SPEC = DatasetSpec( ) +MTRAG_POOLED_SPEC = DatasetSpec( + key="mtrag_pooled", + db_filename="mtrag_pooled_unused.lancedb", + document_loader=_unused_document_loader, + document_mapper=map_mtrag_document, + qa_loader=_unused_document_loader, + qa_case_builder=build_mtrag_case, + retrieval_loader=lambda: Dataset.from_list(load_pooled_queries("lastturn")), + retrieval_mapper=map_mtrag_retrieval, + retrieval_evaluators=[ + RecallEvaluator(5), + RecallEvaluator(10), + NDCGEvaluator(5), + MAPEvaluator(), + ], + citation_evaluator=CitationMAPEvaluator(), + retrieval_limit=5, + ingest_batch_size=INGEST_BATCH_SIZE, +) + + async def main() -> None: parser = argparse.ArgumentParser( description=( @@ -283,6 +508,17 @@ async def main() -> None: 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( + "--pooled", + action="store_true", + help="build the four-domain pooled corpus instead of clapnq alone", + ) + parser.add_argument( + "--alpha", + type=float, + default=0.0, + help="pooled only: 0 keeps a collection to one domain, 1 shards across all", + ) parser.add_argument( "--out", type=Path, @@ -291,14 +527,26 @@ async def main() -> None: ) args = parser.parse_args() - settings = emitted_config(args.config, args.n, args.seed) + if args.pooled: + settings = load_yaml_config(args.config) + lancedb = dict(settings.get("lancedb") or {}) + lancedb.pop("uri", None) + lancedb["databases"] = pooled_database_paths(args.n, args.alpha, args.seed) + settings["lancedb"] = lancedb + else: + 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) + if args.pooled: + written = await build_pooled_databases( + config, args.n, args.alpha, args.seed, args.budget + ) + else: + written = await build_databases(config, args.n, args.seed, args.budget) for name, count in written.items(): print(f"{name}: {count} passages") diff --git a/evaluations/tests/test_mtrag_federated.py b/evaluations/tests/test_mtrag_federated.py index 593ec390..f37943ea 100644 --- a/evaluations/tests/test_mtrag_federated.py +++ b/evaluations/tests/test_mtrag_federated.py @@ -7,12 +7,17 @@ import pytest from evaluations.datasets import DATASETS from evaluations.datasets.mtrag_federated import ( DEFAULT_BUDGET, + DOMAINS, FTSIndexNotCoveringRows, GOLD_TITLE_FLOOR, MTRAG_FEDERATED_SPEC, assert_fts_covers_rows, collection_names, collection_of, + domain_files, + partition_pooled, + pooled_collection_names, + pooled_database_paths, partition_records, pool_composition, sample_records, @@ -296,3 +301,123 @@ class TestFTSCoverageAssertion: async def test_rejects_a_missing_index(self) -> None: with pytest.raises(FTSIndexNotCoveringRows, match="no content_fts_idx"): await assert_fts_covers_rows(self._Table(100, None), "clapnq_0") + + +class TestDomains: + def test_names_the_four_upstream_domains(self) -> None: + assert DOMAINS == ("clapnq", "cloud", "fiqa", "govt") + + def test_paths_follow_the_upstream_layout(self) -> None: + assert domain_files("govt") == ( + "corpora/passage_level/govt.jsonl.zip", + "mtrag-human/retrieval_tasks/govt/qrels/dev.tsv", + "mtrag-human/retrieval_tasks/govt/govt_lastturn.jsonl", + ) + + +class TestDomainPartition: + """With four real domains, alpha finally means something: 0 keeps a + collection to one topic, 1 shards titles across all of them.""" + + def test_alpha_zero_keeps_a_domain_together_when_n_matches(self) -> None: + for domain_index, domain in enumerate(DOMAINS): + assigned = { + collection_of(f"{domain} title {i}", 4, alpha=0.0, domain=domain) + for i in range(50) + } + assert assigned == {domain_index} + + def test_alpha_zero_subdivides_within_a_domain_when_n_exceeds_it(self) -> None: + for domain_index, domain in enumerate(DOMAINS): + assigned = { + collection_of(f"{domain} title {i}", 8, alpha=0.0, domain=domain) + for i in range(200) + } + assert assigned == {domain_index * 2, domain_index * 2 + 1} + + def test_alpha_zero_groups_domains_when_n_is_below_it(self) -> None: + assigned = { + (domain, collection_of(f"t{i}", 2, alpha=0.0, domain=domain)) + for domain in DOMAINS + for i in range(20) + } + by_collection: dict[int, set[str]] = {} + for domain, collection in assigned: + by_collection.setdefault(collection, set()).add(domain) + assert set(by_collection) == {0, 1} + assert all(len(v) == 2 for v in by_collection.values()) + + def test_alpha_one_ignores_the_domain(self) -> None: + """The shard endpoint: a title's collection must not depend on its domain.""" + titles = [f"title {i}" for i in range(200)] + as_clapnq = [collection_of(t, 8, alpha=1.0, domain="clapnq") for t in titles] + as_govt = [collection_of(t, 8, alpha=1.0, domain="govt") for t in titles] + assert as_clapnq == as_govt + + def test_alpha_one_spreads_a_single_domain_across_every_collection(self) -> None: + assigned = { + collection_of(f"title {i}", 8, alpha=1.0, domain="clapnq") + for i in range(400) + } + assert assigned == set(range(8)) + + def test_intermediate_alpha_moves_some_titles_off_their_domain(self) -> None: + titles = [f"title {i}" for i in range(400)] + home = [collection_of(t, 4, alpha=0.0, domain="fiqa") for t in titles] + mixed = [collection_of(t, 4, alpha=0.5, domain="fiqa") for t in titles] + moved = sum(1 for a, b in zip(home, mixed) if a != b) + assert 0 < moved < len(titles), f"alpha=0.5 moved {moved} of {len(titles)}" + + def test_default_alpha_is_the_domain_partition(self) -> None: + for domain in DOMAINS: + assert collection_of("t", 4, domain=domain) == collection_of( + "t", 4, alpha=0.0, domain=domain + ) + + +class TestPooledPartition: + def test_names_are_distinct_from_the_single_domain_set(self) -> None: + """The two datasets must never share database paths.""" + assert not set(pooled_collection_names(4)) & set(collection_names(4)) + + def test_database_paths_separate_alpha_and_n(self) -> None: + a = pooled_database_paths(4, 0.0) + b = pooled_database_paths(4, 1.0) + c = pooled_database_paths(8, 0.0) + assert not set(a.values()) & set(b.values()) + assert not set(a.values()) & set(c.values()) + + def test_routes_each_record_by_its_own_domain(self) -> None: + records = [ + { + "_id": f"{domain}-{i}", + "title": f"{domain} t{i}", + "text": "x", + "domain": domain, + } + for domain in DOMAINS + for i in range(20) + ] + grouped = partition_pooled(records, 4, alpha=0.0) + for name, rows in grouped.items(): + domains = {row["domain"] for row in rows} + assert len(domains) == 1, f"{name} mixes domains at alpha=0: {domains}" + + def test_alpha_one_mixes_domains_in_every_collection(self) -> None: + records = [ + {"_id": f"{domain}-{i}", "title": f"t{i}", "text": "x", "domain": domain} + for domain in DOMAINS + for i in range(60) + ] + grouped = partition_pooled(records, 4, alpha=1.0) + assert all(len({r["domain"] for r in rows}) > 1 for rows in grouped.values()) + + def test_keeps_every_record(self) -> None: + records = [ + {"_id": f"{d}-{i}", "title": f"{d} t{i}", "text": "x", "domain": d} + for d in DOMAINS + for i in range(15) + ] + for alpha in (0.0, 0.5, 1.0): + grouped = partition_pooled(records, 8, alpha=alpha) + assert sum(len(v) for v in grouped.values()) == len(records)