wehub-resource-sync
155 lines
4.4 KiB
Python
155 lines
4.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.
|
|
|
|
This script is incorrect for pseudo test cases since we cannot ensure the correctness of each test case, but it is appropriate for ground-truth test cases,
|
|
serving as hard limit, i.e., if there is one completion for some prefix has passed all test cases, then it is a gold prefix.
|
|
|
|
"""
|
|
|
|
|
|
def counting_partial_response_value(res):
|
|
return sum([1 if x else 0 for x in res])
|
|
|
|
|
|
def parse_value(v, binary: bool):
|
|
if binary:
|
|
return 1 if v > 0 else 0
|
|
return v
|
|
|
|
|
|
def _process_trajectories_worker(item, top_k: int, binary: bool):
|
|
item_id, trajectories = item
|
|
outputs = {
|
|
"idx": item_id
|
|
}
|
|
for i in range(top_k):
|
|
level_trajectories = [(traj["vs"][i], traj["prefix"]) for traj in trajectories if len(traj["vs"]) > i]
|
|
if len(level_trajectories) == 0:
|
|
continue
|
|
prefix_vis = set()
|
|
level_values = []
|
|
level_prefixes = []
|
|
for v, prefix in level_trajectories:
|
|
if prefix in prefix_vis:
|
|
continue
|
|
level_values.append(parse_value(v, binary))
|
|
level_prefixes.append(prefix)
|
|
prefix_vis.add(prefix)
|
|
outputs[f"traj_level_{i}_values"] = level_values
|
|
outputs[f"traj_level_{i}_prefixes"] = level_prefixes
|
|
return outputs
|
|
|
|
|
|
def _annotate(file):
|
|
return json.load(open(file, encoding="utf-8"))
|
|
|
|
|
|
def multiprocessing_loading(files, num_workers: int = 8):
|
|
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 += json.load(open(file, encoding="utf-8"))
|
|
# return data
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
parser.add_argument("--binary", default=False, action="store_true")
|
|
parser.add_argument("--num_workers", type=int, default=8)
|
|
args = parser.parse_args()
|
|
|
|
print("Collecting data...")
|
|
if os.path.exists(args.input_file):
|
|
data = json.load(open(args.input_file))
|
|
else:
|
|
# data = []
|
|
# for file in glob(args.input_file):
|
|
# print(file)
|
|
# data += json.load(open(file))
|
|
files = glob(args.input_file)
|
|
files = sorted(files)
|
|
print(len(files))
|
|
print(files)
|
|
data = multiprocessing_loading(files)
|
|
|
|
print(len(data))
|
|
|
|
num_prefixes = 0
|
|
val_cnt = collections.Counter()
|
|
outputs = []
|
|
preference_pairs = dict()
|
|
missing = 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
|
|
v = counting_partial_response_value(item["res"])
|
|
v = parse_value(v, args.binary)
|
|
|
|
outputs.append({
|
|
"problem_id": problem_id,
|
|
"prefix": prefix,
|
|
"value": v,
|
|
})
|
|
num_prefixes += 1
|
|
val_cnt[v] += 1
|
|
|
|
if problem_id not in preference_pairs:
|
|
preference_pairs[problem_id] = {
|
|
"pos": [],
|
|
"neg": [],
|
|
}
|
|
|
|
if v > 0:
|
|
preference_pairs[problem_id]["pos"].append(prefix)
|
|
else:
|
|
preference_pairs[problem_id]["neg"].append(prefix)
|
|
|
|
preference_pairs = [
|
|
{
|
|
"problem_id": problem_id,
|
|
"pos": pair["pos"],
|
|
"neg": pair["neg"],
|
|
}
|
|
for problem_id, pair in preference_pairs.items()
|
|
]
|
|
|
|
print(f"Missing: {missing}")
|
|
print(val_cnt)
|
|
print(f"Processed {num_prefixes} prefixes.")
|
|
print(f"Averaged {num_prefixes / len(data)} prefixes per problem.")
|
|
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", "_pairs.json"), "w", encoding="utf-8"), indent=2, ensure_ascii=False)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
|
"""
|
|
>>>
|
|
"""
|