Add the pooled four-domain retrieval dataset
`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 that hurts a real deployment is never exercised. Every fusion conclusion measured on it is therefore provisional. `mtrag_pooled` pools all four MTRAG domains, so a query belongs to one and the rest are genuinely off-topic. `collection_of` gains `alpha`, which now means something: 0 keeps a collection to one domain, 1 ignores the domain and shards titles uniformly. Domains map onto collections proportionally, subdividing by title where there are more collections than domains and grouping where there are fewer. Passage ids are checked for collisions across domains, since gold is uri-keyed and a shared id would make it ambiguous. Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc
This commit is contained in:
Yiorgis Gozadinos
parent
3183cc0ad7
commit
16a6881454
5 changed files with 433 additions and 9 deletions
|
|
@ -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
|
||||
|
|
|
|||
49
evaluations/configs/mtrag_pooled.yaml
Normal file
49
evaluations/configs/mtrag_pooled.yaml
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue