OpenCompass/opencompass/datasets/apps.py
2024-04-19 20:49:46 +08:00

878 lines
32 KiB
Python

import faulthandler
import itertools
import json
import multiprocessing
import os
import platform
import re
# to run the solution files we're using a timing based approach
import signal
import sys
from collections import defaultdict
# used for debugging to time steps
from datetime import datetime
from enum import Enum
# for capturing the stdout
from io import StringIO
# used for testing the code that reads from input
from unittest.mock import mock_open, patch
import numpy as np
from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
try:
from pyext import RuntimeModule
except ImportError:
RuntimeModule = None
from opencompass.openicl.icl_evaluator import BaseEvaluator
from opencompass.registry import ICL_EVALUATORS, LOAD_DATASET
from .base import BaseDataset
TIMEOUT = 10
@LOAD_DATASET.register_module()
class APPSDataset(BaseDataset):
@staticmethod
def load(path: str, num_repeats: int = 1):
dataset = load_dataset(path)
new_dataset = DatasetDict()
# add new column "starter" in the prompt
for split in dataset.keys():
new_samples = []
for _, sample in enumerate(dataset[split]):
starter_code = None if len(
sample['starter_code']) == 0 else sample['starter_code']
try:
input_output = json.loads(sample['input_output'])
fn_name = (None if not input_output.get('fn_name') else
input_output['fn_name'])
except ValueError:
fn_name = None
starter = ''
if starter_code:
starter += starter_code
if (not fn_name) and (not starter_code):
call_format = '\\nUse Standard Input format'
starter += call_format
else:
call_format = '\\nUse Call-Based format'
starter += call_format
# Add the new column "starter" to the sample
sample['starter'] = starter
new_samples.append(sample)
new_data = {
key: [sample[key] for sample in new_samples]
for key in new_samples[0].keys()
}
new_dataset[split] = Dataset.from_dict(new_data)
# num_repeats duplicate
# train_repeated = []
test_repeated = []
# for sample in new_dataset['train']:
# train_repeated.extend([sample] * num_repeats)
for sample in new_dataset['test']:
test_repeated.extend([sample] * num_repeats)
# dataset_train_repeated = new_dataset['train'].from_list(
# train_repeated
# )
dataset_test_repeated = new_dataset['test'].from_list(test_repeated)
return DatasetDict({
# 'train': dataset_train_repeated,
'test': dataset_test_repeated
})
@LOAD_DATASET.register_module()
class APPS_miniDataset(BaseDataset):
@staticmethod
def load(path: str, num_repeats: int = 1):
dataset = load_from_disk(path)
new_dataset = DatasetDict()
# add new column "starter" in the prompt
for split in dataset.keys():
new_samples = []
for _, sample in enumerate(dataset[split]):
starter_code = None if len(
sample['starter_code']) == 0 else sample['starter_code']
try:
input_output = json.loads(sample['input_output'])
fn_name = (None if not input_output.get('fn_name') else
input_output['fn_name'])
except ValueError:
fn_name = None
starter = ''
if starter_code:
starter += starter_code
if (not fn_name) and (not starter_code):
call_format = '\\nUse Standard Input format'
starter += call_format
else:
call_format = '\\nUse Call-Based format'
starter += call_format
# Add the new column "starter" to the sample
sample['starter'] = starter
new_samples.append(sample)
new_data = {
key: [sample[key] for sample in new_samples]
for key in new_samples[0].keys()
}
new_dataset[split] = Dataset.from_dict(new_data)
# num_repeats duplicate
# train_repeated = []
test_repeated = []
# for sample in new_dataset['train']:
# train_repeated.extend([sample] * num_repeats)
for sample in new_dataset['test']:
test_repeated.extend([sample] * num_repeats)
# dataset_train_repeated = new_dataset['train'].from_list(
# train_repeated
# )
dataset_test_repeated = new_dataset['test'].from_list(test_repeated)
return DatasetDict({
# 'train': dataset_train_repeated,
'test': dataset_test_repeated
})
EOF_STRINGS = ['\nQUESTION', '\n---', '\nANSWER', '<|endoftext|>']
@ICL_EVALUATORS.register_module()
class APPSEvaluator(BaseEvaluator):
def post_process(self, text):
if '```' in text:
blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
if len(blocks) == 0:
text = text.split('```')[1] # fall back to default strategy
else:
text = blocks[0] # fetch the first code block
if not text.startswith('\n'): # starting with ```python
text = text[max(text.find('\n') + 1, 0):]
return text
TIMEOUT = 10
def check_correctness(self, 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`
"""
def _temp_run(sample, generation, debug, result):
result.append(run_test(sample, test=generation, debug=debug))
manager = multiprocessing.Manager()
result = manager.list()
p = multiprocessing.Process(target=_temp_run,
args=(sample, generation, debug, result))
p.start()
p.join(timeout=timeout + 1)
if p.is_alive():
p.kill()
if not result:
in_outs = json.loads(sample['input_output'])
# consider that all tests failed
result = [[-1 for i in range(len(in_outs['inputs']))]]
if debug:
print('global timeout')
return result[0]
def evaluate_generations(self,
generations,
samples,
idx=None,
debug=False):
assert len(generations.keys()) == len(samples)
results = {}
idx = 0
for task_id, problem_generations in generations.items():
sample = samples.iloc[idx]
res = []
# loop over the generations
for o_idx, o in enumerate(problem_generations):
curr_res = [-2]
try:
curr_res = self.check_correctness(sample,
o,
timeout=TIMEOUT,
debug=debug)
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
if not np.all(curr_res):
if debug:
print('Results were not True for all test cases')
except Exception as e:
if debug:
print(
f'Compilation failed, test framework exception = {repr(e)}{e}\n' # noqa: E501
)
break
finally:
assert isinstance(curr_res, list)
res.append(curr_res)
results[task_id] = res
idx += 1
return results
def estimate_pass_at_k(self, 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))
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(self, results, k_list=[1, 10, 100]):
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}': self.estimate_pass_at_k(total, correct, k).tolist()
for k in ks if (total >= k).all()
}
pass_at_k = {
f'pass@{k}':
100 * (self.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 score(self, predictions, references, test_set):
assert len(predictions) == len(references)
generations = defaultdict(list)
for refer, pred in zip(references, predictions):
pred = self.post_process(pred)
generations[refer].append(pred)
# convert to non-duplicated version
test_set = test_set.to_pandas()
test_set_origin = test_set.drop_duplicates(subset='problem_id')
results = self.evaluate_generations(generations, test_set_origin)
metrics = self.compute_metrics(results)
return metrics
class CODE_TYPE(Enum):
call_based = 0
standard_input = 1
# stuff for setting up signal timer
class TimeoutException(Exception):
pass
def timeout_handler(signum, frame):
print('alarm went off')
raise TimeoutException
try:
signal.signal(signal.SIGALRM, timeout_handler)
except AttributeError:
print('signal.SIGALRM is not available on this platform')
timeout = 4 # seconds
# used to capture stdout as a list
# from https://stackoverflow.com/a/16571630/6416660
# alternative use redirect_stdout() from contextlib
class Capturing(list):
def __enter__(self):
self._stdout = sys.stdout
sys.stdout = self._stringio = StringIO()
# Make closing the StringIO a no-op
self._stringio.close = lambda x: 1
return self
def __exit__(self, *args):
self.extend(self._stringio.getvalue().splitlines())
del self._stringio # free up some memory
sys.stdout = self._stdout
def run_test(sample, test=None, debug=False):
"""if test(generated_code) is not None it'll try to run the code.
otherwise it'll just return an input and output pair.
"""
# Disable functionalities that can make destructive changes to the test.
reliability_guard()
if debug:
print(f'start = {datetime.now().time()}')
try:
in_outs = json.loads(sample['input_output'])
except ValueError:
in_outs = None
if in_outs:
if in_outs.get('fn_name') is None:
which_type = CODE_TYPE.standard_input # Standard input
method_name = None
else:
which_type = CODE_TYPE.call_based # Call-based
method_name = in_outs['fn_name']
if debug:
print(f'loaded input_output = {datetime.now().time()}')
if test is None:
return in_outs
elif test is not None:
results = []
sol = 'import sys\nimport time\nimport itertools\nfrom itertools import accumulate, product, permutations, combinations\nimport collections\nfrom collections import Counter, OrderedDict, deque, defaultdict, ChainMap\nfrom functools import lru_cache\nimport math\nfrom math import sqrt, sin, cos, tan, ceil, fabs, floor, gcd, exp, log, log2\nimport fractions\nfrom typing import List, Tuple\nimport numpy as np\nimport random\nimport heapq\nfrom heapq import *\n' # noqa: E501
if debug:
print(f'loading test code = {datetime.now().time()}')
if which_type == CODE_TYPE.call_based:
sol += test
if debug:
print(f'sol = {sol}')
signal.alarm(timeout)
try:
tmp_sol = RuntimeModule.from_string('tmp_sol', '', sol)
if 'class Solution' not in test:
tmp = tmp_sol
else:
tmp = tmp_sol.Solution()
signal.alarm(0)
except Exception as e:
signal.alarm(0)
if debug:
print(f'type 0 compilation error = {e}')
results.append(-2)
return results
signal.alarm(0)
elif which_type == CODE_TYPE.standard_input:
# sol
tmp_test = test.split('\n')
new_test = []
for x in tmp_test:
if (not x.startswith('from ')) and (
not x.startswith('import ')):
new_test.append('\t' + x + '\n')
else:
new_test.append(x + '\n')
tmp_test = new_test
new_test = ''
started = False
for i in tmp_test:
if i.startswith('\t') and not started:
new_test += 'stdin = sys.stdin\nstdout = sys.stdout\n'
new_test += 'def code():\n'
new_test += i
started = True
elif started and ((i.startswith('from ')) or
(i.startswith('import '))):
new_test += '\t' + i
else:
new_test += i
tmp_test = new_test
sol += tmp_test
if debug:
print(f'sol = {sol}')
method_name = 'code'
signal.alarm(timeout)
try:
tmp_sol = RuntimeModule.from_string('tmp_sol', '', sol)
tmp = tmp_sol
signal.alarm(0)
except Exception as e:
signal.alarm(0)
if debug:
print(f'type 1 compilation error = {e}')
results.append(-2)
return results
signal.alarm(0)
if debug:
print(f'get method = {datetime.now().time()}')
try:
method = getattr(tmp,
method_name) # get_attr second arg must be str
except: # noqa: E722
signal.alarm(0)
e = sys.exc_info()
print(f'unable to get function error = {e}')
results.append(-2)
return results
for index, inputs in enumerate(in_outs['inputs']):
# JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) # noqa: E501
try:
if isinstance(inputs[0], dict):
inputs = [{int(k): v for k, v in inputs[0].items()}]
except: # noqa: E722
True
try:
if isinstance(in_outs['outputs'][index], dict):
in_outs['outputs'][index] = [{
int(k): v
for k, v in in_outs['outputs'][index].items()
}]
except: # noqa: E722
True
try:
if isinstance(in_outs['outputs'][index][0], dict):
in_outs['outputs'][index] = [{
int(k): v
for k, v in in_outs['outputs'][index][0].items()
}]
except: # noqa: E722
True
if debug:
print(
f'time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}' # noqa: E501
)
if which_type == CODE_TYPE.call_based: # Call-based
signal.alarm(timeout)
faulthandler.enable()
try:
output = method(*inputs)
# ground truth sequences are not tuples
if isinstance(output, tuple):
output = list(output)
tmp_result = output == in_outs['outputs'][index]
if isinstance(in_outs['outputs'][index],
list) and in_outs['outputs'][index]:
tmp_result = tmp_result or (
output == in_outs['outputs'][index][0])
# ground truth sequences are not tuples
try:
if isinstance(output[0], tuple):
tmp_result = tmp_result or ([
list(x) for x in output
] == in_outs['outputs'][index][0])
except: # noqa: E722
True
results.append(tmp_result)
# reset the alarm
signal.alarm(0)
except Exception as e:
signal.alarm(0)
faulthandler.disable()
if debug:
print(
f'Standard input runtime error or time limit exceeded error = {e}' # noqa: E501
)
results.append(-1)
continue
faulthandler.disable()
signal.alarm(0)
if debug:
print(
f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
elif which_type == CODE_TYPE.standard_input: # Standard input
faulthandler.enable()
signal.alarm(timeout)
passed = False
if isinstance(inputs, list):
inputs = '\n'.join(inputs)
if isinstance(in_outs['outputs'][index], list):
in_outs['outputs'][index] = '\n'.join(
in_outs['outputs'][index])
with Capturing() as output:
try:
call_method(method, inputs)
# reset the alarm
signal.alarm(0)
passed = True
except Exception as e:
# runtime error or took too long
signal.alarm(0)
print(
f'Call-based runtime error or time limit exceeded error = {repr(e)}{e}' # noqa: E501
)
results.append(-1)
signal.alarm(0)
if not passed:
if debug:
nl = '\n'
if not isinstance(inputs, list):
print(
f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
else:
print(
f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
continue
if passed and debug:
print(
f"==> output = {output}, test outputs = {in_outs['outputs'][index]}" # noqa: E501
)
if custom_compare_(output, in_outs['outputs'][index]):
tmp_result = True
results.append(tmp_result)
continue
# ground truth sequences are expressed as lists not tuples
if isinstance(output, tuple):
output = list(output)
tmp_result = False
try:
tmp_result = (output == [in_outs['outputs'][index]])
if isinstance(in_outs['outputs'][index], list):
tmp_result = tmp_result or (
output == in_outs['outputs'][index])
if isinstance(output[0], str):
tmp_result = tmp_result or ([
e.strip() for e in output
] == in_outs['outputs'][index])
except Exception as e:
if debug:
print(f'Failed check1 exception = {e}')
pass
if tmp_result == True: # noqa: E712
results.append(tmp_result)
continue
# try one more time without \n
if isinstance(in_outs['outputs'][index], list):
for tmp_index, i in enumerate(in_outs['outputs'][index]):
in_outs['outputs'][index][tmp_index] = i.split('\n')
in_outs['outputs'][index][tmp_index] = [
x.strip()
for x in in_outs['outputs'][index][tmp_index] if x
]
else:
in_outs['outputs'][index] = in_outs['outputs'][
index].split('\n')
in_outs['outputs'][index] = list(
filter(len, in_outs['outputs'][index]))
in_outs['outputs'][index] = list(
map(lambda x: x.strip(), in_outs['outputs'][index]))
try:
tmp_result = (output == [in_outs['outputs'][index]])
if isinstance(in_outs['outputs'][index], list):
tmp_result = tmp_result or (
output == in_outs['outputs'][index])
except Exception as e:
if debug:
print(f'Failed check2 exception = {e}')
pass
if tmp_result == True: # noqa: E712
results.append(tmp_result)
continue
# try by converting the output into a split up list too
if isinstance(output, list):
output = list(filter(len, output))
if debug:
nl = '\n'
if not isinstance(inputs, list):
print(
f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
else:
print(
f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
if tmp_result == True: # noqa: E712
results.append(tmp_result)
continue
try:
tmp_result = (output == [in_outs['outputs'][index]])
if isinstance(in_outs['outputs'][index], list):
tmp_result = tmp_result or (
output == in_outs['outputs'][index])
except Exception as e:
if debug:
print(f'Failed check3 exception = {e}')
pass
try:
output_float = [float(e) for e in output]
gt_float = [float(e) for e in in_outs['outputs'][index]]
tmp_result = tmp_result or (
(len(output_float) == len(gt_float))
and np.allclose(output_float, gt_float))
except Exception:
pass
try:
if isinstance(output[0], list):
output_float = [float(e) for e in output[0]]
gt_float = [
float(e) for e in in_outs['outputs'][index][0]
]
tmp_result = tmp_result or (
(len(output_float) == len(gt_float))
and np.allclose(output_float, gt_float))
except Exception:
pass
if tmp_result == True: # noqa: E712
results.append(tmp_result)
continue
# try by converting the stuff into split up list
if isinstance(in_outs['outputs'][index], list):
for tmp_index, i in enumerate(in_outs['outputs'][index]):
in_outs['outputs'][index][tmp_index] = set(i.split())
else:
in_outs['outputs'][index] = set(
in_outs['outputs'][index].split())
try:
tmp_result = (output == in_outs['outputs'][index])
except Exception as e:
if debug:
print(f'Failed check4 exception = {e}')
continue
if tmp_result == True: # noqa: E712
results.append(tmp_result)
continue
# try by converting the output into a split up list too
if isinstance(output, list):
for tmp_index, i in enumerate(output):
output[tmp_index] = i.split()
output = list(filter(len, output))
for tmp_index, i in enumerate(output):
output[tmp_index] = set(i)
else:
output = output.split()
output = list(filter(len, output))
output = set(output)
try:
tmp_result = (set(frozenset(s) for s in output) == set(
frozenset(s) for s in in_outs['outputs'][index]))
except Exception as e:
if debug:
print(f'Failed check5 exception = {e}')
# if they are all numbers, round so that similar numbers are treated as identical # noqa: E501
try:
tmp_result = tmp_result or (
set(
frozenset(round(float(t), 3) for t in s)
for s in output) == # noqa: E502, E128
set(
frozenset(round(float(t), 3) for t in s)
for s in in_outs['outputs'][index])) # noqa: E501
except Exception as e:
if debug:
print(f'Failed check6 exception = {e}')
if tmp_result == True and debug: # noqa: E712
print('PASSED')
results.append(tmp_result)
if debug:
nl = '\n'
if not isinstance(inputs, list):
print(
f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
else:
print(
f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" # noqa: E501
)
return results
def custom_compare_(output, ground_truth):
if isinstance(output, list):
output_1 = '\n'.join(output)
if stripped_string_compare(output_1, ground_truth):
return True
if isinstance(output, list):
output_2 = [o.lstrip().rstrip() for o in output]
output_2 = '\n'.join(output_2)
if stripped_string_compare(output_2, ground_truth):
return True
return False
def stripped_string_compare(s1, s2):
s1 = s1.lstrip().rstrip()
s2 = s2.lstrip().rstrip()
return s1 == s2
def call_method(method, inputs):
if isinstance(inputs, list):
inputs = '\n'.join(inputs)
inputs_line_iterator = iter(inputs.split('\n'))
# sys.setrecursionlimit(10000)
# @patch('builtins.input', side_effect=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)
# @patch('sys.stdout.write', print)
def _inner_call_method(_method):
try:
return _method()
except SystemExit:
pass
finally:
pass
return _inner_call_method(method)
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.exit = None
builtins.quit = None
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 # type: ignore
__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