Spaces:
Sleeping
Sleeping
File size: 3,856 Bytes
d3d0e0e | 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 | from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from data_agent_baseline.benchmark.schema import PublicTask, TaskAssets, TaskRecord
TASK_DIR_PREFIX = "task_"
def _task_number(task_id: str) -> int:
try:
if not task_id.startswith(TASK_DIR_PREFIX):
# raise ValueError(f"Invalid task id: {task_id}")
return float("inf")
return int(task_id.removeprefix(TASK_DIR_PREFIX))
except ValueError:
return float("inf")
def _load_task_record(task_json_path: Path) -> TaskRecord:
payload = json.loads(task_json_path.read_text())
required_keys = {"task_id", "question"}
actual_keys = set(payload)
missing_required = required_keys - actual_keys
if missing_required:
raise ValueError(
f"Missing required task.json keys for {task_json_path.parent.name}: "
f"{sorted(missing_required)}"
)
task_id = str(payload.get("task_id", "")).strip()
question = str(payload.get("question", "")).strip()
raw_difficulty = payload.get("difficulty", "unknown")
difficulty = str(raw_difficulty).strip() or "unknown"
if not task_id:
raise ValueError(f"Invalid task_id in {task_json_path}: expected non-empty string")
if not question:
raise ValueError(f"Invalid question in {task_json_path}: expected non-empty string")
return TaskRecord(
task_id=task_id,
difficulty=difficulty,
question=question,
)
@dataclass(frozen=True, slots=True)
class DABenchPublicDataset:
root_dir: Path
@property
def exists(self) -> bool:
return self.root_dir.is_dir()
def task_dirs(self) -> list[Path]:
if not self.exists:
return []
task_dirs = [
path
for path in self.root_dir.iterdir()
if path.is_dir() and path.name.startswith(TASK_DIR_PREFIX)
]
task_dirs.sort(key=lambda path: _task_number(path.name))
return task_dirs
def list_task_ids(self) -> list[str]:
return [path.name for path in self.task_dirs()]
def get_task(self, task_id: str) -> PublicTask:
task_dir = self.root_dir / task_id
task_json_path = task_dir / "task.json"
if not task_json_path.exists():
raise FileNotFoundError(f"Missing task.json: {task_json_path}")
record = _load_task_record(task_json_path)
if record.task_id != task_dir.name:
raise ValueError(f"task_id mismatch for {task_dir}: task.json has {record.task_id}")
context_dir = task_dir / "context"
if not context_dir.is_dir():
raise FileNotFoundError(f"Missing context dir: {context_dir}")
assets = TaskAssets(task_dir=task_dir, context_dir=context_dir)
return PublicTask(record=record, assets=assets)
def iter_tasks(
self,
*,
task_ids: list[str] | None = None,
difficulty: str | None = None,
difficulties: list[str] | None = None,
) -> list[PublicTask]:
selected_ids = set(task_ids or [])
selected_difficulties = set(difficulties or [])
if difficulty is not None:
selected_difficulties.add(difficulty)
tasks: list[PublicTask] = []
for task_dir in self.task_dirs():
if selected_ids and task_dir.name not in selected_ids:
continue
task = self.get_task(task_dir.name)
if selected_difficulties and task.difficulty not in selected_difficulties:
continue
tasks.append(task)
return tasks
def task_counts(self) -> dict[str, int]:
counts: dict[str, int] = {}
for task in self.iter_tasks():
counts[task.difficulty] = counts.get(task.difficulty, 0) + 1
return counts
|