19d5b2e7f6
benchmark.py was 1220 lines holding six unrelated jobs: populating a database, running retrieval, running QA, resolving datasets, moving databases to and from HuggingFace, and wiring the Typer CLI. qa.py takes both QA runners with their live summary, refusal metrics and target resolution. population.py takes populate_db and the batched ingest. retrieval.py takes run_retrieval_benchmark. artifacts.py takes HF_REPO_ID and the download/upload bodies. experiment.py takes DEFAULT_JUDGE_MODEL and build_experiment_metadata, which retrieval and QA both record. benchmark.py keeps the CLI at 258 lines: the Typer app, config and case-id loading, dataset resolution, evaluate_dataset, and three commands whose bodies are now a loop over specs. The module-level side effects stay with it — load_dotenv before configure_telemetry, so credentials and LOGFIRE_TOKEN are in the environment before telemetry and model setup read them — so no importable module carries one. Test patch targets follow the code. get_model, run_capability_question, run_capability_conversation, set_eval_attribute and HaikuRAG are patched inside moved code, so they move with it; run_qa_benchmark, run_retrieval_benchmark and find_config_file stay patchable on evaluations.benchmark because evaluate_dataset and _load_config still look them up there. One assertion got stronger: a QA test patched benchmark.HaikuRAG to prove the QA path does not open its own client. qa.py has no HaikuRAG reference at all now, so the test asserts that instead.
106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
"""Pre-built evaluation databases on HuggingFace."""
|
|
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import HfApi, snapshot_download
|
|
from rich.console import Console
|
|
|
|
from evaluations.config import DatasetSpec
|
|
|
|
console = Console()
|
|
|
|
HF_REPO_ID = "ggozad/haiku-rag-eval-dbs"
|
|
|
|
|
|
def download_dataset_db(spec: DatasetSpec, force: bool = False) -> None:
|
|
"""Fetch one dataset's database from HuggingFace into its local path."""
|
|
db = spec.db_path()
|
|
if db.exists() and not force:
|
|
console.print(
|
|
f"[yellow]Skipping {spec.key}: database already exists at {db}[/yellow]"
|
|
)
|
|
console.print("Use --force to overwrite.")
|
|
return
|
|
|
|
console.print(f"[blue]Downloading {spec.key}...[/blue]")
|
|
|
|
try:
|
|
downloaded_path = snapshot_download(
|
|
repo_id=HF_REPO_ID,
|
|
repo_type="dataset",
|
|
allow_patterns=f"{spec.db_filename}/*",
|
|
)
|
|
except Exception as e:
|
|
console.print(f"[red]Failed to download {spec.key}: {e}[/red]")
|
|
return
|
|
|
|
source_path = Path(downloaded_path) / spec.db_filename
|
|
if not source_path.exists():
|
|
console.print(f"[red]Database {spec.key} not found in HuggingFace repo.[/red]")
|
|
console.print(
|
|
f"[yellow]The database may not have been uploaded yet. "
|
|
f"Try running 'evaluations build {spec.key}' to create it locally.[/yellow]"
|
|
)
|
|
return
|
|
|
|
if db.exists():
|
|
shutil.rmtree(db)
|
|
|
|
db.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copytree(source_path, db)
|
|
|
|
console.print(f"[green]Downloaded {spec.key} to {db}[/green]")
|
|
|
|
|
|
def upload_dataset_db(spec: DatasetSpec) -> None:
|
|
"""Push one dataset's database to HuggingFace (maintainer only).
|
|
|
|
Uses ``upload_large_folder`` for resumable, parallel transfer — important
|
|
for the multi-GB ORB databases which would otherwise abort on any transient
|
|
network failure under plain ``upload_folder``.
|
|
|
|
``upload_large_folder`` has no ``path_in_repo`` — it ships the contents of
|
|
``folder_path`` to the repo root. Stage the db under a temp parent with
|
|
hardlinks so the basename becomes the remote path, leaving everything else
|
|
at the root undisturbed.
|
|
"""
|
|
db = spec.db_path()
|
|
if not db.exists():
|
|
console.print(f"[red]Database not found at {db}[/red]")
|
|
return
|
|
|
|
api = HfApi()
|
|
|
|
# Wipe the existing remote path so we don't accumulate orphaned files from
|
|
# prior uploads. upload_large_folder doesn't accept delete_patterns, so we
|
|
# do this as a separate commit. Safe to run if the path is missing.
|
|
try:
|
|
api.delete_folder(
|
|
path_in_repo=spec.db_filename,
|
|
repo_id=HF_REPO_ID,
|
|
repo_type="dataset",
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as staging:
|
|
target = Path(staging) / spec.db_filename
|
|
target.mkdir()
|
|
for src in db.rglob("*"):
|
|
if not src.is_file():
|
|
continue
|
|
dest = target / src.relative_to(db)
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
os.link(src, dest)
|
|
|
|
console.print(f"[blue]Uploading {spec.key} ({db})...[/blue]")
|
|
api.upload_large_folder(
|
|
folder_path=staging,
|
|
repo_id=HF_REPO_ID,
|
|
repo_type="dataset",
|
|
)
|
|
|
|
console.print(f"[green]Uploaded {spec.key} to {HF_REPO_ID}[/green]")
|