Spaces:
Running
Running
File size: 3,622 Bytes
9c84f9d | 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 | import csv
import json
import logging
from pathlib import Path
from huggingface_hub import hf_hub_download, list_repo_files
from solar_eval.config import get_provider_config
logger = logging.getLogger(__name__)
class DatasetLoader:
"""Loads datasets from HuggingFace repos or local filesystem."""
def __init__(self, hf_token: str | None = None, base_dir: Path | None = None) -> None:
self.hf_token = hf_token or get_provider_config().hf_token or None
self.base_dir = base_dir
def download_file(self, repo: str, path: str, force: bool = False) -> Path:
"""Download a single file from HF repo. Returns local cache path."""
local_path = hf_hub_download(
repo_id=repo,
filename=path,
repo_type="dataset",
token=self.hf_token,
force_download=force,
)
return Path(local_path)
def _resolve_local_path(self, repo: str, path: str) -> Path:
"""Resolve a local file path from repo dir and dataset path."""
repo_path = Path(repo)
if repo_path.is_absolute():
return repo_path / path
if self.base_dir:
return self.base_dir / repo / path
return repo_path / path
def load_jsonl(
self, repo: str, path: str, source: str = "huggingface", force: bool = False
) -> list[dict]:
"""Load and parse a JSONL file from HuggingFace or local filesystem."""
if source == "local":
local_path = self._resolve_local_path(repo, path)
if not local_path.exists():
raise FileNotFoundError(f"Local dataset not found: {local_path}")
else:
local_path = self.download_file(repo, path, force=force)
records = []
with open(local_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
records.append(json.loads(line))
logger.info(f"Loaded {len(records)} records from {local_path}")
return records
def load_csv(
self, repo: str, path: str, source: str = "huggingface", force: bool = False
) -> list[dict]:
"""Download and parse a CSV file.
Supports both DictReader (with headers) and headerless CSV.
For headerless CSV, columns are mapped to 'wrong' and 'correct' keys.
"""
if source == "local":
local_path = self._resolve_local_path(repo, path)
if not local_path.exists():
raise FileNotFoundError(f"Local CSV not found: {local_path}")
else:
local_path = self.download_file(repo, path, force=force)
with open(local_path, encoding="utf-8") as f:
# Peek first line to detect headers
first_line = f.readline().strip()
f.seek(0)
if first_line and ("wrong" in first_line or "before" in first_line):
reader = csv.DictReader(f)
records = list(reader)
else:
# Headerless CSV: treat as (wrong, correct) pairs
reader = csv.reader(f)
records = []
for row in reader:
if len(row) >= 2 and row[0].strip():
records.append({"wrong": row[0].strip(), "correct": row[1].strip()})
logger.info(f"Loaded {len(records)} records from {local_path}")
return records
def list_files(self, repo: str) -> list[str]:
"""List all files in HF repo."""
return list(list_repo_files(repo, repo_type="dataset", token=self.hf_token))
|