Spaces:
Sleeping
Sleeping
| """Genuine OpenEnv adapter around one SpreadsheetBench core instance.""" | |
| from __future__ import annotations | |
| from collections.abc import Callable | |
| import os | |
| from pathlib import Path | |
| from typing import Any | |
| from uuid import uuid4 | |
| from openenv.core.env_server.interfaces import Environment | |
| from openenv.core.env_server.types import EnvironmentMetadata | |
| from spreadsheetbench_env.catalog import get_split | |
| from spreadsheetbench_env.core import CoreObservation, SpreadsheetBenchCore | |
| from spreadsheetbench_env.executor import ExecutorLimits, ScriptExecutor | |
| from spreadsheetbench_env.interface import environment_description | |
| from spreadsheetbench_env.models import ( | |
| SpreadsheetBenchAction, | |
| SpreadsheetBenchObservation, | |
| SpreadsheetBenchState, | |
| ) | |
| _CoreFactory = Callable[[], SpreadsheetBenchCore] | |
| def _optional_integer(name: str) -> int | None: | |
| raw = os.getenv(name) | |
| if raw is None or not raw.strip(): | |
| return None | |
| try: | |
| return int(raw) | |
| except ValueError as error: | |
| raise ValueError(f"{name} must be an integer, got {raw!r}") from error | |
| def _positive_number(name: str, default: str) -> float: | |
| raw = os.getenv(name, default) | |
| try: | |
| value = float(raw) | |
| except ValueError as error: | |
| raise ValueError(f"{name} must be numeric, got {raw!r}") from error | |
| if value <= 0: | |
| raise ValueError(f"{name} must be positive") | |
| return value | |
| def _default_core_factory() -> SpreadsheetBenchCore: | |
| data_root = Path(os.getenv("SPREADSHEETBENCH_DATA", "/opt/spreadsheetbench")) | |
| work_root = Path(os.getenv("SPREADSHEETBENCH_WORK_ROOT", "/tmp")) | |
| limits = ExecutorLimits( | |
| timeout_s=_positive_number("SCRIPT_TIMEOUT_S", "120"), | |
| cpu_seconds=int(_positive_number("SCRIPT_CPU_SECONDS", "120")), | |
| memory_bytes=int(_positive_number("SCRIPT_MEMORY_BYTES", str(2 * 1024**3))), | |
| file_bytes=int(_positive_number("SCRIPT_FILE_BYTES", str(256 * 1024**2))), | |
| process_count=int(_positive_number("SCRIPT_PROCESS_COUNT", "64")), | |
| open_files=int(_positive_number("SCRIPT_OPEN_FILES", "128")), | |
| run_uid_base=_optional_integer("SPREADSHEETBENCH_RUNNER_UID_BASE"), | |
| run_gid=_optional_integer("SPREADSHEETBENCH_RUNNER_GID"), | |
| ) | |
| return SpreadsheetBenchCore( | |
| data_root, | |
| work_root=work_root, | |
| executor_factory=lambda root: ScriptExecutor(root, limits=limits), | |
| ) | |
| class SpreadsheetBenchEnvironment( | |
| Environment[ | |
| SpreadsheetBenchAction, | |
| SpreadsheetBenchObservation, | |
| SpreadsheetBenchState, | |
| ] | |
| ): | |
| """One connection-owned, typed SpreadsheetBench episode.""" | |
| SUPPORTS_CONCURRENT_SESSIONS = True | |
| REQUIRES_SINGLE_THREAD_EXECUTOR = False | |
| def __init__( | |
| self, | |
| core_factory: _CoreFactory = _default_core_factory, | |
| *, | |
| default_task_id: int | None = None, | |
| ) -> None: | |
| super().__init__() | |
| self._core_factory = core_factory | |
| self._core: SpreadsheetBenchCore | None = None | |
| self._default_task_id = default_task_id | |
| self._episode_id: str | None = None | |
| self._closed = False | |
| def _get_core(self) -> SpreadsheetBenchCore: | |
| if self._closed: | |
| raise RuntimeError("SpreadsheetBench environment is closed") | |
| if self._core is None: | |
| self._core = self._core_factory() | |
| return self._core | |
| def get_metadata(self) -> EnvironmentMetadata: | |
| return EnvironmentMetadata( | |
| name="SpreadsheetBench", | |
| description=environment_description(), | |
| version="0.1.0", | |
| documentation_url="https://arxiv.org/abs/2406.14991", | |
| ) | |
| def list_splits(self) -> list[dict[str, Any]]: | |
| return [ | |
| { | |
| "name": name, | |
| "role": get_split(name).role, | |
| "num_tasks": len(get_split(name).tasks), | |
| "source_revision": get_split(name).source_revision, | |
| "sha256": get_split(name).sha256, | |
| } | |
| for name in ("development", "test") | |
| ] | |
| def list_tasks(self, split: str) -> list[dict[str, Any]]: | |
| task_split = get_split(split) | |
| return [ | |
| { | |
| "id": task.task_id, | |
| "source_id": task.source_id, | |
| "instruction": task.instruction, | |
| "task_type": task.task_type, | |
| "split": task_split.role, | |
| "source_revision": task_split.source_revision, | |
| "catalog_sha256": task_split.sha256, | |
| } | |
| for task in task_split.tasks | |
| ] | |
| def num_tasks(self, split: str) -> int: | |
| return len(get_split(split).tasks) | |
| def get_task(self, split: str, index: int) -> dict[str, Any]: | |
| tasks = self.list_tasks(split) | |
| try: | |
| return tasks[index] | |
| except IndexError as error: | |
| raise IndexError( | |
| f"SpreadsheetBench {split} task index out of range: {index}" | |
| ) from error | |
| def get_task_range( | |
| self, | |
| split: str, | |
| start: int | None = None, | |
| stop: int | None = None, | |
| ) -> list[dict[str, Any]]: | |
| return self.list_tasks(split)[slice(start, stop)] | |
| def reset( | |
| self, | |
| seed: int | None = None, | |
| episode_id: str | None = None, | |
| *, | |
| task_id: int | None = None, | |
| training_feedback: bool | None = None, | |
| **kwargs: Any, | |
| ) -> SpreadsheetBenchObservation: | |
| del seed, kwargs | |
| resolved = self._default_task_id if task_id is None else int(task_id) | |
| if resolved is None: | |
| raise ValueError("SpreadsheetBench reset requires task_id") | |
| if not isinstance(training_feedback, bool): | |
| # Chosen per run by the client: SkillOpt training rollouts receive the gold-based | |
| # wrong-output feedback, evaluation rollouts never do. | |
| raise ValueError("SpreadsheetBench reset requires training_feedback (true or false)") | |
| core_observation = self._get_core().reset(resolved, training_feedback=training_feedback) | |
| self._episode_id = episode_id or str(uuid4()) | |
| return self._observation( | |
| core_observation, reward=0.0, terminated=False, truncated=False | |
| ) | |
| def step( | |
| self, | |
| action: SpreadsheetBenchAction, | |
| timeout_s: float | None = None, | |
| **kwargs: Any, | |
| ) -> SpreadsheetBenchObservation: | |
| del timeout_s, kwargs | |
| transition = self._get_core().step(action.code) | |
| return self._observation( | |
| transition.observation, | |
| reward=transition.reward, | |
| terminated=transition.terminated, | |
| truncated=transition.truncated, | |
| ) | |
| def state(self) -> SpreadsheetBenchState: | |
| if self._core is None: | |
| return SpreadsheetBenchState( | |
| episode_id=self._episode_id, | |
| step_count=0, | |
| closed=self._closed, | |
| ) | |
| state = self._core.state() | |
| return SpreadsheetBenchState( | |
| episode_id=self._episode_id, | |
| step_count=state.step_count, | |
| task_id=state.task_id, | |
| goal=state.goal, | |
| reward=state.reward, | |
| success=state.success, | |
| progress=state.progress, | |
| terminated=state.terminated, | |
| truncated=state.truncated, | |
| step_limit=state.step_limit, | |
| diagnostics=state.diagnostics, | |
| closed=state.closed, | |
| ) | |
| def close(self) -> None: | |
| if self._closed: | |
| return | |
| if self._core is not None: | |
| self._core.close() | |
| self._closed = True | |
| def _observation( | |
| self, | |
| observation: CoreObservation, | |
| *, | |
| reward: float, | |
| terminated: bool, | |
| truncated: bool, | |
| ) -> SpreadsheetBenchObservation: | |
| return SpreadsheetBenchObservation( | |
| task_id=observation.task_id, | |
| text=observation.text, | |
| goal=observation.goal, | |
| task_type=observation.task_type, | |
| case_count=observation.case_count, | |
| answer_position=observation.answer_position, | |
| success=observation.success, | |
| progress=observation.progress, | |
| terminated=terminated, | |
| truncated=truncated, | |
| done=terminated or truncated, | |
| reward=reward, | |
| step_limit=observation.step_limit, | |
| ) | |