190 lines
7.4 KiB
Python
190 lines
7.4 KiB
Python
import json
|
|
import argparse
|
|
from glob import glob
|
|
import os
|
|
import sys
|
|
import collections
|
|
from multiprocessing import Pool
|
|
from tqdm import tqdm
|
|
from functools import partial
|
|
|
|
sys.set_int_max_str_digits(0)
|
|
|
|
"""
|
|
For self-consistency based input-output pairs, first copy the pseudo test cases into the prefix data, run `prefix_fail_extract_pseudo_label.py`,
|
|
and the run this script.
|
|
|
|
|
|
Note:
|
|
1. Sometimes the pseudo test cases can include some extremely large cases, leading to out-of-memory error.
|
|
"""
|
|
|
|
|
|
def counting_partial_response_value(full_res):
|
|
pass_num = []
|
|
for pred_res in full_res:
|
|
if len(pred_res) == 0:
|
|
pass_num.append(0)
|
|
# elif pred_res[0] == -1: # FIXME: This line is commented since 2024/09/25. Seems no difference.
|
|
# pass_num.append(0)
|
|
else:
|
|
pass_num.append(sum([1 for x in pred_res if x is True]))
|
|
|
|
return pass_num
|
|
|
|
|
|
def annotate(file, exclude: str = ""):
|
|
exclude = exclude.split(",")
|
|
if any([e and e in file for e in exclude]):
|
|
print(f"Excluding {file}")
|
|
return []
|
|
return json.load(open(file, encoding="utf-8"))
|
|
|
|
|
|
def multiprocessing_loading(files, exclude: str = "", num_workers: int = 8):
|
|
_annotate = partial(annotate, exclude=exclude)
|
|
|
|
# with Pool(num_workers) as p:
|
|
# data = list(tqdm(p.imap(_annotate, files), total=len(files)))
|
|
# all_data = []
|
|
# for d in data:
|
|
# all_data.extend(d)
|
|
# return all_data
|
|
data = []
|
|
for file in tqdm(files):
|
|
data += _annotate(file)
|
|
return data
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
parser.add_argument("--pass_case_margin", type=float, default=1)
|
|
parser.add_argument("--pass_case_lower_bound", type=float, default=0.5)
|
|
parser.add_argument("--test_case_field", type=str, default="pseudo_input_output")
|
|
parser.add_argument("--num_workers", type=int, default=8)
|
|
parser.add_argument("--reduction", type=str, default="max")
|
|
parser.add_argument("--exclude", type=str, default="")
|
|
args = parser.parse_args()
|
|
|
|
print("Collecting data...")
|
|
if os.path.exists(args.input_file):
|
|
data = json.load(open(args.input_file))
|
|
else:
|
|
files = glob(args.input_file)
|
|
files = sorted(files)
|
|
print(len(files))
|
|
print(files)
|
|
data = multiprocessing_loading(files, args.exclude)
|
|
|
|
print(len(data))
|
|
|
|
num_prefixes = 0
|
|
val_cnt = collections.Counter()
|
|
outputs = []
|
|
p_id2prefixes = collections.defaultdict(list)
|
|
preference_pairs = []
|
|
missing = 0
|
|
missing_test_cases = 0
|
|
for item in tqdm(data):
|
|
problem_id, resp_id, prefix_id = item["prefix_id"].split("_")
|
|
prefix = item["prefix"]
|
|
# problem_id = int(problem_id)
|
|
|
|
if "res" not in item:
|
|
missing += 1
|
|
continue
|
|
|
|
if args.test_case_field not in item or (not isinstance(item[args.test_case_field], dict)) or "inputs" not in item[args.test_case_field]:
|
|
missing_test_cases += 1
|
|
continue
|
|
|
|
test_case_num = len(item[args.test_case_field]["inputs"])
|
|
if test_case_num == 0:
|
|
missing_test_cases += 1
|
|
continue
|
|
|
|
pass_num = counting_partial_response_value(item["full_res"])
|
|
if args.reduction == "max":
|
|
max_pass_num = max(pass_num)
|
|
elif args.reduction == "avg":
|
|
max_pass_num = sum(pass_num) / len(pass_num)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
outputs.append({
|
|
"problem_id": problem_id,
|
|
"prefix": prefix,
|
|
"pass_num": pass_num,
|
|
"max_pass_num": max_pass_num,
|
|
"test_case_num": test_case_num,
|
|
})
|
|
num_prefixes += 1
|
|
val_cnt.update(pass_num)
|
|
|
|
p_id2prefixes[problem_id].append(outputs[-1])
|
|
|
|
for problem_id, all_prefixes in tqdm(p_id2prefixes.items()):
|
|
max_pass_num2prefixes = collections.defaultdict(list)
|
|
for prefix in all_prefixes:
|
|
max_pass_num2prefixes[prefix["max_pass_num"]].append(prefix)
|
|
max_pass_num2prefixes = sorted(max_pass_num2prefixes.items(), key=lambda x: x[0])
|
|
|
|
pos_prefixes = []
|
|
neg_prefixes = []
|
|
for p in all_prefixes:
|
|
pass_ratio = p["max_pass_num"] / p["test_case_num"]
|
|
if pass_ratio < args.pass_case_lower_bound:
|
|
continue
|
|
neg_upper_pass_num = p["max_pass_num"] - args.pass_case_margin
|
|
|
|
target_neg = []
|
|
for pass_num, prefixes in max_pass_num2prefixes:
|
|
if pass_num < neg_upper_pass_num:
|
|
target_neg.extend([x["prefix"] for x in prefixes])
|
|
else:
|
|
break
|
|
|
|
pos_prefixes.append(p["prefix"])
|
|
neg_prefixes.append(target_neg)
|
|
|
|
preference_pairs.append({
|
|
"problem_id": problem_id,
|
|
"pos": pos_prefixes,
|
|
"neg": neg_prefixes
|
|
})
|
|
|
|
print(f"Missing: {missing}")
|
|
print(f"Missing test cases: {missing_test_cases}")
|
|
print(val_cnt)
|
|
print(f"Processed {num_prefixes} prefixes.")
|
|
print(f"Averaged {num_prefixes / len(data)} prefixes per problem.")
|
|
print(f"Processed {len(preference_pairs)} problems.")
|
|
json.dump(outputs, open(args.output_file, "w", encoding="utf-8"), indent=2, ensure_ascii=False)
|
|
json.dump(preference_pairs, open(args.output_file.replace(".json", f"_low{args.pass_case_lower_bound}_m{args.pass_case_margin}_{args.reduction}.json"),
|
|
"w", encoding="utf-8"), indent=2, ensure_ascii=False)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
|
"""
|
|
>>> python scripts/apps/prm/construct_process_rm_sample_fix.py \
|
|
--input_file "/mnt/fangkai_blob/reward_modeling//experiments/deepseek-coder-v1.5-ins.7b.apps.r2c.gpt4o.distil.V100.w8.v3.1.dp4.tp4.s42/apps/checkpoint-200/train.tem1.0.n10.prefix.upper0.8.r0.3.completion.tem1.0.n5.v2.0.[0-9]*-of-256.pseudo_test_case.exec.json" \
|
|
--output_file /mnt/fangkai_blob/reward_modeling//experiments/deepseek-coder-v1.5-ins.7b.apps.r2c.gpt4o.distil.V100.w8.v3.1.dp4.tp4.s42/apps/checkpoint-200/train.tem1.0.n10.prefix.upper0.8.r0.3.completion.tem1.0.n5.v2.0.pseudo_test_case.prefix_pass_num.fix.json \
|
|
--pass_case_lower_bound 0.8 --pass_case_margin 4 --test_case_field pseudo_input_output
|
|
|
|
|
|
>>> python scripts/apps/prm/construct_process_rm_sample_fix.py \
|
|
--input_file "${OUTPUT_PREFIX_PATH}/experiments/deepseek-coder-v1.5-ins.7b.apps.r2c.sft_ps_test_case.process-dpo.V100.tp8dp16.v4.9.s42/oss-instruct-apps-train/checkpoint-700/train.tem1.0.n10.prefix.upper0.8.r0.3.sample20_per.completion.tem1.0.n3.pseudo_input_output.exec.*-of-256.json" \
|
|
--output_file ${OUTPUT_PREFIX_PATH}/experiments/deepseek-coder-v1.5-ins.7b.apps.r2c.sft_ps_test_case.process-dpo.V100.tp8dp16.v4.9.s42/oss-instruct-apps-train/checkpoint-700/train.tem1.0.n10.prefix.upper0.8.r0.3.sample20_per.completion.tem1.0.n3.pseudo_input_output.prefix_pass_num.fix.json \
|
|
--pass_case_lower_bound 0.5 --pass_case_margin 4 --test_case_field pseudo_input_output --reduction avg --test_case_field input_output --exclude "204-of-256,225-of-256"
|
|
|
|
Missing: 0
|
|
Missing test cases: 0
|
|
Counter({10: 448824, 0: 153053, 9: 25420, 1: 23469, 8: 16718, 5: 15115, 2: 14803, 3: 12821, 4: 12253, 7: 12053, 6: 11947, 11: 724, 20: 112, 16: 3}) # TODO: This should be a problem. Why there are some problems have more than 10 test cases?
|
|
Processed 249105 prefixes.
|
|
Averaged 1.0 prefixes per problem.
|
|
"""
|