danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add the federated ClapNQ retrieval dataset

Browse source 
Measures whether cross-database fusion reaches what a query needs, scored
on retrieval alone so no model or judge sits between the fusion and the
number.

The corpus is MTRAG ClapNQ partitioned by article title, whole titles to a
collection, so an article's passages never split and a query's gold stays
concentrated in one collection, which is the condition a per-collection
depth quota punishes. collection_of keys on sha256 rather than hash(),
which is salted per process: the partition is never stored, and scoring
recomputes it in a different process than the one that ingested.

The 148 titles holding a gold passage carry 10,723 passages between them,
so a budget near that floor leaves no cross-topic distractors and inflates
recall. The default is 40,000 and the build reports the gold/distractor
split, warning when there are none.

build_databases opens each collection by configured name with a scope of
one, since populate_db writes to a single database. The operator entry
point emits the config for the partition it just built, so a config cannot
search a differently-partitioned build.

Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc
This commit is contained in:
Yiorgis Gozadinos 2026-08-31 13:37:52 +03:00
parent b0818f4d3d
commit daa6629879
No known key found for this signature in database
5 changed files with 536 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -18,6 +18,9 @@
- `dot` removed from `search.vector_index_metric`; switch to `cosine` or `l2`
and rerun `create-index`.
### Added
- `mtrag_federated` evaluation dataset and its reference config `evaluations/configs/mtrag_federated.yaml`: MTRAG ClapNQ partitioned by article title into `n` collections, scored on retrieval only with Recall@5/@10, nDCG@5 and MAP. `python -m evaluations.datasets.mtrag_federated --config REF --n N --out PATH` builds the partition and emits the config that searches it.
### Fixed

56
evaluations/configs/mtrag_federated.yaml Normal file
View 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

2
evaluations/evaluations/datasets/__init__.py
View file

