diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 0f58de66..b24668b1 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -406,6 +406,16 @@ async def run_qa_benchmark( skill_factory = _skill_factory_for_target(target) resolved_skill_model = get_model(skill_config, config) + from haiku.rag.tools.filters import build_multi_document_filter + + question_to_filter: dict[str, str | None] = {} + for case in cases: + meta = case.metadata or {} + target_uri = meta.get("target_doc_uri") + question_to_filter[case.inputs] = ( + build_multi_document_filter([target_uri]) if target_uri else None + ) + async def answer_question(question: str) -> str: result = await run_skill_question( skill_factory=skill_factory, @@ -413,6 +423,7 @@ async def run_qa_benchmark( config=config, question=question, skill_model=resolved_skill_model, + document_filter=question_to_filter.get(question), ) set_eval_attribute("cited_uris", result.cited_uris) return result.answer diff --git a/evaluations/evaluations/datasets/mmlongbench.py b/evaluations/evaluations/datasets/mmlongbench.py index 120cba90..436f9e1f 100644 --- a/evaluations/evaluations/datasets/mmlongbench.py +++ b/evaluations/evaluations/datasets/mmlongbench.py @@ -105,6 +105,7 @@ def build_mmlb_case( metadata: dict[str, str] = { "case_index": str(index), "doc_id": doc["doc_id"], + "target_doc_uri": doc["doc_id"], "doc_type": doc.get("doc_type", ""), "answer_format": doc.get("answer_format", ""), "evidence_pages": str(evidence_pages), diff --git a/evaluations/tests/test_datasets.py b/evaluations/tests/test_datasets.py index 82e792c8..66f9a844 100644 --- a/evaluations/tests/test_datasets.py +++ b/evaluations/tests/test_datasets.py @@ -240,6 +240,7 @@ class TestMMLongBenchDoc: assert case.expected_output == "42" assert case.metadata is not None assert case.metadata["doc_id"] == "report.pdf" + assert case.metadata["target_doc_uri"] == "report.pdf" assert case.metadata["doc_type"] == "Financial report" assert case.metadata["answer_format"] == "Int" assert case.metadata["evidence_pages"] == "[5]"