SwiftSage / code_executor.py
yuchenlin's picture
Upload 14 files
1a0cf07 verified
raw
history blame
3.3 kB
"""
Source and credits: https://github.com/ZubinGou/math-evaluation-harness/blob/main/python_executor.py
We modified it to be more simple.
"""
import io
import pickle
import traceback
from concurrent.futures import ProcessPoolExecutor, TimeoutError
from contextlib import redirect_stdout
class GenericRuntime:
GLOBAL_DICT = {}
LOCAL_DICT = None
HEADERS = []
def __init__(self):
self._global_vars = self.GLOBAL_DICT.copy()
self._local_vars = self.LOCAL_DICT.copy() if self.LOCAL_DICT else None
for c in self.HEADERS:
self.exec_code(c)
def exec_code(self, code_piece: str) -> None:
exec(code_piece, self._global_vars)
def eval_code(self, expr: str) -> any:
return eval(expr, self._global_vars)
def inject(self, var_dict):
self._global_vars.update(var_dict)
@property
def answer(self):
return self._global_vars['answer']
class PythonExecutor:
def __init__(
self,
runtime=None,
get_answer_symbol=None,
get_answer_expr=None,
get_answer_from_stdout=False,
timeout_length=5,
):
self.runtime = runtime if runtime else GenericRuntime()
self.answer_symbol = get_answer_symbol
self.get_answer_expr = get_answer_expr
self.get_answer_from_stdout = get_answer_from_stdout
self.timeout_length = timeout_length
def execute(self, code):
try:
if self.get_answer_from_stdout:
program_io = io.StringIO()
with redirect_stdout(program_io):
self.runtime.exec_code('\n'.join(code))
program_io.seek(0)
result = program_io.read()
elif self.answer_symbol:
self.runtime.exec_code('\n'.join(code))
result = self.runtime._global_vars[self.answer_symbol]
elif self.get_answer_expr:
self.runtime.exec_code('\n'.join(code))
result = self.runtime.eval_code(self.get_answer_expr)
else:
self.runtime.exec_code('\n'.join(code[:-1]))
result = self.runtime.eval_code(code[-1])
report = "Done"
pickle.dumps(result) # Serialization check
except Exception as e:
result = ''
report = str(e)
return result, report
def apply(self, code):
code_snippet = code.split('\n')
# Use ProcessPoolExecutor to enforce timeout
with ProcessPoolExecutor() as executor:
future = executor.submit(self.execute, code_snippet)
try:
result, report = future.result(timeout=self.timeout_length)
except TimeoutError:
result, report = "", "Timeout Error"
return result.strip(), report.strip()
# Example usage
if __name__ == "__main__":
executor = PythonExecutor(get_answer_from_stdout=True)
code = """
from sympy import Matrix
def null_space_basis():
A = Matrix([[3, 3, -1, -6], [9, -1, -8, -1], [7, 4, -2, -9]])
basis = A.nullspace()
return [v.evalf(3) for v in basis]
result = null_space_basis()
print(result)
"""
result, report = executor.apply(code)
print("Result:", result)
print("Report:", report)