From 22d3a7c421fac9724c6697a4b637998616aa123f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 17:20:23 +0300 Subject: [PATCH] Sample and partition the pooled corpus at passage level `title` is the empty string for every cloud and fiqa passage, so two of the four domains have exactly one title covering 72,442 and 61,022 passages, and govt's titles are web-scrape artifacts with one 10,192-passage bucket. Only clapnq has titles that identify a document. Keeping whole titles therefore put the pooled gold floor at 146,543 passages: a budget of 120,000 yielded zero distractors, and 58 gold titles alone accounted for 135,479 passages. Passage level costs nothing the heterogeneous comparison needs. At alpha=0 the domain places a collection, so a query's gold is concentrated by construction rather than by the atom, and a titleless domain now spreads across its own collections instead of collapsing into one. Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc --- .../evaluations/datasets/mtrag_federated.py | 46 ++++++++++++++- evaluations/tests/test_mtrag_federated.py | 57 +++++++++++++++++++ 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/evaluations/evaluations/datasets/mtrag_federated.py b/evaluations/evaluations/datasets/mtrag_federated.py index 801af935..51b30a70 100644 --- a/evaluations/evaluations/datasets/mtrag_federated.py +++ b/evaluations/evaluations/datasets/mtrag_federated.py @@ -301,7 +301,43 @@ def pooled_gold_ids(variant: str = "lastturn") -> set[str]: 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) + return sample_pooled_records(load_pooled_records(), pooled_gold_ids(), budget, seed) + + +def sample_pooled_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 at passage level, keeping every gold passage. + + The single-domain dataset keeps whole titles, which cannot work here: `title` + is the empty string for every cloud and fiqa passage, so two of the four + domains have exactly one title covering 72,442 and 61,022 passages. Whole + titles put the gold floor at 146,543 passages, leaving no distractors at any + budget below the entire corpus. + + Passage level costs nothing this comparison needs: at alpha=0 the domain + places a collection, so a query's gold is concentrated by construction rather + than by the atom. + """ + wanted = set(gold_ids) + by_id = {row["_id"]: row for row in records} + missing = sorted(wanted - set(by_id)) + if missing: + raise ValueError( + f"{len(missing)} gold passages do not resolve to the pooled corpus, " + f"first few: {missing[:3]}" + ) + kept = set(wanted) + others = [row["_id"] for row in records if row["_id"] not in wanted] + random.Random(seed).shuffle(others) + for passage_id in others: + if len(kept) >= budget: + break + kept.add(passage_id) + return [row for row in records if row["_id"] in kept] def partition_pooled( @@ -310,11 +346,15 @@ def partition_pooled( alpha: float, seed: int = DEFAULT_SEED, ) -> dict[str, list[Mapping[str, Any]]]: - """Route pooled records to collections, honouring each record's domain.""" + """Route pooled records to collections, honouring each record's domain. + + Keyed on the passage id rather than the title, because two of the four + domains have no titles. See `sample_pooled_records`. + """ 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"]) + index = collection_of(row["_id"], n, seed, alpha=alpha, domain=row["domain"]) grouped[names[index]].append(row) return grouped diff --git a/evaluations/tests/test_mtrag_federated.py b/evaluations/tests/test_mtrag_federated.py index f37943ea..c9de74b8 100644 --- a/evaluations/tests/test_mtrag_federated.py +++ b/evaluations/tests/test_mtrag_federated.py @@ -18,6 +18,7 @@ from evaluations.datasets.mtrag_federated import ( partition_pooled, pooled_collection_names, pooled_database_paths, + sample_pooled_records, partition_records, pool_composition, sample_records, @@ -421,3 +422,59 @@ class TestPooledPartition: 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) + + +class TestPooledSampler: + """Two of the four domains have no titles at all, so the pooled corpus is + sampled and partitioned at passage level rather than by title.""" + + @staticmethod + def _pool(per_domain: int = 40) -> list[dict[str, str]]: + return [ + { + "_id": f"{domain}-{i}", + # cloud and fiqa carry an empty title upstream. + "title": "" if domain in ("cloud", "fiqa") else f"{domain} t{i}", + "text": "x", + "domain": domain, + } + for domain in DOMAINS + for i in range(per_domain) + ] + + def test_keeps_every_gold_passage(self) -> None: + pool = self._pool() + gold = {"cloud-3", "fiqa-7", "clapnq-1", "govt-39"} + kept = sample_pooled_records(pool, gold, budget=20) + assert gold <= {row["_id"] for row in kept} + + def test_respects_the_budget_above_the_gold_floor(self) -> None: + pool = self._pool() + kept = sample_pooled_records(pool, {"cloud-3"}, budget=25) + assert len(kept) == 25 + + def test_a_titleless_domain_does_not_drag_in_its_whole_corpus(self) -> None: + """The failure this replaces: whole-title keeping pulled all 72,442 cloud + passages in because they share one empty title.""" + pool = self._pool() + kept = sample_pooled_records(pool, {"cloud-3"}, budget=10) + cloud = [row for row in kept if row["domain"] == "cloud"] + assert len(cloud) < 40, f"kept {len(cloud)} of 40 cloud passages" + + def test_rejects_gold_the_pool_does_not_hold(self) -> None: + with pytest.raises(ValueError, match="do not resolve"): + sample_pooled_records(self._pool(), {"nope-1"}, budget=10) + + def test_partition_is_passage_level_not_title_level(self) -> None: + """A titleless domain must still spread across its own collections.""" + pool = self._pool(per_domain=200) + grouped = partition_pooled(pool, 8, alpha=0.0) + cloud_collections = { + name + for name, rows in grouped.items() + if any(row["domain"] == "cloud" for row in rows) + } + assert len(cloud_collections) == 2, ( + f"cloud landed in {len(cloud_collections)} collections; with one empty " + "title a title-keyed partition would give 1" + )