MCPcopy Create free account
hub / github.com/PRIME-RL/PRIME / estimate_pass_at_k

Function estimate_pass_at_k

eval/utils/evaluation_leetcode.py:85–108  ·  view source on GitHub ↗

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
)

Source from the content-addressed store, hash-verified

83 return dataset
84
85def estimate_pass_at_k(
86 num_samples: Union[int, List[int], np.ndarray],
87 num_correct: Union[List[int], np.ndarray],
88 k: int
89) -> np.ndarray:
90 """
91 Estimates pass@k of each problem and returns them in an array.
92 """
93
94 def estimator(n: int, c: int, k: int) -> float:
95 """
96 Calculates 1 - comb(n - c, k) / comb(n, k).
97 """
98 if n - c < k:
99 return 1.0
100 return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
101
102 if isinstance(num_samples, int):
103 num_samples_it = itertools.repeat(num_samples, len(num_correct))
104 else:
105 assert len(num_samples) == len(num_correct)
106 num_samples_it = iter(num_samples)
107
108 return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
109
110def process_humaneval_test(sample, problems, example_test=False, is_mbpp=False, language="python"):
111 task_id = sample["task_id"]

Callers 1

Calls 1

estimatorFunction · 0.70

Tested by

no test coverage detected