Spaces:
Running
Running
File size: 2,878 Bytes
ed65693 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Dataset loaders for benchmark-style retrieval evaluation."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
SCIFACT_DIR = Path("data/scifact")
def iter_jsonl(path: Path):
with open(path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
yield json.loads(line)
def scifact_available(data_dir: Path = SCIFACT_DIR) -> bool:
return (
(data_dir / "corpus.jsonl").exists()
and (data_dir / "queries.jsonl").exists()
and (data_dir / "qrels" / "test.tsv").exists()
)
def load_scifact_corpus(data_dir: Path = SCIFACT_DIR) -> list[dict[str, str]]:
"""Load SciFact corpus records as document dictionaries."""
records: list[dict[str, str]] = []
for item in iter_jsonl(data_dir / "corpus.jsonl"):
title = item.get("title", "").strip()
body = item.get("text", "").strip()
text = f"{title}\n\n{body}".strip()
if text:
records.append(
{
"doc_id": str(item["_id"]),
"title": title,
"text": text,
}
)
return records
def load_scifact_queries(data_dir: Path = SCIFACT_DIR) -> dict[str, str]:
return {
str(item["_id"]): item.get("text", "").strip()
for item in iter_jsonl(data_dir / "queries.jsonl")
if item.get("text", "").strip()
}
def load_scifact_qrels(data_dir: Path = SCIFACT_DIR, split: str = "test") -> dict[str, list[str]]:
qrels_path = data_dir / "qrels" / f"{split}.tsv"
qrels: dict[str, list[str]] = {}
with open(qrels_path, encoding="utf-8") as f:
next(f, None) # skip header
for line in f:
parts = line.strip().split("\t")
if len(parts) != 3:
continue
query_id, corpus_id, score = parts
if int(score) <= 0:
continue
qrels.setdefault(str(query_id), []).append(str(corpus_id))
return qrels
def load_scifact_golden_set(
data_dir: Path = SCIFACT_DIR, split: str = "test"
) -> list[dict[str, Any]]:
"""Return BEIR-style golden rows compatible with scripts/evaluate.py."""
queries = load_scifact_queries(data_dir)
qrels = load_scifact_qrels(data_dir, split=split)
golden_set: list[dict[str, Any]] = []
for query_id in sorted(qrels, key=lambda value: int(value) if value.isdigit() else value):
query = queries.get(query_id)
if not query:
continue
golden_set.append(
{
"query_id": query_id,
"query": query,
"relevant_sources": sorted(set(qrels[query_id])),
"relevant_keywords": [],
"query_type": "scifact",
}
)
return golden_set
|