Spaces:
Running
Running
File size: 5,865 Bytes
ba270e9 af304a5 ba270e9 2da981f ba270e9 | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """Genuine OpenEnv adapter around one MATH core instance."""
from __future__ import annotations
from collections.abc import Callable
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 math_env.catalog import get_split
from math_env.core import CoreObservation, MathCore
from math_env.interface import environment_description
from math_env.models import (
MathAction,
MathObservation,
MathState,
)
_CoreFactory = Callable[[], MathCore]
class MathEnvironment(
Environment[MathAction, MathObservation, MathState]
):
"""One connection-owned, typed MATH episode."""
SUPPORTS_CONCURRENT_SESSIONS = True
REQUIRES_SINGLE_THREAD_EXECUTOR = False
def __init__(
self,
core_factory: _CoreFactory = MathCore,
*,
default_task_id: int | None = None,
) -> None:
super().__init__()
self._core_factory = core_factory
self._core: MathCore | None = None
self._default_task_id = default_task_id
self._episode_id: str | None = None
self._closed = False
def _get_core(self) -> MathCore:
if self._closed:
raise RuntimeError("MATH environment is closed")
if self._core is None:
self._core = self._core_factory()
return self._core
def get_metadata(self) -> EnvironmentMetadata:
return EnvironmentMetadata(
name="MATH",
description=environment_description(),
version="0.1.0",
documentation_url="https://huggingface.co/datasets/HuggingFaceH4/MATH",
)
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.source_class,
"source_class": task.source_class,
"level": task.level,
"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"MATH {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,
**kwargs: Any,
) -> MathObservation:
del seed, kwargs
resolved = self._default_task_id if task_id is None else int(task_id)
if resolved is None:
raise ValueError("MATH reset requires task_id")
core_observation = self._get_core().reset(resolved)
self._episode_id = episode_id or str(uuid4())
return self._observation(core_observation, reward=0.0, terminated=False)
def step(
self,
action: MathAction,
timeout_s: float | None = None,
**kwargs: Any,
) -> MathObservation:
del timeout_s, kwargs
transition = self._get_core().step(action.answer)
return self._observation(
transition.observation,
reward=transition.reward,
terminated=transition.terminated,
)
@property
def state(self) -> MathState:
if self._core is None:
return MathState(
episode_id=self._episode_id,
step_count=0,
closed=self._closed,
)
state = self._core.state()
return MathState(
episode_id=self._episode_id,
step_count=state.step_count,
task_id=state.task_id,
goal=state.goal,
source_class=state.source_class,
level=state.level,
reward=state.reward,
success=state.success,
progress=state.progress,
terminated=state.terminated,
truncated=False,
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,
) -> MathObservation:
return MathObservation(
task_id=observation.task_id,
text=observation.text,
goal=observation.goal,
success=observation.success,
progress=observation.progress,
terminated=terminated,
truncated=False,
done=terminated,
reward=reward,
step_limit=observation.step_limit,
) |