feat: add search_filter to evaluations
This commit is contained in:
parent
981353eb98
commit
93c21272d1
4 changed files with 244 additions and 6 deletions
|
|
@ -65,6 +65,7 @@ evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.l
|
|||
- `--name NAME` - Override the evaluation name
|
||||
- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
|
||||
- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
|
||||
- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)).
|
||||
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
|
||||
|
|
@ -86,6 +87,31 @@ evaluations:
|
|||
enable_thinking: true
|
||||
```
|
||||
|
||||
### Restricting the corpus
|
||||
|
||||
When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`, `db_source`):
|
||||
|
||||
```bash
|
||||
evaluations run mqf --skip-db --config haiku.rag.s3.yaml \
|
||||
--filter "db_source in ('dataset')"
|
||||
```
|
||||
|
||||
The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `search_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results.
|
||||
|
||||
A dataset can declare its own default in its `DatasetSpec`, so runs need no flag:
|
||||
|
||||
```python
|
||||
MQF_SPEC = DatasetSpec(
|
||||
key="mqf",
|
||||
...
|
||||
search_filter="db_source in ('dataset')",
|
||||
)
|
||||
```
|
||||
|
||||
`--filter` overrides that default; passing an empty string (`--filter ""`) clears it and searches the whole database.
|
||||
|
||||
Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus.
|
||||
|
||||
## Methodology
|
||||
|
||||
### Retrieval Metrics
|
||||
|
|
@ -132,21 +158,21 @@ Two approaches are benchmarked separately:
|
|||
| Embedding Model | Reranker | Cases | MAP |
|
||||
|------------------------------------------|------------------------------------------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | none | 3045 | 0.9774 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9798 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9709 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `nvidia/llama-nemotron-rerank-vl-1b-v2` (multimodal) | 3045 | 0.9913 |
|
||||
|
||||
*The nemotron row without a reranker is measured on this release. The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text, measured on haiku.rag main post-v0.67.3.*
|
||||
*The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text. Measured on haiku.rag main post-v0.67.3 (multimodal reranking ships in the next release).*
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Embedding Model | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3039 | 0.9263 | 0.9761 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3040 | 0.9362 | 0.9343 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
|
||||
|
||||
*Both nemotron `Gemma-4` rows are measured on this release, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` with thinking on, and exclude the cases that errored (6 of 3045 for `rag-capability`, 5 for `analysis-capability`). The `rag-capability` row cites at 99.64% with a mean of 1.08 citations per case, at a median 4.7s per case against 5.0s for `analysis-capability`. Citation coverage is what moved on this release: 4.9% of analysis cases register no citation, against 26.3% before, at unchanged searches and code executions per case. The remaining rows are from haiku.rag v0.52.0, where Qwen3-VL covered 1409 / 3045 cases.*
|
||||
*Measured on haiku.rag v0.52.0, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases.*
|
||||
|
||||
#### Text embedder + VLM picture descriptions
|
||||
|
||||
|
|
|
|||
|
|
@ -57,6 +57,16 @@ configure_telemetry(service_name="evals", scrubbing=False)
|
|||
configure_cli_logging()
|
||||
console = Console()
|
||||
|
||||
def resolve_search_filter(spec: DatasetSpec, override: str | None) -> str | None:
|
||||
"""Pick the document filter for a run: `--filter` wins over the dataset's.
|
||||
|
||||
An empty `--filter ""` is honoured as "no filter", so a dataset that
|
||||
declares one can still be run against the whole database.
|
||||
"""
|
||||
if override is None:
|
||||
return spec.search_filter
|
||||
return override or None
|
||||
|
||||
|
||||
def build_experiment_metadata(
|
||||
dataset_key: str,
|
||||
|
|
@ -65,6 +75,7 @@ def build_experiment_metadata(
|
|||
judge_config: ModelConfig | None = None,
|
||||
target: Target = "rag-capability",
|
||||
capability_config: ModelConfig | None = None,
|
||||
search_filter: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build experiment metadata for Logfire tracking."""
|
||||
metadata: dict[str, Any] = {
|
||||
|
|
@ -88,6 +99,7 @@ def build_experiment_metadata(
|
|||
"qa_enable_thinking": config.qa.model.enable_thinking,
|
||||
"qa_extra_body": config.qa.model.extra_body,
|
||||
"qa_max_searches": config.qa.max_searches,
|
||||
"search_filter": search_filter,
|
||||
}
|
||||
if judge_config is not None:
|
||||
metadata.update(
|
||||
|
|
@ -191,6 +203,7 @@ async def run_retrieval_benchmark(
|
|||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
multimodal_only: bool = False,
|
||||
search_filter: str | 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.")
|
||||
|
|
@ -246,7 +259,9 @@ async def run_retrieval_benchmark(
|
|||
async with HaikuRAG(db, config=config, read_only=True) as rag:
|
||||
|
||||
async def retrieval_target(question: str) -> list[str]:
|
||||
chunks = await rag.search(query=question, limit=5, include_images=False)
|
||||
chunks = await rag.search(
|
||||
query=question, limit=5, include_images=False, filter=search_filter
|
||||
)
|
||||
|
||||
seen = set()
|
||||
identifiers = []
|
||||
|
|
@ -264,6 +279,7 @@ async def run_retrieval_benchmark(
|
|||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
search_filter=search_filter,
|
||||
)
|
||||
|
||||
report = await dataset.evaluate(
|
||||
|
|
@ -364,6 +380,7 @@ async def run_qa_benchmark(
|
|||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
search_filter: str | None = None,
|
||||
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
|
||||
corpus = spec.qa_loader()
|
||||
corpus = _filter_qa_corpus(corpus, case_ids)
|
||||
|
|
@ -419,6 +436,7 @@ async def run_qa_benchmark(
|
|||
judge_config=judge_config,
|
||||
target=target,
|
||||
capability_config=capability_config,
|
||||
search_filter=search_filter,
|
||||
)
|
||||
|
||||
async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
|
||||
|
|
@ -440,6 +458,7 @@ async def run_qa_benchmark(
|
|||
config=config,
|
||||
question=question,
|
||||
capability_model=resolved_capability_model,
|
||||
document_filter=search_filter,
|
||||
)
|
||||
set_eval_attribute("cited_uris", result.cited_uris)
|
||||
set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
|
||||
|
|
@ -530,7 +549,13 @@ async def evaluate_dataset(
|
|||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
search_filter: str | None = None,
|
||||
) -> None:
|
||||
# Resolved once so both phases score the same subset of the database.
|
||||
resolved_filter = resolve_search_filter(spec, search_filter)
|
||||
if resolved_filter is not None:
|
||||
console.print(f"Document filter: {resolved_filter}", style="dim")
|
||||
|
||||
if not skip_db:
|
||||
console.print(f"Using dataset: {spec.key}", style="bold magenta")
|
||||
await populate_db(
|
||||
|
|
@ -546,6 +571,7 @@ async def evaluate_dataset(
|
|||
name=name,
|
||||
db_path=db_path,
|
||||
multimodal_only=multimodal_only,
|
||||
search_filter=resolved_filter,
|
||||
)
|
||||
|
||||
if not skip_qa:
|
||||
|
|
@ -562,6 +588,7 @@ async def evaluate_dataset(
|
|||
target=target,
|
||||
capability_model=capability_model,
|
||||
case_ids=case_ids,
|
||||
search_filter=resolved_filter,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -651,6 +678,17 @@ def run(
|
|||
"analysis.model when --target is analysis-capability) from the config."
|
||||
),
|
||||
),
|
||||
search_filter: str | None = typer.Option(
|
||||
None,
|
||||
"--filter",
|
||||
"-f",
|
||||
help=(
|
||||
"SQL WHERE clause over document columns (id, uri, title, metadata, "
|
||||
"db_source, ...) restricting every benchmark search, e.g. "
|
||||
"\"db_source in ('dataset')\". Overrides the dataset's own "
|
||||
"filter; pass an empty string to search the whole database."
|
||||
),
|
||||
),
|
||||
filter_ids: Path | None = typer.Option(
|
||||
None,
|
||||
"--filter-ids",
|
||||
|
|
@ -688,6 +726,7 @@ def run(
|
|||
target=target_value,
|
||||
capability_model=capability_model_config,
|
||||
case_ids=_load_case_ids(filter_ids),
|
||||
search_filter=search_filter,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ class DatasetSpec:
|
|||
retrieval_evaluator: Evaluator | None = None
|
||||
qa_evaluator: Evaluator | None = None
|
||||
document_limit: int | None = None
|
||||
search_filter: str | None = None
|
||||
|
||||
def db_path(self, override_path: Path | None = None) -> Path:
|
||||
"""Get the database path.
|
||||
|
|
|
|||
|
|
@ -581,6 +581,178 @@ class TestRetrievalTarget:
|
|||
assert result["map"] == 0.5
|
||||
|
||||
|
||||
class TestResolveSearchFilter:
|
||||
def test_dataset_filter_used_when_no_override(self) -> None:
|
||||
from evaluations.benchmark import resolve_search_filter
|
||||
|
||||
spec = _stub_spec(search_filter="db_source in ('dataset')")
|
||||
assert resolve_search_filter(spec, None) == "db_source in ('dataset')"
|
||||
|
||||
def test_override_wins(self) -> None:
|
||||
from evaluations.benchmark import resolve_search_filter
|
||||
|
||||
spec = _stub_spec(search_filter="db_source in ('dataset')")
|
||||
assert resolve_search_filter(spec, "uri LIKE '%.pdf'") == "uri LIKE '%.pdf'"
|
||||
|
||||
def test_empty_override_clears_dataset_filter(self) -> None:
|
||||
"""`--filter ""` runs a filtered dataset against the whole database."""
|
||||
from evaluations.benchmark import resolve_search_filter
|
||||
|
||||
spec = _stub_spec(search_filter="db_source in ('dataset')")
|
||||
assert resolve_search_filter(spec, "") is None
|
||||
|
||||
def test_none_when_neither_is_set(self) -> None:
|
||||
from evaluations.benchmark import resolve_search_filter
|
||||
|
||||
assert resolve_search_filter(_stub_spec(), None) is None
|
||||
|
||||
|
||||
class TestSearchFilterThreading:
|
||||
"""The resolved filter must reach both benchmark phases, so retrieval and
|
||||
QA score the same subset of the database."""
|
||||
|
||||
def test_metadata_records_filter(self) -> None:
|
||||
result = build_experiment_metadata(
|
||||
dataset_key="test",
|
||||
test_cases=1,
|
||||
config=AppConfig(),
|
||||
search_filter="db_source in ('dataset')",
|
||||
)
|
||||
assert result["search_filter"] == "db_source in ('dataset')"
|
||||
|
||||
def test_metadata_filter_is_none_when_unset(self) -> None:
|
||||
result = build_experiment_metadata(
|
||||
dataset_key="test", test_cases=1, config=AppConfig()
|
||||
)
|
||||
assert result["search_filter"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieval_search_receives_filter(self, tmp_path: Path) -> None:
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
|
||||
from evaluations.benchmark import run_retrieval_benchmark
|
||||
from evaluations.config import RetrievalSample
|
||||
from evaluations.evaluators import MAPEvaluator
|
||||
|
||||
searches: list[dict] = []
|
||||
|
||||
class FakeRag:
|
||||
async def search(self, **kwargs) -> list[SearchResult]:
|
||||
searches.append(kwargs)
|
||||
return [SearchResult(content="x", score=1.0, document_uri="uri-x")]
|
||||
|
||||
spec = _stub_spec(
|
||||
retrieval_loader=lambda: [{"q": "What is X?", "uris": ("uri-x",)}],
|
||||
retrieval_mapper=lambda d: RetrievalSample(
|
||||
question=d["q"], expected_uris=d["uris"]
|
||||
),
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
)
|
||||
|
||||
with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
|
||||
mock_haiku.return_value.__aenter__.return_value = FakeRag()
|
||||
await run_retrieval_benchmark(
|
||||
spec,
|
||||
AppConfig(),
|
||||
db_path=tmp_path / "test.lancedb",
|
||||
search_filter="db_source in ('dataset')",
|
||||
)
|
||||
|
||||
assert searches[0]["filter"] == "db_source in ('dataset')"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_qa_capability_run_receives_filter(self, tmp_path: Path) -> None:
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.capability_runner import CapabilityRunResult
|
||||
from evaluations.evaluators import NumberMatchEvaluator
|
||||
|
||||
# A deterministic evaluator, so no judge model is constructed.
|
||||
spec = DatasetSpec(
|
||||
key="test",
|
||||
db_filename="test.lancedb",
|
||||
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
||||
document_mapper=lambda doc: None,
|
||||
qa_loader=lambda: [{"question": "What is X?", "answer": "42"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
||||
qa_case_builder=lambda idx, doc: Case(
|
||||
name=f"case-{idx}",
|
||||
inputs=doc["question"],
|
||||
expected_output=doc["answer"],
|
||||
),
|
||||
qa_evaluator=NumberMatchEvaluator(),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"evaluations.benchmark.run_capability_question",
|
||||
new_callable=AsyncMock,
|
||||
return_value=CapabilityRunResult(answer="ANSWER: 42"),
|
||||
) as mock_run:
|
||||
await run_qa_benchmark(
|
||||
spec,
|
||||
AppConfig(),
|
||||
db_path=tmp_path / "test.lancedb",
|
||||
search_filter="db_source in ('dataset')",
|
||||
)
|
||||
|
||||
mock_run.assert_awaited_once()
|
||||
assert mock_run.await_args[1]["document_filter"] == "db_source in ('dataset')"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evaluate_dataset_resolves_once_for_both_phases(self) -> None:
|
||||
"""The dataset's own filter reaches retrieval and QA without a flag."""
|
||||
spec = _stub_spec(search_filter="db_source in ('dataset','other')")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.benchmark.run_retrieval_benchmark", new_callable=AsyncMock
|
||||
) as mock_retrieval,
|
||||
patch(
|
||||
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
|
||||
) as mock_qa,
|
||||
):
|
||||
await evaluate_dataset(
|
||||
spec=spec,
|
||||
config=AppConfig(),
|
||||
skip_db=True,
|
||||
skip_retrieval=False,
|
||||
skip_qa=False,
|
||||
limit=None,
|
||||
name=None,
|
||||
db_path=None,
|
||||
)
|
||||
|
||||
expected = "db_source in ('dataset','other')"
|
||||
assert mock_retrieval.call_args[1]["search_filter"] == expected
|
||||
assert mock_qa.call_args[1]["search_filter"] == expected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evaluate_dataset_override_reaches_both_phases(self) -> None:
|
||||
spec = _stub_spec(search_filter="db_source in ('dataset','other')")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.benchmark.run_retrieval_benchmark", new_callable=AsyncMock
|
||||
) as mock_retrieval,
|
||||
patch(
|
||||
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
|
||||
) as mock_qa,
|
||||
):
|
||||
await evaluate_dataset(
|
||||
spec=spec,
|
||||
config=AppConfig(),
|
||||
skip_db=True,
|
||||
skip_retrieval=False,
|
||||
skip_qa=False,
|
||||
limit=None,
|
||||
name=None,
|
||||
db_path=None,
|
||||
search_filter="db_source in ('other')",
|
||||
)
|
||||
|
||||
assert mock_retrieval.call_args[1]["search_filter"] == "db_source in ('other')"
|
||||
assert mock_qa.call_args[1]["search_filter"] == "db_source in ('other')"
|
||||
|
||||
|
||||
class TestEvaluateDatasetCaseIds:
|
||||
def _spec(self) -> DatasetSpec:
|
||||
return _stub_spec()
|
||||
|
|
|
|||
Loading…
Reference in a new issue