mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-25 17:10:42 +00:00
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
from modal.functions import Function
|
|
from evaluator import OptimizerEvaluator
|
|
|
|
def score_optimizer(optimizer_code: str, architecture: str):
|
|
"""
|
|
Test an optimizer implementation by sending it to the Modal sandbox environment.
|
|
|
|
Args:
|
|
optimizer_code (str): The optimizer code to test
|
|
|
|
Returns:
|
|
dict: Results containing stdout, stderr, code and filename
|
|
"""
|
|
send_code = Function.from_name("optimizer-test", "send_code")
|
|
|
|
response_obj = send_code.remote(optimizer_code)
|
|
|
|
print(response_obj)
|
|
|
|
evaluator = OptimizerEvaluator()
|
|
|
|
validity = evaluator.check_validity(optimizer_code=optimizer_code, stdout=response_obj["stdout"], stderr=response_obj["stderr"])
|
|
|
|
if not validity:
|
|
return 0
|
|
else:
|
|
score = evaluator.score_optimizer(optimizer_code=optimizer_code, architecture=architecture)
|
|
return score
|
|
|
|
if __name__ == "__main__":
|
|
test_code = """
|
|
import torch
|
|
import math
|
|
|
|
class CustomOptimizer(torch.optim.Optimizer):
|
|
def __init__(self, params, lr=0.01, beta1=0.9, beta2=0.999, eps=1e-8):
|
|
defaults = dict(lr=lr, beta1=beta1, beta2=beta2, eps=eps)
|
|
"""
|
|
result = score_optimizer(test_code, "mnist")
|
|
print(result)
|