Compare commits
9 commits
main
...
archive/ev
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a60d0f45d3 | ||
|
|
33bd0be702 | ||
|
|
058120c05a | ||
|
|
ae5c798af7 | ||
|
|
cc344fb205 | ||
|
|
73578a1198 | ||
|
|
b46e8a4491 | ||
|
|
409f60e1bf | ||
|
|
daa6629879 |
12 changed files with 1281 additions and 5 deletions
|
|
@ -13,6 +13,9 @@
|
|||
text untruncated. `format_citations_rich` takes a `full` argument.
|
||||
- `doctor` fails when the chunks FTS index covers no rows.
|
||||
- FTS and hybrid searches log a warning when the FTS index covers no rows.
|
||||
- `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.
|
||||
|
||||
### Removed
|
||||
|
||||
|
|
@ -21,6 +24,9 @@
|
|||
|
||||
### Fixed
|
||||
|
||||
- MTRAG retrieval queries have their `|speaker|: ` tag stripped. It reached the embedder, the BM25 query and the reranker's query, costing about 3.6pp recall on the reranker and nothing on the fusion path, measured paired over 777 queries. Affects `mtrag_clapnq`, `mtrag_clapnq_rewrite`, `mtrag_federated` and `mtrag_pooled` retrieval scores; the QA path never carried it.
|
||||
- Batched evaluation ingest converts inline content as text instead of letting `HaikuRAG.convert` disambiguate it, so a passage beginning with a URL is stored rather than fetched over HTTP. 187 MTRAG cloud and fiqa passages start with one; no clapnq passage does, so no existing dataset's numbers change.
|
||||
- `mtrag_federated` builds vacuum each collection after ingest and assert the chunks FTS index covers every row. Without the vacuum the index stays at zero rows, and full-text search returns near-arbitrary rows while still returning results.
|
||||
- FTS and hybrid search on a database whose FTS index covers no rows. Chunk
|
||||
writes now build the index and rebuild it if it covers none; `haiku-rag
|
||||
vacuum` also repairs it.
|
||||
|
|
|
|||
56
evaluations/configs/mtrag_federated.yaml
Normal file
56
evaluations/configs/mtrag_federated.yaml
Normal file
|
|
@ -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
|
||||
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
|
||||
|
|
@ -41,6 +41,7 @@ async def evaluate_dataset(
|
|||
vacuum_interval: int = 100,
|
||||
multimodal_only: bool = False,
|
||||
judge_model: ModelConfig | None = None,
|
||||
retrieval_limit: int | None = None,
|
||||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
|
|
@ -72,6 +73,7 @@ async def evaluate_dataset(
|
|||
db_path=db_path,
|
||||
multimodal_only=multimodal_only,
|
||||
document_filter=document_filter,
|
||||
retrieval_limit=retrieval_limit,
|
||||
)
|
||||
|
||||
if not skip_qa:
|
||||
|
|
@ -169,6 +171,14 @@ def run(
|
|||
None, "--limit", help="Limit number of test cases for both retrieval and QA."
|
||||
),
|
||||
name: str | None = typer.Option(None, "--name", help="Override evaluation name."),
|
||||
retrieval_limit: int | None = typer.Option(
|
||||
None,
|
||||
"--retrieval-limit",
|
||||
help=(
|
||||
"Candidates each database fetches, overriding the dataset's. "
|
||||
"Sets how deep hybrid search looks before its results are scored."
|
||||
),
|
||||
),
|
||||
vacuum_interval: int = typer.Option(
|
||||
100, "--vacuum-interval", help="Vacuum every N documents during DB population."
|
||||
),
|
||||
|
|
@ -235,6 +245,7 @@ def run(
|
|||
vacuum_interval=vacuum_interval,
|
||||
multimodal_only=multimodal_only,
|
||||
judge_model=judge_model_config,
|
||||
retrieval_limit=retrieval_limit,
|
||||
target=target_value,
|
||||
capability_model=capability_model_config,
|
||||
case_ids=_load_case_ids(filter_ids),
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from .mtrag import (
|
|||
MTRAG_CLAPNQ_REWRITE_SPEC,
|
||||
MTRAG_CLAPNQ_SPEC,
|
||||
)
|
||||
from .mtrag_federated import MTRAG_FEDERATED_SPEC, MTRAG_POOLED_SPEC
|
||||
from .open_rag_bench import (
|
||||
ORB_MULTIMODAL_NEMOTRON_SPEC,
|
||||
ORB_MULTIMODAL_SPEC,
|
||||
|
|
@ -24,6 +25,8 @@ DATASETS: dict[str, DatasetSpec] = {
|
|||
MTRAG_CLAPNQ_REWRITE_SPEC,
|
||||
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,4 +1,5 @@
|
|||
import json
|
||||
import re
|
||||
import zipfile
|
||||
from collections.abc import Iterable, Mapping
|
||||
from functools import partial
|
||||
|
|
@ -85,6 +86,20 @@ def _validate_qrels_resolve(
|
|||
)
|
||||
|
||||
|
||||
_SPEAKER_TAG = re.compile(r"^\|[^|]+\|:\s*")
|
||||
|
||||
|
||||
def strip_speaker_markup(text: str) -> str:
|
||||
"""Drop the leading `|speaker|: ` tag MTRAG's query files encode.
|
||||
|
||||
Anchored, so only a leading tag is markup: a pipe later in the question is
|
||||
content, and a second tag survives. The tag reaches the embedder, the BM25
|
||||
query and the reranker's query, and on the reranker it costs about 3.6pp
|
||||
recall, measured paired over 777 queries.
|
||||
"""
|
||||
return _SPEAKER_TAG.sub("", text)
|
||||
|
||||
|
||||
def _join_queries_qrels(
|
||||
queries: Iterable[Mapping[str, Any]], qrels: Mapping[str, list[str]]
|
||||
) -> list[dict[str, Any]]:
|
||||
|
|
@ -97,7 +112,7 @@ def _join_queries_qrels(
|
|||
records.append(
|
||||
{
|
||||
"query_id": query_id,
|
||||
"question": query["text"],
|
||||
"question": strip_speaker_markup(query["text"]),
|
||||
"expected_uris": expected,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
599
evaluations/evaluations/datasets/mtrag_federated.py
Normal file
599
evaluations/evaluations/datasets/mtrag_federated.py
Normal file
|
|
@ -0,0 +1,599 @@
|
|||
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
|
||||
|
||||
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,
|
||||
strip_speaker_markup,
|
||||
)
|
||||
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"
|
||||
# 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
|
||||
# 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
|
||||
FTS_INDEX_NAME = "content_fts_idx"
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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")
|
||||
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, ...]:
|
||||
"""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)
|
||||
|
||||
|
||||
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": strip_speaker_markup(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_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(
|
||||
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.
|
||||
|
||||
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["_id"], 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"
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class FTSIndexNotCoveringRows(AssertionError):
|
||||
"""The chunks FTS index does not cover every row, so full-text search is
|
||||
dead while still returning results."""
|
||||
|
||||
|
||||
async def assert_fts_covers_rows(table: Any, name: str) -> None:
|
||||
rows = await table.count_rows()
|
||||
indices = {index.name for index in await table.list_indices()}
|
||||
if FTS_INDEX_NAME not in indices:
|
||||
raise FTSIndexNotCoveringRows(f"{name}: no {FTS_INDEX_NAME} on {rows} rows")
|
||||
stats = await table.index_stats(FTS_INDEX_NAME)
|
||||
indexed = getattr(stats, "num_indexed_rows", 0) or 0
|
||||
if indexed < rows:
|
||||
raise FTSIndexNotCoveringRows(
|
||||
f"{name}: {FTS_INDEX_NAME} covers {indexed} of {rows} rows; "
|
||||
"full-text search would return near-arbitrary rows"
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
# The chunks FTS index is built once when the table is created, over
|
||||
# zero rows, and nothing folds later rows into it but an optimize.
|
||||
# `auto_vacuum` is false here, as in every reference config, so
|
||||
# without this the index covers nothing and full-text search returns
|
||||
# near-arbitrary rows while still looking like it works.
|
||||
await client.store.vacuum(retention_seconds=0)
|
||||
await assert_fts_covers_rows(client.store.chunks_table, name)
|
||||
written[name] = len(grouped[name])
|
||||
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
|
||||
# Passage level, not `pool_composition`: that counts gold-bearing titles,
|
||||
# which is meaningless here since cloud and fiqa have one empty title each.
|
||||
gold = pooled_gold_ids()
|
||||
gold_kept = sum(1 for row in pool if row["_id"] in gold)
|
||||
print(
|
||||
f"pool: {len(pool)} passages, {gold_kept} gold, "
|
||||
f"{len(pool) - gold_kept} 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
|
||||
# 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,
|
||||
)
|
||||
|
||||
|
||||
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=(
|
||||
"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(
|
||||
"--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,
|
||||
required=True,
|
||||
help="where to write the emitted config for this partition",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
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))
|
||||
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")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
|
@ -39,6 +39,9 @@ async def _ingest_batched(
|
|||
row["uri"]: row["id"] for row in uri_rows if row["id"] not in chunked_ids
|
||||
}
|
||||
|
||||
from haiku.rag.converters import get_converter
|
||||
|
||||
converter = get_converter(rag._config)
|
||||
batch: list[DocumentImport] = []
|
||||
for doc in corpus:
|
||||
payload = spec.document_mapper(cast(Mapping[str, Any], doc))
|
||||
|
|
@ -48,7 +51,13 @@ async def _ingest_batched(
|
|||
if payload.uri in chunkless:
|
||||
await rag.delete_document(chunkless[payload.uri])
|
||||
assert payload.content is not None, "batched ingest requires inline content"
|
||||
docling_document = await rag.convert(payload.content, format=payload.format)
|
||||
# Convert as text explicitly. `rag.convert` disambiguates a str by
|
||||
# parsing it, and a passage beginning with a URL (187 of them across
|
||||
# MTRAG's cloud and fiqa corpora) is then fetched over HTTP instead of
|
||||
# stored. Batched ingest has already asserted the content is inline.
|
||||
docling_document = await converter.convert_text(
|
||||
payload.content, format=payload.format
|
||||
)
|
||||
chunks = await rag.chunk(docling_document)
|
||||
batch.append(
|
||||
DocumentImport(
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ async def run_retrieval_benchmark(
|
|||
db_path: Path | None = None,
|
||||
multimodal_only: bool = False,
|
||||
document_filter: str | None = None,
|
||||
retrieval_limit: int | None = None,
|
||||
) -> dict[str, float] | None:
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
console.print("Skipping retrieval benchmark; no retrieval config.")
|
||||
|
|
@ -72,6 +73,7 @@ async def run_retrieval_benchmark(
|
|||
evaluators=list(spec.retrieval_evaluators),
|
||||
)
|
||||
|
||||
fetch = retrieval_limit or spec.retrieval_limit
|
||||
db = (
|
||||
None
|
||||
if spec.uses_configured_databases(config, db_path)
|
||||
|
|
@ -82,7 +84,7 @@ async def run_retrieval_benchmark(
|
|||
async def retrieval_target(question: str) -> list[str]:
|
||||
chunks = await rag.search(
|
||||
query=question,
|
||||
limit=spec.retrieval_limit,
|
||||
limit=fetch,
|
||||
include_images=False,
|
||||
filter=document_filter,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -888,7 +888,11 @@ class TestBatchedIngest:
|
|||
rag.store.chunks_table = _table(
|
||||
[{"document_id": f"id-{uri}"} for uri in complete_uris]
|
||||
)
|
||||
rag.convert = AsyncMock(side_effect=lambda content, **kw: f"docling:{content}")
|
||||
# Batched ingest converts text through the configured converter, the
|
||||
# same path create_document uses, so the double needs a real config.
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
rag._config = AppConfig()
|
||||
rag.chunk = AsyncMock(return_value=[])
|
||||
rag.import_documents = AsyncMock()
|
||||
rag.delete_document = AsyncMock()
|
||||
|
|
@ -933,8 +937,11 @@ class TestBatchedIngest:
|
|||
await _ingest_batched(rag, self._spec(), corpus, batch_size=10)
|
||||
|
||||
(batch,), _ = rag.import_documents.call_args
|
||||
# Conversion goes through the configured converter now, not rag.convert,
|
||||
# so the batch contents are the assertion: exactly the incomplete uris,
|
||||
# which is stricter than counting conversions.
|
||||
assert [imp.uri for imp in batch] == ["u1", "u3"]
|
||||
assert rag.convert.await_count == 2
|
||||
assert len(batch) == 2
|
||||
rag.delete_document.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -13,8 +13,10 @@ from evaluations.datasets.mtrag import (
|
|||
_validate_qrels_resolve,
|
||||
build_mtrag_case,
|
||||
build_mtrag_live_case,
|
||||
load_clapnq_retrieval,
|
||||
map_mtrag_document,
|
||||
map_mtrag_retrieval,
|
||||
strip_speaker_markup,
|
||||
)
|
||||
from evaluations.evaluators import (
|
||||
CitationMAPEvaluator,
|
||||
|
|
@ -276,3 +278,40 @@ class TestLiveConversations:
|
|||
"mtrag_mode": "live_session",
|
||||
"compaction": False,
|
||||
}
|
||||
|
||||
|
||||
class TestSpeakerMarkup:
|
||||
"""MTRAG's retrieval query files encode the speaker into the query text.
|
||||
It reaches the embedder, the BM25 query and the reranker's query; on the
|
||||
reranker it costs about 3.6pp recall, measured paired over 777 queries.
|
||||
"""
|
||||
|
||||
def test_strips_a_leading_speaker_tag(self) -> None:
|
||||
assert (
|
||||
strip_speaker_markup("|user|: How many teams are in the NFL?")
|
||||
== "How many teams are in the NFL?"
|
||||
)
|
||||
|
||||
def test_strips_any_speaker_not_just_user(self) -> None:
|
||||
assert strip_speaker_markup("|agent|: Twelve of them.") == "Twelve of them."
|
||||
|
||||
def test_leaves_an_unmarked_question_alone(self) -> None:
|
||||
assert (
|
||||
strip_speaker_markup("How many teams are in the NFL?")
|
||||
== "How many teams are in the NFL?"
|
||||
)
|
||||
|
||||
def test_leaves_a_pipe_mid_sentence_alone(self) -> None:
|
||||
"""Only a leading tag is markup; a pipe in the question is content."""
|
||||
text = "What does the |> operator do?"
|
||||
assert strip_speaker_markup(text) == text
|
||||
|
||||
def test_does_not_strip_a_second_tag(self) -> None:
|
||||
"""One tag is the encoding; a second is content and must survive."""
|
||||
assert strip_speaker_markup("|user|: |agent|: nested") == "|agent|: nested"
|
||||
|
||||
def test_retrieval_samples_arrive_clean(self) -> None:
|
||||
rows = list(load_clapnq_retrieval("lastturn"))
|
||||
assert rows, "no retrieval rows"
|
||||
marked = [r for r in rows if r["question"].startswith("|")]
|
||||
assert not marked, f"{len(marked)} of {len(rows)} still carry markup"
|
||||
|
|
|
|||
480
evaluations/tests/test_mtrag_federated.py
Normal file
480
evaluations/tests/test_mtrag_federated.py
Normal file
|
|
@ -0,0 +1,480 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
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,
|
||||
sample_pooled_records,
|
||||
partition_records,
|
||||
pool_composition,
|
||||
sample_records,
|
||||
)
|
||||
|
||||
|
||||
def _smoke_config():
|
||||
"""A config placing two databases, so the run resolves a federated client."""
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
return AppConfig.model_validate(
|
||||
{
|
||||
"lancedb": {
|
||||
"databases": {
|
||||
"clapnq_0": "/tmp/a.lancedb",
|
||||
"clapnq_1": "/tmp/b.lancedb",
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class TestRetrievalLimitOverride:
|
||||
async def test_override_replaces_the_spec_value(self, monkeypatch) -> None:
|
||||
"""Fetch depth is a run knob: hybrid search degenerates below roughly 50
|
||||
candidates, so every regime would otherwise need its own dataset."""
|
||||
seen: list[int | None] = []
|
||||
|
||||
async def fake_search(self, query, limit=None, **kwargs): # noqa: ANN001
|
||||
seen.append(limit)
|
||||
return []
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
monkeypatch.setattr(HaikuRAG, "search", fake_search)
|
||||
from evaluations.retrieval import run_retrieval_benchmark
|
||||
|
||||
await run_retrieval_benchmark(
|
||||
MTRAG_FEDERATED_SPEC,
|
||||
_smoke_config(),
|
||||
limit=1,
|
||||
retrieval_limit=77,
|
||||
)
|
||||
assert seen and set(seen) == {77}
|
||||
|
||||
async def test_spec_value_is_the_default(self, monkeypatch) -> None:
|
||||
seen: list[int | None] = []
|
||||
|
||||
async def fake_search(self, query, limit=None, **kwargs): # noqa: ANN001
|
||||
seen.append(limit)
|
||||
return []
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
monkeypatch.setattr(HaikuRAG, "search", fake_search)
|
||||
from evaluations.retrieval import run_retrieval_benchmark
|
||||
|
||||
await run_retrieval_benchmark(MTRAG_FEDERATED_SPEC, _smoke_config(), limit=1)
|
||||
assert seen and set(seen) == {MTRAG_FEDERATED_SPEC.retrieval_limit}
|
||||
|
||||
|
||||
class TestFTSCoverageAssertion:
|
||||
"""The chunks FTS index is built once over zero rows and only an optimize
|
||||
folds later rows in, so a build that skips it ships dead full-text search
|
||||
that still returns results."""
|
||||
|
||||
class _Index:
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
|
||||
class _Stats:
|
||||
def __init__(self, indexed: int) -> None:
|
||||
self.num_indexed_rows = indexed
|
||||
|
||||
class _Table:
|
||||
def __init__(self, rows: int, indexed: int | None) -> None:
|
||||
self._rows = rows
|
||||
self._indexed = indexed
|
||||
|
||||
async def count_rows(self) -> int:
|
||||
return self._rows
|
||||
|
||||
async def list_indices(self):
|
||||
if self._indexed is None:
|
||||
return []
|
||||
return [TestFTSCoverageAssertion._Index("content_fts_idx")]
|
||||
|
||||
async def index_stats(self, name: str):
|
||||
assert name == "content_fts_idx"
|
||||
return TestFTSCoverageAssertion._Stats(self._indexed or 0)
|
||||
|
||||
async def test_passes_when_the_index_covers_every_row(self) -> None:
|
||||
await assert_fts_covers_rows(self._Table(100, 100), "clapnq_0")
|
||||
|
||||
async def test_rejects_a_zero_row_index(self) -> None:
|
||||
with pytest.raises(FTSIndexNotCoveringRows, match="covers 0 of 100"):
|
||||
await assert_fts_covers_rows(self._Table(100, 0), "clapnq_0")
|
||||
|
||||
async def test_rejects_a_partially_covering_index(self) -> None:
|
||||
with pytest.raises(FTSIndexNotCoveringRows, match="covers 60 of 100"):
|
||||
await assert_fts_covers_rows(self._Table(100, 60), "clapnq_0")
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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"
|
||||
)
|
||||
Loading…
Reference in a new issue