@ -8,6 +8,7 @@ from .mtrag import (
MTRAG_CLAPNQ_REWRITE_SPEC,
MTRAG_CLAPNQ_SPEC,
)
from .mtrag_federated import MTRAG_FEDERATED_SPEC
from .open_rag_bench import (
ORB_MULTIMODAL_NEMOTRON_SPEC,
ORB_MULTIMODAL_SPEC,
@ -24,6 +25,7 @@ DATASETS: dict[str, DatasetSpec] = {
MTRAG_CLAPNQ_REWRITE_SPEC,
MTRAG_CLAPNQ_LIVE_SPEC,
MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC,
MTRAG_FEDERATED_SPEC,
ORB_TEXT_SPEC,
ORB_MULTIMODAL_SPEC,
ORB_MULTIMODAL_NEMOTRON_SPEC,

280
evaluations/evaluations/datasets/mtrag_federated.py Normal file
View file

@ -0,0 +1,280 @@
import argparse
import asyncio
import hashlib
import random
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any
import yaml
from datasets import Dataset
from evaluations.config import DatasetSpec
from evaluations.datasets.mtrag import (
_load_qrels,
build_mtrag_case,
load_clapnq_corpus,
load_clapnq_retrieval,
map_mtrag_document,
map_mtrag_retrieval,
)
from evaluations.evaluators import (
CitationMAPEvaluator,
MAPEvaluator,
NDCGEvaluator,
RecallEvaluator,
)
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_default_data_dir
COLLECTION_PREFIX = "clapnq"
DEFAULT_SEED = 20260831
# Whole titles are kept, and the 148 titles holding a gold passage carry 10,723
# passages between them, so that is the floor. A budget near it leaves no
# cross-topic distractors at all and inflates recall; this default leaves about
# 29,000, a gold-title share near a quarter.
GOLD_TITLE_FLOOR = 10_723
DEFAULT_BUDGET = 40_000
INGEST_BATCH_SIZE = 512
def collection_of(title: str, n: int, seed: int = DEFAULT_SEED) -> int:
"""Which of `n` collections holds a title's passages.
Keyed on the title, so an article's passages never split. sha256 rather than
`hash()`, which is salted per process: the partition is never stored, and
scoring recomputes it in a different process than the one that ingested.
"""
if n < 1:
raise ValueError("a partition needs at least one collection")
digest = hashlib.sha256(f"{seed}/{title}".encode()).digest()
return int.from_bytes(digest[:8], "big") % n
def collection_names(n: int) -> tuple[str, ...]:
"""The collection names in declaration order.
Order is load-bearing: fusion resolves equal ranks to configured order, so
permuting these names is an arm rather than a cosmetic change.
"""
return tuple(f"{COLLECTION_PREFIX}_{index}" for index in range(n))
def database_paths(n: int, seed: int = DEFAULT_SEED) -> dict[str, str]:
"""Where each collection's database lives.
The partition is in the filename, so a build at one `(n, seed)` can never
overwrite another's databases or be searched by the wrong config.
"""
root = get_default_data_dir() / "evaluations" / "dbs"
return {
name: str(root / f"mtrag_federated_s{seed}_n{n}_{index}.lancedb")
for index, name in enumerate(collection_names(n))
}
def sample_records(
records: Sequence[Mapping[str, Any]],
gold_ids: Iterable[str],
budget: int = DEFAULT_BUDGET,
seed: int = DEFAULT_SEED,
) -> list[Mapping[str, Any]]:
"""A fixed sub-corpus: every gold passage, plus seeded distractor titles.
Whole titles are kept or dropped together. Gold is mandatory, so a budget
below the gold floor yields the gold titles alone rather than an incomplete
corpus that would score as missing retrievals.
"""
by_title: dict[str, list[Mapping[str, Any]]] = {}
title_of: dict[str, str] = {}
for row in records:
by_title.setdefault(row["title"], []).append(row)
title_of[row["_id"]] = row["title"]
wanted = set(gold_ids)
missing = sorted(wanted - set(title_of))
if missing:
raise ValueError(
f"{len(missing)} gold passages do not resolve to the corpus, "
f"first few: {missing[:3]}"
)
gold_titles = {title_of[passage_id] for passage_id in wanted}
kept = {title for title in by_title if title in gold_titles}
total = sum(len(by_title[title]) for title in kept)
distractors = [title for title in by_title if title not in gold_titles]
random.Random(seed).shuffle(distractors)
for title in distractors:
size = len(by_title[title])
if total + size > budget:
continue
kept.add(title)
total += size
return [row for row in records if row["title"] in kept]
def partition_records(
records: Sequence[Mapping[str, Any]],
n: int,
seed: int = DEFAULT_SEED,
) -> dict[str, list[Mapping[str, Any]]]:
"""Route every record to its collection, naming all `n` even when empty."""
names = collection_names(n)
grouped: dict[str, list[Mapping[str, Any]]] = {name: [] for name in names}
for row in records:
grouped[names[collection_of(row["title"], n, seed)]].append(row)
return grouped
def pool_composition(
records: Sequence[Mapping[str, Any]], gold_ids: Iterable[str]
) -> tuple[int, int]:
"""Passages in gold-bearing titles, and passages in distractor titles.
A pool with no distractors scores as an easy retrieval task and says
nothing, so the build reports this rather than leaving it to be inferred
from the budget.
"""
wanted = set(gold_ids)
gold_titles = {row["title"] for row in records if row["_id"] in wanted}
gold_side = sum(1 for row in records if row["title"] in gold_titles)
return gold_side, len(records) - gold_side
def gold_passage_ids() -> set[str]:
"""Every corpus id the ClapNQ qrels reference."""
return {passage_id for ids in _load_qrels().values() for passage_id in ids}
def load_pool(
budget: int = DEFAULT_BUDGET, seed: int = DEFAULT_SEED
) -> list[Mapping[str, Any]]:
records = [dict(row) for row in load_clapnq_corpus()]
return sample_records(records, gold_passage_ids(), budget, seed)
def _unused_document_loader() -> Dataset:
raise RuntimeError(
"the federated corpus is built by build_databases(); run with --skip-db"
)
def emitted_config(reference: Path, n: int, seed: int = DEFAULT_SEED) -> dict[str, Any]:
"""The reference config with this partition's databases placed in it."""
settings = load_yaml_config(reference)
lancedb = dict(settings.get("lancedb") or {})
lancedb.pop("uri", None)
lancedb["databases"] = database_paths(n, seed)
settings["lancedb"] = lancedb
return settings
async def build_databases(
config: AppConfig,
n: int,
seed: int = DEFAULT_SEED,
budget: int = DEFAULT_BUDGET,
) -> dict[str, int]:
"""Ingest the partition into one database per collection.
Each member is opened by configured name with a scope of one, which is what
makes `create=True` legal on a client whose config places several.
"""
from haiku.rag.client import HaikuRAG
from evaluations.population import _ingest_batched
names = collection_names(n)
configured = set(config.lancedb.databases or {})
missing = sorted(set(names) - configured)
if missing:
raise ValueError(
f"lancedb.databases must place every collection; missing {missing}"
)
pool = load_pool(budget, seed)
gold_side, distractors = pool_composition(pool, gold_passage_ids())
print(
f"pool: {len(pool)} passages, {gold_side} in gold-bearing titles, "
f"{distractors} distractors"
)
if not distractors:
print(
" WARNING: no distractor titles, so every passage belongs to an "
f"answer-bearing article; raise --budget above {GOLD_TITLE_FLOOR}"
)
grouped = partition_records(pool, n, seed)
written: dict[str, int] = {}
for name in names:
async with HaikuRAG(config=config, sources=[name], create=True) as client:
await _ingest_batched(
client, MTRAG_FEDERATED_SPEC, grouped[name], INGEST_BATCH_SIZE
)
written[name] = len(grouped[name])
return written
MTRAG_FEDERATED_SPEC = DatasetSpec(
key="mtrag_federated",
# Never read: the run searches the configured set. Present because the spec
# requires one, and pointed at the first collection so a stray --db is
# obviously wrong rather than silently plausible.
db_filename="mtrag_federated_unused.lancedb",
document_loader=_unused_document_loader,
document_mapper=map_mtrag_document,
# The QA phase is not wired yet, and the loader is what run_qa_benchmark
# reaches first, so a forgotten --skip-qa fails loudly there. The builder is
# the one the generation tasks will need when QA arrives.
qa_loader=_unused_document_loader,
qa_case_builder=build_mtrag_case,
retrieval_loader=lambda: load_clapnq_retrieval("lastturn"),
retrieval_mapper=map_mtrag_retrieval,
retrieval_evaluators=[
RecallEvaluator(5),
RecallEvaluator(10),
NDCGEvaluator(5),
MAPEvaluator(),
],
citation_evaluator=CitationMAPEvaluator(),
# The product default, which is where the fusion depth quota bites.
retrieval_limit=5,
ingest_batch_size=INGEST_BATCH_SIZE,
)
async def main() -> None:
parser = argparse.ArgumentParser(
description=(
"Build the federated ClapNQ partition and emit the config that "
"searches exactly it."
)
)
parser.add_argument("--config", type=Path, required=True, help="reference config")
parser.add_argument("--n", type=int, required=True, help="collection count")
parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
parser.add_argument("--budget", type=int, default=DEFAULT_BUDGET)
parser.add_argument(
"--out",
type=Path,
required=True,
help="where to write the emitted config for this partition",
)
args = parser.parse_args()
settings = emitted_config(args.config, args.n, args.seed)
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(yaml.safe_dump(settings, sort_keys=False))
print(f"wrote {args.out}")
# Reload from disk, so the config that builds is the file that will search.
config = AppConfig.model_validate(load_yaml_config(args.out))
written = await build_databases(config, args.n, args.seed, args.budget)
for name, count in written.items():
print(f"{name}: {count} passages")
if __name__ == "__main__":
asyncio.run(main())

195
evaluations/tests/test_mtrag_federated.py Normal file
View file

@ -0,0 +1,195 @@
import os
import subprocess
import sys
import pytest
from evaluations.datasets import DATASETS
from evaluations.datasets.mtrag_federated import (
DEFAULT_BUDGET,
GOLD_TITLE_FLOOR,
MTRAG_FEDERATED_SPEC,
collection_names,
collection_of,
partition_records,
pool_composition,
sample_records,
)
def record(passage_id: str, title: str) -> dict[str, str]:
return {"_id": passage_id, "title": title, "text": f"text of {passage_id}"}
def corpus(titles: dict[str, int]) -> list[dict[str, str]]:
"""One record per passage, `titles` mapping a title to its passage count."""
return [
record(f"{title}_{index}", title)
for title, count in titles.items()
for index in range(count)
]
class TestCollectionOf:
def test_assigns_within_range(self) -> None:
for n in (2, 4, 8):
assigned = {collection_of(f"title {i}", n) for i in range(200)}
assert assigned <= set(range(n))
def test_uses_every_collection(self) -> None:
"""A partition that leaves a collection empty is not a partition."""
for n in (2, 4, 8):
assigned = {collection_of(f"title {i}", n) for i in range(200)}
assert assigned == set(range(n))
def test_is_stable_across_processes(self) -> None:
"""Salted `hash()` would make a build unreproducible between runs.
The partition is never stored, so scoring recomputes it in a different
process than the one that ingested.
"""
code = (
"from evaluations.datasets.mtrag_federated import collection_of;"
"print([collection_of(f'title {i}', 8) for i in range(12)])"
)
runs = {
subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
check=True,
env={**os.environ, "PYTHONHASHSEED": seed},
).stdout.strip()
for seed in ("0", "1", "12345")
}
assert len(runs) == 1, f"assignment varies with PYTHONHASHSEED: {runs}"
def test_seed_changes_the_assignment(self) -> None:
titles = [f"title {i}" for i in range(200)]
one = [collection_of(t, 8, seed=1) for t in titles]
two = [collection_of(t, 8, seed=2) for t in titles]
assert one != two
def test_rejects_a_collection_count_below_one(self) -> None:
with pytest.raises(ValueError, match="at least one collection"):
collection_of("title", 0)
class TestCollectionNames:
def test_names_one_per_collection(self) -> None:
assert collection_names(3) == (
"clapnq_0",
"clapnq_1",
"clapnq_2",
)
def test_declaration_order_is_the_name_order(self) -> None:
"""Fusion resolves ties to configured order, so the order is load-bearing."""
names = collection_names(4)
assert list(names) == sorted(names, key=lambda name: int(name.split("_")[1]))
class TestPartitionRecords:
def test_keeps_every_record(self) -> None:
records = corpus({"a": 3, "b": 2, "c": 4})
grouped = partition_records(records, 2)
assert sum(len(rows) for rows in grouped.values()) == len(records)
def test_never_splits_a_title(self) -> None:
"""A title is the atom: its passages must share a collection, or a
query's gold spreads for reasons the partition never intended."""
records = corpus({f"title {i}": 5 for i in range(40)})
grouped = partition_records(records, 4)
holders: dict[str, set[str]] = {}
for name, rows in grouped.items():
for row in rows:
holders.setdefault(row["title"], set()).add(name)
split = {title: names for title, names in holders.items() if len(names) > 1}
assert not split, f"titles split across collections: {split}"
def test_names_every_collection_even_when_one_is_empty(self) -> None:
"""The config declares n databases, so the build must create n."""
records = corpus({"only": 2})
grouped = partition_records(records, 4)
assert set(grouped) == set(collection_names(4))
class TestSampleRecords:
def test_keeps_every_gold_passage(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
gold = {"title 3_1", "title 17_4", "title 42_9"}
sampled = sample_records(records, gold, budget=60)
assert gold <= {row["_id"] for row in sampled}
def test_keeps_whole_titles_holding_gold(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
sampled = sample_records(records, {"title 3_1"}, budget=0)
assert sorted(row["_id"] for row in sampled) == sorted(
f"title 3_{i}" for i in range(10)
)
def test_respects_the_budget(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
sampled = sample_records(records, {"title 3_1"}, budget=100)
assert len(sampled) <= 100
def test_budget_below_the_gold_floor_still_keeps_gold(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
gold = {f"title {i}_0" for i in range(20)}
sampled = sample_records(records, gold, budget=5)
assert len(sampled) == 200
assert gold <= {row["_id"] for row in sampled}
def test_is_stable_for_a_seed(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
first = sample_records(records, {"title 0_0"}, budget=100, seed=7)
second = sample_records(records, {"title 0_0"}, budget=100, seed=7)
assert [row["_id"] for row in first] == [row["_id"] for row in second]
def test_seed_changes_the_distractors(self) -> None:
records = corpus({f"title {i}": 10 for i in range(50)})
first = sample_records(records, {"title 0_0"}, budget=100, seed=7)
second = sample_records(records, {"title 0_0"}, budget=100, seed=8)
assert {row["_id"] for row in first} != {row["_id"] for row in second}
def test_rejects_gold_the_corpus_does_not_hold(self) -> None:
records = corpus({"a": 2})
with pytest.raises(ValueError, match="do not resolve"):
sample_records(records, {"missing"}, budget=10)
class TestSpec:
def test_registers_under_its_key(self) -> None:
assert DATASETS[MTRAG_FEDERATED_SPEC.key] is MTRAG_FEDERATED_SPEC
def test_opts_out_of_the_shared_population(self) -> None:
"""The databases are built by build_databases, not populate_db."""
with pytest.raises(RuntimeError, match="build_databases"):
MTRAG_FEDERATED_SPEC.document_loader()
def test_retrieval_limit_matches_the_product_default(self) -> None:
"""5 is config's search.limit, the setting the depth quota bites at."""
assert MTRAG_FEDERATED_SPEC.retrieval_limit == 5
def test_scores_retrieval_without_a_judge(self) -> None:
assert MTRAG_FEDERATED_SPEC.retrieval_evaluators
assert MTRAG_FEDERATED_SPEC.retrieval_loader is not None
assert MTRAG_FEDERATED_SPEC.retrieval_mapper is not None
class TestPoolComposition:
def test_separates_gold_bearing_titles_from_distractors(self) -> None:
records = corpus({"answers": 4, "filler": 6})
gold_side, distractors = pool_composition(records, {"answers_2"})
assert (gold_side, distractors) == (4, 6)
def test_reports_no_distractors_when_the_budget_is_at_the_floor(self) -> None:
"""A pool of only answer-bearing articles scores as an easy task and
says nothing, so the build has to be able to see it."""
records = corpus({f"title {i}": 10 for i in range(5)})
gold = {f"title {i}_0" for i in range(5)}
sampled = sample_records(records, gold, budget=1)
assert pool_composition(sampled, gold) == (50, 0)
def test_the_default_budget_clears_the_gold_floor(self) -> None:
assert DEFAULT_BUDGET > GOLD_TITLE_FLOOR