mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
|
|
def estimator(num_samples: int, num_correct: int, k: int) -> float:
|
|
"""Unbiased Pass@k estimator: 1 - C(n-c, k) / C(n, k).
|
|
|
|
(The Codex/HumanEval estimator.)
|
|
* for k=1 it degenerates to c/n, the mean per-sample accuracy, i.e. the usual avg@n / Pass@1;
|
|
* when n-c < k (too few wrong samples for a k-subset to be all-correct) it returns 1.0.
|
|
"""
|
|
assert 0 <= num_correct <= num_samples, f"invalid num_correct={num_correct} num_samples={num_samples}"
|
|
assert 0 <= k <= num_samples, f"invalid k={k} num_samples={num_samples}"
|
|
if num_samples - num_correct < k:
|
|
return 1.0
|
|
return float(1.0 - math.prod(1.0 - k / i for i in range(num_samples - num_correct + 1, num_samples + 1)))
|
|
|
|
|
|
def pass_at_k_by_task(
|
|
results: list[dict[str, Any]],
|
|
k_list: list[int],
|
|
*,
|
|
task_key: str = "task_id",
|
|
correct_key: str = "correct",
|
|
) -> dict[str, float]:
|
|
"""Mean Pass@k over tasks, with tasks grouped by `task_key`.
|
|
|
|
When a task has fewer than k samples (judging failure, truncated sampling), approximate with
|
|
min(k, n) instead of dropping the task: dropping it would compute Pass@k only over the "clean"
|
|
tasks, and those are usually the easier ones, which biases the result upward.
|
|
"""
|
|
grouped: dict[str, list[bool]] = defaultdict(list)
|
|
for row in results:
|
|
grouped[row[task_key]].append(bool(row[correct_key]))
|
|
|
|
metrics: dict[str, float] = {}
|
|
for k in k_list:
|
|
per_task = [estimator(len(flags), sum(flags), min(k, len(flags))) for flags in grouped.values() if flags]
|
|
metrics[f"Pass@{k}"] = sum(per_task) / len(per_task) if per_task else 0.0
|
|
return metrics
|
|
|
|
|
|
def default_k_list(n_sample: int) -> list[int]:
|
|
"""Default k values: 1, 2, 4, ... (powers of two), with n_sample appended."""
|
|
k_list = [2**i for i in range(max(math.ceil(math.log2(n_sample)), 1))]
|
|
if n_sample not in k_list:
|
|
k_list.append(n_sample)
|
|
return k_list
|
|
|
|
|
|
def summarize(
|
|
results: list[dict[str, Any]],
|
|
*,
|
|
n_sample: int,
|
|
k_list: list[int] | None = None,
|
|
task_key: str = "task_id",
|
|
extra: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""One rollout per row: {task_id, correct, ...}."""
|
|
k_list = k_list or default_k_list(n_sample)
|
|
grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
for row in results:
|
|
grouped[row[task_key]].append(row)
|
|
|
|
summary: dict[str, Any] = {
|
|
"num_tasks": len(grouped),
|
|
"num_rollouts": len(results),
|
|
"n_sample": n_sample,
|
|
"score": sum(1 for row in results if row["correct"]) / len(results) if results else 0.0,
|
|
"task_accuracy": {
|
|
task_id: sum(1 for row in task_rows if row["correct"]) / len(task_rows) for task_id, task_rows in sorted(grouped.items())
|
|
},
|
|
}
|
|
summary.update(pass_at_k_by_task(results, k_list, task_key=task_key))
|
|
if extra:
|
|
summary.update(extra)
|
|
return summary
|
|
|
|
|
|
def print_report(summary: dict[str, Any], *, title: str = "") -> None:
|
|
print("\n" + "=" * 72)
|
|
if title:
|
|
print(title)
|
|
print(f"tasks={summary['num_tasks']} rollouts={summary['num_rollouts']} n_sample={summary['n_sample']}")
|
|
pass_keys = sorted(
|
|
(metric for metric in summary if metric.startswith("Pass@")),
|
|
key=lambda k: int(k.split("@")[1]),
|
|
)
|
|
for key in pass_keys:
|
|
print(f" {key:<10s} = {summary[key] * 100:6.2f}")
|
|
note = "" if abs(summary["score"] - summary.get("Pass@1", 0.0)) < 1e-9 else " (ragged sample counts, differs from Pass@1)"
|
|
print(
|
|
f" {'score':<10s} = {summary['score'] * 100:6.2f} (= mean per-sample accuracy, equals Pass@1 when sample counts are even){note}"
|
|
)
|
|
print("-" * 72)
|
|
print("per-task accuracy:")
|
|
for tid, acc in summary["task_accuracy"].items():
|
|
bar = "#" * int(round(acc * 20))
|
|
print(f" {tid:<20s} {acc * 100:6.2f} {bar}")
|
|
print("=" * 72)
|