refactor & fix order bug & add early stop option

code_eval.py

@@ -20,11 +20,14 @@ import itertools
 import os
 from collections import Counter, defaultdict
 from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
+from typing import List, Optional
+import time
 
 import datasets
 import evaluate
 import numpy as np
 from tqdm import tqdm
+from pydantic import BaseModel
 
 from .execute import check_correctness
 
@@ -155,9 +158,11 @@ class CodeEval(evaluate.Metric):
         self,
         predictions,
         references,
+        task_ids=None,
         k=[1, 10, 100],
         num_workers=4,
         timeout=3.0,
+        early_stop=False,
     ):
         """Returns the scores"""
 
@@ -169,69 +174,43 @@ class CodeEval(evaluate.Metric):
                 "This metric is currently not supported on Windows."
             )
 
+        task_ids = task_ids or list(range(len(predictions)))
+
         with ThreadPoolExecutor(max_workers=num_workers) as executor:
-            futures = []
-            future_dict = defaultdict(lambda: defaultdict(list))
-            completion_id = Counter()
-            results = defaultdict(list)
-
-            for task_id, (candidates, test_case) in enumerate(
-                zip(predictions, references)
-            ):
-                for candidate in candidates:
-                    for _test_case in test_case:
-                        assert isinstance(_test_case, str)
-                        test_program = candidate + "\n" + _test_case
-                        args = (
-                            test_program,
-                            timeout,
-                            task_id,
-                            completion_id[task_id],
-                        )
+            results = {}
+            for tid, pred, ref in zip(task_ids, predictions, references):
+                results[tid] = []
+                for candidate in pred:
+                    result = Result(task_id=tid, completion_id=len(results))
+                    for test_case in ref:
+                        assert isinstance(test_case, str)
+                        test_program = candidate + "\n" + test_case
+                        args = (test_program, timeout, tid)
                         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):
-                try:
-                    result = future.result()
-                except CancelledError:
-                    pbar.update(1)
-                    continue
-
-                results[result["task_id"]].append((result["completion_id"], result))
-                pbar.update(1)
-
-                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 = []
-            result.sort(key=lambda x: x[0])
-            for completion_id, group in itertools.groupby(result, key=lambda x: x[0]):
-                group = list(group)
-                new_result.append(
-                    (
-                        completion_id,
-                        dict(
-                            task_id=key,
-                            passed=all(r[1]["passed"] for r in group),
-                            result=[r[1]["result"] for r in group],
-                            completion_id=completion_id,
-                        ),
-                    )
-                )
-            new_results[key] = new_result
-        results = new_results
+                        result.add(future)
+                    results[tid].append(result)
+
+            pbar = tqdm(total=len(results))
+            prev_done_count = 0
+            while not all(r.done() for r in results.values()):
+                cur_done_count = 0
+                for result in results.values():
+                    for r in result:
+                        if not r.done():
+                            r.refresh(early_stop)
+                        else:
+                            cur_done_count += 1
+                pbar.update(cur_done_count - prev_done_count)
+                prev_done_count = cur_done_count
+                time.sleep(1)
+
+        results = {
+            task_id: [(r.completion_id, r.dict(exclude={"futures"})) for r in result]
+            for task_id, result in results.items()
+        }
 
         total, correct = [], []
         for result in results.values():
-            result.sort(key=lambda x: x[0])
             passed = [r[1]["passed"] for r in result]
             total.append(len(passed))
             correct.append(sum(passed))
@@ -266,3 +245,37 @@ def estimate_pass_at_k(num_samples, num_correct, k):
     return np.array(
         [estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]
     )
+
+
+class Result(BaseModel):
+    task_id: int
+    completion_id: int
+
+    passed: Optional[bool] = None
+    result: List[str] = []
+    futures: List[object] = []
+
+    def add(self, future):
+        self.futures.append(future)
+        self.result.append(None)
+
+    def refresh(self, early_stop=False):
+        for i, future in enumerate(self.futures):
+            if self.result[i] is None and future.done():
+                try:
+                    self.result[i] = future.result()
+                except CancelledError:
+                    self.result[i] = "Early Stopped"
+                except Exception as e:
+                    self.result[i] = str(e)
+
+                    if early_stop:
+                        # cancel all other futures
+                        for future in self.futures[i + 1 :]:
+                            future.cancel()
+
+        if all(r is not None for r in self.result):
+            self.passed = all(r["passed"] for r in self.result)
+
+    def done(self):
+        return self.passed is not None

requirements.txt

@@ -1 +1,2 @@
-git+https://github.com/huggingface/evaluate@af3c30561d840b83e54fc5f7150ea58046d6af69
+pydantic
+numpy