separate timelimits

code_eval.py

@@ -22,13 +22,12 @@ from collections import Counter, defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import datasets
-import numpy as np
-
 import evaluate
+import numpy as np
+from tqdm import tqdm
 
 from .execute import check_correctness
 
-
 _CITATION = """\
 @misc{chen2021evaluating,
       title={Evaluating Large Language Models Trained on Code},
@@ -143,7 +142,7 @@ class CodeEval(evaluate.Metric):
             features=datasets.Features(
                 {
                     "predictions": datasets.Sequence(datasets.Value("string")),
-                    "references": datasets.Value("string"),
+                    "references": datasets.Sequence(datasets.Value("string")),
                 }
             ),
             homepage="https://github.com/openai/human-eval",
@@ -152,33 +151,65 @@ class CodeEval(evaluate.Metric):
             license=_LICENSE,
         )
 
-    def _compute(self, predictions, references, k=[1, 10, 100], num_workers=4, timeout=3.0):
+    def _compute(
+        self,
+        predictions,
+        references,
+        k=[1, 10, 100],
+        num_workers=4,
+        timeout=3.0,
+    ):
         """Returns the scores"""
 
         if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1":
             raise ValueError(_WARNING)
 
         if os.name == "nt":
-            raise NotImplementedError("This metric is currently not supported on Windows.")
+            raise NotImplementedError(
+                "This metric is currently not supported on Windows."
+            )
 
         with ThreadPoolExecutor(max_workers=num_workers) as executor:
             futures = []
             completion_id = Counter()
-            n_samples = 0
             results = defaultdict(list)
 
-            for task_id, (candidates, test_case) in enumerate(zip(predictions, references)):
+            for task_id, (candidates, test_case) in enumerate(
+                zip(predictions, references)
+            ):
                 for candidate in candidates:
-                    test_program = candidate + "\n" + test_case
-                    args = (test_program, timeout, task_id, completion_id[task_id])
-                    future = executor.submit(check_correctness, *args)
-                    futures.append(future)
+                    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],
+                        )
+                        future = executor.submit(check_correctness, *args)
+                        futures.append(future)
                     completion_id[task_id] += 1
-                    n_samples += 1
 
+            pbar = tqdm(total=len(futures))
             for future in as_completed(futures):
                 result = future.result()
                 results[result["task_id"]].append((result["completion_id"], result))
+                pbar.update(1)
+
+        for result in results.values():
+            new_result = []
+            for completion_id, group in itertools.groupby(results, key=lambda x: x[0]):
+                group = list(group)
+                new_result.append(
+                    dict(
+                        task_id=group[0]["task_id"],
+                        passed=all(r["passed"] for r in group),
+                        result=[r["result"] for r in group],
+                        completion_id=completion_id,
+                    )
+                )
+            result = new_result
 
         total, correct = [], []
         for result in results.values():
@@ -190,7 +221,11 @@ class CodeEval(evaluate.Metric):
         correct = np.array(correct)
 
         ks = k
-        pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() for k in ks if (total >= k).all()}
+        pass_at_k = {
+            f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
+            for k in ks
+            if (total >= k).all()
+        }
 
         return pass_at_k, results
 
@@ -210,4 +245,6 @@ def estimate_pass_at_k(num_samples, num_correct, k):
         assert len(num_samples) == len(num_correct)
         num_samples_it = iter(num_samples)
 
-    return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
+    return np.array(
+        [estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]
+    )