Evaluates the functional correctness of generated samples, and writes results to f"{sample_file}_results.jsonl.gz"
(
sample_file: str,
k: List[int] = [1, 10, 100],
n_workers: int = 4,
timeout: float = 3.0,
problem_file: str = HUMAN_EVAL,
)
| 37 | |
| 38 | |
| 39 | def evaluate_functional_correctness( |
| 40 | sample_file: str, |
| 41 | k: List[int] = [1, 10, 100], |
| 42 | n_workers: int = 4, |
| 43 | timeout: float = 3.0, |
| 44 | problem_file: str = HUMAN_EVAL, |
| 45 | ): |
| 46 | """ |
| 47 | Evaluates the functional correctness of generated samples, and writes |
| 48 | results to f"{sample_file}_results.jsonl.gz" |
| 49 | """ |
| 50 | |
| 51 | problems = read_problems(problem_file) |
| 52 | |
| 53 | # Check the generated samples against test suites. |
| 54 | with ThreadPoolExecutor(max_workers=n_workers) as executor: |
| 55 | |
| 56 | futures = [] |
| 57 | completion_id = Counter() |
| 58 | n_samples = 0 |
| 59 | results = defaultdict(list) |
| 60 | |
| 61 | print("Reading samples...") |
| 62 | for sample in tqdm.tqdm(stream_jsonl(sample_file)): |
| 63 | task_id = sample["task_id"] |
| 64 | completion = sample["completion"] |
| 65 | args = (problems[task_id], completion, timeout, completion_id[task_id]) |
| 66 | future = executor.submit(check_correctness, *args) |
| 67 | futures.append(future) |
| 68 | completion_id[task_id] += 1 |
| 69 | n_samples += 1 |
| 70 | |
| 71 | assert len(completion_id) == len(problems), "Some problems are not attempted." |
| 72 | |
| 73 | print("Running test suites...") |
| 74 | for future in tqdm.tqdm(as_completed(futures), total=len(futures)): |
| 75 | result = future.result() |
| 76 | results[result["task_id"]].append((result["completion_id"], result)) |
| 77 | |
| 78 | # Calculate pass@k. |
| 79 | total, correct = [], [] |
| 80 | for result in results.values(): |
| 81 | result.sort() |
| 82 | passed = [r[1]["passed"] for r in result] |
| 83 | total.append(len(passed)) |
| 84 | correct.append(sum(passed)) |
| 85 | total = np.array(total) |
| 86 | correct = np.array(correct) |
| 87 | |
| 88 | ks = k |
| 89 | pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() |
| 90 | for k in ks if (total >= k).all()} |
| 91 | |
| 92 | # Finally, save the results in one file: |
| 93 | def combine_results(): |
| 94 | for sample in stream_jsonl(sample_file): |
| 95 | task_id = sample["task_id"] |
| 96 | result = results[task_id].pop(0) |
no test coverage detected