Skip to content

codecritic
codecritic
Commit eb972b4c by nzy
Options

transfer to llmkit

parent 432936ac
Showing with 0 additions and 904 deletions
codecritic/cli/run_hf_score.py deleted 100644 → 0
import argparse
import os
from pathlib import Path
import pprint
from tqdm import tqdm
import torch
import torch.nn.functional as F
import transformers
import accelerate
from codecritic.utils.json import load_jsonl, save_jsonl
from codecritic.data.verify import get_score_token_id
from codecritic.utils.metric import group_results, score_pass_at_k
@torch.inference_mode()
def hf_score(accelerator, model, tokenizer, prompts):
score_token = get_score_token_id(tokenizer)
with accelerator.split_between_processes(prompts) as partial_prompts:
results = []
for item in tqdm(partial_prompts):
input_ids = tokenizer.apply_chat_template(
item["messages"], add_generation_prompt=True, return_tensors="pt"
).to("cuda")
output = model(input_ids)
next_token_logits = output.logits[0, -1, :]
score = F.softmax(next_token_logits, dim=0)[score_token].item()
results.append({**item, "score": score})
return accelerate.utils.gather_object(results)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str)
parser.add_argument("--test", type=str)
parser.add_argument("--apps", type=str)
args = parser.parse_args()
home_path = Path(args.model).parent
result_dir = home_path / "hf_eval"
result_dir.mkdir(exist_ok=True)
prompts = load_jsonl(args.test)
# os.environ["TOKENIZERS_PARALLELISM"] = "false"
accelerator = accelerate.Accelerator()
tokenizer = transformers.AutoTokenizer.from_pretrained(args.model)
model = transformers.AutoModelForCausalLM.from_pretrained(args.model, device_map="auto")
# model, tokenizer = accelerator.prepare(model, tokenizer)
for name, param in model.named_parameters():
print(f"{name}: {param.device}")
# model.generation_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id else tokenizer.eos_token_id
# model.generation_config.eos_token_id = tokenizer.eos_token_id
model.eval()
accelerator.wait_for_everyone()
results = hf_score(accelerator, model, tokenizer, prompts)
if accelerator.is_main_process:
score_path = result_dir / "scores.jsonl"
save_jsonl(results, score_path)
# compute pass@k
eval_result_path = result_dir / "passk.jsonl"
# results = load_jsonl(score_path)
groups = group_results(results, args.apps)
eval_results = [score_pass_at_k(groups, k, home_path.stem) for k in range(1, 16)]
save_jsonl(eval_results, eval_result_path)
pprint.pp(eval_results)
if __name__ == "__main__":
main()
codecritic/cli/sample_apps.py deleted 100644 → 0
import argparse
from pathlib import Path
from codecritic.utils.json import save_json, save_jsonl
from codecritic.utils.vllm import vllm_chatcomplete
from codecritic.sampling.sample_apps import mk_sample_prompt
from codecritic.sampling.evaluate_code import evaluate
from codecritic.sampling.sort_split_dataset import sort_and_split_dataset
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str)
parser.add_argument("--apps", type=str)
parser.add_argument("--output_dir", type=str)
args = parser.parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(exist_ok=True)
prompts = mk_sample_prompt(args.model, args.apps)
sampling_params = dict(n=50, temperature=0.6, max_new_tokens=2048)
save_json(sampling_params, output_dir / "sampling_params.json")
codes = vllm_chatcomplete(args.model, prompts, sampling_params)
save_jsonl(codes, output_dir / "sample.jsonl")
dataset = evaluate(codes, args.apps)
save_jsonl(dataset, output_dir / "dataset.jsonl")
train, test, min_test = sort_and_split_dataset(dataset, sampling_params["n"])
save_jsonl(train, output_dir / "train.jsonl")
save_jsonl(test, output_dir / "test.jsonl")
save_jsonl(min_test, output_dir / "min_test.jsonl")
codecritic/sampling/apps_test.py deleted 100644 → 0
# copy from codeparrot/apps_metric/testing_util.py
# https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/testing_util.py
import json
import sys
import faulthandler
import platform
# used for debugging to time steps
from datetime import datetime
# to run the solution files we're using a timing based approach
import signal
import numpy as np
# for capturing the stdout
from io import StringIO
# used for testing the code that reads from input
from unittest.mock import patch, mock_open
from pyext import RuntimeModule
from enum import Enum
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")
#return
raise TimeoutException
signal.signal(signal.SIGALRM, timeout_handler)
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"
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:
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)
try:
if isinstance(inputs[0], dict):
inputs = [{int(k): v for k,v in inputs[0].items()}]
except:
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:
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:
True
if debug:
print(f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}")
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:
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}")
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]]}")
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}")
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]]}")
else:
print(f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
continue
if passed and debug:
print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}")
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:
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:
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]]}")
else:
print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
if tmp_result == True:
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 as e:
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 as e:
pass
if tmp_result == True:
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:
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
try:
tmp_result = tmp_result or (set(frozenset(round(float(t),3) for t in s) for s in output) ==\
set(frozenset(round(float(t),3) for t in s) for s in in_outs["outputs"][index]))
except Exception as e:
if debug:
print(f"Failed check6 exception = {e}")
if tmp_result == True and debug:
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]]}")
else:
print(f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}")
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 as e:
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
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 # 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
\ No newline at end of file
codecritic/sampling/evaluate_code.py deleted 100644 → 0
# copy from codeparrot/apps_metric/utils.py
# https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py
import json
import multiprocessing
import numpy as np
from datasets import load_dataset
from tqdm.contrib.concurrent import process_map
from codecritic.sampling.apps_test import run_test
from codecritic.utils.json import save_jsonl
from codecritic.data.code import extract_code
TIMEOUT = 10
def check_correctness(sample, generation, timeout, debug=False):
"""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(f"global timeout")
return result[0]
def test_generation(args, debug=False):
apps_item, code_sample = args
message = code_sample["messages"][-1]
assert message["role"] == "assistant"
code = extract_code(message["content"])
curr_res = [-2]
try:
curr_res = check_correctness(apps_item, code, timeout=TIMEOUT, debug=debug)
if debug:
print(f"\nSuccessful compilation of task {code}!")
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(curr_res)
print(f"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")
finally:
assert isinstance(curr_res, list)
problem_result = np.asarray(curr_res)
code_sample["eval_result"] = bool(np.all(problem_result > 0))
code_sample["testcase"] = curr_res
return code_sample
def get_apps_item(item, apps):
problem_id = item["problem_id"]
split, idx = problem_id.split("_")
# get corresponding samples from APPS dataset
return apps[split][int(idx)]
def evaluate_code_samples(code_samples, apps):
args = [(get_apps_item(sample, apps), sample) for sample in code_samples]
cpu_num = multiprocessing.cpu_count()
chunksize = max(len(code_samples) // (cpu_num * 5), 1)
results = process_map(
test_generation, args, max_workers=cpu_num, chunksize=chunksize
)
return results
def evaluate_incorrect_code_samples_again(results, apps, loop_num):
"""
There are some strange bugs in apps evaluation that cannot be reproduced.
The observable issue is that the same code will yield different 'eval_result' values.
Typically, the test framework may encounter an exception or decide that the code has timed out unreasonably.
This function is an ugly workaround to address this problem:
If the function returns a timeout result or raises an exception, it will be run twice to verify if the result is consistent.
The 'loop_num' parameter controls the number of times the function will be retried until the test framework obtains a consistent result.
"""
maybe_incorrect_lst, correct_lst = [], []
for item in results:
if any(x in item["testcase"] for x in (-1, -2)):
maybe_incorrect_lst.append(item)
else:
correct_lst.append(item)
for _ in range(loop_num):
if len(maybe_incorrect_lst) == 0:
break
new_results = evaluate_code_samples(maybe_incorrect_lst, apps)
print(f"maybe incorrect lst size: {len(maybe_incorrect_lst)}")
check_lst = []
for i in range(len(new_results)):
old_item, new_item = maybe_incorrect_lst[i], new_results[i]
old_eval, new_eval = old_item["eval_result"], new_item["eval_result"]
if old_eval == new_eval:
correct_lst.append(old_item)
else:
check_lst.append(new_item)
print(old_item["problem_id"], old_eval, new_item["problem_id"], new_eval)
maybe_incorrect_lst = check_lst
if len(results) != len(correct_lst):
save_jsonl(maybe_incorrect_lst, "debug.jsonl")
# raise ValueError("cannot correctly evaluate codes")
print("cannot correctly evalute code. see debug.jsonl")
if len(maybe_incorrect_lst) < 5:
correct_lst.extend(maybe_incorrect_lst)
return correct_lst
def evaluate(code_samples, dataset_path):
apps = load_dataset(dataset_path)
results = evaluate_code_samples(code_samples, apps)
for item in results:
item["testcase"] = item["eval_result"]
item["eval_result"] = bool(np.all(np.asarray(item["testcase"]) > 0))
results = evaluate_incorrect_code_samples_again(results, apps, 10)
return results
codecritic/sampling/sample_apps.py deleted 100644 → 0
from datasets import load_dataset
import json
from transformers import AutoTokenizer
def mk_prompt(doc) -> str:
prompt = "Write Python code to solve competitive programming problems in a markdown code block."
starter_code = None if len(doc["starter_code"]) == 0 else doc["starter_code"]
try:
input_outpout = json.loads(doc["input_output"])
fn_name = None if not input_outpout.get("fn_name") else input_outpout["fn_name"]
except ValueError:
fn_name = None
prompt += "\nQUESTION:\n"
prompt += doc["question"]
if starter_code:
prompt += starter_code
if not fn_name:
prompt += "\nUse Standard Input format"
else:
prompt += "\nUse Call-Based format"
prompt += "\nPlease generate the code in a ```python markdown block, ensuring to include the closing ``` at the end."
conversation = [{"role": "user", "content": prompt}]
return conversation
def mk_sample_prompt(model_path, apps_path):
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token
prompts = []
for split in ["train", "test"]:
ds = load_dataset(apps_path, split=split)
for sample in ds:
problem_id = split + "_" + str(sample["problem_id"])
# Filter problems without input_output
try:
json.loads(sample["input_output"])
except ValueError:
print(f"Skipping {problem_id}: Invalid JSON in input_output")
continue
prompt = mk_prompt(sample)
# Filter long prompts
chat_text = tokenizer.apply_chat_template(prompt, tokenize=False)
tokenized_prompt = tokenizer.encode(chat_text)
if len(tokenized_prompt) > (4096 - 512):
print(
f"Skipping {problem_id}: Token length {len(tokenized_prompt)} exceeds limit"
)
continue
prompts.append(dict(problem_id=problem_id, messages=prompt))
return prompts
codecritic/sampling/sort_split_dataset.py deleted 100644 → 0
def mk_key_for_sort(item):
problem_id = item['problem_id']
prefix, idx = problem_id.split('_')
prefix_weight = 0 if prefix == 'train' else 1
return (prefix_weight, int(idx))
def sort_and_drop_key(dataset, key):
dataset = sorted(dataset, key=lambda x: x[key])
for item in dataset:
item.pop(key)
return dataset
TEST_RANGES = [(0, 300), (3000, 3100), (4000, 4100)]
def is_in_test_range(prefix_weight, idx):
if prefix_weight == 1:
for start, end in TEST_RANGES:
if start <= idx < end:
return True
return False
def sort_and_split_dataset(dataset, n):
"""
The dataset will be divided into two parts: Train and Test.
From the Test set, 10% of items across varying difficulties will be selected.
Among these, only those items for which the LLM can generate correct solutions will be included in the minimal testset.
This approach reduces the test time by approximately 1/5.
"""
# add `key_for_sort`
new_train, new_test = [], []
for item in dataset:
item["key_for_sort"] = mk_key_for_sort(item)
if is_in_test_range(*item["key_for_sort"]):
new_test.append(item)
else:
new_train.append(item)
new_train = sort_and_drop_key(new_train, "key_for_sort")
new_test = sort_and_drop_key(new_test, "key_for_sort")
minimal_test = []
assert len(new_test) % n == 0
for i in range(len(new_test) // n):
problem = new_test[i * n : (i + 1) * n]
has_correct_solution = any(d["eval_result"] for d in problem)
if has_correct_solution:
minimal_test.extend(problem)
return new_train, new_test, minimal_test
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment