Estimates pass@k of each problem and returns them in an array.
(
num_samples: Union[int, List[int], np.ndarray],
num_correct: Union[List[int], np.ndarray],
k: int
)
| 25 | |
| 26 | |
| 27 | def estimate_pass_at_k( |
| 28 | num_samples: Union[int, List[int], np.ndarray], |
| 29 | num_correct: Union[List[int], np.ndarray], |
| 30 | k: int |
| 31 | ) -> np.ndarray: |
| 32 | """ |
| 33 | Estimates pass@k of each problem and returns them in an array. |
| 34 | """ |
| 35 | |
| 36 | def estimator(n: int, c: int, k: int) -> float: |
| 37 | """ |
| 38 | Calculates 1 - comb(n - c, k) / comb(n, k). |
| 39 | """ |
| 40 | if n - c < k: |
| 41 | return 1.0 |
| 42 | return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) |
| 43 | |
| 44 | if isinstance(num_samples, int): |
| 45 | num_samples_it = itertools.repeat(num_samples, len(num_correct)) |
| 46 | else: |
| 47 | assert len(num_samples) == len(num_correct) |
| 48 | num_samples_it = iter(num_samples) |
| 49 | |
| 50 | return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]) |
no test coverage detected