Add Bayes@N task variants for HumanEval and MBPP

bigcode_eval/evaluator.py

@@ -42,6 +42,20 @@ def __init__(self, accelerator, model, tokenizer, args):
 
     def generate_text(self, task_name, intermediate_generations=None):
         task = tasks.get_task(task_name, self.args)
+        if getattr(task, "bayes_at_n", False) and intermediate_generations is not None:
+            raise ValueError(
+                "Bayes@N does not support intermediate checkpoints; "
+                "use a complete saved-generation file."
+            )
+        if (
+            getattr(task, "bayes_at_n", False)
+            and self.args.load_generations_path
+            and self.args.limit_start != 0
+        ):
+            raise ValueError(
+                "Bayes@N does not support limit_start with loaded generations; "
+                "the existing loader does not apply that offset."
+            )
         dataset = task.get_dataset()
         # if args.limit is None, use all samples
         # if args.limit is used, make sure args.limit_start + args.limit <= len(dataset)
@@ -80,6 +94,8 @@ def generate_text(self, task_name, intermediate_generations=None):
             intermediate_save_generations_path=intermediate_save_generations_path,
         )
 
+        if getattr(task, "bayes_at_n", False) and not generations:
+            raise ValueError("Bayes@N needs at least one problem with generations")
         if len(generations[0]) > self.args.n_samples:
             generations = [l[: self.args.n_samples] for l in generations]
             warnings.warn(

bigcode_eval/generation.py

@@ -56,6 +56,21 @@ def parallel_generations(
         # load generated code
         with open(args.load_generations_path) as fp:
             generations = json.load(fp)
+            if getattr(task, "bayes_at_n", False):
+                if (
+                    not isinstance(generations, list)
+                    or n_tasks < 1
+                    or len(generations) < n_tasks
+                    or any(
+                        not isinstance(row, list)
+                        or not row
+                        or not all(isinstance(program, str) for program in row)
+                        for row in generations[:n_tasks]
+                    )
+                ):
+                    raise ValueError(
+                        "Bayes@N needs a nonempty program list for every selected problem."
+                    )
             if accelerator.is_main_process:
                 print(
                     f"generations loaded, {n_tasks} selected from {len(generations)} with {len(generations[0])} candidates"

bigcode_eval/tasks/__init__.py

@@ -20,6 +20,7 @@
     **humanevalplus.create_all_tasks(),
     **humanevalpack.create_all_tasks(),
     "mbpp": mbpp.MBPP,
+    "mbpp-bayes-at-n": mbpp.MBPPBayesAtN,
     "mbppplus": mbppplus.MBPPPlus,
     "parity": parity.Parity,
     "python_bugs": python_bugs.PythonBugs,
@@ -43,6 +44,8 @@ def get_task(task_name, args=None):
             kwargs["prompt"] = args.prompt
         if "load_data_path" in inspect.signature(TASK_REGISTRY[task_name]).parameters:
             kwargs["load_data_path"] = args.load_data_path
+        if "n_samples" in inspect.signature(TASK_REGISTRY[task_name]).parameters:
+            kwargs["n_samples"] = getattr(args, "n_samples", None)
         return TASK_REGISTRY[task_name](**kwargs)
     except KeyError:
         print("Available tasks:")

bigcode_eval/tasks/custom_metrics/bayes_at_n.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Binary Bayes@N moments from https://arxiv.org/abs/2510.04265."""
+
+import math
+from numbers import Integral
+
+
+def estimate_bayes_at_n(num_samples, num_correct):
+    """Return posterior mean and SD of average problem success probability.
+
+    Each problem has an independent Beta(1, 1) prior. All problems must have
+    the same positive number of independently generated programs.
+
+    Args:
+        num_samples: Positive integer sample counts, one per problem.
+        num_correct: Integer passing-program counts, one per problem.
+
+    Returns:
+        A pair of floats, posterior mean and standard deviation on [0, 1].
+
+    Raises:
+        ValueError: Counts are empty, unequal, noninteger, or out of range.
+    """
+    if len(num_samples) == 0 or len(num_samples) != len(num_correct):
+        raise ValueError("Bayes@N needs matching, nonempty problem counts")
+    for n, c in zip(num_samples, num_correct):
+        if (
+            isinstance(n, bool)
+            or isinstance(c, bool)
+            or not isinstance(n, Integral)
+            or not isinstance(c, Integral)
+            or n <= 0
+            or not 0 <= c <= n
+        ):
+            raise ValueError("Bayes@N needs integer counts with N > 0 and 0 <= c <= N")
+    if any(n != num_samples[0] for n in num_samples):
+        raise ValueError("Bayes@N needs the same sample count for every problem")
+    n = int(num_samples[0])
+    total = n + 2
+    means = [(int(c) + 1) / total for c in num_correct]
+    variances = [
+        ((int(c) + 1) / total) * ((n - int(c) + 1) / total) / (total + 1)
+        for c in num_correct
+    ]
+    problems = len(num_samples)
+    return math.fsum(means) / problems, math.sqrt(math.fsum(variances)) / problems

bigcode_eval/tasks/custom_metrics/code_eval.py

@@ -20,9 +20,11 @@
 import os
 from collections import Counter, defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from numbers import Integral
 
 import numpy as np
 
+from .bayes_at_n import estimate_bayes_at_n
 from .execute import check_correctness
 
 
@@ -126,8 +128,57 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE."""
 
-def compute_code_eval(predictions, references, k=[1, 10, 100], num_workers=4, timeout=3.0):
-    """Returns the scores"""
+def compute_code_eval(
+    predictions,
+    references,
+    k=[1, 10, 100],
+    num_workers=4,
+    timeout=3.0,
+    bayes_at_n=False,
+    expected_n=None,
+):
+    """Evaluate programs and return aggregate scores plus per-program verdicts.
+
+    Args:
+        predictions: Lists of generated programs, one list per problem.
+        references: One test program per problem.
+        k: Sample budgets for pass@k.
+        num_workers: Number of concurrent program checks.
+        timeout: Timeout in seconds per program.
+        bayes_at_n: Add binary, uniform-prior posterior mean and SD fields.
+        expected_n: Required trials per problem when Bayes@N is enabled.
+            If omitted, infer N from the first row and require every row to match.
+
+    Returns:
+        The score dictionary and per-problem completion verdicts. Optional
+        fields are ``bayes_at_n`` and ``bayes_at_n_std`` on the [0, 1] scale.
+    """
+
+    if bayes_at_n:
+        if (
+            not isinstance(predictions, (list, tuple))
+            or not predictions
+            or not isinstance(references, (list, tuple))
+            or len(predictions) != len(references)
+        ):
+            raise ValueError("Bayes@N needs matching, nonempty predictions and references")
+        if expected_n is not None and (
+            isinstance(expected_n, bool)
+            or not isinstance(expected_n, Integral)
+            or expected_n < 1
+        ):
+            raise ValueError("Bayes@N expected_n must be a positive integer")
+        for candidates, reference in zip(predictions, references):
+            if (
+                not isinstance(candidates, (list, tuple))
+                or not candidates
+                or not all(isinstance(candidate, str) for candidate in candidates)
+                or not isinstance(reference, str)
+            ):
+                raise ValueError("Bayes@N needs nonempty program lists and string references")
+        expected_n = len(predictions[0]) if expected_n is None else expected_n
+        if any(len(row) != expected_n for row in predictions):
+            raise ValueError("Bayes@N needs exactly expected_n programs for every problem")
 
     if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1":
         raise ValueError(_WARNING)
@@ -155,7 +206,15 @@ def compute_code_eval(predictions, references, k=[1, 10, 100], num_workers=4, ti
             results[result["task_id"]].append((result["completion_id"], result))
 
     total, correct = [], []
+    if bayes_at_n and set(results) != set(range(len(predictions))):
+        raise ValueError("Bayes@N is missing evaluated problems")
     for result in results.values():
+        if bayes_at_n and (
+            len(result) != expected_n
+            or {item[0] for item in result} != set(range(expected_n))
+            or any(not isinstance(item[1].get("passed"), (bool, np.bool_)) for item in result)
+        ):
+            raise ValueError("Bayes@N needs a complete set of unique binary verdicts")
         result.sort()
         passed = [r[1]["passed"] for r in result]
         total.append(len(passed))
@@ -168,6 +227,10 @@ def compute_code_eval(predictions, references, k=[1, 10, 100], num_workers=4, ti
         ks = [ks]
     pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() for k in ks if (total >= k).all()}
 
+    if bayes_at_n:
+        mean, std = estimate_bayes_at_n(total, correct)
+        pass_at_k.update(bayes_at_n=mean, bayes_at_n_std=std)
+
     return pass_at_k, results
 
 

bigcode_eval/tasks/humaneval.py

@@ -29,13 +29,20 @@ def create_all_tasks():
     :return: {task_name: task}
         e.g. {multiple-py: Task, multiple-java: Task}
     """
-    return {"humaneval": create_task(True), "humaneval-unstripped": create_task(False)}
+    return {
+        "humaneval": create_task(True),
+        "humaneval-unstripped": create_task(False),
+        "humaneval-bayes-at-n": create_task(True, bayes_at_n=True),
+        "humaneval-unstripped-bayes-at-n": create_task(False, bayes_at_n=True),
+    }
 
 
-def create_task(strip_prompt):
+def create_task(strip_prompt, bayes_at_n=False):
     class HumanEval(GeneralHumanEval):
-        def __init__(self, **kwargs):
-            super().__init__(strip_prompt, **kwargs)
+        def __init__(self, n_samples=None, **kwargs):
+            super().__init__(
+                strip_prompt, bayes_at_n=bayes_at_n, n_samples=n_samples, **kwargs
+            )
 
     return HumanEval
 
@@ -47,7 +54,15 @@ class GeneralHumanEval(Task):
 
     DATASET_PATH = "openai_humaneval"
 
-    def __init__(self, strip_prompt, k=[1, 10, 100], num_workers=16, timeout=3.0):
+    def __init__(
+        self,
+        strip_prompt,
+        k=[1, 10, 100],
+        num_workers=16,
+        timeout=3.0,
+        bayes_at_n=False,
+        n_samples=None,
+    ):
         super().__init__(
             stop_words=["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif", "\n```", "<file_sep>"],
             requires_execution=True,
@@ -56,6 +71,8 @@ def __init__(self, strip_prompt, k=[1, 10, 100], num_workers=16, timeout=3.0):
         self.k = k
         self.num_workers = num_workers
         self.timeout = timeout
+        self.bayes_at_n = bayes_at_n
+        self.n_samples = n_samples
 
     def get_dataset(self):
         """Returns dataset for the task or an iterable of any object, that get_prompt can handle"""
@@ -101,5 +118,7 @@ def process_results(self, generations, references):
             k=self.k,
             num_workers=self.num_workers,
             timeout=self.timeout,
+            bayes_at_n=self.bayes_at_n,
+            expected_n=self.n_samples,
         )
         return results

bigcode_eval/tasks/mbpp.py

@@ -30,11 +30,13 @@ class MBPP(Task):
 
     DATASET_PATH = "mbpp"
 
-    def __init__(self):
+    def __init__(self, bayes_at_n=False, n_samples=None):
         super().__init__(
             stop_words=["\nclass", "\nassert", '\n"""', "\nprint", "\nif", "\n<|/", "\n```"],
             requires_execution=True,
         )
+        self.bayes_at_n = bayes_at_n
+        self.n_samples = n_samples
 
     def get_dataset(self):
         """Returns dataset for the task or an iterable of any object, that get_prompt can handle"""
@@ -82,5 +84,14 @@ def process_results(self, generations, references):
         results, _ = compute_code_eval(
             references=references,
             predictions=generations,
+            bayes_at_n=self.bayes_at_n,
+            expected_n=self.n_samples,
         )
         return results
+
+
+class MBPPBayesAtN(MBPP):
+    """MBPP with binary Bayes@N posterior summaries alongside pass@k."""
+
+    def __init__(self, n_samples=None):
+        super().__init__(bayes_at_n=True, n_samples=n_samples)

docs/README.md

@@ -554,3 +554,8 @@ The complete prompt with 8 shot examples (as used in [PAL](https://github.com/re
 ## How to add a new benchmark
 
 We welcome contributions to add new code benchmarks to this evaluation harness. You can find a step-by-step guide in [`guide.md`](https://github.com/bigcode-project/bigcode-evaluation-harness/blob/main/docs/guide.md).
+
+## Repeated-trial posterior summaries
+
+The [Bayes@N guide](bayes_at_n.md) describes optional posterior mean and standard
+deviation reporting for the HumanEval and MBPP tasks.

docs/bayes_at_n.md

@@ -0,0 +1,53 @@
+# Bayes@N for HumanEval and MBPP
+
+Use `humaneval-bayes-at-n`, `humaneval-unstripped-bayes-at-n`, or
+`mbpp-bayes-at-n` to report the posterior mean and posterior standard deviation
+of average problem success probability alongside pass@k. These task variants
+keep the original prompts, references, postprocessing and execution checks.
+
+For example, evaluate 16 previously sampled programs per problem:
+
+```sh
+accelerate launch main.py --tasks humaneval-bayes-at-n \
+    --load_generations_path generations.json --n_samples 16 \
+    --allow_code_execution --metric_output_path evaluation_results.json
+```
+
+The generation file uses the existing list-of-lists format, in dataset order.
+Set `--n_samples` to the intended number of retained programs per problem.
+The normal generation path can use the same task names. Use independent
+stochastic generations from a fixed model and sampling configuration.
+
+Each task result gains two fields on the [0, 1] scale: `bayes_at_n` is the
+posterior mean, and `bayes_at_n_std` is its posterior standard deviation.
+For two problems with outcomes `[[0, 0], [1, 1]]`, they are `0.5` and about
+`0.136931`. The original task names continue to report only their existing
+metrics.
+
+This implements the binary, uniform-prior case of
+[Bayes@N](https://arxiv.org/abs/2510.04265), Algorithm 1 and Appendix A. Each
+problem has its own Beta(1, 1) prior. If `c` of `N` programs pass the entire test
+reference, its posterior is Beta(`c + 1`, `N - c + 1`). The corpus mean averages
+problem posterior means. The corpus SD is `sqrt(sum(problem variances)) / M`,
+where `M` is the number of selected problems. SD measures uncertainty about the
+latent success probability of this fixed problem set, rather than variation
+across benchmark problems or uncertainty in pass@k.
+
+Every selected problem needs the same positive number of programs and a test
+reference. Invalid or incomplete inputs raise `ValueError` before execution;
+missing or duplicate execution verdicts also fail. Programs that fail tests or
+time out are failures under the existing evaluator's grading rules. Historical
+observations, categorical rubrics and credible intervals are outside this
+integration's scope.
+
+For evaluation-only files, use `--limit_start 0`. The existing loader does not
+apply an offset to saved generation rows, so these variants reject a nonzero
+offset to avoid pairing programs with the wrong references. Device-padding
+candidates are cropped to `--n_samples` by the existing generation pipeline.
+Intermediate checkpoints are unsupported for these variants: the existing
+resume path compacts missing rows and cannot preserve their question identity.
+Use a complete saved-generation file instead.
+The standalone `compute_code_eval(..., bayes_at_n=True, expected_n=N)` API also
+supports the calculation; if `expected_n` is omitted, it infers a fixed N from
+the first row. Callers must retain every intended problem. A standalone matrix
+cannot reveal a problem omitted entirely.