add early termination

code_eval.py

@@ -19,7 +19,7 @@ described in the paper "Evaluating Large Language Models Trained on Code"
 import itertools
 import os
 from collections import Counter, defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
 
 import datasets
 import evaluate
@@ -171,6 +171,7 @@ class CodeEval(evaluate.Metric):
 
         with ThreadPoolExecutor(max_workers=num_workers) as executor:
             futures = []
+            future_dict = defaultdict(lambda: defaultdict(list))
             completion_id = Counter()
             results = defaultdict(list)
 
@@ -189,31 +190,47 @@ class CodeEval(evaluate.Metric):
                         )
                         future = executor.submit(check_correctness, *args)
                         futures.append(future)
+                        future_dict[task_id][completion_id[task_id]].append(future)
                     completion_id[task_id] += 1
 
             pbar = tqdm(total=len(futures))
             for future in as_completed(futures):
-                result = future.result()
+                try:
+                    result = future.result()
+                except CancelledError:
+                    pbar.update(1)
+                    continue
+
                 results[result["task_id"]].append((result["completion_id"], result))
                 pbar.update(1)
 
-        for result in results.values():
+                if not result["passed"]:
+                    future_list = future_dict[result["task_id"]][result["completion_id"]]
+                    for future in future_list:
+                        future.cancel()
+
+        new_results = {}
+        for key, result in results.items():
             new_result = []
             for completion_id, group in itertools.groupby(result, key=lambda x: x[0]):
                 group = list(group)
                 new_result.append(
-                    dict(
-                        task_id=group[0][0],
-                        passed=all(r[1]["passed"] for r in group),
-                        result=[r[1]["result"] for r in group],
-                        completion_id=completion_id,
+                    (
+                        group[0][0],
+                        dict(
+                            task_id=group[0][0],
+                            passed=all(r[1]["passed"] for r in group),
+                            result=[r[1]["result"] for r in group],
+                            completion_id=completion_id,
+                        ),
                     )
                 )
-            result = new_result
+            new_results[key] = new_result
+        results = new_results
 
         total, correct = [], []
         for result in results.values():
-            result.sort()
+            result.sort(key=lambda x: x[0])
             passed = [r[1]["passed"] for r in result]
             total.append(len(passed))
             correct.append(sum(passed))