|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import io |
|
|
import sys |
|
|
import ast |
|
|
import time |
|
|
import json |
|
|
import signal |
|
|
import tempfile |
|
|
import platform |
|
|
import contextlib |
|
|
import faulthandler |
|
|
import numpy as np |
|
|
import resource |
|
|
import multiprocessing |
|
|
from enum import Enum |
|
|
from io import StringIO |
|
|
from decimal import Decimal |
|
|
from types import ModuleType |
|
|
from datetime import datetime |
|
|
from collections import defaultdict |
|
|
|
|
|
from unittest.mock import patch, mock_open |
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed |
|
|
|
|
|
import_string = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(50000)\n" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TimeoutException(Exception): |
|
|
pass |
|
|
|
|
|
|
|
|
class CODE_TYPE(Enum): |
|
|
call_based = 0 |
|
|
standard_input = 1 |
|
|
|
|
|
|
|
|
def timeout_handler(signum, frame): |
|
|
print("timeout occured: alarm went off") |
|
|
raise TimeoutException |
|
|
|
|
|
|
|
|
def reliability_guard(maximum_memory_bytes=None): |
|
|
""" |
|
|
This disables various destructive functions and prevents the generated code |
|
|
from interfering with the test (e.g. fork bomb, killing other processes, |
|
|
removing filesystem files, etc.) |
|
|
WARNING |
|
|
This function is NOT a security sandbox. Untrusted code, including, model- |
|
|
generated code, should not be blindly executed outside of one. See the |
|
|
Codex paper for more information about OpenAI's code sandbox, and proceed |
|
|
with caution. |
|
|
""" |
|
|
|
|
|
if maximum_memory_bytes is not None: |
|
|
import resource |
|
|
|
|
|
resource.setrlimit( |
|
|
resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes) |
|
|
) |
|
|
resource.setrlimit( |
|
|
resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes) |
|
|
) |
|
|
if not platform.uname().system == "Darwin": |
|
|
resource.setrlimit( |
|
|
resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes) |
|
|
) |
|
|
|
|
|
faulthandler.disable() |
|
|
|
|
|
import builtins |
|
|
|
|
|
|
|
|
builtins.quit = None |
|
|
|
|
|
import os |
|
|
|
|
|
os.environ["OMP_NUM_THREADS"] = "1" |
|
|
|
|
|
os.kill = None |
|
|
os.system = None |
|
|
os.putenv = None |
|
|
os.remove = None |
|
|
os.removedirs = None |
|
|
os.rmdir = None |
|
|
os.fchdir = None |
|
|
os.setuid = None |
|
|
os.fork = None |
|
|
os.forkpty = None |
|
|
os.killpg = None |
|
|
os.rename = None |
|
|
os.renames = None |
|
|
os.truncate = None |
|
|
os.replace = None |
|
|
os.unlink = None |
|
|
os.fchmod = None |
|
|
os.fchown = None |
|
|
os.chmod = None |
|
|
os.chown = None |
|
|
os.chroot = None |
|
|
os.fchdir = None |
|
|
os.lchflags = None |
|
|
os.lchmod = None |
|
|
os.lchown = None |
|
|
os.getcwd = None |
|
|
os.chdir = None |
|
|
|
|
|
import shutil |
|
|
|
|
|
shutil.rmtree = None |
|
|
shutil.move = None |
|
|
shutil.chown = None |
|
|
|
|
|
import subprocess |
|
|
|
|
|
subprocess.Popen = None |
|
|
|
|
|
__builtins__["help"] = None |
|
|
|
|
|
import sys |
|
|
|
|
|
sys.modules["ipdb"] = None |
|
|
sys.modules["joblib"] = None |
|
|
sys.modules["resource"] = None |
|
|
sys.modules["psutil"] = None |
|
|
sys.modules["tkinter"] = None |
|
|
|
|
|
|
|
|
def get_function(compiled_sol, fn_name: str): |
|
|
try: |
|
|
assert hasattr(compiled_sol, fn_name) |
|
|
return getattr(compiled_sol, fn_name) |
|
|
except Exception: |
|
|
return |
|
|
|
|
|
|
|
|
def compile_code(code: str, timeout: int): |
|
|
signal.alarm(timeout) |
|
|
try: |
|
|
tmp_sol = ModuleType("tmp_sol", "") |
|
|
exec(code, tmp_sol.__dict__) |
|
|
if "class Solution" in code: |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
compiled_sol = tmp_sol.Solution() |
|
|
else: |
|
|
|
|
|
compiled_sol = tmp_sol |
|
|
|
|
|
assert compiled_sol is not None |
|
|
finally: |
|
|
signal.alarm(0) |
|
|
|
|
|
return compiled_sol |
|
|
|
|
|
|
|
|
def truncatefn(s, length=300): |
|
|
if isinstance(s, str): |
|
|
pass |
|
|
else: |
|
|
s = str(s) |
|
|
if len(s) <= length: |
|
|
return s |
|
|
|
|
|
return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :] |
|
|
|
|
|
|
|
|
def grade_call_based( |
|
|
code: str, all_inputs: list, all_outputs: list, fn_name: str, timeout: int |
|
|
): |
|
|
|
|
|
|
|
|
code = import_string + "\n\n" + code |
|
|
compiled_sol = compile_code(code, timeout) |
|
|
|
|
|
if compiled_sol is None: |
|
|
return |
|
|
|
|
|
method = get_function(compiled_sol, fn_name) |
|
|
|
|
|
if method is None: |
|
|
return |
|
|
|
|
|
all_inputs = [ |
|
|
[json.loads(line) for line in inputs.split("\n")] for inputs in all_inputs |
|
|
] |
|
|
|
|
|
all_outputs = [json.loads(output) for output in all_outputs] |
|
|
|
|
|
total_execution = 0 |
|
|
all_results = [] |
|
|
for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)): |
|
|
signal.alarm(timeout) |
|
|
faulthandler.enable() |
|
|
try: |
|
|
|
|
|
start = time.time() |
|
|
prediction = method(*gt_inp) |
|
|
total_execution += time.time() - start |
|
|
signal.alarm(0) |
|
|
|
|
|
|
|
|
|
|
|
if isinstance(prediction, tuple): |
|
|
prediction = list(prediction) |
|
|
|
|
|
tmp_result = prediction == gt_out |
|
|
|
|
|
|
|
|
|
|
|
all_results.append(tmp_result) |
|
|
|
|
|
if not tmp_result: |
|
|
return all_results, { |
|
|
"output": truncatefn(prediction), |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
"error_code": -2, |
|
|
"error_message": "Wrong Answer", |
|
|
} |
|
|
except Exception as e: |
|
|
signal.alarm(0) |
|
|
if "timeoutexception" in repr(e).lower(): |
|
|
all_results.append(-3) |
|
|
return all_results, { |
|
|
"error": repr(e), |
|
|
"error_code": -3, |
|
|
"error_message": "Time Limit Exceeded", |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
} |
|
|
else: |
|
|
all_results.append(-4) |
|
|
return all_results, { |
|
|
"error": repr(e), |
|
|
"error_code": -4, |
|
|
"error_message": "Runtime Error", |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
} |
|
|
|
|
|
finally: |
|
|
signal.alarm(0) |
|
|
faulthandler.disable() |
|
|
|
|
|
return all_results, {"execution time": total_execution} |
|
|
|
|
|
|
|
|
def clean_if_name(code: str) -> str: |
|
|
try: |
|
|
astree = ast.parse(code) |
|
|
last_block = astree.body[-1] |
|
|
if isinstance(last_block, ast.If): |
|
|
condition = last_block.test |
|
|
if ast.unparse(condition).strip() == "__name__ == '__main__'": |
|
|
code = ( |
|
|
ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body) |
|
|
) |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
return code |
|
|
|
|
|
|
|
|
def make_function(code: str) -> str: |
|
|
try: |
|
|
import_stmts = [] |
|
|
all_other_stmts = [] |
|
|
astree = ast.parse(code) |
|
|
for stmt in astree.body: |
|
|
if isinstance(stmt, (ast.Import, ast.ImportFrom)): |
|
|
import_stmts.append(stmt) |
|
|
else: |
|
|
all_other_stmts.append(stmt) |
|
|
|
|
|
function_ast = ast.FunctionDef( |
|
|
name="wrapped_function", |
|
|
args=ast.arguments( |
|
|
posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[] |
|
|
), |
|
|
body=all_other_stmts, |
|
|
decorator_list=[], |
|
|
lineno=-1, |
|
|
) |
|
|
main_code = ( |
|
|
import_string |
|
|
+ "\n" |
|
|
+ ast.unparse(import_stmts) |
|
|
+ "\n" |
|
|
+ ast.unparse(function_ast) |
|
|
) |
|
|
return main_code |
|
|
except Exception: |
|
|
return code |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Capturing(list): |
|
|
def __enter__(self): |
|
|
self._stdout = sys.stdout |
|
|
sys.stdout = self._stringio = StringIO() |
|
|
|
|
|
self._stringio.close = lambda x: 1 |
|
|
return self |
|
|
|
|
|
def __exit__(self, *args): |
|
|
self.append(self._stringio.getvalue()) |
|
|
del self._stringio |
|
|
sys.stdout = self._stdout |
|
|
|
|
|
|
|
|
def call_method(method, inputs): |
|
|
|
|
|
if isinstance(inputs, list): |
|
|
inputs = "\n".join(inputs) |
|
|
|
|
|
inputs_line_iterator = iter(inputs.split("\n")) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@patch("builtins.open", mock_open(read_data=inputs)) |
|
|
@patch("sys.stdin", StringIO(inputs)) |
|
|
@patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) |
|
|
@patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) |
|
|
@patch("sys.stdin.read", lambda *args: inputs) |
|
|
|
|
|
def _inner_call_method(_method): |
|
|
try: |
|
|
return _method() |
|
|
except SystemExit: |
|
|
pass |
|
|
finally: |
|
|
pass |
|
|
|
|
|
return _inner_call_method(method) |
|
|
|
|
|
|
|
|
def get_stripped_lines(val: str): |
|
|
|
|
|
val = val.strip() |
|
|
|
|
|
return [val_line.strip() for val_line in val.split("\n")] |
|
|
|
|
|
|
|
|
def convert_line_to_decimals(line: str) -> tuple[bool, list[Decimal]]: |
|
|
try: |
|
|
decimal_line = [Decimal(elem) for elem in line.split()] |
|
|
except Exception: |
|
|
return False, [] |
|
|
return True, decimal_line |
|
|
|
|
|
|
|
|
def grade_stdio( |
|
|
code: str, |
|
|
all_inputs: list, |
|
|
all_outputs: list, |
|
|
timeout: int, |
|
|
): |
|
|
|
|
|
code = clean_if_name(code) |
|
|
|
|
|
|
|
|
code = make_function(code) |
|
|
|
|
|
compiled_sol = compile_code(code, timeout) |
|
|
if compiled_sol is None: |
|
|
return |
|
|
|
|
|
method = get_function(compiled_sol, "wrapped_function") |
|
|
|
|
|
if method is None: |
|
|
return |
|
|
|
|
|
all_results = [] |
|
|
total_execution_time = 0 |
|
|
for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)): |
|
|
signal.alarm(timeout) |
|
|
faulthandler.enable() |
|
|
|
|
|
signal.alarm(timeout) |
|
|
with Capturing() as captured_output: |
|
|
try: |
|
|
start = time.time() |
|
|
call_method(method, gt_inp) |
|
|
total_execution_time += time.time() - start |
|
|
|
|
|
signal.alarm(0) |
|
|
except Exception as e: |
|
|
signal.alarm(0) |
|
|
if "timeoutexception" in repr(e).lower(): |
|
|
all_results.append(-3) |
|
|
return all_results, { |
|
|
"error": repr(e), |
|
|
"error_code": -3, |
|
|
"error_message": "Time Limit Exceeded", |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
} |
|
|
else: |
|
|
all_results.append(-4) |
|
|
return all_results, { |
|
|
"error": repr(e), |
|
|
"error_code": -4, |
|
|
"error_message": "Runtime Error", |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
} |
|
|
|
|
|
finally: |
|
|
signal.alarm(0) |
|
|
faulthandler.disable() |
|
|
|
|
|
prediction = captured_output[0] |
|
|
|
|
|
stripped_prediction_lines = get_stripped_lines(prediction) |
|
|
stripped_gt_out_lines = get_stripped_lines(gt_out) |
|
|
|
|
|
|
|
|
|
|
|
WA_send_args = { |
|
|
"output": truncatefn(prediction), |
|
|
"inputs": truncatefn(gt_inp), |
|
|
"expected": truncatefn(gt_out), |
|
|
"error_code": -2, |
|
|
} |
|
|
|
|
|
if len(stripped_prediction_lines) != len(stripped_gt_out_lines): |
|
|
all_results.append(-2) |
|
|
WA_send_args["error_message"] = "Wrong answer: mismatched output length" |
|
|
return all_results, WA_send_args |
|
|
|
|
|
for output_line_idx, ( |
|
|
stripped_prediction_line, |
|
|
stripped_gt_out_line, |
|
|
) in enumerate(zip(stripped_prediction_lines, stripped_gt_out_lines)): |
|
|
WA_send_args["error_message"] = ( |
|
|
f"Wrong answer at {output_line_idx=}: {truncatefn(stripped_prediction_line)} != {truncatefn(stripped_gt_out_line)}" |
|
|
) |
|
|
|
|
|
|
|
|
if stripped_prediction_line == stripped_gt_out_line: |
|
|
continue |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
success, decimal_prediction_line = convert_line_to_decimals( |
|
|
stripped_prediction_line |
|
|
) |
|
|
if not success: |
|
|
all_results.append(-2) |
|
|
return all_results, WA_send_args |
|
|
success, decimal_gtout_line = convert_line_to_decimals(stripped_gt_out_line) |
|
|
if not success: |
|
|
all_results.append(-2) |
|
|
return all_results, WA_send_args |
|
|
|
|
|
if decimal_prediction_line == decimal_gtout_line: |
|
|
continue |
|
|
|
|
|
all_results.append(-2) |
|
|
return all_results, WA_send_args |
|
|
all_results.append(True) |
|
|
|
|
|
return all_results, {"execution time": total_execution_time} |
|
|
|
|
|
|
|
|
def run_test(sample, test=None, debug=False, timeout=6): |
|
|
""" |
|
|
if test(generated_code) is not None it'll try to run the code. |
|
|
otherwise it'll just return an input and output pair. |
|
|
""" |
|
|
signal.signal(signal.SIGALRM, timeout_handler) |
|
|
|
|
|
|
|
|
|
|
|
reliability_guard(4 * 1024 * 1024 * 1024) |
|
|
|
|
|
if debug: |
|
|
print(f"start = {datetime.now().time()}") |
|
|
|
|
|
try: |
|
|
in_outs = json.loads(sample["input_output"]) |
|
|
except ValueError as e: |
|
|
raise e |
|
|
in_outs = None |
|
|
|
|
|
if in_outs: |
|
|
if in_outs.get("fn_name") is None: |
|
|
which_type = CODE_TYPE.standard_input |
|
|
method_name = None |
|
|
|
|
|
else: |
|
|
which_type = CODE_TYPE.call_based |
|
|
method_name = in_outs["fn_name"] |
|
|
|
|
|
if debug: |
|
|
print(f"loaded input_output = {datetime.now().time()}") |
|
|
|
|
|
if test is None: |
|
|
assert False, "should not happen: test code is none" |
|
|
return in_outs, {"error": "no test code provided"} |
|
|
elif test is not None: |
|
|
results = [] |
|
|
|
|
|
if debug: |
|
|
print(f"loading test code = {datetime.now().time()}") |
|
|
|
|
|
if which_type == CODE_TYPE.call_based: |
|
|
signal.alarm(timeout) |
|
|
try: |
|
|
results, metadata = grade_call_based( |
|
|
code=test, |
|
|
all_inputs=in_outs["inputs"], |
|
|
all_outputs=in_outs["outputs"], |
|
|
fn_name=method_name, |
|
|
timeout=timeout, |
|
|
) |
|
|
return results, metadata |
|
|
except Exception as e: |
|
|
return [-4], { |
|
|
"error_code": -4, |
|
|
"error_message": f"Error during testing: {e}", |
|
|
} |
|
|
finally: |
|
|
signal.alarm(0) |
|
|
elif which_type == CODE_TYPE.standard_input: |
|
|
|
|
|
|
|
|
|
|
|
signal.alarm(timeout) |
|
|
try: |
|
|
results, metadata = grade_stdio( |
|
|
code=test, |
|
|
all_inputs=in_outs["inputs"], |
|
|
all_outputs=in_outs["outputs"], |
|
|
timeout=timeout, |
|
|
) |
|
|
return results, metadata |
|
|
except Exception as e: |
|
|
return [-4], { |
|
|
"error_code": -4, |
|
|
"error_message": f"Error during testing: {e}", |
|
|
} |
|
|
finally: |
|
|
signal.alarm(0) |
|
|
|
|
|
|
|
|
def _temp_run(sample, generation, debug, result, metadata_list, timeout): |
|
|
res, metadata = run_test(sample, test=generation, debug=debug, timeout=timeout) |
|
|
result.append(res) |
|
|
metadata_list.append(metadata) |
|
|
|
|
|
|
|
|
def check_correctness(sample, generation, timeout, debug=True): |
|
|
"""Check correctness of code generation with a global timeout. |
|
|
The global timeout is to catch some extreme/rare cases not handled by the timeouts |
|
|
inside `run_test`""" |
|
|
|
|
|
manager = multiprocessing.Manager() |
|
|
result = manager.list() |
|
|
metadata_list = manager.list() |
|
|
p = multiprocessing.Process( |
|
|
target=_temp_run, |
|
|
args=(sample, generation, debug, result, metadata_list, timeout), |
|
|
) |
|
|
p.start() |
|
|
p.join( |
|
|
timeout=(timeout + 1) * len(json.loads(sample["input_output"])["inputs"]) + 5 |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if p.is_alive(): |
|
|
p.kill() |
|
|
if not result: |
|
|
in_outs = json.loads(sample["input_output"]) |
|
|
|
|
|
result = [[-1 for i in range(len(in_outs["inputs"]))]] |
|
|
if debug: |
|
|
print("global timeout") |
|
|
return result[0], metadata_list[0] |
|
|
|
|
|
import signal |
|
|
from contextlib import contextmanager |
|
|
|
|
|
class TimeoutException(Exception): pass |
|
|
|
|
|
@contextmanager |
|
|
def time_limit(seconds): |
|
|
def signal_handler(signum, frame): |
|
|
raise TimeoutException("Timed out!") |
|
|
signal.signal(signal.SIGALRM, signal_handler) |
|
|
signal.alarm(seconds) |
|
|
try: |
|
|
yield |
|
|
finally: |
|
|
signal.alarm(0) |
|
|
def evaluate_generations_by_problem(args): |
|
|
problem_generations: list[str] = args[0] |
|
|
sample = args[1] |
|
|
debug: bool = args[2] |
|
|
timeout: int = args[3] |
|
|
|
|
|
res = [] |
|
|
metadata = [] |
|
|
reliability_guard(200*1024*1024*1024) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for o_idx, o in enumerate(problem_generations): |
|
|
print("we are testing!") |
|
|
curr_res = [-2] |
|
|
curr_res, curr_metadata = check_correctness( |
|
|
sample, o, timeout=timeout, debug=debug |
|
|
) |
|
|
|
|
|
print(curr_res, curr_metadata) |
|
|
if debug: |
|
|
print(f"\nSuccessful compilation of task {o_idx}!") |
|
|
fixed = [] |
|
|
for e in curr_res: |
|
|
if isinstance(e, np.ndarray): |
|
|
e = e.item(0) |
|
|
if isinstance(e, np.bool_): |
|
|
e = bool(e) |
|
|
fixed.append(e) |
|
|
curr_res = fixed |
|
|
res.append(curr_res) |
|
|
metadata.append(curr_metadata) |
|
|
return res, metadata |
|
|
|
|
|
|
|
|
def evaluate_generations( |
|
|
samples_list: list, |
|
|
generations_list: list[list[str]], |
|
|
debug: bool = False, |
|
|
num_process_evaluate: int = 16, |
|
|
timeout=6, |
|
|
): |
|
|
"""We take the list of code generations and try to compile them |
|
|
and the run their corresponding unit tests which are retrieved from the APPS dataset. |
|
|
|
|
|
Args: |
|
|
generations: list of code generations (same order as samples in APPS dataset) |
|
|
level: difficulty level used in the generation, can be "all", "introductory", "interview" or "competition" |
|
|
|
|
|
Returns: |
|
|
results: dictionary of results, key is the problem index, value is a list of results for each generation |
|
|
""" |
|
|
|
|
|
|
|
|
|
|
|
inputs = [ |
|
|
[(generations_list[index], samples_list[index], debug, timeout), index] |
|
|
for index in range(len(generations_list)) |
|
|
] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with ProcessPoolExecutor( |
|
|
max_workers=1 if debug else num_process_evaluate |
|
|
) as executor: |
|
|
futures = { |
|
|
executor.submit(evaluate_generations_by_problem, arg): index |
|
|
for arg, index in inputs |
|
|
} |
|
|
|
|
|
results = {} |
|
|
metadata = {} |
|
|
for future in as_completed(futures): |
|
|
index = futures[future] |
|
|
results[index], metadata[index] = future.result() |
|
|
|
|
|
assert len(results) == len( |
|
|
inputs |
|
|
), f"results = {len(results)} inputs = {len(inputs)} {results=}" |
|
|
|
|
|
|
|
|
return results, metadata |
|
|
|
|
|
|
|
|
def estimate_pass_at_k(num_samples, num_correct, k): |
|
|
"""Estimates pass@k of each problem and returns them in an array.""" |
|
|
|
|
|
def estimator(n: int, c: int, k: int) -> float: |
|
|
"""Calculates 1 - comb(n - c, k) / comb(n, k).""" |
|
|
if n - c < k: |
|
|
return 1.0 |
|
|
return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) |
|
|
|
|
|
import itertools |
|
|
|
|
|
if isinstance(num_samples, int): |
|
|
num_samples_it = itertools.repeat(num_samples, len(num_correct)) |
|
|
else: |
|
|
assert len(num_samples) == len(num_correct) |
|
|
num_samples_it = iter(num_samples) |
|
|
|
|
|
return np.array( |
|
|
[estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)] |
|
|
) |
|
|
|
|
|
|
|
|
def compute_metrics_from_results(results, k_list=[1, 5]): |
|
|
total = [] |
|
|
correct = [] |
|
|
task_ids = [] |
|
|
for task_id, res in results.items(): |
|
|
all_correct = [] |
|
|
for generation in res: |
|
|
gen = np.array(generation) |
|
|
all_correct.append(np.all(gen > 0)) |
|
|
task_ids.append(task_id) |
|
|
total.append(len(all_correct)) |
|
|
correct.append(sum(all_correct)) |
|
|
total = np.array(total) |
|
|
correct = np.array(correct) |
|
|
ks = k_list |
|
|
detail_pass_at_k = { |
|
|
f"pass@{k}": estimate_pass_at_k(total, correct, k).tolist() |
|
|
for k in ks |
|
|
if (total >= k).all() |
|
|
} |
|
|
pass_at_k = { |
|
|
f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() |
|
|
for k in ks |
|
|
if (total >= k).all() |
|
|
} |
|
|
detail_metrics = {k: dict(zip(task_ids, v)) for k, v in detail_pass_at_k.items()} |
|
|
pass_at_k["detail"] = detail_metrics |
|
|
return pass_at_k |
|
|
|
|
|
|
|
|
def codegen_metrics( |
|
|
samples_list, |
|
|
generations_list, |
|
|
k_list=[1, 5, 10, 20, 40, 50, 75, 100, 125, 150, 200, 500, 1000], |
|
|
num_process_evaluate=4, |
|
|
timeout=6, |
|
|
debug=False, |
|
|
): |
|
|
|
|
|
|
|
|
samples_linear = [] |
|
|
generations_linear = [] |
|
|
remap_index = [] |
|
|
results = defaultdict(list) |
|
|
metadatas = defaultdict(list) |
|
|
for idx, (sample, generation_list) in enumerate( |
|
|
zip(samples_list, generations_list) |
|
|
): |
|
|
assert isinstance(generation_list, list), generations_list[0] |
|
|
for generation in generation_list: |
|
|
assert isinstance(generation, str), generations_list[0] |
|
|
samples_linear.append(sample) |
|
|
generations_linear.append([generation]) |
|
|
remap_index.append(idx) |
|
|
|
|
|
|
|
|
results_linear, metadatas_linear = evaluate_generations( |
|
|
samples_linear, |
|
|
generations_linear, |
|
|
debug=debug, |
|
|
num_process_evaluate=num_process_evaluate, |
|
|
timeout=timeout, |
|
|
) |
|
|
|
|
|
for idx, sub_results in sorted(results_linear.items(), key=lambda x: x[0]): |
|
|
results[remap_index[idx]].append(sub_results[0]) |
|
|
|
|
|
for idx, sub_metadatas in sorted(metadatas_linear.items(), key=lambda x: x[0]): |
|
|
metadatas[remap_index[idx]].append(sub_metadatas[0]) |
|
|
|
|
|
metrics = compute_metrics_from_results(results, k_list=k_list) |
|
|
|
|
|
final_metadata = [] |
|
|
for key in sorted(list(metadatas.keys())): |
|
|
final_metadata.append(metadatas[key]) |
|
|
|
|
|
for i in range(len(final_metadata)): |
|
|
if type(final_metadata[i]) is not list: |
|
|
final_metadata[i] = [json.dumps(final_metadata[i])] |
|
|
else: |
|
|
final_metadata[i] = [json.dumps(x) for x in final_metadata[i]] |
|
|
|
|
|
assert len(final_metadata[i]) == len( |
|
|
generations_list[0] |
|
|
), f"{len(final_metadata[i])=}" |
|
|
|
|
|
return [metrics, results, final_metadata] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def parse_assert_statement(statement): |
|
|
""" |
|
|
Parse a Python assert statement and extract the expected output |
|
|
from the right side of the '==' operator as a string. |
|
|
|
|
|
:param statement: A string containing the assert statement. |
|
|
:return: The expected output from the assert statement as a string. |
|
|
""" |
|
|
try: |
|
|
parsed = ast.parse(statement, mode="exec") |
|
|
except SyntaxError: |
|
|
return "Invalid syntax" |
|
|
|
|
|
if len(parsed.body) == 0: |
|
|
return "Empty statement" |
|
|
|
|
|
if not isinstance(parsed.body[0], ast.Assert): |
|
|
return "Not an assert statement" |
|
|
|
|
|
comparison = parsed.body[0].test |
|
|
|
|
|
if not isinstance(comparison, ast.Compare) or not isinstance( |
|
|
comparison.ops[0], ast.Eq |
|
|
): |
|
|
return "Not an equality assertion" |
|
|
|
|
|
|
|
|
return ast.get_source_segment(statement, comparison.comparators[0]) |
|
|
|
|
|
|
|
|
def check_testcase_output(testcase_str, expected_output): |
|
|
|
|
|
if len(testcase_str.splitlines()) > 1: |
|
|
for line in testcase_str.splitlines(): |
|
|
if line.startswith("#"): |
|
|
continue |
|
|
if "assert" in line: |
|
|
testcase_str = line |
|
|
break |
|
|
|
|
|
testcase_str = testcase_str.strip() |
|
|
|
|
|
if "assert" in testcase_str: |
|
|
testcase_output_str = str(parse_assert_statement(testcase_str)) |
|
|
|
|
|
else: |
|
|
testcase_output_str = testcase_str |
|
|
|
|
|
global_result = None |
|
|
|
|
|
try: |
|
|
testcase_output_eval = eval(testcase_output_str) |
|
|
except Exception: |
|
|
global_result = False |
|
|
|
|
|
|
|
|
|
|
|
try: |
|
|
expected_output_eval = json.loads(expected_output) |
|
|
except Exception: |
|
|
global_result = False |
|
|
print("Failed to eval expected testcase output", expected_output) |
|
|
|
|
|
if global_result is None: |
|
|
global_result = testcase_output_eval == expected_output_eval |
|
|
|
|
|
return global_result |
|
|
|
|
|
|
|
|
def test_output_metrics( |
|
|
samples, |
|
|
generations, |
|
|
k_list=[1, 5], |
|
|
): |
|
|
num_samples = len(samples) |
|
|
results = [] |
|
|
|
|
|
for idx in range(num_samples): |
|
|
idx_results = [] |
|
|
sample = samples[idx] |
|
|
extracted_generation_list = generations[idx] |
|
|
for extracted_generation in extracted_generation_list: |
|
|
global_result = check_testcase_output( |
|
|
extracted_generation, sample["output"] |
|
|
) |
|
|
idx_results.append([global_result]) |
|
|
results.append(idx_results) |
|
|
|
|
|
results = {result_idx: results[result_idx] for result_idx in range(len(results))} |
|
|
|
|
|
metrics = compute_metrics_from_results(results, k_list=k_list) |
|
|
|
|
|
return [metrics, results] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
BASE_IMPORTS = """from itertools import accumulate, chain, combinations, count, permutations, product, groupby, islice, repeat |
|
|
from copy import deepcopy |
|
|
from string import ascii_lowercase |
|
|
from math import floor, log2, log10, sqrt, comb, gcd, ceil, inf, isqrt |
|
|
from collections import defaultdict, deque, Counter |
|
|
from bisect import bisect, bisect_left, bisect_right, insort |
|
|
from heapq import heappush, heappop, heapify, merge |
|
|
from functools import reduce, cache, lru_cache |
|
|
from random import randrange, shuffle |
|
|
from operator import itemgetter, sub |
|
|
from re import search as re_search # Assuming 're' refers to a regex search |
|
|
from os.path import commonprefix |
|
|
from typing import List, Tuple, Dict, Set, Optional, Union, Any, Callable, Iterable, Iterator, Generator |
|
|
import copy |
|
|
import string |
|
|
import math |
|
|
import collections |
|
|
import bisect |
|
|
import heapq |
|
|
import functools |
|
|
import random |
|
|
import itertools |
|
|
import operator |
|
|
import re |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
from math import log, prod # 'log' and 'prod' are functions in the math module |
|
|
from collections import deque, defaultdict, Counter, OrderedDict |
|
|
from itertools import accumulate, permutations, combinations, product, groupby, islice, chain, repeat, zip_longest, cycle |
|
|
from functools import lru_cache, reduce, partial |
|
|
# from sortedcontainers import SortedList, SortedDict, SortedSet |
|
|
# import sortedcontainers |
|
|
from operator import iand |
|
|
import sys |
|
|
""" |
|
|
|
|
|
@contextlib.contextmanager |
|
|
def chdir(root): |
|
|
if root == ".": |
|
|
yield |
|
|
return |
|
|
cwd = os.getcwd() |
|
|
os.chdir(root) |
|
|
try: |
|
|
yield |
|
|
except BaseException as exc: |
|
|
raise exc |
|
|
finally: |
|
|
os.chdir(cwd) |
|
|
|
|
|
@contextlib.contextmanager |
|
|
def create_tempdir(): |
|
|
with tempfile.TemporaryDirectory() as dirname: |
|
|
with chdir(dirname): |
|
|
yield dirname |
|
|
|
|
|
|
|
|
@contextlib.contextmanager |
|
|
def time_limit(seconds): |
|
|
def signal_handler(signum, frame): |
|
|
raise TimeoutException("Timed out!") |
|
|
|
|
|
signal.setitimer(signal.ITIMER_REAL, seconds) |
|
|
signal.signal(signal.SIGALRM, signal_handler) |
|
|
try: |
|
|
yield |
|
|
finally: |
|
|
signal.setitimer(signal.ITIMER_REAL, 0) |
|
|
|
|
|
|
|
|
class WriteOnlyStringIO(io.StringIO): |
|
|
"""StringIO that throws an exception when it's read from""" |
|
|
|
|
|
def read(self, *args, **kwargs): |
|
|
raise OSError |
|
|
|
|
|
def readline(self, *args, **kwargs): |
|
|
raise OSError |
|
|
|
|
|
def readlines(self, *args, **kwargs): |
|
|
raise OSError |
|
|
|
|
|
def readable(self, *args, **kwargs): |
|
|
"""Returns True if the IO object can be read.""" |
|
|
return False |
|
|
|
|
|
|
|
|
class redirect_stdin(contextlib._RedirectStream): |
|
|
_stream = "stdin" |
|
|
|
|
|
|
|
|
@contextlib.contextmanager |
|
|
def swallow_io(): |
|
|
stream = WriteOnlyStringIO() |
|
|
with contextlib.redirect_stdout(stream): |
|
|
with contextlib.redirect_stderr(stream): |
|
|
with redirect_stdin(stream): |
|
|
yield |
|
|
|
|
|
|
|
|
def unsafe_execute(check_program, result, timeout): |
|
|
|
|
|
with create_tempdir(): |
|
|
|
|
|
|
|
|
import os |
|
|
import shutil |
|
|
|
|
|
rmtree = shutil.rmtree |
|
|
rmdir = os.rmdir |
|
|
chdir = os.chdir |
|
|
|
|
|
|
|
|
reliability_guard(4 * 1024 * 1024 * 1024) |
|
|
|
|
|
|
|
|
try: |
|
|
exec_globals = {} |
|
|
with swallow_io(): |
|
|
with time_limit(timeout): |
|
|
exec(check_program, exec_globals) |
|
|
result.append("passed") |
|
|
except TimeoutException: |
|
|
result.append("timed out") |
|
|
except BaseException as e: |
|
|
result.append(f"failed: {e}") |
|
|
|
|
|
|
|
|
shutil.rmtree = rmtree |
|
|
os.rmdir = rmdir |
|
|
os.chdir = chdir |
|
|
|
|
|
|
|
|
def check_execution_correctness(check_program, timeout=3): |
|
|
""" |
|
|
Evaluates the functional correctness of a completion by running the test |
|
|
suite provided in the problem. |
|
|
|
|
|
:param completion_id: an optional completion ID so we can match |
|
|
the results later even if execution finishes asynchronously. |
|
|
""" |
|
|
manager = multiprocessing.Manager() |
|
|
result = manager.list() |
|
|
|
|
|
p = multiprocessing.Process(target=unsafe_execute, args=(check_program, result, timeout)) |
|
|
p.start() |
|
|
p.join(timeout=timeout + 1) |
|
|
if p.is_alive(): |
|
|
p.kill() |
|
|
|
|
|
if not result: |
|
|
result.append("timed out") |
|
|
|
|
|
return result[0] == "passed" |
|
|
|
|
|
|
|
|
def evaluate_score(args) -> list[bool]: |
|
|
gs, (c, i, o) = args |
|
|
|
|
|
execution_results = [] |
|
|
for g in gs: |
|
|
if i in g: |
|
|
pass |
|
|
else: |
|
|
code_to_execute = f"{BASE_IMPORTS}\n{c}\nassert {o} == {g}" |
|
|
execution_results.append(check_execution_correctness(code_to_execute, 3)) |
|
|
if len(execution_results) == 0: |
|
|
execution_results = [False] * len(gs) |
|
|
return execution_results |
|
|
|
|
|
|
|
|
def pass_at_k(n, c, k): |
|
|
if n - c < k: |
|
|
return 1.0 |
|
|
return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) |
|
|
|
|
|
|
|
|
def code_execution_metrics( |
|
|
samples, |
|
|
generations, |
|
|
): |
|
|
|
|
|
references = [(doc["code"], doc["input"], doc["output"]) for doc in samples] |
|
|
with ProcessPoolExecutor() as executor: |
|
|
args_list = zip(generations, references) |
|
|
results = executor.map(evaluate_score, args_list) |
|
|
all_results = list(results) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
pass_at_1s = [] |
|
|
for execution_result in all_results: |
|
|
c, n = execution_result.count(True), len(execution_result) |
|
|
pass_at_1s.append(pass_at_k(n, c, 1)) |
|
|
metrics = {"pass@1": sum(pass_at_1s) / len(pass_at_1s)} |
|
|
|
|
|
results = {} |
|
|
for i, r in enumerate(all_results): |
|
|
r_new = [] |
|
|
for _r in r: |
|
|
r_new.append([_r]) |
|
|
results[i] = r_new |
|
|
return [metrics, results] |
|
|
|
