414 lines
16 KiB
Python
414 lines
16 KiB
Python
"""Raw record loading and the FieldSpec-based default conversion.
|
|
|
|
Supported sources:
|
|
- local file: .jsonl / .json / .csv / .tsv / .parquet
|
|
- local dir: looks for ``{subset}_{split}.jsonl`` etc. (and csv/tsv/parquet)
|
|
- hub id: HuggingFace ``datasets`` (optional dependency, imported on demand)
|
|
- modelscope: ``params={'hub': 'modelscope'}`` — native ModelScope raw-file
|
|
HTTP download (no ``modelscope`` package needed, immune to its datasets
|
|
version pinning). File selection mirrors the local-dir convention and
|
|
understands HF-style shards (``test-00000-of-00001.parquet``).
|
|
"""
|
|
|
|
import csv
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import urllib.request
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List, Optional, Union
|
|
|
|
from .sample import Sample
|
|
from .spec import DatasetSpec, FieldSpec
|
|
|
|
_SUPPORTED_EXTS = ['.jsonl', '.json', '.csv', '.tsv', '.parquet']
|
|
# params keys consumed by the loader itself; never forwarded to load_dataset()
|
|
_RESERVED_PARAMS = {'hub', 'ms_files', 'hf_files', 'filter_column'}
|
|
_MS_API = 'https://www.modelscope.cn/api/v1/datasets'
|
|
|
|
|
|
def get_cache_root() -> Path:
|
|
"""Canonical evalharness cache root ($EVALHARNESS_CACHE or ~/.cache/evalharness)."""
|
|
return Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser()
|
|
|
|
|
|
def set_cache_root(path) -> None:
|
|
"""Override the cache root at runtime (also moves the .raw blob store)."""
|
|
os.environ['EVALHARNESS_CACHE'] = str(path)
|
|
|
|
|
|
def load_raw_records(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
|
|
"""Load raw records (list of dicts) from the spec's source.
|
|
|
|
When ``raw_dir`` is given, a copy of the *native* source data is placed
|
|
there (downloaded blob / copy of the local file / record dump for hub
|
|
sources) so every cache entry is self-contained: raw in, samples out.
|
|
"""
|
|
source = spec.source
|
|
if os.path.exists(source):
|
|
records, native_files = _load_local(source, spec)
|
|
if raw_dir is not None and native_files:
|
|
_link_or_copy_all(native_files, raw_dir)
|
|
elif spec.params.get('hub') == 'modelscope':
|
|
records = _load_from_modelscope(spec, raw_dir)
|
|
elif spec.params.get('hub') == 'hf_raw':
|
|
records = _load_from_hf_raw(spec, raw_dir)
|
|
else:
|
|
records = _load_from_hub(spec)
|
|
if raw_dir is not None:
|
|
# hub-native packaging lives in the HF cache; keep an exact,
|
|
# pre-conversion record dump so the entry is self-contained
|
|
raw_dir.mkdir(parents=True, exist_ok=True)
|
|
with open(raw_dir / 'records.jsonl', 'w', encoding='utf-8') as f:
|
|
for r in records:
|
|
f.write(json.dumps(r, ensure_ascii=False, default=str) + '\n')
|
|
records = _apply_subset_filter(spec, records)
|
|
return records
|
|
|
|
|
|
def _load_local(path: str, spec: DatasetSpec) -> tuple:
|
|
"""Load from a local path; also return the native file(s) to preserve."""
|
|
if os.path.isfile(path):
|
|
return _read_file(path), [Path(path)]
|
|
# directory: follow the <subset>_<split>.<ext> / <subset>.<ext> convention
|
|
for ext in _SUPPORTED_EXTS:
|
|
for candidate in (f'{spec.subset}_{spec.split}{ext}', f'{spec.subset}{ext}', f'{spec.split}{ext}'):
|
|
full = os.path.join(path, candidate)
|
|
if os.path.exists(full):
|
|
return _read_file(full), [Path(full)]
|
|
expected = [os.path.join(path, f'{spec.subset}_{spec.split}{e}') for e in _SUPPORTED_EXTS]
|
|
available = sorted(f for f in os.listdir(path) if os.path.splitext(f)[1] in _SUPPORTED_EXTS)
|
|
raise FileNotFoundError(
|
|
f'no dataset file found for subset={spec.subset!r} split={spec.split!r} in {path!r}.\n'
|
|
f'Expected one of:\n - ' + '\n - '.join(expected) + '\n'
|
|
f'Available: {available or "(none)"}'
|
|
)
|
|
|
|
|
|
def _link_or_copy_all(files: List[Path], raw_dir: Path) -> None:
|
|
"""Place native files in raw_dir (hardlink when possible, else copy)."""
|
|
raw_dir.mkdir(parents=True, exist_ok=True)
|
|
for src in files:
|
|
dst = raw_dir / src.name
|
|
if dst.exists() or src.parent.resolve() == raw_dir.resolve():
|
|
continue
|
|
try:
|
|
os.link(src, dst)
|
|
except OSError:
|
|
shutil.copy2(src, dst)
|
|
|
|
|
|
def _read_file(path: str) -> List[Dict[str, Any]]:
|
|
ext = os.path.splitext(path)[1]
|
|
if ext == '.parquet':
|
|
try:
|
|
import pyarrow.parquet as pq
|
|
except ImportError:
|
|
raise ImportError(
|
|
f'{path} is parquet; install the reader first: pip install "evalharness[all]"'
|
|
)
|
|
return pq.read_table(path).to_pylist()
|
|
with open(path, encoding='utf-8') as f:
|
|
if ext == '.jsonl':
|
|
return [json.loads(line) for line in f if line.strip()]
|
|
if ext == '.json':
|
|
data = json.load(f)
|
|
return data if isinstance(data, list) else [data]
|
|
if ext in ('.csv', '.tsv'):
|
|
return list(csv.DictReader(f, delimiter='\t' if ext == '.tsv' else ','))
|
|
raise ValueError(f'unsupported file format: {path}')
|
|
|
|
|
|
def _apply_subset_filter(spec: DatasetSpec, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""When ``params['filter_column']`` is set, subset selects that column's value.
|
|
|
|
Used when a mirror packs all subsets into one file with a category column
|
|
(e.g. evalscope/cmmlu: one parquet, ``category`` = subject).
|
|
"""
|
|
col = spec.params.get('filter_column')
|
|
if not col or spec.subset in ('default', 'all'):
|
|
return records
|
|
return [r for r in records if r.get(col) == spec.subset]
|
|
|
|
|
|
# ---------------- ModelScope raw-file loading (no modelscope package) ----------------
|
|
|
|
|
|
def _ms_api_json(url: str) -> Dict[str, Any]:
|
|
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
|
|
with urllib.request.urlopen(req, timeout=60) as resp:
|
|
return json.loads(resp.read().decode('utf-8'))
|
|
|
|
|
|
def _ms_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
|
|
"""Recursively list blob paths under a ModelScope dataset repo directory."""
|
|
url = f'{_MS_API}/{repo}/repo/tree?Revision=master' + (f'&Root={root}' if root else '')
|
|
try:
|
|
data = _ms_api_json(url).get('Data') or {}
|
|
except Exception:
|
|
return []
|
|
files: List[str] = []
|
|
for entry in data.get('Files') or []:
|
|
if entry.get('Type') == 'blob':
|
|
files.append(entry['Path'])
|
|
elif entry.get('Type') == 'tree' and depth < 4:
|
|
files.extend(_ms_list_files(repo, entry['Path'], depth + 1))
|
|
return files
|
|
|
|
|
|
def _ms_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
|
|
"""Pick repo files for this subset/split, mirroring the local-dir convention.
|
|
|
|
Priority: explicit ``params['ms_files']`` > exact name match > HF-style
|
|
shard match (``{split}-00000-of-00001`` / ``{subset}_{split}-...`` /
|
|
``{subset}_0``) > single-data-file repos. When the subset names a parent
|
|
directory, matches inside it win (e.g. livecodebench ``release_latest/``).
|
|
"""
|
|
explicit = spec.params.get('ms_files')
|
|
if explicit:
|
|
# paths support {subset}/{split} templating (e.g. tau2 domains)
|
|
return [p.format(subset=spec.subset, split=spec.split) for p in explicit]
|
|
|
|
def match_one(path: str) -> bool:
|
|
base = os.path.basename(path)
|
|
name = os.path.splitext(base)[0]
|
|
for ext in _SUPPORTED_EXTS:
|
|
if not base.endswith(ext):
|
|
continue
|
|
stem = base[: -len(ext)]
|
|
if stem in (f'{spec.subset}_{spec.split}', spec.subset, spec.split):
|
|
return True
|
|
if stem == f'{spec.subset}{spec.split}':
|
|
return True
|
|
shard = (
|
|
rf'{re.escape(spec.split)}-\d+-of-\d+$'
|
|
rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*'
|
|
rf'|{re.escape(spec.subset)}[-_]\d+$'
|
|
)
|
|
if re.fullmatch(shard, stem, flags=re.ASCII):
|
|
return True
|
|
return False
|
|
|
|
candidates = [f for f in files if match_one(f)]
|
|
if not candidates:
|
|
data_files = [f for f in files if os.path.splitext(f)[1] in _SUPPORTED_EXTS]
|
|
if len(data_files) == 1:
|
|
return data_files
|
|
return []
|
|
in_subset_dir = [f for f in candidates if os.path.dirname(f) == spec.subset]
|
|
return sorted(in_subset_dir or candidates)
|
|
|
|
|
|
def _parquet_ok(path: Path) -> bool:
|
|
"""Cheap integrity check: parquet files end with magic 'PAR1'."""
|
|
try:
|
|
with open(path, 'rb') as f:
|
|
f.seek(-4, 2)
|
|
return f.read(4) == b'PAR1'
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _ms_download(repo: str, path: str, dest_dir: Path) -> Path:
|
|
"""Download one repo file into the raw cache (content-addressed, reused)."""
|
|
dest = dest_dir / os.path.basename(path)
|
|
if dest.exists() and dest.stat().st_size > 0 \
|
|
and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
|
|
return dest
|
|
if dest.exists(): # truncated/corrupt (e.g. an interrupted download)
|
|
dest.unlink()
|
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
url = f'{_MS_API}/{repo}/repo?Revision=master&FilePath={path}'
|
|
tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
|
|
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
|
|
with urllib.request.urlopen(req, timeout=600) as resp, open(tmp, 'wb') as out:
|
|
while True:
|
|
chunk = resp.read(1 << 20)
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
os.replace(tmp, dest)
|
|
return dest
|
|
|
|
|
|
def _load_from_modelscope(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
|
|
import hashlib
|
|
|
|
files = _ms_list_files(spec.source)
|
|
if not files:
|
|
raise FileNotFoundError(f'no files found on ModelScope dataset {spec.source!r}')
|
|
selected = _ms_match_files(spec, files)
|
|
if not selected:
|
|
raise FileNotFoundError(
|
|
f'no data file for subset={spec.subset!r} split={spec.split!r} in '
|
|
f'modelscope {spec.source!r}. Available (first 10): {files[:10]}'
|
|
)
|
|
# shared blob store: download once per repo, hardlink into each cache entry
|
|
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
|
|
records: List[Dict[str, Any]] = []
|
|
blobs: List[Path] = []
|
|
for path in selected:
|
|
blobs.append(_ms_download(spec.source, path, blob_dir))
|
|
records.extend(_read_file(str(blobs[-1])))
|
|
if raw_dir is not None:
|
|
_link_or_copy_all(blobs, raw_dir)
|
|
return records
|
|
|
|
|
|
# ---------------- HuggingFace raw-file loading (no auth, exact bytes) ----------------
|
|
# For repos whose layout datasets.load_dataset cannot express (raw json trees,
|
|
# shard dirs without configs) or that only exist as plain files. Honors
|
|
# $HF_ENDPOINT (e.g. https://hf-mirror.com).
|
|
|
|
|
|
def _hf_base() -> str:
|
|
return os.environ.get('HF_ENDPOINT', 'https://huggingface.co').rstrip('/')
|
|
|
|
|
|
def _hf_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
|
|
"""Recursively list file paths under an HF dataset repo directory."""
|
|
url = f'{_hf_base()}/api/datasets/{repo}/tree/main' + (f'/{root}' if root else '') + '?limit=1000'
|
|
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=60) as resp:
|
|
entries = json.loads(resp.read().decode('utf-8'))
|
|
except Exception:
|
|
return []
|
|
files: List[str] = []
|
|
for entry in entries:
|
|
if entry.get('type') == 'file':
|
|
files.append(entry['path'])
|
|
elif entry.get('type') == 'directory' and depth < 4:
|
|
files.extend(_hf_list_files(repo, entry['path'], depth + 1))
|
|
return files
|
|
|
|
|
|
def _hf_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
|
|
"""Pick repo files for this subset/split (see module docstring conventions)."""
|
|
explicit = spec.params.get('hf_files')
|
|
if explicit:
|
|
return [p.format(subset=spec.subset, split=spec.split) for p in explicit]
|
|
|
|
def ext_ok(p: str) -> bool:
|
|
return os.path.splitext(p)[1] in _SUPPORTED_EXTS
|
|
|
|
data = [f for f in files if ext_ok(f)]
|
|
if spec.subset != 'default':
|
|
# subset wins exclusively: never ALSO match bare-split shards in the
|
|
# same dir (repos like sam-paech LCB mix test-*.parquet and
|
|
# release_v*-*.parquet in one data/ folder)
|
|
in_dir = [f for f in data if os.path.dirname(f) == spec.subset]
|
|
if in_dir:
|
|
return sorted(in_dir)
|
|
exact = []
|
|
for f in data:
|
|
stem = os.path.splitext(os.path.basename(f))[0]
|
|
if stem in (f'{spec.subset}_{spec.split}', spec.subset):
|
|
exact.append(f)
|
|
if exact:
|
|
return sorted(exact)
|
|
shards = [
|
|
f for f in data
|
|
if re.fullmatch(rf'{re.escape(spec.subset)}[-_]\d+(-of-\d+)?'
|
|
rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*',
|
|
os.path.splitext(os.path.basename(f))[0])
|
|
]
|
|
if shards:
|
|
return sorted(shards)
|
|
else:
|
|
exact = [
|
|
f for f in data
|
|
if os.path.splitext(os.path.basename(f))[0] == spec.split
|
|
]
|
|
if exact:
|
|
return sorted(exact)
|
|
shards = [
|
|
f for f in data
|
|
if re.fullmatch(rf'{re.escape(spec.split)}[-_]\d+(-of-\d+)?',
|
|
os.path.splitext(os.path.basename(f))[0])
|
|
]
|
|
if shards:
|
|
return sorted(shards)
|
|
if len(data) == 1:
|
|
return data
|
|
return []
|
|
|
|
|
|
def _hf_download(repo: str, path: str, dest_dir: Path) -> Path:
|
|
"""Download one repo file (follows the CDN redirect) into the blob store."""
|
|
dest = dest_dir / os.path.basename(path)
|
|
if dest.exists() and dest.stat().st_size > 0 \
|
|
and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
|
|
return dest
|
|
if dest.exists(): # truncated/corrupt (e.g. an interrupted download)
|
|
dest.unlink()
|
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
url = f'{_hf_base()}/datasets/{repo}/resolve/main/{path}'
|
|
tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
|
|
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
|
|
with urllib.request.urlopen(req, timeout=1800) as resp, open(tmp, 'wb') as out:
|
|
while True:
|
|
chunk = resp.read(1 << 20)
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
os.replace(tmp, dest)
|
|
return dest
|
|
|
|
|
|
def _load_from_hf_raw(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
|
|
import hashlib
|
|
|
|
files = _hf_list_files(spec.source)
|
|
if not files:
|
|
raise FileNotFoundError(f'no files found on HF dataset {spec.source!r}')
|
|
selected = _hf_match_files(spec, files)
|
|
if not selected:
|
|
raise FileNotFoundError(
|
|
f'no data file for subset={spec.subset!r} split={spec.split!r} in '
|
|
f'HF {spec.source!r}. Available (first 10): {files[:10]}'
|
|
)
|
|
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
|
|
records: List[Dict[str, Any]] = []
|
|
blobs: List[Path] = []
|
|
for path in selected:
|
|
blobs.append(_hf_download(spec.source, path, blob_dir))
|
|
records.extend(_read_file(str(blobs[-1])))
|
|
if raw_dir is not None:
|
|
_link_or_copy_all(blobs, raw_dir)
|
|
return records
|
|
|
|
|
|
def _load_from_hub(spec: DatasetSpec) -> List[Dict[str, Any]]:
|
|
try:
|
|
import datasets
|
|
except ImportError:
|
|
raise ImportError(
|
|
f'dataset {spec.name!r} lives on a hub ({spec.source!r}); '
|
|
'install it first: pip install evalharness (light deps are default)'
|
|
)
|
|
kwargs = {k: v for k, v in spec.params.items() if k not in _RESERVED_PARAMS}
|
|
subset = None if spec.subset == 'default' else spec.subset
|
|
ds = datasets.load_dataset(spec.source, subset, split=spec.split, revision=spec.version, **kwargs)
|
|
return [dict(r) for r in ds]
|
|
|
|
|
|
def field_spec_to_record_fn(fields: FieldSpec) -> Callable[[Dict[str, Any]], Sample]:
|
|
"""Build the default record->Sample converter from a FieldSpec."""
|
|
|
|
def convert(record: Dict[str, Any]) -> Sample:
|
|
target = record.get(fields.target, '')
|
|
if isinstance(target, (int, float)):
|
|
target = str(target)
|
|
metadata = {k: record.get(k) for k in fields.metadata}
|
|
return Sample(
|
|
input=record.get(fields.input, ''),
|
|
choices=record.get(fields.choices) if fields.choices in record else None,
|
|
target=target,
|
|
id=record.get(fields.id) if fields.id else None,
|
|
metadata=metadata,
|
|
)
|
|
|
|
return convert
|