| import re |
| import json |
| from pathlib import Path |
| from collections import defaultdict |
| from human_eval.evaluation import evaluate_functional_correctness |
|
|
| version = "20240121-Jul" |
|
|
| DATA_DIR = Path(__file__).parent / "data" |
|
|
| def extract_python_code(generation: str): |
| generation = generation.replace("[PYTHON]", '```python').replace("[/PYTHON]", '```') |
| if '```python' in generation: |
| p_code = re.compile(r'```python\n(.*?)\n```', flags=re.DOTALL) |
| code_block = p_code.findall(generation)[0] |
| return code_block |
| else: |
| codelist = re.split("\ndef|\nclass|\nif|\n#|\nprint", generation) |
| return codelist[0] |
| |
| def evaluate_main(generation_path: str, result_path: str, temp_dir: str): |
| problem_path = (DATA_DIR / f"{version}.jsonl").as_posix() |
|
|
| print(problem_path) |
| problems = [json.loads(line) for line in open(problem_path, 'r')] |
|
|
| id2problems = { x['task_id']: x for x in problems } |
|
|
| results = [json.loads(line) for line in open(generation_path, 'r')] |
| for result in results: |
| if 'task_id' not in result: |
| result['task_id'] = problems[result['index']]['task_id'] |
|
|
| if 'generation' not in result: |
| try: |
| if 'output' not in result: |
| result['output'] = result['response'] |
| if result['output'].startswith("\n "): |
| func_code = extract_python_code(result['prompt_sft']).strip() |
| result['generation'] = func_code + '\n' + result['output'] |
| else: |
| result['generation'] = extract_python_code(result['output']) |
| except: |
| result['generation'] = result['output'] |
| |
| with open(result_path, 'w') as fr: |
| for result in results: |
| fr.write(json.dumps(result) + "\n") |
|
|
| score = evaluate_functional_correctness( |
| input_file=result_path, |
| tmp_dir=temp_dir, |
| problem_file=problem_path, |
| result_path=result_path |
| ) |
|
|
| hardness_results = defaultdict(int) |
| for result in [json.loads(line) for line in open(result_path, 'r')]: |
| problem = id2problems[result['task_id']] |
|
|
| hardness = problem['meta']['difficulty'] |
| hardness_results[hardness] += 1 |
| hardness_results[hardness + "_correct"] += result['passed'] |
|
|
| print("="*100) |
| print("Evaluate {} over.".format(generation_path)) |
| print("Pass@1: {:.3f}".format(score["pass@1"])) |
| for key in ["Easy", "Medium", "Hard"]: |
| if key.endswith("_correct"): |
| continue |
| acc = hardness_results[key+"_correct"] / hardness_results[key] |
| print("{}: {:.3f}({}/{})".format(key, acc, hardness_results[key+"_correct"], hardness_results[key])) |
|
|
| if __name__ == '__main__': |
| import argparse |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--generation_path", type=str, required=True) |
| parser.add_argument("--result_path", type=str) |
| parser.add_argument("--temp_dir", type=str, default="output/temp") |
| args = parser.parse_args() |
|
|
| if args.result_path is None: |
| args.result_path = args.generation_path.replace(".jsonl", "_result.jsonl") |
| |
| evaluate_main(args.generation_path, args.result_path, temp_dir=args.temp_dir) |
| pass |
|
|
