File size: 3,993 Bytes
5e27996 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import os
from dataclasses import dataclass
from enum import Enum, auto
from typing import ClassVar
from huggingface_hub import hf_hub_download
class Category(Enum):
"""Benchmark categories with associated root directories and path patterns."""
RAG = auto()
RAG_0108 = auto()
LENGTH_SCALE = auto()
# ============================================================================
# HuggingFace config & local data root
# ============================================================================
HF_REPO_ID = "Anoy123423123/MSA-RAG-BENCHMARKS"
_DATA_ROOT = os.path.join(os.getcwd(), "data")
@dataclass(frozen=True)
class BenchmarkSpec:
"""Immutable specification for a single benchmark's file layout."""
bench_name: str # benchmark name, also the HF subdirectory
query_file: str
memory_file: str
def _resolve(self, filename: str) -> str:
"""Return local path if cached, otherwise download from HF into data/."""
local_path = os.path.join(_DATA_ROOT, self.bench_name, filename)
if os.path.exists(local_path):
return local_path
os.makedirs(os.path.dirname(local_path), exist_ok=True)
return hf_hub_download(
repo_id=HF_REPO_ID,
filename=f"{self.bench_name}/{filename}",
repo_type="dataset",
local_dir=_DATA_ROOT,
)
@property
def query_path(self) -> str:
return self._resolve(self.query_file)
@property
def memory_path(self) -> str:
return self._resolve(self.memory_file)
def get_bench_files(self) -> tuple[str, str]:
return self.query_path, self.memory_path
# ============================================================================
# Registry: benchmark name -> spec
# ============================================================================
def _rag(name: str) -> BenchmarkSpec:
return BenchmarkSpec(name, f"qdata_{name}.pkl", f"mdata_{name}.pkl")
def _rag_0108(name: str) -> BenchmarkSpec:
return BenchmarkSpec(name, f"qdata_{name}.pkl", f"mdata_{name}.pkl")
_REGISTRY: dict[str, BenchmarkSpec] = {
# --- Length-scale benchmarks ---
"ms_100M": BenchmarkSpec("ms_100M", "qdata_msmarco_16K.pkl", "mdata_msmarco_100M.pkl"),
# --- Multi-hop QA ---
"2wikimultihopqa": _rag("2wikimultihopqa"),
"hotpotqa": _rag("hotpotqa"),
"musique": _rag("musique"),
# --- HippoRAG ---
"hipporag_narrative": _rag_0108("hipporag_narrative"),
"hipporag_popqa": _rag_0108("hipporag_popqa"),
# --- Single-hop QA ---
"nature_questions": _rag("nature_questions"),
"triviaqa_06M": _rag("triviaqa_06M"),
"triviaqa_10M": _rag("triviaqa_10M"),
# --- Multilingual / Passage retrieval ---
"dureader": _rag("dureader"),
"msmarco_v1": _rag("msmarco_v1"),
}
ALL_BENCH_NAMES: list[str] = list(_REGISTRY)
# ============================================================================
# Public API
# ============================================================================
class BenchMarks:
"""Resolve benchmark name to query / memory file paths.
Usage:
bench = BenchMarks("hotpotqa")
query_file, memory_file = bench.get_bench_files()
"""
AVAILABLE: ClassVar[list[str]] = ALL_BENCH_NAMES
def __init__(self, bench_name: str) -> None:
if bench_name not in _REGISTRY:
raise ValueError(
f"Unknown benchmark: {bench_name!r}. "
f"Available: {', '.join(ALL_BENCH_NAMES)}"
)
self._spec = _REGISTRY[bench_name]
self.name = bench_name
self.bench_name = self._spec.bench_name
self.query_file_name = self._spec.query_file
self.memory_file_name = self._spec.memory_file
def get_bench_files(self) -> tuple[str, str]:
return self._spec.get_bench_files()
def __repr__(self) -> str:
return f"BenchMarks({self.name!r})"
|