153 lines
5 KiB
Python
153 lines
5 KiB
Python
"""Adapter for loading local GSM1K benchmark files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from evals.generalization.item_schema import GeneralizationAuditItem
|
|
|
|
|
|
def load_gsm1k_items(
|
|
*,
|
|
local_cache: Path,
|
|
split: str,
|
|
max_items: int | None = None,
|
|
) -> tuple[GeneralizationAuditItem, ...]:
|
|
"""Load and normalize GSM1K items from a local cache directory or file.
|
|
|
|
Args:
|
|
local_cache: Path to the cache directory or file.
|
|
split: The requested dataset split (e.g. 'test').
|
|
max_items: Optional maximum number of items to load.
|
|
|
|
Returns:
|
|
A tuple of loaded GeneralizationAuditItem records.
|
|
"""
|
|
if not local_cache.exists():
|
|
raise FileNotFoundError(
|
|
f"Local cache path does not exist: {local_cache}"
|
|
)
|
|
|
|
target_file = None
|
|
if local_cache.is_file():
|
|
target_file = local_cache
|
|
elif local_cache.is_dir():
|
|
# Check standard file naming conventions for local cache
|
|
candidate_jsonl = local_cache / f"{split}.jsonl"
|
|
candidate_json = local_cache / f"{split}.json"
|
|
if candidate_jsonl.is_file():
|
|
target_file = candidate_jsonl
|
|
elif candidate_json.is_file():
|
|
target_file = candidate_json
|
|
else:
|
|
# Fallback search for files containing the split name in their filename
|
|
all_files = sorted(local_cache.glob("*"))
|
|
for f in all_files:
|
|
if (
|
|
f.is_file()
|
|
and split in f.name
|
|
and f.suffix in (".jsonl", ".json")
|
|
):
|
|
target_file = f
|
|
break
|
|
|
|
if not target_file:
|
|
raise FileNotFoundError(
|
|
f"No JSON or JSONL file found for split {split!r} in {local_cache}"
|
|
)
|
|
|
|
try:
|
|
content = target_file.read_text(encoding="utf-8")
|
|
except OSError as exc:
|
|
raise ValueError(
|
|
f"Failed to read local GSM1K file {target_file}: {exc}"
|
|
) from exc
|
|
|
|
# Parse as JSONL first, fallback to JSON array
|
|
raw_records = []
|
|
lines = content.strip().splitlines()
|
|
is_jsonl = True
|
|
for idx, line in enumerate(lines):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
parsed = json.loads(line)
|
|
# JSONL lines must be objects (mappings)
|
|
if not isinstance(parsed, dict):
|
|
is_jsonl = False
|
|
break
|
|
raw_records.append(parsed)
|
|
except Exception:
|
|
is_jsonl = False
|
|
break
|
|
|
|
if not is_jsonl or not raw_records:
|
|
try:
|
|
data = json.loads(content)
|
|
if isinstance(data, list):
|
|
raw_records = data
|
|
# Since it parsed successfully as a JSON array, it's not JSONL format
|
|
is_jsonl = False
|
|
elif isinstance(data, dict):
|
|
raw_records = [data]
|
|
is_jsonl = False
|
|
else:
|
|
raise ValueError("JSON root must be a list or dictionary.")
|
|
except Exception as exc:
|
|
raise ValueError(
|
|
f"Failed to parse GSM1K file {target_file} as JSON/JSONL: {exc}"
|
|
) from exc
|
|
|
|
items: list[GeneralizationAuditItem] = []
|
|
for idx, rec in enumerate(raw_records):
|
|
if max_items is not None and len(items) >= max_items:
|
|
break
|
|
|
|
# Extract question (query)
|
|
question = rec.get("question") or rec.get("prompt")
|
|
if question is None:
|
|
raise ValueError(
|
|
f"GSM1K record at index {idx} in {target_file.name} is missing 'question' or 'prompt' field."
|
|
)
|
|
|
|
# Extract answer
|
|
answer = rec.get("answer") or rec.get("grade") or rec.get("label")
|
|
if answer is None:
|
|
raise ValueError(
|
|
f"GSM1K record at index {idx} in {target_file.name} is missing 'answer', 'grade', or 'label' field."
|
|
)
|
|
|
|
# Resolve stable item ID
|
|
item_id = str(rec.get("id") or rec.get("item_id") or idx)
|
|
|
|
# Build stable opaque prompt reference
|
|
prompt_ref = f"gsm1k:{split}:{item_id}"
|
|
|
|
q_sha256 = hashlib.sha256(str(question).encode("utf-8")).hexdigest()
|
|
a_sha256 = hashlib.sha256(str(answer).encode("utf-8")).hexdigest()
|
|
|
|
metadata_list = [
|
|
("source_format", "jsonl" if is_jsonl else "json"),
|
|
("source_index", str(idx)),
|
|
("source_file_name", target_file.name),
|
|
("source_record_id", item_id),
|
|
("question_sha256", q_sha256),
|
|
("answer_sha256", a_sha256),
|
|
("question_length", str(len(str(question)))),
|
|
("answer_kind", "numeric_text"),
|
|
]
|
|
|
|
items.append(
|
|
GeneralizationAuditItem(
|
|
dataset="GSM1K",
|
|
split=split,
|
|
item_id=item_id,
|
|
prompt_ref=prompt_ref,
|
|
answer_kind="numeric_text",
|
|
metadata=tuple(metadata_list),
|
|
)
|
|
)
|
|
|
|
return tuple(items)
|