46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import json
|
|
import argparse
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str, required=True)
|
|
parser.add_argument("--output_file", type=str, default=None)
|
|
args = parser.parse_args()
|
|
|
|
data = json.load(open(args.input_file))
|
|
|
|
success = 0
|
|
success_at_k = 0
|
|
for item in data:
|
|
if isinstance(item["res"], bool):
|
|
assert isinstance(item["response"], str)
|
|
if item["pred"] is None:
|
|
item["res"] = False
|
|
|
|
if item["res"]:
|
|
success += 1
|
|
success_at_k += 1
|
|
else:
|
|
if len(item["res"]):
|
|
assert isinstance(item["response"], list)
|
|
for i in range(len(item["response"])):
|
|
if item["pred"][i] is None:
|
|
item["res"][i] = False
|
|
item["full_res"][i] = [[False] * len(item["test_cases"]["inputs"]) if item["test_cases"] else 1]
|
|
|
|
if any(item["res"]):
|
|
success_at_k += 1
|
|
if item["res"][0]:
|
|
success += 1
|
|
|
|
metrics = {"acc": success / len(data), "pass@k": success_at_k / len(data), "correct": success, "total": len(data)}
|
|
print(json.dumps(metrics, indent=2))
|
|
if args.output_file is not None:
|
|
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()
|