wehub-resource-sync
84 lines
3.5 KiB
Python
84 lines
3.5 KiB
Python
import argparse
|
|
import json
|
|
import os.path
|
|
import sys
|
|
from glob import glob
|
|
from collections import Counter
|
|
from datasets import load_dataset
|
|
|
|
sys.set_int_max_str_digits(0)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
parser.add_argument("--split", type=str, default="test")
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.exists(f'apps_difficulty_{args.split}.json'):
|
|
_dataset = load_dataset("codeparrot/apps", split=args.split).to_list()
|
|
problem_id2difficulty = {item["problem_id"]: item["difficulty"] for item in _dataset}
|
|
all_difficulties = Counter(problem_id2difficulty.values())
|
|
json.dump(problem_id2difficulty, open(f'apps_difficulty_{args.split}.json', "w"), ensure_ascii=False)
|
|
json.dump(all_difficulties, open(f'apps_difficulty_{args.split}_all.json', "w"), ensure_ascii=False)
|
|
else:
|
|
problem_id2difficulty = json.load(open(f'apps_difficulty_{args.split}.json'))
|
|
all_difficulties = json.load(open(f'apps_difficulty_{args.split}_all.json'))
|
|
problem_id2difficulty = {int(k): v for k, v in problem_id2difficulty.items()}
|
|
|
|
if os.path.exists(args.input_file):
|
|
if args.input_file.endswith(".json"):
|
|
data = json.load(open(args.input_file))
|
|
else:
|
|
data = [json.loads(line) for line in open(args.input_file).readlines()]
|
|
else:
|
|
data = []
|
|
for file in glob(args.input_file):
|
|
if file.endswith(".json"):
|
|
# data += json.load(open(file))
|
|
f = open(file, "r")
|
|
data += json.load(f)
|
|
else:
|
|
# data += [json.loads(line) for line in open(file).readlines()]
|
|
lines = list(open(file).readlines())
|
|
for line in lines:
|
|
data.append(json.loads(line))
|
|
|
|
success = 0
|
|
success_at_k = 0
|
|
successes_at_difficulty = {difficulty: 0 for difficulty in all_difficulties}
|
|
successes_at_k_at_difficulty = {difficulty: 0 for difficulty in all_difficulties}
|
|
for item in data:
|
|
p_id = item["id"]
|
|
if isinstance(item["res"], bool):
|
|
if item["res"]:
|
|
success += 1
|
|
success_at_k += 1
|
|
successes_at_difficulty[problem_id2difficulty[p_id]] += 1
|
|
successes_at_k_at_difficulty[problem_id2difficulty[p_id]] += 1
|
|
else:
|
|
if len(item["res"]):
|
|
if any(item["res"]):
|
|
success_at_k += 1
|
|
successes_at_k_at_difficulty[problem_id2difficulty[p_id]] += 1
|
|
if item["res"][0]:
|
|
success += 1
|
|
successes_at_difficulty[problem_id2difficulty[p_id]] += 1
|
|
|
|
metrics = {"acc": success / len(data), "pass@k": success_at_k / len(data), "correct": success, "total": len(data)}
|
|
for difficulty in all_difficulties:
|
|
metrics[f"acc_{difficulty}"] = successes_at_difficulty[difficulty] / all_difficulties[difficulty]
|
|
metrics[f"pass@k_{difficulty}"] = successes_at_k_at_difficulty[difficulty] / all_difficulties[difficulty]
|
|
metrics[f"correct_{difficulty}"] = successes_at_difficulty[difficulty]
|
|
metrics[f"total_{difficulty}"] = all_difficulties[difficulty]
|
|
print(json.dumps(metrics, indent=2))
|
|
|
|
if args.output_file:
|
|
json.dump(data, open(args.output_file, "w"), ensure_ascii=False)
|
|
json.dump(metrics, open(args.output_file.replace(".json", ".metrics"), "w"), ensure_ascii=False)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|