mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-23 16:54:56 +00:00
24 lines
755 B
Python
24 lines
755 B
Python
from atropos.environments.optimizer.wrapper import score_optimizer
|
|
from atroposlib.envs.base import BaseEnv, BaseEnvConfig
|
|
|
|
|
|
class OptimizerBenchmarkEnvConfig(BaseEnvConfig):
|
|
architecture: str = "mnist"
|
|
|
|
|
|
class OptimizerBenchmarkEnvironment(BaseEnv):
|
|
env_config_cls = OptimizerBenchmarkEnvConfig
|
|
|
|
def __init__(
|
|
self,
|
|
config: OptimizerBenchmarkEnvConfig,
|
|
server_configs=None,
|
|
slurm=False,
|
|
testing=False,
|
|
):
|
|
super().__init__(config, server_configs, slurm, testing)
|
|
self.architecture = config.architecture
|
|
|
|
async def evaluate(self, optimizer_code: str, *args, **kwargs):
|
|
# This method is required by BaseEnv
|
|
return score_optimizer(optimizer_code, self.architecture)
|