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

Score retrieval evals by document URI

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-08-06 11:06:13 +03:00
parent a4839c0ad2
commit 788a2fe731
No known key found for this signature in database
2 changed files with 14 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
evaluations/evaluations/benchmark.py
View file

@ -235,11 +235,10 @@ async def run_retrieval_benchmark(
seen = set()
identifiers = []
for result in chunks:
# Use arxiv_id from metadata if present, otherwise use URI
doc_id = result.document_meta.get("arxiv_id") or result.document_uri
if doc_id and doc_id not in seen:
identifiers.append(doc_id)
seen.add(doc_id)
uri = result.document_uri
if uri and uri not in seen:
identifiers.append(uri)
seen.add(uri)
return identifiers

19
evaluations/tests/test_benchmark.py
View file

@ -515,21 +515,21 @@ class TestRetrievalTarget:
assert searches[0]["include_images"] is False
@pytest.mark.asyncio
async def test_prefers_metadata_identifier_over_uri(self, tmp_path: Path) -> None:
async def test_ranks_each_document_once(self, tmp_path: Path) -> None:
from haiku.rag.store.models.chunk import SearchResult
from evaluations.benchmark import run_retrieval_benchmark
def _result(uri: str, score: float) -> SearchResult:
return SearchResult(content="x", score=score, document_uri=uri)
class FakeRag:
async def search(self, **kwargs) -> list[SearchResult]:
return [
SearchResult(
content="x",
score=1.0,
document_id="doc-1",
document_uri="uri-other",
document_meta={"arxiv_id": "uri-x"},
)
_result("uri-other", 1.0),
_result("uri-x", 0.9),
_result("uri-other", 0.8),
_result("uri-x", 0.7),
]
with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
@ -538,8 +538,9 @@ class TestRetrievalTarget:
self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
)
# uri-x is the only relevant document and ranks second of two
assert result is not None
assert result["map"] == 1.0
assert result["map"] == 0.5
class TestEvaluateDatasetCaseIds: